import copy
import logging
import random
import time

import numpy as np
import torch
import wandb

from .utils import transform_list_to_tensor


class FedAVGAggregator(object):

    def __init__(self, train_global, test_global, all_train_data_num,
                 train_data_local_dict, test_data_local_dict, train_data_local_num_dict, worker_num, device,
                 args, model_trainer):
        self.trainer = model_trainer

        self.args = args
        self.train_global = train_global
        self.test_global = test_global
        self.val_global = self._generate_validation_set()
        self.all_train_data_num = all_train_data_num

        self.train_data_local_dict = train_data_local_dict
        self.test_data_local_dict = test_data_local_dict
        self.train_data_local_num_dict = train_data_local_num_dict
        # to take avg of all clients
        self.test_acc_all_clients = dict()


        self.worker_num = worker_num
        self.device = device
        self.model_dict = dict()


        self.sample_num_dict = dict()
        self.averaged_loss_dict = dict()

        # for personalization
        self.per_acc_dict = dict()
        self.best_acc_prev_round = 0.0
        self.acc_avg_personalized_model_on_all_clients = 0.0

        self.flag_client_model_uploaded_dict = dict()
        for idx in range(self.worker_num):
            self.flag_client_model_uploaded_dict[idx] = False

    def get_global_model_params(self):
        return self.trainer.get_model_params()

    def set_global_model_params(self, model_parameters):
        self.trainer.set_model_params(model_parameters)

    def add_local_trained_result(self, index, model_params, sample_num, averaged_loss, personalized_accuracy, client_index, round_idx):
        logging.info("add_model. index = %d" % index)
        self.model_dict[index] = model_params
        self.sample_num_dict[index] = sample_num

        self.averaged_loss_dict[index] = averaged_loss

        if round_idx % self.args.frequency_of_the_test == 0:
            # for personalization
            self.per_acc_dict[index] = personalized_accuracy
            # to keep record of all clients
            self.test_acc_all_clients[client_index] = personalized_accuracy

        self.flag_client_model_uploaded_dict[index] = True

    def check_whether_all_receive(self):
        for idx in range(self.worker_num):
            if not self.flag_client_model_uploaded_dict[idx]:
                return False
        for idx in range(self.worker_num):
            self.flag_client_model_uploaded_dict[idx] = False
        return True

    def aggregate(self):
        start_time = time.time()
        model_list = []
        training_num = 0

        for idx in range(self.worker_num):
            if self.args.is_mobile == 1:
                self.model_dict[idx] = transform_list_to_tensor(self.model_dict[idx])
            model_list.append((self.sample_num_dict[idx], self.model_dict[idx]))
            training_num += self.sample_num_dict[idx]

        logging.info("len of self.model_dict[idx] = " + str(len(self.model_dict)))

        # logging.info("################aggregate: %d" % len(model_list))
        (num0, averaged_params) = model_list[0]
        for k in averaged_params.keys():
            for i in range(0, len(model_list)):
                local_sample_number, local_model_params = model_list[i]
                w = local_sample_number / training_num
                if i == 0:
                    averaged_params[k] = local_model_params[k] * w
                else:
                    averaged_params[k] += local_model_params[k] * w

        # update the global model which is cached at the server side
        self.set_global_model_params(averaged_params)

        end_time = time.time()
        logging.info("aggregate time cost: %d" % (end_time - start_time))
        return averaged_params

    def client_sampling(self, round_idx, client_num_in_total, client_num_per_round):
        if client_num_in_total == client_num_per_round:
            client_indexes = [client_index for client_index in range(client_num_in_total)]
        else:
            num_clients = min(client_num_per_round, client_num_in_total)
            np.random.seed(round_idx)  # make sure for each comparison, we are selecting the same clients each round
            client_indexes = np.random.choice(range(client_num_in_total), num_clients, replace=False)
        logging.info("client_indexes = %s" % str(client_indexes))
        return client_indexes

    def _generate_validation_set(self, num_samples=10000):
        if self.args.dataset.startswith("stackoverflow"):
            test_data_num = len(self.test_global.dataset)
            sample_indices = random.sample(range(test_data_num), min(num_samples, test_data_num))
            subset = torch.utils.data.Subset(self.test_global.dataset, sample_indices)
            sample_testset = torch.utils.data.DataLoader(subset, batch_size=self.args.batch_size)
            return sample_testset
        else:
            return self.test_global

    def test_on_server_for_all_clients(self, round_idx):
        if self.trainer.test_on_the_server(self.train_data_local_dict, self.test_data_local_dict, self.device,
                                           self.args):
            return

        if round_idx % self.args.frequency_of_the_test == 0 or round_idx == self.args.comm_round - 1:
            logging.info("################test_on_server_for_all_clients. Round = {}".format(round_idx))

            """
                personalization
            """
            # averaged training loss
            averaged_training_loss = 0.0
            for idx in range(self.worker_num):
                averaged_training_loss += self.averaged_loss_dict[idx]
            averaged_training_loss = averaged_training_loss / self.worker_num
            logging.info("averaged_training_loss = %f" % averaged_training_loss)
            wandb.log({"SSL-Train/Loss": averaged_training_loss, "round": round_idx})

            # averaged personalized accuracy
            averaged_personalized_acc = 0.0
            for idx in range(self.worker_num):
                averaged_personalized_acc += self.per_acc_dict[idx]
            averaged_personalized_acc = averaged_personalized_acc / self.worker_num
            if averaged_personalized_acc > 0.0:
                output_json = {"Personalized Model Test/Averaged Acc": averaged_personalized_acc, "round": round_idx}
                wandb.log(output_json)
                logging.info(output_json)

            # test the global model on each client's local dataset
            test_num_samples = []
            test_tot_corrects = []
            test_losses = []

            test_acc_list = []
            for client_idx in self.test_data_local_dict.keys():
                test_data = self.test_data_local_dict[client_idx]
                metrics = self.trainer.test(test_data, self.device, self.args)

                test_tot_correct, test_num_sample, test_loss = metrics['test_correct'], metrics['test_total'], metrics[
                    'test_loss']
                test_tot_corrects.append(copy.deepcopy(test_tot_correct))
                test_num_samples.append(copy.deepcopy(test_num_sample))
                test_losses.append(copy.deepcopy(test_loss))

                # test on test dataset
                test_acc = sum(test_tot_corrects) / sum(test_num_samples)
                test_acc_list.append(test_acc)
            averaged_acc = sum(test_acc_list) / len(test_acc_list)
            output_json = {"Global Model Test/Averaged Acc": averaged_acc, "round": round_idx}
            wandb.log(output_json)
            logging.info(output_json)

    def local_statistics(self, round_idx):
        # Average validation accuracy (personalized) of all client's
        val_acc_list_all_clients = self.test_acc_all_clients.values()
        logging.info("Client accuracies of all clients "+str(val_acc_list_all_clients))
        self.acc_avg_personalized_model_on_all_clients = sum(val_acc_list_all_clients) / len(val_acc_list_all_clients)
        logging.info('Round {:3d}, Averaged Validation Accuracy on Clients of all round {:.3f}'.format(round_idx,
                                                                                                           self.acc_avg_personalized_model_on_all_clients))
        wandb.log(
            {"Averaged Validation Accuracy (All Clients)": self.acc_avg_personalized_model_on_all_clients,
             "Round": round_idx})


        # Average validation accuracy (personalized) of this round's client's
        val_acc_list_this_round = self.per_acc_dict.values()
        logging.info("Client accuracies of this round "+str(val_acc_list_this_round))
        avg_valid_acc_this_round = sum(val_acc_list_this_round) / len(val_acc_list_this_round)
        logging.info('Round {:3d}, Averaged Validation Accuracy on Clients of this round {:.3f}'.format(round_idx,
                                                                                                           avg_valid_acc_this_round))
        wandb.log(
            {"Averaged Validation Accuracy (clients of this round) ": avg_valid_acc_this_round,
             "Round": round_idx})

        if self.acc_avg_personalized_model_on_all_clients >= self.best_acc_prev_round:
            self.best_acc_prev_round = self.acc_avg_personalized_model_on_all_clients
            wandb_table_p = wandb.Table(columns=["Client's Validation Accuracy of Best round (all Clients)"])
            wandb_table_p.add_data(str(self.test_acc_all_clients))
            wandb.log({"Validation Accuracy (Personalized Model (all clients)": wandb_table_p})

        wandb_table_pe = wandb.Table(columns=["Client's Validation Accuracy of Best round (this round)"])
        wandb_table_pe.add_data(str(self.per_acc_dict))
        wandb.log({"Validation Accuracy (Personalized Model (this round's clients))": wandb_table_pe})

