import logging

from fedml_core import RunConfig
from fedml_core.distributed.communication.message import Message
from fedml_core.distributed.server.server_manager import ServerManager
from .message_define import MyMessage
from .utils import transform_tensor_to_list


class FedAVGServerManager(ServerManager):
    def __init__(self, aggregator, config: RunConfig, comm=None, rank=0, size=0, is_preprocessed=False,
                 preprocessed_client_lists=None, grpc_ipconfig_path=None, trpc_master_config_path=None):
        super().__init__(config, comm, rank, size, grpc_ipconfig_path, trpc_master_config_path)
        self.aggregator = aggregator
        self.config = config
        self.round_idx = 0
        self.is_preprocessed = is_preprocessed
        self.preprocessed_client_lists = preprocessed_client_lists

    def send_init_msg(self):
        # sampling clients
        client_indexes = self.aggregator.client_sampling(self.round_idx, self.config.client_num_in_total,
                                                         self.config.client_num_per_round)
        global_model_params = self.aggregator.get_global_model_params()
        if self.config.is_mobile:
            global_model_params = transform_tensor_to_list(global_model_params)
        for process_id in range(1, self.size):
            self.send_message_init_config(process_id, global_model_params, client_indexes[process_id - 1])

    def register_message_receive_handlers(self):
        self.register_message_receive_handler(MyMessage.MSG_TYPE_C2S_SEND_MODEL_TO_SERVER,
                                              self.handle_message_receive_model_from_client)

    def handle_message_receive_model_from_client(self, msg_params):
        sender_id = msg_params.get(MyMessage.MSG_ARG_KEY_SENDER)
        model_params = msg_params.get(MyMessage.MSG_ARG_KEY_MODEL_PARAMS)
        local_sample_number = msg_params.get(MyMessage.MSG_ARG_KEY_NUM_SAMPLES)

        self.aggregator.add_local_trained_result(sender_id - 1, model_params, local_sample_number)
        b_all_received = self.aggregator.check_whether_all_receive()
        logging.info(f"b_all_received = {b_all_received}")
        if b_all_received:
            global_model_params = self.aggregator.aggregate()
            self.aggregator.test_on_server_for_all_clients(self.round_idx)

            # start the next round
            self.round_idx += 1
            if self.round_idx == self.config.comm_round:
                # post_complete_message_to_sweep_process(self.args)
                self.finish()
                logging.info('Finishing server manager')
                return
            if self.is_preprocessed:
                if self.preprocessed_client_lists is None:
                    # sampling has already been done in data preprocessor
                    client_indexes = [self.round_idx] * self.config.client_num_per_round
                else:
                    client_indexes = self.preprocessed_client_lists[self.round_idx]
            else:
                # sampling clients
                client_indexes = self.aggregator.client_sampling(self.round_idx, self.config.client_num_in_total,
                                                                 self.config.client_num_per_round)

            logging.info(f'indexes of clients: {client_indexes}')
            logging.info(f'size = {self.size:d}')
            if self.config.is_mobile:
                global_model_params = transform_tensor_to_list(global_model_params)

            for receiver_id in range(1, self.size):
                self.send_message_sync_model_to_client(receiver_id, global_model_params,
                                                       client_indexes[receiver_id - 1])

    def send_message_init_config(self, receive_id, global_model_params, client_index):
        message = Message(MyMessage.MSG_TYPE_S2C_INIT_CONFIG, self.get_sender_id(), receive_id)
        message.add_params(MyMessage.MSG_ARG_KEY_MODEL_PARAMS, global_model_params)
        message.add_params(MyMessage.MSG_ARG_KEY_CLIENT_INDEX, str(client_index))
        self.send_message(message)

    def send_message_sync_model_to_client(self, receive_id, global_model_params, client_index):
        logging.info(f"send_message_sync_model_to_client. receive_id = {receive_id:d}")
        message = Message(MyMessage.MSG_TYPE_S2C_SYNC_MODEL_TO_CLIENT, self.get_sender_id(), receive_id)
        message.add_params(MyMessage.MSG_ARG_KEY_MODEL_PARAMS, global_model_params)
        message.add_params(MyMessage.MSG_ARG_KEY_CLIENT_INDEX, str(client_index))
        self.send_message(message)

import logging

