#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Python version: 3.6

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import copy
import numpy as np
from torchvision import datasets, transforms
import torch
import os
import torch.distributed as dist
from utils.sampling import mnist_iid, mnist_noniid, cifar_iid
from utils.options import args_parser
from models.Update import LocalUpdate
from models.Update import LocalUpdateF
from models.Nets import MLP, CNNMnist, CNNCifar
from models.Fed import FedAvg
from models.test import test_img

from torch.multiprocessing import Process
from deep_gradient_compression import DGC
import json

# __name__是内置的变量，在执行当前文件（main_fed.py）时，默认值为__main__
# 但是如果其他.py文件import当前文件（main_fed.py）时，在其他文件中执行main_fed.py中的__name__,此时main_fed.py中的__name__默认值为文件名main_fed.py
if __name__ == '__main__':
    # parse args
    args = args_parser()
    args.device = torch.device('cuda:{}'.format(args.gpu))
    torch.manual_seed(0)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    rank = 0
    device_id = rank
    os.environ['MASTER_ADDR'] = '127.0.0.1'
    os.environ['MASTER_PORT'] = '29500'
    dist.init_process_group(backend='gloo', rank=rank, world_size=args.world_size)



    # if torch.cuda.is_available() and args.gpu != -1 else 'cpu'


    # load dataset and split users
    if args.dataset == 'mnist':
        # ToTensor():归一数据到（0,1），Normalize（）：（date-0.1307）/0.3081,将数据分布到（-1， 1）
        trans_mnist = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
        if trans_mnist is not None:
            print(1)
            print(trans_mnist)
        # 测试（60000）和训练集（10000）
        dataset_train = datasets.MNIST('../data/mnist/', train=True, download=True, transform=trans_mnist)
        dataset_test = datasets.MNIST('../data/mnist/', train=False, download=True, transform=trans_mnist)
        # sample users
        # Noniid数据
        if args.iid:
            dict_users = mnist_iid(dataset_train, args.num_users)
        else:
            dict_users = mnist_noniid(dataset_train, args.num_users)
    elif args.dataset == 'cifar':
        trans_cifar = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        dataset_train = datasets.CIFAR10('../data/cifar', train=True, download=True, transform=trans_cifar)
        dataset_test = datasets.CIFAR10('../data/cifar', train=False, download=True, transform=trans_cifar)
        if args.iid:
            dict_users = cifar_iid(dataset_train, args.num_users)
        else:
            exit('Error: only consider IID setting in CIFAR10')
    else:
        exit('Error: unrecognized dataset')
    img_size = dataset_train[0][0].shape
    # print('df ',img_size) [1,28,28]

    # build model
    # print(args.model)
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_glob = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_glob = CNNMnist(args=args).to(args.device)

    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            # print('x取值',x)
            len_in *= x
        net_glob = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)
        # add
        control_global = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')

    # 设置为训练模型
    net_glob.train()
    print(net_glob)

    control_weights =control_global.state_dict()
    # copy weights
    # 初始化全局权重
    w_glob = net_glob.state_dict()
    c_glob = copy.deepcopy(net_glob.state_dict())

    # print(w_glob)
    # training
    loss_train = []
    accuracy = []
    cv_loss, cv_acc = [], []
    val_loss_pre, counter = 0, 0
    net_best = None
    best_loss = None
    val_acc_list, net_list = [], []
    count = 0, 0
    test_acc_list = []

    if args.all_clients: 
        print("Aggregation over all clients")
        w_locals = [w_glob for i in range(args.num_users)]
        # add



    else:
        # 初始化本地权重
        c_local = [MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device) for i in
                   range(args.num_users)]
        for net in c_local:
            net.load_state_dict(control_weights)
        delta_c = copy.deepcopy(net_glob.state_dict())
        # delta_x = copy.deepcopy(net_glob.state_dict())
        # with open("test.txt", "w") as f:
        #     for i in range(0, len(c_local)):
        #         for k,v in c_local[i].state_dict().items():
        #             f.write(f"{k},{v}\n".format(k,v))
        # with open("test.txt", "a") as f:
        #     for i in range(0, len(c_local)):
        #         for k, v in w_locals[i].items():
        #             f.write(f"{k},{v}\n".format(k, v))
            # add 初始化变化量

        # print("why?")



    for iter in range(args.epochs):
        # 初始换控制变量
        for i in delta_c:
            delta_c[i] = 0.0
        # for i in delta_x:
        #     delta_x[i] = 0.0

        loss_locals = []

        if not args.all_clients:
            w_locals = []
        m = max(int(args.frac * args.num_users), 1)

        # 每次随机十位幸运观众
        idxs_users = np.random.choice(range(args.num_users), m, replace=False)
        for idx in idxs_users:
            # momentum法SGD
            local = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
            w, loss, local_delta_c, local_delta, control_local_w= local.train(net=copy.deepcopy(net_glob).to(args.device), control_local
                = c_local[idx], control_global=control_global, rank=rank, device_id=device_id, size=args.world_size)

            # add
            if iter != 0:
                c_local[idx].load_state_dict(control_local_w)


            if args.all_clients:
                w_locals[idx] = copy.deepcopy(w)
            else:
                w_locals.append(copy.deepcopy(w))

            # add
            loss_locals.append(copy.deepcopy(loss))

            # add
            for i in delta_c:
                if iter != 0:
                    delta_c[i] += w[i]
                else:
                    delta_c[i] += local_delta_c[i]
                    # delta_x[i] += local_delta[i]

        # add
        # update the delta C
        for i in delta_c:
            delta_c[i] /= m
            # delta_x[i] /= m

        # update global weights
        w_glob = FedAvg(w_locals)
        # add 更新全局c，w
        # w_glob = net_glob.state_dict()
        control_global_w = control_global.state_dict()
        for i in control_global_w:
            if iter !=0:
            #     w_glob[i] = delta_x[i]
            # else:
            #     w_glob[i] += delta_x[i]
                control_global_w[i] += (m / args.num_users) * delta_c[i]


        # copy weight to net_glob
        net_glob.load_state_dict(w_glob)
        # add
        control_global.load_state_dict(control_global_w)


        # print loss
        loss_avg = sum(loss_locals) / len(loss_locals)
        print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avg))
        loss_train.append(loss_avg)

        # acc_train, loss_train = test_img(net_glob, dataset_train, args)
        acc_test, loss_test = test_img(net_glob, dataset_test, args)
        accuracy.append(acc_test)


        # add
        for c in range(args.num_users):
            local_model = LocalUpdate(args=args, dataset=dataset_train, idxs=dict_users[idx])
            torch.cuda.empty_cache()

        # net_glob.eval()

        # print("Training accuracy: {:.2f}".format(acc_train))
        # print("Testing accuracy: {:.2f}".format(acc_test))

