import torch
import time
import copy
import re
import random
import numpy as np

from sim.algorithms.FL_base import SFLClient, SFLServer
from sim.data.datasets import build_dataset
from sim.data.partition import build_partition
from sim.models.build_models import build_model
from sim.utils.record_utils import logconfig, add_log, record_exp_result2
from sim.utils.utils import setup_seed, replace_with_samples, label_distribution, distance, relu, client_probabilities, initialize_state_probabilities, compute_alpha, compute_weights
from sim.utils.optim_utils import OptimKit, LrUpdater
from torch.utils.data import Dataset, DataLoader, Subset
from sim.utils.options import args_parser


args = args_parser()

torch.set_num_threads(4)
setup_seed(args.seed)
device = torch.device("cuda:{}".format(args.device) if torch.cuda.is_available() else "cpu")
args.beta = [int(args.beta[0]), args.beta[1]] if args.partition == 'exdir' else args.beta

def customize_record_name(args):
    if args.partition == 'exdir':
        partition = f'{args.partition}{args.beta[0]},{args.beta[1]}'
    elif args.partition == 'iid':
        partition = f'{args.partition}'
    elif args.partition == 'dir':
        partition = f'{args.partition}{args.beta[0]}'
    record_name = f'FedProx_N{args.N}_M{args.M}_S{args.S}_R{args.R}_T{args.T}_K{args.K}_{args.m}_{args.d}_{partition}'\
                + f'_seed{args.seed}_{args.capacity}_{args.distance}_alpha{args.alpha}_set{args.set}_dynamic{args.dynamic}_a{args.a}'
    return record_name
record_name = customize_record_name(args)

def main():
    global args, record_name, device
    logconfig(name=record_name, flag=args.log)
    add_log('{}'.format(args), flag=args.log)
    add_log('record_name: {}'.format(record_name), flag=args.log)
    
    client = SFLClient()
    server = SFLServer()

    state_dataidx_map = build_partition(args.d, args.M, args.partition, args.beta)
    
    train_dataset, test_dataset = build_dataset(args.d)

    client.setup_test_dataset(test_dataset)
    train_labels = torch.tensor([label for _, label in train_dataset])
    uniform_distribution = np.ones(args.num_classes) / args.num_classes
    class_probabilities_list = [label_distribution(train_labels, state_dataidx_map[i]) for i in range(args.M)]
    distance_list = [distance(args.distance, class_probabilities_list[i], uniform_distribution) for i in range(args.M)]
    add_log('distance_list: {}'.format(distance_list), flag=args.log)
    
    global_model = build_model(model=args.m, dataset=args.d)
    server.setup_model(global_model.to(device))
    add_log('{}'.format(global_model), flag=args.log)

    # construct optim kit
    client_optim_kit = OptimKit(optim_name=args.optim, batch_size=args.batch_size, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    client_optim_kit.setup_lr_updater(LrUpdater.exponential_lr_updater, mul=args.lr_decay)
    client.setup_optim_kit(client_optim_kit)
    client.setup_criterion(torch.nn.CrossEntropyLoss())
    server.setup_optim_settings(lr=args.global_lr)

    pi_list, data_map_list, nonzero_indices_list = initialize_state_probabilities(args.N, args.M, args.S, args.set)
    add_log('pi_list: {}'.format(pi_list), flag=args.log)

    gamma = 0
    for i in range(1, args.T + 1):
        gamma += (1 - args.alpha) ** i / args.T

    if args.dynamic == 'on':
        alpha_matrix = compute_alpha(args.N, args.M, pi_list, distance_list, nonzero_indices_list, args.a1, args.b1, args.alpha)
    elif args.dynamic == 'off':
        alpha_matrix = np.zeros((args.N, args.M))
        for n in range(args.N):
            for i in range(args.M):
                if pi_list[n][i] != 0:
                    alpha_matrix[n][i] = args.alpha
                else:
                    alpha_matrix[n][i] = 0

    add_log('alpha_matrix: {}'.format(alpha_matrix), flag=args.log)

    client_probs = client_probabilities(args.N)

    client_weights = compute_weights(args.N, args.M, pi_list, distance_list, alpha_matrix, gamma, client_probs, args.a2, args.b2, args.alpha, args.distance)
    add_log('client_weights: {}'.format(client_weights), flag=args.log)

    record_exp_result2(record_name, {'round':0})
    for round in range(0, args.R):
        selected_clients = server.select_clients_prob(args.N, client_probs)
        add_log('selected clients: {}'.format(selected_clients), flag=args.log)
        local_param_list = []
        weight_list = []
        for c_id in selected_clients:
            server.setup_temp_model(copy.deepcopy(server.global_model.to(device)))
            state_list = np.random.choice(args.M, args.T, p=pi_list[c_id], replace=True)
            add_log('state_list: {}'.format(state_list), flag=args.log)

            for state in state_list:
                data_map_list[c_id] = replace_with_samples(data_map_list[c_id], state_dataidx_map[state], alpha_matrix[c_id][state], args.capacity)

                local_dataset = Subset(train_dataset, data_map_list[c_id])
                local_param, local_delta = client.local_update_step_prox(local_dataset=local_dataset, model=copy.deepcopy(server.temp_model), g_model=copy.deepcopy(server.global_model), num_steps=args.K, device=device, a=args.a, clip=args.clip)
                torch.nn.utils.vector_to_parameters(local_param, server.temp_model.parameters())

            local_param_list.append(local_param)
            weight_list.append(relu(client_weights[c_id]))

        if args.dynamic == 'on':
            param_after_FL = sum([local_param_list[i] * weight_list[i] for i in range(len(local_param_list))]) / sum(weight_list)
        else:
            param_after_FL = server.aggregate_update(local_param_list)
        torch.nn.utils.vector_to_parameters(param_after_FL, server.global_model.parameters())
        
        client.optim_kit.update_lr()

        test_losses, test_top1, test_top5 = client.evaluate_dataset(model=server.global_model, dataset=client.test_dataset, device=args.device)
        add_log("Round {}'s server  test  acc: {:6.2f}%, test  loss: {:.4f}".format(round+1, test_top1.avg, test_losses.avg), 'red', flag=args.log)

        record_exp_result2(record_name, {'round':round+1, 'test_loss'  : test_losses.avg, 'test_top1'  : test_top1.avg})
    
    if args.save_model == 1:
        torch.save({'model': torch.nn.utils.parameters_to_vector(server.global_model.parameters())}, './save_model/{}.pt'.format(record_name))

if __name__ == '__main__':
    main()