import copy
import logging
import random
from typing import List, Tuple

import numpy as np
import torch
import wandb
from torch.utils.data import ConcatDataset

from fedml_api.standalone.fedavg.my_model_trainer import MyModelTrainer
from fedml_api.standalone.federated_sgan.ac_gan_model_trainer import ACGANModelTrainer
from fedml_api.standalone.federated_sgan.client import FedSSGANClient
from fedml_api.standalone.federated_sgan.model_trainer import FedSSGANModelTrainer
from fedml_api.standalone.utils.HeterogeneousModelBaseTrainerAPI import HeterogeneousModelBaseTrainerAPI


class FedSSGANAPI(HeterogeneousModelBaseTrainerAPI):
    def __init__(self, dataset, device, args, adapter_model, client_models: List[Tuple[torch.nn.Module, int]]):
        """
        Args:
            dataset: Dataset presplit into data loaders
            device: Device to run training on
            args: Additional args
            client_models: List of client models and their frequency participating (assuming a stateful algorithm for simplicity)
        """
        super().__init__(dataset, device, args)

        self.global_model = MyModelTrainer(adapter_model)

        self._setup_clients(self.train_data_local_num_dict, self.train_data_local_dict, self.test_data_local_dict,
                            client_models)

        self._plot_client_training_data_distribution()

    def _setup_clients(self, train_data_local_num_dict, train_data_local_dict, test_data_local_dict,
                       client_models):
        logging.info("############setup_clients (START)#############")

        c_idx = 0
        for local_model, freq in client_models:
            for i in range(freq):
                model_trainer = ACGANModelTrainer(
                    copy.deepcopy(self.global_model.model),
                    copy.deepcopy(local_model)
                )
                c = FedSSGANClient(c_idx, train_data_local_dict[c_idx], test_data_local_dict[c_idx],
                                   train_data_local_num_dict[c_idx], self.test_global, self.args, self.device,
                                   model_trainer)
                c_idx += 1
                self.client_list.append(c)

        logging.info("############setup_clients (END)#############")

    def train(self):
        logging.info('\n###############Pre-Training clients#############\n')
        for i, c in enumerate(self.client_list):
            logging.info(f'Pre=training client: {i}')
            c.pre_train()
        logging.info('###############Pre-Training clients (END)###########\n')

        unlabelled_synthesised_data = None
        w_global = self.global_model.get_model_params()
        for round_idx in range(self.args.comm_round):

            logging.info("################Communication round : {}".format(round_idx))

            w_locals = []
            synthesised_data_locals = []
            client_synthesised_data_lens = {'round': round_idx}

            client: FedSSGANClient
            for idx, client in enumerate(self.client_list):
                # Update client synthetic datasets
                # client.set_synthetic_dataset(unlabelled_synthesised_data)

                # Local round
                w = client.train(copy.deepcopy(w_global), round_idx)
                # self.logger.info("local weights = " + str(w))
                w_locals.append((client.get_sample_number(), copy.deepcopy(w)))

            #     synthetic_data = client.generate_synthetic_dataset()
            #     if synthetic_data is not None:
            #         synthesised_data_locals.append(synthetic_data)
            #         client_synthesised_data_lens[f'Client_{idx}: Synthetic Dataset Size'] = len(synthetic_data)
            #     else:
            #         client_synthesised_data_lens[f'Client_{idx}: Synthetic Dataset Size'] = 0
            #
            # if len(synthesised_data_locals) > 0:
            #     unlabelled_synthesised_data = ConcatDataset(synthesised_data_locals)
            #     logging.info(f'\n Synthetic Unlabelled Dataset Size: {len(unlabelled_synthesised_data)}\n')
            #     client_synthesised_data_lens['Total Synthetic Dataset Size'] = len(unlabelled_synthesised_data)
            # else:
            #     unlabelled_synthesised_data = None
            #     client_synthesised_data_lens['Total Synthetic Dataset Size'] = 0

            # wandb.log(client_synthesised_data_lens)

            # update global weights
            w_global = self._aggregate(w_locals)
            self.global_model.set_model_params(w_global)

            # test results
            # at last round
            if round_idx == self.args.comm_round - 1:
                self._local_test_on_all_clients(round_idx)
            # per {frequency_of_the_test} round
            elif round_idx % self.args.frequency_of_the_test == 0:
                if self.args.dataset.startswith("stackoverflow"):
                    self._local_test_on_validation_set(round_idx)
                else:
                    self._local_test_on_all_clients(round_idx)

[((29, 28, 29, 57), 'fedml_api.standalone.fedavg.my_model_trainer.MyModelTrainer', 'MyModelTrainer', ({(29, 43, 29, 56): 'adapter_model'}, {}), '(adapter_model)', False, 'from fedml_api.standalone.fedavg.my_model_trainer import MyModelTrainer\n'), ((38, 8, 38, 70), 'logging.info', 'logging.info', ({(38, 21, 38, 69): '"""############setup_clients (START)#############"""'}, {}), "('############setup_clients (START)#############')", False, 'import logging\n'), ((53, 8, 53, 68), 'logging.info', 'logging.info', ({(53, 21, 53, 67): '"""############setup_clients (END)#############"""'}, {}), "('############setup_clients (END)#############')", False, 'import logging\n'), ((56, 8, 56, 76), 'logging.info', 'logging.info', ({(56, 21, 56, 75): '"""\n###############Pre-Training clients#############\n"""'}, {}), '("""\n###############Pre-Training clients#############\n""")', False, 'import logging\n'), ((60, 8, 60, 78), 'logging.info', 'logging.info', ({(60, 21, 60, 77): '"""###############Pre-Training clients (END)###########\n"""'}, {}), "('###############Pre-Training clients (END)###########\\n')", False, 'import logging\n'), ((58, 12, 58, 53), 'logging.info', 'logging.info', ({(58, 25, 58, 52): 'f"""Pre=training client: {i}"""'}, {}), "(f'Pre=training client: {i}')", False, 'import logging\n'), ((47, 20, 49, 49), 'fedml_api.standalone.federated_sgan.client.FedSSGANClient', 'FedSSGANClient', ({(47, 35, 47, 40): 'c_idx', (47, 42, 47, 70): 'train_data_local_dict[c_idx]', (47, 72, 47, 99): 'test_data_local_dict[c_idx]', (48, 35, 48, 67): 'train_data_local_num_dict[c_idx]', (48, 69, 48, 85): 'self.test_global', (48, 87, 48, 96): 'self.args', (48, 98, 48, 109): 'self.device', (49, 35, 49, 48): 'model_trainer'}, {}), '(c_idx, train_data_local_dict[c_idx], test_data_local_dict[\n    c_idx], train_data_local_num_dict[c_idx], self.test_global, self.args,\n    self.device, model_trainer)', False, 'from fedml_api.standalone.federated_sgan.client import FedSSGANClient\n'), ((44, 20, 44, 58), 'copy.deepcopy', 'copy.deepcopy', ({(44, 34, 44, 57): 'self.global_model.model'}, {}), '(self.global_model.model)', False, 'import copy\n'), ((45, 20, 45, 46), 'copy.deepcopy', 'copy.deepcopy', ({(45, 34, 45, 45): 'local_model'}, {}), '(local_model)', False, 'import copy\n'), ((78, 33, 78, 56), 'copy.deepcopy', 'copy.deepcopy', ({(78, 47, 78, 55): 'w_global'}, {}), '(w_global)', False, 'import copy\n'), ((80, 61, 80, 77), 'copy.deepcopy', 'copy.deepcopy', ({(80, 75, 80, 76): 'w'}, {}), '(w)', False, 'import copy\n')]