#######################################################################################################################
#######################################################################################################################
#######################################################################################################################
#######################################################################################################################
    # Fedavg
    # build model
    if args.model == 'cnn' and args.dataset == 'cifar':
        net_globF = CNNCifar(args=args).to(args.device)
    elif args.model == 'cnn' and args.dataset == 'mnist':
        net_globF = CNNMnist(args=args).to(args.device)
    elif args.model == 'mlp':
        len_in = 1
        for x in img_size:
            len_in *= x
        net_globF = MLP(dim_in=len_in, dim_hidden=200, dim_out=args.num_classes).to(args.device)
    else:
        exit('Error: unrecognized model')
    print(net_globF)
    net_globF.train()

    # copy weights
    w_globF = net_globF.state_dict()

    # training
    loss_trainF = []
    accuracyF = []
    cv_loss, cv_acc = [], []
    val_loss_pre, counter = 0, 0
    net_best = None
    best_loss = None
    val_acc_list, net_list = [], []

    if args.all_clients:
        print("Aggregation over all clients")
        w_localsF = [w_globF for i in range(args.num_users)]
    for iter in range(args.epochs):
        loss_locals = []
        if not args.all_clients:
            w_localsF = []
        m = max(int(args.frac * args.num_users), 1)
        idxs_users = np.random.choice(range(args.num_users), m, replace=False)
        for idx in idxs_users:
            localF = LocalUpdateF(args=args, dataset=dataset_train, idxs=dict_users[idx])
            w, loss = localF.train(net=copy.deepcopy(net_globF).to(args.device))
            if args.all_clients:
                w_localsF[idx] = copy.deepcopy(w)
            else:
                w_localsF.append(copy.deepcopy(w))
            loss_locals.append(copy.deepcopy(loss))
        # update global weights
        w_globF = FedAvg(w_localsF)

        # copy weight to net_globF
        net_globF.load_state_dict(w_globF)

        # print loss
        loss_avgF = sum(loss_locals) / len(loss_locals)
        print('Round {:3d}, Average loss {:.3f}'.format(iter, loss_avgF))
        loss_trainF.append(loss_avgF)

        acc_test, loss_test = test_img(net_globF, dataset_test, args)
        accuracyF.append(acc_test)



    # plot loss curve
    plt.figure()
    print(loss_train, loss_trainF)
    plt.plot(range(len(loss_train)), loss_train, label='Scaffold', zorder=2)
    plt.plot(range(len(loss_trainF)), loss_trainF, 'r', label='FedAvg',zorder=1)
    plt.ylabel('train_loss')
    plt.xlabel('epochs')
    plt.legend(loc='best')
    plt.savefig('./save/fed_{}_{}_{}_{}_iid{}.png'.format(args.dataset, args.model, args.epochs, 'train_loss', args.iid))


    # testing
    net_glob.eval()
    acc_train, loss_train = test_img(net_glob, dataset_train, args)
    acc_test, loss_test = test_img(net_glob, dataset_test, args)
    print("Training accuracy: {:.2f}".format(acc_train))
    print("Testing accuracy: {:.2f}".format(acc_test))

    # plot loss curve
    plt.figure()
    # plt.plot((np.arange(1, len(accuracy)), 1), accuracy, 'r')
    plt.plot(range(len(accuracy)), accuracy, label='Scaffold', zorder=2)
    plt.plot(range(len(accuracyF)), accuracyF, 'r', label='FedAvg', zorder=1)
    plt.ylabel('test_acc')
    plt.xlabel('epochs')
    plt.legend(loc='best')
    plt.savefig('./save/fed_{}_{}_{}_{}_iid{}.png'.format(args.dataset, args.model, args.epochs, 'acc_test', args.iid))


