import logging
import time


from .message_define import MyMessage
from .utils import transform_list_to_tensor
from ....core.distributed.fedml_comm_manager import FedMLCommManager
from ....core.distributed.communication.message import Message


class FedNovaClientManager(FedMLCommManager):
    def __init__(
        self,
        args,
        trainer,
        comm=None,
        rank=0,
        size=0,
        backend="MPI",
    ):
        super().__init__(args, comm, rank, size, backend)
        self.trainer = trainer
        self.num_rounds = args.comm_round
        self.round_idx = 0
        self.worker_id = self.rank - 1

    def run(self):
        super().run()

    def register_message_receive_handlers(self):
        self.register_message_receive_handler(
            MyMessage.MSG_TYPE_S2C_INIT_CONFIG, self.handle_message_init
        )
        self.register_message_receive_handler(
            MyMessage.MSG_TYPE_S2C_SYNC_MODEL_TO_CLIENT,
            self.handle_message_receive_model_from_server,
        )

    def handle_message_init(self, msg_params):
        global_model_params = msg_params.get(MyMessage.MSG_ARG_KEY_MODEL_PARAMS)
        # client_index = msg_params.get(MyMessage.MSG_ARG_KEY_CLIENT_INDEX)

        average_weight_dict = msg_params.get(MyMessage.MSG_ARG_KEY_AVG_WEIGHTS)
        client_schedule = msg_params.get(MyMessage.MSG_ARG_KEY_CLIENT_SCHEDULE)
        client_indexes = client_schedule[self.worker_id]

        self.round_idx = 0
        self.__train(global_model_params, client_indexes, average_weight_dict)

    def start_training(self):
        self.round_idx = 0
        # self.__train()

    def handle_message_receive_model_from_server(self, msg_params):
        logging.info("handle_message_receive_model_from_server.")
        global_model_params = msg_params.get(MyMessage.MSG_ARG_KEY_MODEL_PARAMS)
        # client_index = msg_params.get(MyMessage.MSG_ARG_KEY_CLIENT_INDEX)

        average_weight_dict = msg_params.get(MyMessage.MSG_ARG_KEY_AVG_WEIGHTS)
        client_schedule = msg_params.get(MyMessage.MSG_ARG_KEY_CLIENT_SCHEDULE)
        client_indexes = client_schedule[self.worker_id]

        self.round_idx += 1
        self.__train(global_model_params, client_indexes, average_weight_dict)
        if self.round_idx == self.num_rounds - 1:
            # post_complete_message_to_sweep_process(self.args)
            self.finish()


    def send_result_to_server(self, receive_id, weights, client_runtime_info):
        message = Message(
            MyMessage.MSG_TYPE_C2S_SEND_MODEL_TO_SERVER,
            self.get_sender_id(),
            receive_id,
        )
        message.add_params(MyMessage.MSG_ARG_KEY_MODEL_PARAMS, weights)
        # message.add_params(MyMessage.MSG_ARG_KEY_NUM_SAMPLES, local_sample_num)
        message.add_params(MyMessage.MSG_ARG_KEY_CLIENT_RUNTIME_INFO, client_runtime_info)
        self.send_message(message)


    def add_client_model(self, local_agg_model_params, client_index, grad, t_eff, weight=1.0):
        # Add params that needed to be reduces from clients
        # for name, param in model_params.items():
        #     if name not in local_agg_model_params:
        #         local_agg_model_params[name] = param * weight
        #     else:
        #         local_agg_model_params[name] += param * weight
        # local_agg_model_params[client_index]["grad"] = grad
        # local_agg_model_params[client_index]["t_eff"] = t_eff
        local_agg_model_params.append({
            "grad": grad, "t_eff": t_eff,
        })


    def __train(self, global_model_params, client_indexes, average_weight_dict):
        logging.info("#######training########### round_id = %d" % self.round_idx)

        # local_agg_model_params = {}
        local_agg_model_params = []
        client_runtime_info = {}
        for client_index in client_indexes:
            logging.info("#######training########### Simulating client_index = %d, average weight: %f " % \
                (client_index, average_weight_dict[client_index]))
            start_time = time.time()
            self.trainer.update_model(global_model_params)
            self.trainer.update_dataset(int(client_index))
            # weights, local_sample_num = self.trainer.train(self.round_idx)
            loss, grad, t_eff = self.trainer.train(self.round_idx)
            self.add_client_model(local_agg_model_params, client_index, grad, t_eff,
                                weight=average_weight_dict[client_index])

            end_time = time.time()
            client_runtime = end_time - start_time
            client_runtime_info[client_index] = client_runtime
            logging.info("#######training########### End Simulating client_index = %d, consuming time: %f" % \
                (client_index, client_runtime))
        self.send_result_to_server(0, local_agg_model_params, client_runtime_info)