from fedml_core import RunConfig
from fedml_core.distributed.communication.message import Message
from fedml_core.distributed.server.server_manager import ServerManager
from .message_define import MyMessage
from .utils import transform_tensor_to_list


class FedAVGServerManager(ServerManager):
    def __init__(self, aggregator, config: RunConfig, comm=None, rank=0, size=0, is_preprocessed=False,
                 preprocessed_client_lists=None, grpc_ipconfig_path=None, trpc_master_config_path=None):
        super().__init__(config, comm, rank, size, grpc_ipconfig_path, trpc_master_config_path)
        self.aggregator = aggregator
        self.config = config
        self.round_idx = 0
        self.is_preprocessed = is_preprocessed
        self.preprocessed_client_lists = preprocessed_client_lists

    def send_init_msg(self):
        # sampling clients
        client_indexes = self.aggregator.client_sampling(self.round_idx, self.config.client_num_in_total,
                                                         self.config.client_num_per_round)
        global_model_params = self.aggregator.get_global_model_params()
        if self.config.is_mobile:
            global_model_params = transform_tensor_to_list(global_model_params)
        for process_id in range(1, self.size):
            self.send_message_init_config(process_id, global_model_params, client_indexes[process_id - 1])

    def register_message_receive_handlers(self):
        self.register_message_receive_handler(MyMessage.MSG_TYPE_C2S_SEND_MODEL_TO_SERVER,
                                              self.handle_message_receive_model_from_client)

    def handle_message_receive_model_from_client(self, msg_params):
        sender_id = msg_params.get(MyMessage.MSG_ARG_KEY_SENDER)
        model_params = msg_params.get(MyMessage.MSG_ARG_KEY_MODEL_PARAMS)
        local_sample_number = msg_params.get(MyMessage.MSG_ARG_KEY_NUM_SAMPLES)

        self.aggregator.add_local_trained_result(sender_id - 1, model_params, local_sample_number)
        b_all_received = self.aggregator.check_whether_all_receive()
        logging.info(f"b_all_received = {b_all_received}")
        if b_all_received:
            global_model_params = self.aggregator.aggregate()
            self.aggregator.test_on_server_for_all_clients(self.round_idx)

            # start the next round
            self.round_idx += 1
            if self.round_idx == self.config.comm_round:
                # post_complete_message_to_sweep_process(self.args)
                self.finish()
                logging.info('Finishing server manager')
                return
            if self.is_preprocessed:
                if self.preprocessed_client_lists is None:
                    # sampling has already been done in data preprocessor
                    client_indexes = [self.round_idx] * self.config.client_num_per_round
                else:
                    client_indexes = self.preprocessed_client_lists[self.round_idx]
            else:
                # sampling clients
                client_indexes = self.aggregator.client_sampling(self.round_idx, self.config.client_num_in_total,
                                                                 self.config.client_num_per_round)

            logging.info(f'indexes of clients: {client_indexes}')
            logging.info(f'size = {self.size:d}')
            if self.config.is_mobile:
                global_model_params = transform_tensor_to_list(global_model_params)

            for receiver_id in range(1, self.size):
                self.send_message_sync_model_to_client(receiver_id, global_model_params,
                                                       client_indexes[receiver_id - 1])

    def send_message_init_config(self, receive_id, global_model_params, client_index):
        message = Message(MyMessage.MSG_TYPE_S2C_INIT_CONFIG, self.get_sender_id(), receive_id)
        message.add_params(MyMessage.MSG_ARG_KEY_MODEL_PARAMS, global_model_params)
        message.add_params(MyMessage.MSG_ARG_KEY_CLIENT_INDEX, str(client_index))
        self.send_message(message)

    def send_message_sync_model_to_client(self, receive_id, global_model_params, client_index):
        logging.info(f"send_message_sync_model_to_client. receive_id = {receive_id:d}")
        message = Message(MyMessage.MSG_TYPE_S2C_SYNC_MODEL_TO_CLIENT, self.get_sender_id(), receive_id)
        message.add_params(MyMessage.MSG_ARG_KEY_MODEL_PARAMS, global_model_params)
        message.add_params(MyMessage.MSG_ARG_KEY_CLIENT_INDEX, str(client_index))
        self.send_message(message)

# sampling clients # start the next round # post_complete_message_to_sweep_process(self.args) # sampling has already been done in data preprocessor # sampling clients