[((6, 0, 6, 21), 'matplotlib.use', 'matplotlib.use', ({(6, 15, 6, 20): '"""Agg"""'}, {}), "('Agg')", False, 'import matplotlib\n'), ((30, 11, 30, 24), 'utils.options.args_parser', 'args_parser', ({}, {}), '()', False, 'from utils.options import args_parser\n'), ((32, 4, 32, 24), 'torch.manual_seed', 'torch.manual_seed', ({(32, 22, 32, 23): '(0)'}, {}), '(0)', False, 'import torch\n'), ((40, 4, 40, 82), 'torch.distributed.init_process_group', 'dist.init_process_group', (), '', True, 'import torch.distributed as dist\n'), ((298, 4, 298, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((302, 4, 302, 28), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(302, 15, 302, 27): '"""train_loss"""'}, {}), "('train_loss')", True, 'import matplotlib.pyplot as plt\n'), ((303, 4, 303, 24), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(303, 15, 303, 23): '"""epochs"""'}, {}), "('epochs')", True, 'import matplotlib.pyplot as plt\n'), ((304, 4, 304, 26), 'matplotlib.pyplot.legend', 'plt.legend', (), '', True, 'import matplotlib.pyplot as plt\n'), ((310, 28, 310, 67), 'models.test.test_img', 'test_img', ({(310, 37, 310, 45): 'net_glob', (310, 47, 310, 60): 'dataset_train', (310, 62, 310, 66): 'args'}, {}), '(net_glob, dataset_train, args)', False, 'from models.test import test_img\n'), ((311, 26, 311, 64), 'models.test.test_img', 'test_img', ({(311, 35, 311, 43): 'net_glob', (311, 45, 311, 57): 'dataset_test', (311, 59, 311, 63): 'args'}, {}), '(net_glob, dataset_test, args)', False, 'from models.test import test_img\n'), ((316, 4, 316, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((320, 4, 320, 26), 'matplotlib.pyplot.ylabel', 'plt.ylabel', ({(320, 15, 320, 25): '"""test_acc"""'}, {}), "('test_acc')", True, 'import matplotlib.pyplot as plt\n'), ((321, 4, 321, 24), 'matplotlib.pyplot.xlabel', 'plt.xlabel', ({(321, 15, 321, 23): '"""epochs"""'}, {}), "('epochs')", True, 'import matplotlib.pyplot as plt\n'), ((322, 4, 322, 26), 'matplotlib.pyplot.legend', 'plt.legend', (), '', True, 'import matplotlib.pyplot as plt\n'), ((55, 24, 55, 106), 'torchvision.datasets.MNIST', 'datasets.MNIST', (), '', False, 'from torchvision import datasets, transforms\n'), ((56, 23, 56, 106), 'torchvision.datasets.MNIST', 'datasets.MNIST', (), '', False, 'from torchvision import datasets, transforms\n'), ((194, 17, 194, 33), 'models.Fed.FedAvg', 'FedAvg', ({(194, 24, 194, 32): 'w_locals'}, {}), '(w_locals)', False, 'from models.Fed import FedAvg\n'), ((218, 30, 218, 68), 'models.test.test_img', 'test_img', ({(218, 39, 218, 47): 'net_glob', (218, 49, 218, 61): 'dataset_test', (218, 63, 218, 67): 'args'}, {}), '(net_glob, dataset_test, args)', False, 'from models.test import test_img\n'), ((282, 18, 282, 35), 'models.Fed.FedAvg', 'FedAvg', ({(282, 25, 282, 34): 'w_localsF'}, {}), '(w_localsF)', False, 'from models.Fed import FedAvg\n'), ((292, 30, 292, 69), 'models.test.test_img', 'test_img', ({(292, 39, 292, 48): 'net_globF', (292, 50, 292, 62): 'dataset_test', (292, 64, 292, 68): 'args'}, {}), '(net_globF, dataset_test, args)', False, 'from models.test import test_img\n'), ((60, 25, 60, 65), 'utils.sampling.mnist_iid', 'mnist_iid', ({(60, 35, 60, 48): 'dataset_train', (60, 50, 60, 64): 'args.num_users'}, {}), '(dataset_train, args.num_users)', False, 'from utils.sampling import mnist_iid, mnist_noniid, cifar_iid\n'), ((62, 25, 62, 68), 'utils.sampling.mnist_noniid', 'mnist_noniid', ({(62, 38, 62, 51): 'dataset_train', (62, 53, 62, 67): 'args.num_users'}, {}), '(dataset_train, args.num_users)', False, 'from utils.sampling import mnist_iid, mnist_noniid, cifar_iid\n'), ((65, 24, 65, 107), 'torchvision.datasets.CIFAR10', 'datasets.CIFAR10', (), '', False, 'from torchvision import datasets, transforms\n'), ((66, 23, 66, 107), 'torchvision.datasets.CIFAR10', 'datasets.CIFAR10', (), '', False, 'from torchvision import datasets, transforms\n'), ((162, 20, 162, 87), 'models.Update.LocalUpdate', 'LocalUpdate', (), '', False, 'from models.Update import LocalUpdate\n'), ((224, 26, 224, 93), 'models.Update.LocalUpdate', 'LocalUpdate', (), '', False, 'from models.Update import LocalUpdate\n'), ((225, 12, 225, 36), 'torch.cuda.empty_cache', 'torch.cuda.empty_cache', ({}, {}), '()', False, 'import torch\n'), ((274, 21, 274, 89), 'models.Update.LocalUpdateF', 'LocalUpdateF', (), '', False, 'from models.Update import LocalUpdateF\n'), ((50, 42, 50, 63), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import datasets, transforms\n'), ((50, 65, 50, 107), 'torchvision.transforms.Normalize', 'transforms.Normalize', ({(50, 86, 50, 95): '(0.1307,)', (50, 97, 50, 106): '(0.3081,)'}, {}), '((0.1307,), (0.3081,))', False, 'from torchvision import datasets, transforms\n'), ((68, 25, 68, 65), 'utils.sampling.cifar_iid', 'cifar_iid', ({(68, 35, 68, 48): 'dataset_train', (68, 50, 68, 64): 'args.num_users'}, {}), '(dataset_train, args.num_users)', False, 'from utils.sampling import mnist_iid, mnist_noniid, cifar_iid\n'), ((79, 19, 79, 38), 'models.Nets.CNNCifar', 'CNNCifar', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((172, 32, 172, 48), 'copy.deepcopy', 'copy.deepcopy', ({(172, 46, 172, 47): 'w'}, {}), '(w)', False, 'import copy\n'), ((177, 31, 177, 50), 'copy.deepcopy', 'copy.deepcopy', ({(177, 45, 177, 49): 'loss'}, {}), '(loss)', False, 'import copy\n'), ((239, 20, 239, 39), 'models.Nets.CNNCifar', 'CNNCifar', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((277, 33, 277, 49), 'copy.deepcopy', 'copy.deepcopy', ({(277, 47, 277, 48): 'w'}, {}), '(w)', False, 'import copy\n'), ((280, 31, 280, 50), 'copy.deepcopy', 'copy.deepcopy', ({(280, 45, 280, 49): 'loss'}, {}), '(loss)', False, 'import copy\n'), ((64, 42, 64, 63), 'torchvision.transforms.ToTensor', 'transforms.ToTensor', ({}, {}), '()', False, 'from torchvision import datasets, transforms\n'), ((64, 65, 64, 119), 'torchvision.transforms.Normalize', 'transforms.Normalize', ({(64, 86, 64, 101): '(0.5, 0.5, 0.5)', (64, 103, 64, 118): '(0.5, 0.5, 0.5)'}, {}), '((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))', False, 'from torchvision import datasets, transforms\n'), ((81, 19, 81, 38), 'models.Nets.CNNMnist', 'CNNMnist', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((125, 19, 125, 79), 'models.Nets.MLP', 'MLP', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((174, 32, 174, 48), 'copy.deepcopy', 'copy.deepcopy', ({(174, 46, 174, 47): 'w'}, {}), '(w)', False, 'import copy\n'), ((241, 20, 241, 39), 'models.Nets.CNNMnist', 'CNNMnist', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((279, 33, 279, 49), 'copy.deepcopy', 'copy.deepcopy', ({(279, 47, 279, 48): 'w'}, {}), '(w)', False, 'import copy\n'), ((88, 19, 88, 79), 'models.Nets.MLP', 'MLP', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((90, 25, 90, 85), 'models.Nets.MLP', 'MLP', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((246, 20, 246, 80), 'models.Nets.MLP', 'MLP', (), '', False, 'from models.Nets import MLP, CNNMnist, CNNCifar\n'), ((163, 82, 163, 105), 'copy.deepcopy', 'copy.deepcopy', ({(163, 96, 163, 104): 'net_glob'}, {}), '(net_glob)', False, 'import copy\n'), ((275, 39, 275, 63), 'copy.deepcopy', 'copy.deepcopy', ({(275, 53, 275, 62): 'net_globF'}, {}), '(net_globF)', False, 'import copy\n')]