import numpy as np

import torch
import torch.nn.functional as F
import datetime
import time
from utils import *
from defense import *

from tqdm.auto import tqdm

from models.vgg import get_vgg_model
import pandas as pd

from torch.nn.utils import parameters_to_vector, vector_to_parameters
from lira_helper import *
import csv_record

import datasets
from termcolor import colored

#段代码定义了一个指数衰减函数，用于根据初始值、衰减率和时间步数（或者幂次）计算衰减后的值。
def exponential_decay(init_val, decay_rate, t):
    return init_val*(1.0 - decay_rate)**t

class Net(nn.Module):
    def __init__(self, num_classes):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout2d(0.25)
        self.dropout2 = nn.Dropout2d(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        output = self.fc2(x)
        #output = F.log_softmax(x, dim=1)
        return output


def get_results_filename(poison_type, attack_method, model_replacement, project_frequency, defense_method, norm_bound, prox_attack, fixed_pool=False, model_arch="vgg9"):
    filename = "{}_{}_{}".format(poison_type, model_arch, attack_method)
    if fixed_pool:
        filename += "_fixed_pool" 
    
    if model_replacement:
        filename += "_with_replacement"
    else:
        filename += "_without_replacement"
    
    if attack_method == "pgd":
        filename += "_1_{}".format(project_frequency)
    
    if prox_attack:
        filename += "_prox_attack"

    if defense_method in ("norm-clipping", "norm-clipping-adaptive", "weak-dp"):
        filename += "_{}_m_{}".format(defense_method, norm_bound)
    elif defense_method in ("krum", "multi-krum", "rfa"):
        filename += "_{}".format(defense_method)
               
    filename += "_acc_results.csv"

    return filename

#这段代码是用于计算两个模型之间参数的范数差异
def calc_norm_diff(gs_model, vanilla_model, epoch, fl_round, mode="bad"):
    norm_diff = 0
    for p_index, p in enumerate(gs_model.parameters()):
        norm_diff += torch.norm(list(gs_model.parameters())[p_index] - list(vanilla_model.parameters())[p_index]) ** 2
    norm_diff = torch.sqrt(norm_diff).item()
    if mode == "bad":
        #pdb.set_trace()
        logger.info("===> ND `|w_bad-w_g|` in local epoch: {} | FL round: {} |, is {}".format(epoch, fl_round, norm_diff))
    elif mode == "normal":
        logger.info("===> ND `|w_normal-w_g|` in local epoch: {} | FL round: {} |, is {}".format(epoch, fl_round, norm_diff))
    elif mode == "avg":
        logger.info("===> ND `|w_avg-w_g|` in local epoch: {} | FL round: {} |, is {}".format(epoch, fl_round, norm_diff))

    return norm_diff

def vectorize_net(net):
    return torch.cat([p.view(-1) for p in net.parameters()])

def load_model_weight(net, weight):
    index_bias = 0
    for p_index, p in enumerate(net.parameters()):
        p.data =  weight[index_bias:index_bias+p.numel()].view(p.size())
        index_bias += p.numel()

def fed_nova_aggregator(model, net_list, list_ni, device, list_ai):
    # https://github.com/Xtra-Computing/NIID-Bench/blob/692569f790af0f5908dba23a15d4d80ce7e7aec4/experiments.py#L375
    total_n = sum(list_ni)
    print(f"Aggregating models with FedNova with aggregation weight {list_ai}")
    print(f"List of sample size {list_ni}")
    
    weight_accumulator = {}
    for name, params in model.state_dict().items():
        weight_accumulator[name] = torch.zeros(params.size()).to(device)
    
    for net_index, net in enumerate(net_list):
        for name, data in net.state_dict().items():
            # weight_accumulator[name].add_( ( (data - model.state_dict()[name]) / list_ai[net_index] ) * (list_ni[net_index] / total_n))
            weight_accumulator[name].add_( torch.true_divide( (model.state_dict()[name] - data), list_ai[net_index] ) * list_ni[net_index] / total_n)
    
    coeff = 0.0
    for i in range(len(list_ai)):
        coeff += list_ai[i] * list_ni[i] / total_n
        # print(f"ai: {list_ai[i]} ni: {list_ni[i]} coeff: {coeff}")
    print(f"coeff: {coeff}")
    
    for name, params in model.state_dict().items():
        update_per_layer = coeff * weight_accumulator[name]
        # params.add_(coeff * update_per_layer)
        if params.type() != update_per_layer.type():
            params.sub_(update_per_layer.to(torch.int64))
        else:
            params.sub_(update_per_layer)
    return model

def fed_avg_aggregator(init_model, net_list, net_freq, device, model="lenet"):
    # import IPython
    # IPython.embed()
    
    weight_accumulator = {}
    
    for name, params in init_model.state_dict().items():
        weight_accumulator[name] = torch.zeros_like(params).float()
    
    for i in range(0, len(net_list)):
        diff = dict()
        for name, data in net_list[i].state_dict().items():
            # diff[name] = (data - model_server_before_aggregate.state_dict()[name]).cpu().detach().numpy()
            diff[name] = (data - init_model.state_dict()[name])
            try:
                weight_accumulator[name].add_(net_freq[i]  *  diff[name])
                # weight_accumulator[name].add_(0.1  *  diff[name])
                
            except Exception as e:
                print(e)
                import IPython
                IPython.embed()
                exit(0)
        # print(f"diff: {diff}")
    for idl, (name, data) in enumerate(init_model.state_dict().items()):
        update_per_layer = weight_accumulator[name] #  * self.conf["lambda"]
        if data.type() != update_per_layer.type():
            data.add_(update_per_layer.to(torch.int64))            
        else:
            data.add_(update_per_layer)
            
    return init_model

def estimate_wg(model, device, train_loader, optimizer, epoch, log_interval, criterion):
    logger.info("Prox-attack: Estimating wg_hat")
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
    
    if batch_idx % log_interval == 0:
            logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))
  
def train_lira(model, atkmodel, tgtmodel, optimizer, atkmodel_optimizer, train_loader, criterion, atk_eps = 0.001, attack_alpha = 0.5,
               attack_portion = 1.0, clip_image=None, target_transform=None, atk_train_epoch=1, atkmodel_train=False, device=None,
               pgd_attack = False, batch_idx = 1, project_frequency = 1, pgd_eps = None, model_original=None, adv_optimizer=None,
               proj="l_2", mask_grad_list=[], aggregator="fedavg", wg_hat=None, mu=0.1, local_e=0):
    # print(f"attack_alpha: {attack_alpha}")
    # print(f"atkmodel_optimizer: {atkmodel_optimizer}")
    tgtmodel.eval()
    wg_clone = copy.deepcopy(model)
    # atk_copy = copy.deepcopy(tgtmodel)
    # atk_optimizer = None
    loss_fn = nn.CrossEntropyLoss()
    func_fn = loss_fn
    correct = 0

    correct_clean = 0
    correct_poison = 0

    dataset_size = 0
    poison_size = 0
    clean_size = 0
    loss_list = []

    if not atkmodel_train:
        # optimizer = optim.SGD(model.parameters(), lr=0.001)
        model.train()
        # atkmodel.eval()
        # Sub-training phase
        for batch_idx, batch in enumerate(train_loader):
            bs = len(batch)
            data, targets = batch
            # data, target = data.to(device), target.to(device)
            # clean_images, clean_targets, poison_images, poison_targets, poisoning_per_batch = get_poison_batch(batch, attack_portion)
            clean_images, clean_targets = copy.deepcopy(data).to(device), copy.deepcopy(targets).to(device)
            poison_images, poison_targets = copy.deepcopy(data).to(device), copy.deepcopy(targets).to(device)
            # dataset_size += len(data)
            clean_size += len(clean_images)
            optimizer.zero_grad()
            if pgd_attack:
                adv_optimizer.zero_grad()
            output = model(clean_images)
            loss_clean = loss_fn(output, clean_targets)
            if aggregator == "fedprox":
                # print(f"fedprox!!!")
                wg_hat_vec = parameters_to_vector(list(wg_hat.parameters()))
                model_vec = parameters_to_vector(list(model.parameters()))
                prox_term = torch.norm(wg_hat_vec - model_vec)**2
                loss_clean = loss_clean + mu/2*prox_term
            # attack_alpha说明只有干净样本无poison样本
            if attack_alpha == 1.0:
                optimizer.zero_grad()
                # atkmodel_optimizer.zero_grad()
                loss_clean.backward()
                if mask_grad_list:
                    apply_grad_mask(model, mask_grad_list)
                if not pgd_attack:
                    optimizer.step()
                else:
                    if proj == "l_inf":
                        w = list(model.parameters())
                        n_layers = len(w)
                        # adversarial learning rate
                        eta = 0.001
                        for i in range(len(w)):
                            # uncomment below line to restrict proj to some layers
                            if True:#i == 6 or i == 8 or i == 10 or i == 0 or i == 18:
                                w[i].data = w[i].data - eta * w[i].grad.data
                                # projection step
                                m1 = torch.lt(torch.sub(w[i], model_original[i]), -pgd_eps)
                                m2 = torch.gt(torch.sub(w[i], model_original[i]), pgd_eps)
                                w1 = (model_original[i] - pgd_eps) * m1
                                w2 = (model_original[i] + pgd_eps) * m2
                                w3 = (w[i]) * (~(m1+m2))
                                wf = w1+w2+w3
                                w[i].data = wf.data
                    else:
                        # do l2_projection
                        adv_optimizer.step()
                        w = list(model.parameters())
                        w_vec = parameters_to_vector(w)
                        model_original_vec = parameters_to_vector(model_original)
                        # make sure you project on last iteration otherwise, high LR pushes you really far
                        # Start
                        if (batch_idx%project_frequency == 0 or batch_idx == len(train_loader)-1) and (torch.norm(w_vec - model_original_vec) > pgd_eps):
                            # project back into norm ball
                            w_proj_vec = pgd_eps*(w_vec - model_original_vec)/torch.norm(
                                    w_vec-model_original_vec) + model_original_vec
                            # plug w_proj back into model
                            vector_to_parameters(w_proj_vec, w)


                pred = output.data.max(1)[1]  # get the index of the max log-probability
                loss_list.append(loss_clean.item())
                correct_clean += pred.eq(clean_targets.data.view_as(pred)).cpu().sum().item()
            else:
                if attack_alpha < 1.0:
                    poison_size += len(poison_images)
                    # poison_images, poison_targets = poison_images.to(device), poison_targets.to(device)
                    with torch.no_grad():
                        noise = tgtmodel(poison_images) * atk_eps
                        atkdata = clip_image(poison_images + noise)
                        atktarget = target_transform(poison_targets)
                        # atkdata.requires_grad_(False)
                        # atktarget.requires_grad_(False)
                        if attack_portion < 1.0:
                            atkdata = atkdata[:int(attack_portion*bs)]
                            atktarget = atktarget[:int(attack_portion*bs)]
                    # import IPython
                    # IPython.embed()
                    atkoutput = model(atkdata.detach())
                    loss_poison = F.cross_entropy(atkoutput, atktarget.detach())
                else:
                    loss_poison = torch.tensor(0.0).to(device)
                #即将干净样本损失与毒化样本损失按照指定的权重进行加权求和。
                loss2 = loss_clean * attack_alpha  + (1.0 - attack_alpha) * loss_poison

                optimizer.zero_grad()
                loss2.backward()
                if mask_grad_list:
                    apply_grad_mask(model, mask_grad_list)
                if not pgd_attack:
                    optimizer.step()
                else:
                    if proj == "l_inf":
                        w = list(model.parameters())
                        n_layers = len(w)
                        # adversarial learning rate
                        eta = 0.001
                        for i in range(len(w)):
                            # uncomment below line to restrict proj to some layers
                            if True:#i == 6 or i == 8 or i == 10 or i == 0 or i == 18:
                                w[i].data = w[i].data - eta * w[i].grad.data
                                # projection step
                                m1 = torch.lt(torch.sub(w[i], model_original[i]), -pgd_eps)
                                m2 = torch.gt(torch.sub(w[i], model_original[i]), pgd_eps)
                                w1 = (model_original[i] - pgd_eps) * m1
                                w2 = (model_original[i] + pgd_eps) * m2
                                w3 = (w[i]) * (~(m1+m2))
                                wf = w1+w2+w3
                                w[i].data = wf.data
                    else:
                        # do l2_projection
                        adv_optimizer.step()
                        w = list(model.parameters())
                        w_vec = parameters_to_vector(w)
                        model_original_vec = parameters_to_vector(list(model_original.parameters()))
                        # make sure you project on last iteration otherwise, high LR pushes you really far
                        if (local_e%project_frequency == 0 and batch_idx == len(train_loader)-1) and (torch.norm(w_vec - model_original_vec) > pgd_eps):
                            # project back into norm ball
                            w_proj_vec = pgd_eps*(w_vec - model_original_vec)/torch.norm(
                                    w_vec-model_original_vec) + model_original_vec


                            vector_to_parameters(w_proj_vec, w)

                loss_list.append(loss2.item())
                pred = output.data.max(1)[1]  # get the index of the max log-probability
                poison_pred = atkoutput.data.max(1)[1]  # get the index of the max log-probability

                # correct += pred.eq(target.data.view_as(pred)).cpu().sum().item()
                correct_clean += pred.eq(clean_targets.data.view_as(pred)).cpu().sum().item()
                correct_poison += poison_pred.eq(atktarget.data.view_as(poison_pred)).cpu().sum().item()

    else:
        model.eval()
        # atk_optimizer = optim.Adam(atkmodel.parameters(), lr=0.0002)
        atkmodel.train()
        # optimizer.zero_grad()
        for batch_idx, (data, target) in enumerate(train_loader):
            bs = data.size(0)
            data, target = data.to(device), target.to(device)
            # dataset_size += len(data)
            poison_size += len(data)

            ###############################
            #### Update the classifier ####
            ###############################
            # with torch.no_grad():
            noise = atkmodel(data) * atk_eps
            atkdata = clip_image(data + noise)
            atktarget = target_transform(target)
            if attack_portion < 1.0:
                atkdata = atkdata[:int(attack_portion*bs)]
                atktarget = atktarget[:int(attack_portion*bs)]
            # with torch.no_grad():
            # atkoutput = wg_clone(atkdata)
            atkoutput = model(atkdata)
            loss_p = func_fn(atkoutput, atktarget)
            loss2 = loss_p
            # import IPython
            # IPython.embed()
            atkmodel_optimizer.zero_grad()
            loss2.backward()
            atkmodel_optimizer.step()
            pred = atkoutput.data.max(1)[1]  # get the index of the max log-probability
            correct_poison += pred.eq(atktarget.data.view_as(pred)).cpu().sum().item()
            loss_list.append(loss2.item())

    clean_acc = 100.0 * (float(correct_clean)/float(clean_size)) if clean_size else 0.0
    poison_acc = 100.0 * (float(correct_poison)/float(poison_size)) if poison_size else 0.0

    training_avg_loss = sum(loss_list)/len(loss_list)
    # training_avg_loss = 0.0
    if atkmodel_train:
        logger.info(colored("Training loss = {:.2f}, acc = {:.2f} of atk model this epoch".format(training_avg_loss, poison_acc), "yellow"))
    else:
        logger.info(colored("Training loss = {:.2f}, acc = {:.2f} of cls model this epoch".format(training_avg_loss, clean_acc), "yellow"))
        logger.info("Training clean_acc is {:.2f}, poison_acc = {:.2f}".format(clean_acc, poison_acc))
    del wg_clone

def train_baseline(model, optimizer, train_loader, criterion, dataset, device, target_transform=None, local_e=0):
    loss_fn = nn.CrossEntropyLoss()
    func_fn = loss_fn
    total_loss = 0
    correct = 0
    
    correct_clean = 0
    correct_poison = 0
    total_loss = 0
    poison_data_count = 0
    
    dataset_size = 0
    poison_size = 0
    poison_data_count = 0
    clean_size = 0
    loss_list = []
    model.train()

    for batch_idx, batch in enumerate(train_loader):
        bs = len(batch)
        data, targets, poison_num, _, _ = get_poison_batch(batch, dataset, device, target_transform=target_transform)
        # print(f"training targets: {targets}")
        optimizer.zero_grad()
        dataset_size += len(data)
        poison_data_count += poison_num

        output = model(data)
        class_loss = nn.functional.cross_entropy(output, targets)

        # distance_loss = helper.model_dist_norm_var(model, target_params_variables)
        # Lmodel = αLclass + (1 − α)Lano; alpha_loss =1 fixed
        # loss = helper.params['alpha_loss'] * class_loss + \
        #         (1 - helper.params['alpha_loss']) * distance_loss
        loss = class_loss
        loss.backward()

        optimizer.step()
        total_loss += loss.data
        pred = output.data.max(1)[1]  # get the index of the max log-probability
        correct += pred.eq(targets.data.view_as(pred)).sum().item()
        acc = 100.0 * (float(correct) / float(dataset_size))
        total_l = total_loss / dataset_size
    logger.info("Total loss is {:.2f}, training acc is {:.2f}".format(total_l, acc))

def test(model, device, test_loader, test_batch_size, criterion, mode="raw-task", dataset="cifar10", poison_type="fashion"):
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))
    
    if dataset == 'tiny-imagenet':
        class_correct = list(0. for i in range(200))
        class_total = list(0. for i in range(200))
    
    if dataset in ("mnist", "emnist"):
        target_class = 7
        if mode == "raw-task":
            classes = [str(i) for i in range(10)]
        elif mode == "targetted-task":
            if poison_type == 'ardis':
                classes = [str(i) for i in range(10)]
            else: 
                classes = ["T-shirt/top", 
                            "Trouser",
                            "Pullover",
                            "Dress",
                            "Coat",
                            "Sandal",
                            "Shirt",
                            "Sneaker",
                            "Bag",
                            "Ankle boot"]
    
    elif dataset == "cifar10":
        classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
        # target_class = 2 for greencar, 9 for southwest
        if poison_type in ("howto", "greencar-neo"):
            target_class = 2
        else:
            target_class = 9
    
    elif dataset == "tiny-imagenet":
        classes = [str(i) for i in range(200)]
        target_class = 1
      
    model.eval()
    test_loss = 0
    correct = 0
    backdoor_correct = 0
    backdoor_tot = 0
    final_acc = 0
    task_acc = None

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output, 1)
            c = (predicted == target).squeeze()

            #test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
            # check backdoor accuracy
            if poison_type == 'ardis':
                backdoor_index = torch.where(target == target_class)
                target_backdoor = torch.ones_like(target[backdoor_index])
                predicted_backdoor = predicted[backdoor_index]
                backdoor_correct += (predicted_backdoor == target_backdoor).sum().item()
                backdoor_tot = backdoor_index[0].shape[0]
                
            #for image_index in range(test_batch_size):
            for image_index in range(len(target)):
                label = target[image_index]
                class_correct[label] += c[image_index].item()
        
                    
                class_total[label] += 1
    test_loss /= len(test_loader.dataset)
    
    number_class = 10
    if dataset == 'tiny-imagenet':
        number_class = 200
        
    if mode == "raw-task":
        for i in range(number_class):
            # logger.info('Accuracy of %5s : %.2f %%' % (
            #     classes[i], 100 * class_correct[i] / class_total[i]))

            if i == target_class:
                task_acc = 100 * class_correct[i] / class_total[i]

        logger.info('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.2f}%)\n'.format(
            test_loss, correct, len(test_loader.dataset),
            100. * correct / len(test_loader.dataset)))
        final_acc = 100. * correct / len(test_loader.dataset)

    elif mode == "targetted-task":

        if dataset in ("mnist", "emnist"):
            for i in range(10):
                logger.info('Accuracy of %5s : %.2f %%' % (
                    classes[i], 100 * class_correct[i] / class_total[i]))
            if poison_type == 'ardis':
                # ensure 7 is being classified as 1
                logger.info('Backdoor Accuracy of %.2f : %.2f %%' % (
                     target_class, 100 * backdoor_correct / backdoor_tot))
                final_acc = 100 * backdoor_correct / backdoor_tot
            else:
                # trouser acc
                final_acc = 100 * class_correct[1] / class_total[1]
        
        elif dataset == "cifar10":
            logger.info('#### Targetted Accuracy of %5s : %.2f %%' % (classes[target_class], 100 * class_correct[target_class] / class_total[target_class]))
            final_acc = 100 * class_correct[target_class] / class_total[target_class]
    return final_acc, task_acc

def lira_test(args, device, atkmodel, scratchmodel, target_transform, 
         train_loader, test_loader, epoch, trainepoch, clip_image, 
         testoptimizer=None, log_prefix='Internal', epochs_per_test=5):
    #default phase 2 parameters to phase 1 
    
    if args['test_alpha'] is None:
        args['test_alpha'] = args['attack_alpha']
    if args['test_eps'] is None:
        args['test_eps'] = args['eps']
        
    print(f"\n------------------------------------\nStart postraining based on trajectories")
    
    atkmodel.eval()

    if testoptimizer is None:
        testoptimizer = optim.SGD(scratchmodel.parameters(), lr=args['lr'])
    for cepoch in range(trainepoch):
        pbar = tqdm(enumerate(train_loader), total=len(train_loader), position=0, leave=True)
        for batch_idx, (data, target) in pbar:
            bs = data.size(0)
            data, target = data.to(device), target.to(device)
            testoptimizer.zero_grad()
            with torch.no_grad():
                noise = atkmodel(data) * args['test_eps']
                atkdata = clip_image(data + noise)
                atktarget = target_transform(target)
                if args['attack_portion'] < 1.0:
                    atkdata = atkdata[:int(args['attack_portion']*bs)]
                    atktarget = atktarget[:int(args['attack_portion']*bs)]

            atkoutput = scratchmodel(atkdata)
            output = scratchmodel(data)
            
            loss_clean = loss_fn(output, target)
            loss_poison = loss_fn(atkoutput, atktarget)
            
            loss = args['attack_alpha'] * loss_clean + (1-args['test_alpha']) * loss_poison
            
            loss.backward()
            testoptimizer.step()
            
            if batch_idx % 10 == 0 or batch_idx == (len(train_loader)-1):
                pbar.set_description(
                    'Test [{}-{}] Loss: Clean {:.4f} Poison {:.4f} Total {:.5f}'.format(
                        epoch, cepoch+1,
                        loss_clean.item(),
                        loss_poison.item(),
                        loss.item()
                    ))
        print(f"\n------------------------------------Start testing the model") 
        if cepoch % epochs_per_test == 0 or cepoch == trainepoch-1:
            correct = 0    
            correct_transform = 0
            test_loss = 0
            test_transform_loss = 0

            with torch.no_grad():
                for data, target in test_loader:
                    bs = data.size(0)
                    data, target = data.to(device), target.to(device)
                    output = scratchmodel(data)
                    test_loss += loss_fn(output, target).item() * bs  # sum up batch loss
                    pred = output.max(1, keepdim=True)[
                        1]  # get the index of the max log-probability
                    correct += pred.eq(target.view_as(pred)).sum().item()

                    noise = atkmodel(data) * args['test_eps']
                    atkdata = clip_image(data + noise)
                    atkoutput = scratchmodel(atkdata)
                    test_transform_loss += loss_fn(atkoutput, target_transform(target)).item() * bs  # sum up batch loss
                    atkpred = atkoutput.max(1, keepdim=True)[
                        1]  # get the index of the max log-probability
                    correct_transform += atkpred.eq(
                        target_transform(target).view_as(atkpred)).sum().item()

            test_loss /= len(test_loader.dataset)
            test_transform_loss /= len(test_loader.dataset)

            correct /= len(test_loader.dataset)
            correct_transform /= len(test_loader.dataset)

            print(
                '\n{}-Test set [{}]: Loss: clean {:.4f} poison {:.4f}, Accuracy: clean {:.2f} poison {:.2f}'.format(
                    log_prefix, cepoch, 
                    test_loss, test_transform_loss,
                    correct, correct_transform
                ))
        # print(f"\nFinish the post testing------------------------------------") 
    
 
    return correct, correct_transform

def test_updated(args, device, atkmodel, model, target_transform, 
         train_loader, test_loader, epoch, trainepoch, clip_image, 
         testoptimizer=None, log_prefix='Internal', epochs_per_test=3, 
         dataset="cifar10", criterion=None, subpath_saved="", test_eps=None, 
         is_poison=False, ident="global", atk_baseline=False):

    if args['test_alpha'] is None:
        args['test_alpha'] = args['attack_alpha']
    if args['test_eps'] is None:
        args['test_eps'] = args['eps']
    if not test_eps:
        test_eps = args['test_eps']
    # import IPython
    # IPython.embed()
    # start_time = time.time()
    
    model.eval()
    atkmodel.eval()
    test_loss = 0
    correct = 0
    correct_transform = 0
    test_loss = 0
    test_transform_loss = 0
    
    model = model.to(device)
    atkmodel = atkmodel.to(device)
    with torch.no_grad():
        for idx, (data, target) in enumerate(test_loader):
            start_time = time.time()
            bs = data.size(0)
            data, target = data.to(device), target.to(device)
            # batch = (data, target)
            output = model(data)
            
            test_loss += criterion(output, target).item() * bs  # sum up batch loss
            pred = output.max(1, keepdim=True)[1]  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()
            # print("Time for clean image: ", idx, time.time() - start_time)
            noise = atkmodel(data) * test_eps
            atkdata = clip_image(data + noise)
            atktarget = target_transform(target)
            # print("Time for atk image: ", idx, time.time() - start_time)
            
            atkoutput = model(atkdata)
            test_transform_loss += criterion(atkoutput, atktarget).item() * bs  # sum up batch loss
            atkpred = atkoutput.max(1, keepdim=True)[1]  # get the index of the max log-probability
            correct_transform += atkpred.eq(
                target_transform(target).view_as(atkpred)).sum().item()

    # print(f"Finish the post testing------------------------------------")
    # print(f"Time for post testing: {time.time() - start_time}")
    
    test_loss /= len(test_loader.dataset)
    test_transform_loss /= len(test_loader.dataset)

    correct /= len(test_loader.dataset)
    correct_transform /= len(test_loader.dataset)
    # print(f"\nTest result without retraining: ")
    print(
        '{}-Test set [{}]: Loss: clean {:.4f} poison {:.4f}, Accuracy: clean {:.2f} poison {:.2f}'.format(
            log_prefix, 0, 
            test_loss, test_transform_loss,
            correct, correct_transform
        ))
    
    if is_poison:
        csv_record.test_result.append([ident, epoch, test_loss, test_transform_loss, correct, correct_transform])
        
    print(f"Time for end testing: {time.time() - start_time}")
    return correct, correct_transform
    
class FederatedLearningTrainer:
    def __init__(self, *args, **kwargs):
        self.hyper_params = None

    def run(self, client_model, *args, **kwargs):
        raise NotImplementedError()

class FrequencyFederatedLearningTrainer(FederatedLearningTrainer):
    def __init__(self, arguments=None, lira_args=None, *args, **kwargs):
        #self.poisoned_emnist_dataset = arguments['poisoned_emnist_dataset']
        self.vanilla_model = arguments['vanilla_model']
        self.net_avg = arguments['net_avg']
        self.net_dataidx_map = arguments['net_dataidx_map']
        self.num_nets = arguments['num_nets']
        self.part_nets_per_round = arguments['part_nets_per_round']
        self.fl_round = arguments['fl_round']
        self.local_training_period = arguments['local_training_period']
        self.adversarial_local_training_period = arguments['adversarial_local_training_period']
        # self.federated_optimizer = arguments['federated_optimizer'] # default: FedAvg
        self.args_lr = arguments['args_lr']
        self.args_gamma = arguments['args_gamma']
        self.attacking_fl_rounds = arguments['attacking_fl_rounds']
        self.poisoned_emnist_train_loader = arguments['poisoned_emnist_train_loader']
        self.clean_train_loader = arguments['clean_train_loader']
        self.vanilla_emnist_test_loader = arguments['vanilla_emnist_test_loader']
        self.targetted_task_test_loader = arguments['targetted_task_test_loader']
        self.batch_size = arguments['batch_size']
        self.test_batch_size = arguments['test_batch_size']
        self.log_interval = arguments['log_interval']
        self.device = arguments['device']
        self.save_model = arguments['save_model']
        # self.num_dps_poisoned_dataset = int(arguments['num_dps_poisoned_dataset'] * (1.0+lira_args['attack_portion']))
        self.num_dps_poisoned_dataset = arguments['num_dps_poisoned_dataset']
        self.scale_weights_poison = arguments['scale_weights_poison']
        self.defense_technique = arguments["defense_technique"]
        self.norm_bound = arguments["norm_bound"]
        self.attack_method = arguments["attack_method"]
        self.dataset = arguments["dataset"]
        self.model = arguments["model"]
        self.criterion = nn.CrossEntropyLoss()
        self.eps = arguments['eps']
        self.eps_decay = arguments['eps_decay']
        self.poison_type = arguments['poison_type']
        self.model_replacement = arguments['model_replacement']
        self.project_frequency = arguments['project_frequency']
        self.adv_lr = arguments['adv_lr']
        self.atk_lr = arguments['atk_lr']
        self.prox_attack = arguments['prox_attack']
        self.attack_case = arguments['attack_case']
        self.stddev = arguments['stddev']
        self.atkmodel = arguments['atkmodel']
        self.tgtmodel = arguments['tgtmodel']
        self.lira_args = lira_args
        self.baseline = arguments['baseline']
        self.retrain = arguments['retrain']
        self.aggregator = arguments['aggregator'] # default: fedavg
        # self.create_net = arguments['create_net']
        self.scratch_model = arguments['scratch_model']
        self.atk_eps = arguments['atk_eps']
        self.atk_test_eps = arguments['atk_test_eps']
        self.scale_factor = arguments['scale']
        self.mask_ratio = 0.30 #TODO: modify it later, used only for Neurotoxin attack strategy
        self.flatten_net_avg = None
        self.folder_path = arguments['folder_path']
        self.historical_grad_mask = None
        self.atk_baseline = arguments['atk_baseline']

        if self.attack_method == "pgd":
            self.pgd_attack = True
            self.neurotoxin_attack = False
        elif self.attack_method == "neurotoxin":
            self.neurotoxin_attack = True
            self.pgd_attack = False
        else:
            self.pgd_attack = False
            self.neurotoxin_attack = False

        if arguments["defense_technique"] == "no-defense":
            self._defender = None
        elif arguments["defense_technique"] == "norm-clipping" or arguments["defense_technique"] == "norm-clipping-adaptive":
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "weak-dp":
            # doesn't really add noise. just clips
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "krum":
            self._defender = Krum(mode='krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "multi-krum":
            self._defender = Krum(mode='multi-krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "rfa":
            self._defender = RFA()
        elif arguments["defense_technique"] == "crfl":
            self._defender = CRFL()
        elif arguments["defense_technique"] == "rlr":
            pytorch_total_params = sum(p.numel() for p in self.net_avg.parameters())
            args_rlr={
                'aggr':'avg',
                'noise':0,
                'clip': 0,
                'server_lr': self.args_lr,
            }
            # theta = 2
            theta = 1
            self._defender = RLR(n_params=pytorch_total_params, device=self.device, args=args_rlr, robustLR_threshold=theta)
        elif arguments["defense_technique"] == "foolsgold":
            pytorch_total_params = sum(p.numel() for p in self.net_avg.parameters())
            self._defender = FoolsGold(num_clients=self.part_nets_per_round, num_classes=10, num_features=pytorch_total_params)
        else:
            NotImplementedError("Unsupported defense method !")

    def run(self, wandb_ins):
        expected_atk_threshold = 0.85
        cnt_masks = 0
        print(f"self.scale_weights_poison: {self.scale_weights_poison}")
        # if self.dataset == "mnist":
        #     expected_atk_threshold = 0.90
        # elif self.dataset == "cifar10":
        #     expected_atk_threshold = 0.90
            
        main_task_acc = []
        raw_task_acc = []
        backdoor_task_acc = []
        fl_iter_list = []
        adv_norm_diff_list = []
        wg_norm_list = []
        print(colored('Under defense technique = {}'.format(self.defense_technique), 'red'))
        
        # variables for LIRA algorithms only
        trainlosses = []
        best_acc_clean = 0
        best_acc_poison = 0
        clip_image = get_clip_image(self.dataset)
        attack_train_epoch = self.lira_args['train_epoch']
        target_transform = get_target_transform(self.lira_args)

        #模型参数个数
        pytorch_total_params = sum(p.numel() for p in self.net_avg.parameters())
        # The number of previous iterations to use FoolsGold on
        memory_size = 0
        delta_memory = np.zeros((self.num_nets, pytorch_total_params, memory_size))
        summed_deltas = np.zeros((self.num_nets, pytorch_total_params))
        
        basepath, checkpoint_path, bestmodel_path = create_paths(self.lira_args)
        print('========== PATHS ==========')
        print(f'Basepath: {basepath}')
        print(f'Checkpoint Model: {checkpoint_path}')
        print(f'Best Model: {bestmodel_path}')
        
        # LOAD_ATK_MODEL = False
        LOAD_ATK_MODEL = True
        checkpoint_path = bestmodel_path
        if os.path.exists(checkpoint_path) and not self.retrain:
            # Load previously saved models
            checkpoint = torch.load(checkpoint_path)
            print(colored('Load existing attack model from path {}'.format(checkpoint_path), 'red'))
            self.atkmodel.load_state_dict(checkpoint['atkmodel'], strict=True)
            atk_model_loaded = True
        else:
            # Create new model
            print(colored('Create new model from {}'.format(checkpoint_path), 'blue'))
            best_acc_clean = 0
            best_acc_poison = 0
            
        subpath_trigger_saved = f"{self.dataset}baseline_{self.baseline}_atkepc_{attack_train_epoch}_eps_{self.atk_eps}_test_eps_{self.atk_test_eps}"
        if self.baseline:
            subpath_trigger_saved = ""
        
        start_time = time.time()
        
        acc_clean, backdoor_acc = test_updated(self.lira_args, self.device, self.atkmodel, self.net_avg, target_transform, 
                                self.clean_train_loader, self.vanilla_emnist_test_loader, 0, self.lira_args['train_epoch'], clip_image, 
                                testoptimizer=None, log_prefix='External', epochs_per_test=3, dataset=self.dataset, criterion=self.criterion, 
                                subpath_saved=subpath_trigger_saved, is_poison=True, atk_baseline=self.atk_baseline)
        
        print(f"\n----------TEST FOR GLOBAL MODEL BEFORE FEDERATED TRAINING----------------")
        print(f"Main task acc: {round(acc_clean*100, 2)}%")
        print(f"Backdoor task acc: {round(backdoor_acc*100, 2)}%")
        print(f"--------------------------------------------------------------------------\n")
        
        
        # Try to save models at different fl rounds
        current_time = datetime.datetime.now().strftime('%b.%d_%H.%M.%S')
        folder_path = f'saved_models/{self.folder_path}_{"baseline" if self.baseline else "poison"}_{current_time}'
        try:
            os.makedirs(folder_path)
        except FileExistsError:
            logger.info('Folder already exists')
        logger.info(f'current path for saving: {folder_path}')


        local_acc_clean, local_acc_poison = 0.0, 0.0
        best_main_acc = 0.0
        # let's conduct multi-round training
        cur_training_eps = self.atk_eps
        start_decay_r = 0
        start_training_atk_model = True
        tgt_optimizer = optim.Adam(self.tgtmodel.parameters(), lr=self.atk_lr)
        atk_optimizer = optim.Adam(self.atkmodel.parameters(), lr=self.atk_lr)
        self.alternative_training = False
        #使用 flatten_moedel 函数将模型的所有参数展平为一个一维张量
        self.flatten_net_avg = flatten_model(copy.deepcopy(self.net_avg))
        
        # end_time = time.time()
        elapsed_time = time.time() - start_time
        # print time by # seconds
        print(f"Time for TEST FOR GLOBAL MODEL BEFORE FEDERATED TRAINING: {elapsed_time}")
        # print(f"Time for TEST FOR GLOBAL MODEL BEFORE FEDERATED TRAINING: {(end_time - start_time).seconds}")
        
        print("----------START FEDERATED TRAINING----------------")
        
        for flr in range(1, self.fl_round+1):
            print("---"*30)
            print("Start communication round {}".format(flr))
            
            if local_acc_poison >= expected_atk_threshold and self.retrain:
                self.retrain = False
                print(colored('Starting sub-training phase from flr = {}'.format(flr), 'red'))
                self.lira_args["test_eps"] = self.atk_test_eps
                start_decay_r = flr

            #这行代码的作用是确定当前训练时允许的最大扰动大小
            if self.retrain == False:
                cur_training_eps = max(self.atk_test_eps, exponential_decay(self.atk_eps, self.eps_decay, flr-start_decay_r))
                # self.alternative_training = !self.alternative_training
            g_user_indices = []

            if self.defense_technique == "norm-clipping-adaptive":
                # experimental
                norm_diff_collector = []
            
            # tau for FedNova
	        # rho: Parameter controlling the momentum SGD
            list_ni = [] # number of data points for each client
            list_tau = [] # tau for FedNova (number of batch training = batch * epoch)
            list_ai = [] # weight for FedNova

            # if self.neurotoxin_attack and ((backdoor_acc > 0.95 and cur_training_eps <= self.atk_test_eps)):
            if self.neurotoxin_attack and flr >= 400:
                self.baseline = True
                logger.info(colored(f"Stop poisoning from round {flr}!!!", "red"))    
            # if self.atk_baseline and flr >= 400:
            #     self.baseline = True
            #     logger.info(colored(f"Stop poisoning from round {flr}!!!", "red"))  
             
            if flr in self.attacking_fl_rounds and not self.baseline:
                # randomly select participating clients
                selected_node_indices = np.random.choice(self.num_nets, size=self.part_nets_per_round, replace=False)
                num_data_points = [len(self.net_dataidx_map[i]) for i in selected_node_indices]
                #这一行计算了每轮训练中所有被选择的客户端所拥有的总数据点数量
                total_num_dps_per_round = sum(num_data_points)
                
                if not self.retrain:
                    self.alternative_training = not self.alternative_training

                # net_freq = [num_data_points[i]/total_num_dps_per_round for i in range(self.part_nets_per_round)]
                # logger.info("Net freq: {}, FL round: {} without adversary".format(net_freq, flr))

                # we need to reconstruct the net list at the beginning
                #为各个节点分配网络
                net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
                logger.info("################## Starting fl round: {}".format(flr))

                # 假定第一个节点是攻击者
                # 把第一个节点上的数据替换为中毒数据数量
                num_data_points[0] = self.num_dps_poisoned_dataset
                net_freq = [num_data_points[i]/total_num_dps_per_round for i in range(self.part_nets_per_round)]
                logger.info("Net freq: {}, FL round: {} with adversary".format(net_freq, flr))
                # FEDNOVA
                for id_client in range(self.part_nets_per_round):
                    #如果当前客户端是第一个客户端，将计算本地训练轮数 list_tau。
                    if id_client == 0:
                        list_tau.append(np.ceil(num_data_points[id_client] / self.batch_size) * self.adversarial_local_training_period )
                    else:
                        list_tau.append(np.ceil(num_data_points[id_client] / self.batch_size) * self.local_training_period )
                    list_ni.append(num_data_points[id_client])
                #pdb.set_trace()

                # we need to reconstruct the net list at the beginning
                # net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
                # logger.info("################## Starting fl round: {}".format(flr))
                
                model_original = list(self.net_avg.parameters())
                # init_local_model = copy.deepcopy(self.net_avg)
                # super hacky but I'm doing this for the prox-attack
                wg_clone = copy.deepcopy(self.net_avg)
                # wg_hat = None
                # v0 = torch.nn.utils.parameters_to_vector(model_original)
                # wg_norm_list.append(torch.norm(v0).item())
                
                """ Start the FL process """
                for net_idx, net in enumerate(net_list):
                    init_local_model = dict()
                    for name, data in wg_clone.state_dict().items():
                        init_local_model[name] = wg_clone.state_dict()[name].clone()
                    #net  = net_list[net_idx]                
                    # if net_idx == 0:
                    #     global_user_idx = -1 # we assign "-1" as the indices of the attacker in global user indices
                    #     pass
                    # else:
                    #     global_user_idx = selected_node_indices[net_idx-1]
                    #     dataidxs = self.net_dataidx_map[global_user_idx]
                    #     train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size, 
                    #                                     self.test_batch_size, dataidxs) # also get the data loader
                    global_user_idx = selected_node_indices[net_idx]
                    dataidxs = self.net_dataidx_map[global_user_idx] # --> keep the same data for an attacker 

                    # if self.attack_case == "edge-case":
                    train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size, 
                                                    self.test_batch_size, dataidxs) # also get the data loader
                    # else:
                    #     NotImplementedError("Unsupported attack case ...")
                    if net_idx == 0:
                        train_dl_local = copy.deepcopy(self.clean_train_loader)
                    g_user_indices.append(global_user_idx)
                    
                    if net_idx == 0:
                        logger.info("@@@@@@@@ Working on client: {}, which is Attacker".format(net_idx))
                    else:
                        logger.info("@@@@@@@@ Working on client: {}, which is Global user: {}".format(net_idx, global_user_idx))

                    criterion = nn.CrossEntropyLoss()
                    optimizer = optim.SGD(net.parameters(), lr=self.args_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4) # epoch, net, train_loader, optimizer, criterion
                    adv_optimizer = optim.SGD(net.parameters(), lr=self.adv_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4) # looks like adversary needs same lr to hide with others
                    # prox_optimizer = optim.SGD(wg_clone.parameters(), lr=self.args_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4)

                    for param_group in optimizer.param_groups:
                        logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))
                    
                    if net_idx == 0:
                        """ for attacker only """
                        attack_portion = 0.0 if self.retrain else self.lira_args['attack_portion']
                        atk_alpha = self.lira_args['attack_alpha'] if not self.retrain else 1.0
                        
                        print(colored(f"At flr {flr}, test eps is {cur_training_eps}", "green"))

                        mask_grad_list = []
                        # if self.neurotoxin_attack and cur_training_eps != self.atk_eps:
                        pgd_attack = False if cur_training_eps == self.atk_eps else self.pgd_attack
                        clone_local_model = copy.deepcopy(net)

                        
                        for e in range(1, self.adversarial_local_training_period+1):
                            if not self.atk_baseline:
                                pdg_eps = self.eps*self.args_gamma**(flr-start_decay_r) if self.defense_technique in ('krum', 'multi-krum') else self.eps
                                train_lira(net, self.atkmodel, self.tgtmodel, optimizer, atk_optimizer, train_dl_local,
                                        self.criterion, cur_training_eps, atk_alpha, attack_portion, clip_image, 
                                        target_transform, 1, atkmodel_train=False, device=self.device,
                                        pgd_attack=pgd_attack, batch_idx=e, project_frequency=self.project_frequency,
                                        pgd_eps=pdg_eps, model_original=wg_clone, adv_optimizer=adv_optimizer, 
                                        mask_grad_list=mask_grad_list, aggregator=self.aggregator, wg_hat=wg_clone, local_e=e)
                            else:
                                train_baseline(net, optimizer, train_dl_local, self.criterion, self.dataset, self.device, target_transform, e)
                        
                        # mask_flat_all_layer = get_grad_mask_F3BA(clone_local_model, wg_clone, None, None)

                        # Retrain the attack model to avoid catastrophic forgetting            
                        if (self.retrain or self.alternative_training) and not self.atk_baseline:
                            # cp_net = copy.deepcopy(net)
                            # lira_train_eps = self.atk_eps if self.alternative_training else cur_training_eps
                            lira_train_eps = cur_training_eps
                            for e in range(1, attack_train_epoch+1):
                                train_lira(net, self.atkmodel, self.tgtmodel, None, atk_optimizer, train_dl_local,
                                        criterion, lira_train_eps, atk_alpha, 1.0, clip_image, target_transform, 1, 
                                        atkmodel_train=True, device=self.device, pgd_attack=False)
                                self.tgtmodel.load_state_dict(self.atkmodel.state_dict())

                        replacement_flr = 400 if self.dataset == "mnist" else 500
                        if self.model_replacement and cur_training_eps != self.atk_eps and flr >= replacement_flr:
                            v = torch.nn.utils.parameters_to_vector(net.parameters())
                            logger.info("Attacker before scaling : Norm = {}".format(torch.norm(v)))
                            # adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")
                            # logger.info("||w_bad - w_avg|| before scaling = {}".format(adv_norm_diff))
                            #模型替换缩放
                            for idx, param in enumerate(net.parameters()):
                                param.data = (param.data - model_original[idx])*(total_num_dps_per_round/self.num_dps_poisoned_dataset) + model_original[idx]
                            v = torch.nn.utils.parameters_to_vector(net.parameters())
                            logger.info("Attacker after scaling : Norm = {}".format(torch.norm(v)))

                        local_acc_clean, local_acc_poison = test_updated(self.lira_args, self.device, self.tgtmodel, net, target_transform, 
                            self.clean_train_loader, self.vanilla_emnist_test_loader, e, self.lira_args['train_epoch'], clip_image,
                            log_prefix='Internal', dataset=self.dataset, criterion=criterion, subpath_saved=subpath_trigger_saved, test_eps=cur_training_eps)
                        
                        if self.scale_weights_poison != 1.0 and flr%501 == 0:
                            v = torch.nn.utils.parameters_to_vector(net.parameters())
                            logger.info("Attacker before scaling with factor {}: Norm = {}".format(self.scale_weights_poison, torch.norm(v)))
                            # adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")
                            # logger.info("||w_bad - w_avg|| before scaling = {}".format(adv_norm_diff))

                            for idx, param in enumerate(net.parameters()):
                                param.data = (param.data - model_original[idx])*self.scale_weights_poison + model_original[idx]
                            v = torch.nn.utils.parameters_to_vector(net.parameters())
                            logger.info("Attacker after scaling : Norm = {}".format(torch.norm(v)))

                        local_acc_clean, local_acc_poison = test_updated(self.lira_args, self.device, self.tgtmodel, net, target_transform, 
                            self.clean_train_loader, self.vanilla_emnist_test_loader, flr, self.lira_args['train_epoch'], clip_image,
                            log_prefix='Internal', dataset=self.dataset, criterion=criterion, subpath_saved=subpath_trigger_saved, 
                            test_eps=cur_training_eps, is_poison=True, ident=global_user_idx, atk_baseline = self.atk_baseline)
                        
                    else:
                        """ for normal client only """
                        for e in range(1, self.local_training_period+1):
                            train_lira(net, self.atkmodel, self.tgtmodel, optimizer, None, train_dl_local,
                                        self.criterion, cur_training_eps, 1.0, 0.0, clip_image, target_transform, 1, 
                                        atkmodel_train=False, device=self.device, aggregator=self.aggregator, wg_hat=wg_clone)
                    
            else:
                wg_clone = copy.deepcopy(self.net_avg)
                      
                selected_node_indices = np.random.choice(self.num_nets, size=self.part_nets_per_round, replace=False)
                num_data_points = [len(self.net_dataidx_map[i]) for i in selected_node_indices]
                total_num_dps_per_round = sum(num_data_points)
                
                # FEDNOVA
                for id_client in range(self.part_nets_per_round):
                    list_tau.append(np.ceil(num_data_points[id_client] / self.batch_size) * self.local_training_period )
                    list_ni.append(num_data_points[id_client])
                  
                net_freq = [num_data_points[i]/total_num_dps_per_round for i in range(self.part_nets_per_round)]
                logger.info("Net freq: {}, FL round: {} without adversary".format(net_freq, flr))

                # we need to reconstruct the net list at the beginning
                net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
                logger.info("################## Starting fl round: {}".format(flr))

                # start the FL process
                
                # print(f"No attacker in this round {flr} with list client: {net_list}")
                
                for net_idx, net in enumerate(net_list):
                    start_time = time.time()
                    
                    global_user_idx = selected_node_indices[net_idx]
                    dataidxs = self.net_dataidx_map[global_user_idx] # --> keep the same data for an attacker 

                    # if self.attack_case == "edge-case":
                    train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size, 
                                                    self.test_batch_size, dataidxs) # also get the data loader
                    # else:
                    #     NotImplementedError("Unsupported attack case ...")
                    # if net_idx == 0:
                    #     train_dl_local = copy.deepcopy(self.clean_train_loader)
                    g_user_indices.append(global_user_idx)
                    
                    logger.info("@@@@@@@@ Working on client: {}, which is Global user: {}".format(net_idx, global_user_idx))

                    criterion = nn.CrossEntropyLoss()
                    optimizer = optim.SGD(net.parameters(), lr=self.args_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4) # epoch, net, train_loader, optimizer, criterion
                    for param_group in optimizer.param_groups:
                        logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))                    
                  
                    """ for normal client only """
                    for e in range(1, self.local_training_period+1):
                        train_lira(net, self.atkmodel, self.tgtmodel, optimizer, atk_optimizer, train_dl_local,
                                    self.criterion, cur_training_eps, 1.0, 0.0, clip_image, target_transform, 1, 
                                    atkmodel_train=False, device=self.device, aggregator=self.aggregator, wg_hat=wg_clone)
                    
                    honest_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="normal")

                    if self.defense_technique == "norm-clipping-adaptive":
                        # experimental
                        norm_diff_collector.append(honest_norm_diff)   
                    end_time = time.time()
                    logger.info("Time for local training client {}: {}".format(net_idx, end_time-start_time))

                adv_norm_diff_list.append(0)
                model_original = list(self.net_avg.parameters())
                v0 = torch.nn.utils.parameters_to_vector(model_original)
                wg_norm_list.append(torch.norm(v0).item())
            
            
            ### conduct defense here:
            print(colored(f"Conduct defense technique: {self.defense_technique}", "red"))
            if self.defense_technique == "foolsgold":
                delta = np.zeros((self.num_nets, pytorch_total_params))
                if memory_size > 0:
                    for net_idx, global_client_indx in enumerate(selected_node_indices):
                        flatten_local_model = flatten_model(net_list[net_idx])
                        local_update = flatten_local_model - self.flatten_net_avg
                        delta[global_client_indx,:] = local_update
                        # normalize delta
                        if np.linalg.norm(delta[global_client_indx, :]) > 1:
                            delta[global_client_indx, :] = delta[global_client_indx, :] / np.linalg.norm(delta[global_client_indx, :])

                        delta_memory[global_client_indx, :, flr % memory_size] = delta[global_client_indx, :]
                    # Track the total vector from each individual client
                    summed_deltas = np.sum(delta_memory, axis=2)      
                else:
                    for net_idx, global_client_indx in enumerate(selected_node_indices):
                        flatten_local_model = flatten_model(net_list[net_idx])
                        local_update = flatten_local_model - self.flatten_net_avg
                        local_update = local_update.detach().cpu().numpy()
                        delta[global_client_indx,:] = local_update
                        # normalize delta
                        if np.linalg.norm(delta[global_client_indx, :]) > 1:
                            delta[global_client_indx, :] = delta[global_client_indx, :] / np.linalg.norm(delta[global_client_indx, :])
                    # Track the total vector from each individual client
                    # print(f"delta={delta[selected_node_indices,:]}")
                    # print(f"summed_deltas[selected_node_indices,:].shape is: {summed_deltas[selected_node_indices,:].shape}")

                    summed_deltas[selected_node_indices,:] = summed_deltas[selected_node_indices,:] + delta[selected_node_indices,:]
                    # print(f"summed_deltas.shape is: {summed_deltas.shape}")
                    # print(f"summed_deltas={summed_deltas[selected_node_indices,:]}")
            if self.defense_technique == "no-defense":
                pass
            elif self.defense_technique == "norm-clipping":
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "norm-clipping-adaptive":
                # we will need to adapt the norm diff first before the norm diff clipping
                logger.info("#### Let's Look at the Nom Diff Collector : {} ....; Mean: {}".format(norm_diff_collector, 
                    np.mean(norm_diff_collector)))
                self._defender.norm_bound = np.mean(norm_diff_collector)
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "weak-dp":
                # this guy is just going to clip norm. No noise added here
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net,
                                        global_model=self.net_avg,)
            elif self.defense_technique == "krum":
                print("start performing krum...")
                net_list, net_freq = self._defender.exec(client_models=net_list, 
                                                        num_dps=[self.num_dps_poisoned_dataset]+num_data_points,
                                                        g_user_indices=g_user_indices,
                                                        device=self.device)
                print(f"net_list = {len(net_list)}, net_freq={net_freq}")
            elif self.defense_technique == "multi-krum":
                net_list, net_freq = self._defender.exec(client_models=net_list, 
                                                        num_dps=[self.num_dps_poisoned_dataset]+num_data_points,
                                                        g_user_indices=g_user_indices,
                                                        device=self.device)
            elif self.defense_technique == 'rlr':
                print(f"num_data_points: {num_data_points}")
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                        num_dps=num_data_points,
                                                        global_model=copy.deepcopy(self.net_avg))
            elif self.defense_technique == "foolsgold":
                print(f"Performing foolsgold...")
                # net_list, net_freq = self._defender.exec(client_models=net_list, delta = delta[selected_node_indices,:] ,summed_deltas=summed_deltas[selected_node_indices,:], net_avg=self.net_avg, r=flr, device=self.device)
                wv = self._defender.exec(client_models=net_list, delta = delta[selected_node_indices,:] ,summed_deltas=summed_deltas[selected_node_indices,:], net_avg=self.net_avg, r=flr, device=self.device)
                net_freq = wv
                print(f"net_freq: {net_freq}")
            elif self.defense_technique == "rfa":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                        net_freq=net_freq,
                                                        maxiter=500,
                                                        eps=1e-5,
                                                        ftol=1e-7,
                                                        device=self.device)
            elif self.defense_technique == "crfl":
                cp_net_avg = copy.deepcopy(self.net_avg)
                # pseudo_global_net = fed_avg_aggregator(cp_net_avg, net_list, net_freq, device=self.device, model=self.model)
                net_list, net_freq = self._defender.exec(target_model=cp_net_avg,
                                                         epoch=flr,
                                                         sigma_param=0.01,
                                                         dataset_name=self.dataset,
                                                         device=self.device)
            else:
                NotImplementedError("Unsupported defense method !")
            
            
            # after local training periods
            fed_aggregator = self.aggregator
            print(f"\nStart performing fed aggregator: {fed_aggregator}...")
            fed_nova_rho = 0.1
            list_ai = [(tau - fed_nova_rho * (1 - pow(fed_nova_rho, tau)) / (1 - fed_nova_rho)) / (1 - fed_nova_rho) for tau in list_tau]
            
            if fed_aggregator == "fedavg":
                self.net_avg = fed_avg_aggregator(self.net_avg, net_list, net_freq, device=self.device, model=self.model)
            elif fed_aggregator == "fednova":
                self.net_avg = fed_nova_aggregator(self.net_avg, net_list, list_ni, self.device, list_ai)
            else:
                NotImplementedError(f"Unsupported fed agg {fed_aggregator} method !")

            
            acc_clean, backdoor_acc = test_updated(args=self.lira_args, device=self.device, atkmodel=self.tgtmodel, model=self.net_avg, 
                                                   target_transform=target_transform, train_loader=self.clean_train_loader, 
                                                   test_loader=self.vanilla_emnist_test_loader,epoch=flr, trainepoch=self.lira_args['train_epoch'], 
                                                   clip_image=clip_image, testoptimizer=None, log_prefix='External', epochs_per_test=3, 
                                                   dataset=self.dataset, criterion=criterion, subpath_saved=subpath_trigger_saved, test_eps=cur_training_eps,
                                                   is_poison=True, atk_baseline=self.atk_baseline)
            
            
            if self.defense_technique == "weak-dp":
                # add noise to self.net_avg
                noise_adder = AddNoise(stddev=self.stddev)
                noise_adder.exec(client_model=self.net_avg,
                                                device=self.device)

            v = torch.nn.utils.parameters_to_vector(self.net_avg.parameters())
            logger.info("############ Averaged Model : Norm {}".format(torch.norm(v)))

            calc_norm_diff(gs_model=self.net_avg, vanilla_model=self.vanilla_model, epoch=0, fl_round=flr, mode="avg")
            
            logger.info("Measuring the accuracy of the averaged global model, FL round: {} ...".format(flr))

            # START SAVING THE MODEL FOR FURTHER ANALYSIS
            model_name = '{0}/model_last_flr_{1}.pt'.format(folder_path, flr)
            atk_model_name = '{0}/atkmodel_last_flr_{1}.pt'.format(folder_path, flr)
            if flr >= 200 and flr%50 == 0 and self.save_model:
                torch.save(self.net_avg.state_dict(), model_name)
                torch.save(self.tgtmodel.state_dict(), atk_model_name)
            
            wandb_logging_items = {
                'fl_iter': flr,
                'main_task_acc': acc_clean*100.0,
                'backdoor_task_acc': backdoor_acc*100.0, 
                'local_best_acc_clean': best_acc_clean,
                'local_best_acc_poison': best_acc_poison,
                'local_MA': local_acc_clean*100.0,
                'local_BA': local_acc_poison*100.0,
                'cur_training_eps': cur_training_eps,
            }
            if wandb_ins:
                print(f"start logging to wandb")
                wandb_ins.log({"General Information": wandb_logging_items})
            csv_record.save_result_csv(flr, True, folder_path)
            raw_acc = 0
            overall_acc = acc_clean
            fl_iter_list.append(flr)
            main_task_acc.append(overall_acc)
            raw_task_acc.append(raw_acc)
            backdoor_task_acc.append(backdoor_acc)


        saved_path = "checkpoint/last_model"
        torch.save(self.net_avg.state_dict(), f"{saved_path}/lira_{self.dataset}_{self.model}_{self.atk_test_eps}.pt")

        df = pd.DataFrame({'fl_iter': fl_iter_list, 
                            'main_task_acc': main_task_acc, 
                            'backdoor_acc': backdoor_task_acc, 
                            'raw_task_acc':raw_task_acc, 
                            'adv_norm_diff': adv_norm_diff_list, 
                            'wg_norm': wg_norm_list
                            })
       
        if self.poison_type == 'ardis':
            # add a row showing initial accuracies
            df1 = pd.DataFrame({'fl_iter': [0], 'main_task_acc': [88], 'backdoor_acc': [11], 'raw_task_acc': [0], 'adv_norm_diff': [0], 'wg_norm': [0]})
            df = pd.concat([df1, df])

        results_filename = get_results_filename(self.poison_type, self.attack_method, self.model_replacement, self.project_frequency,
                self.defense_technique, self.norm_bound, self.prox_attack, False, self.model)

        df.to_csv(results_filename, index=False)
        logger.info("Wrote accuracy results to: {}".format(results_filename))


class FixedPoolFederatedLearningTrainer(FederatedLearningTrainer):
    def __init__(self, arguments=None, lira_args=None, *args, **kwargs):
        #self.poisoned_emnist_dataset = arguments['poisoned_emnist_dataset']
        self.vanilla_model = arguments['vanilla_model']
        self.net_avg = arguments['net_avg']
        self.net_dataidx_map = arguments['net_dataidx_map']
        self.num_nets = arguments['num_nets']
        self.part_nets_per_round = arguments['part_nets_per_round']
        self.fl_round = arguments['fl_round']
        self.local_training_period = arguments['local_training_period']
        self.adversarial_local_training_period = arguments['adversarial_local_training_period']
        self.args_lr = arguments['args_lr']
        self.args_gamma = arguments['args_gamma']
        self.attacker_pool_size = arguments['attacker_pool_size']
        self.poisoned_emnist_train_loader = arguments['poisoned_emnist_train_loader']
        self.clean_train_loader = arguments['clean_train_loader']
        self.vanilla_emnist_test_loader = arguments['vanilla_emnist_test_loader']
        self.targetted_task_test_loader = arguments['targetted_task_test_loader']
        self.batch_size = arguments['batch_size']
        self.test_batch_size = arguments['test_batch_size']
        self.log_interval = arguments['log_interval']
        self.device = arguments['device']
        self.dataset = arguments["dataset"]
        self.model = arguments["model"]
        self.num_dps_poisoned_dataset = arguments['num_dps_poisoned_dataset']
        self.defense_technique = arguments["defense_technique"]
        self.norm_bound = arguments["norm_bound"]
        self.attack_method = arguments["attack_method"]
        self.criterion = nn.CrossEntropyLoss()
        self.eps = arguments['eps']
        self.eps_decay = arguments['eps_decay']
        self.scale_weights_poison = arguments['scale_weights_poison']
        self.dataset = arguments["dataset"]
        self.poison_type = arguments['poison_type']
        self.model_replacement = arguments['model_replacement']
        self.project_frequency = arguments['project_frequency']
        self.adv_lr = arguments['adv_lr']
        self.atk_lr = arguments['atk_lr']

        self.prox_attack = arguments['prox_attack']
        self.attack_case = arguments['attack_case']
        self.stddev = arguments['stddev']
        self.atkmodel = arguments['atkmodel']
        self.tgtmodel = arguments['tgtmodel']
        self.lira_args = lira_args
        self.baseline = arguments['baseline']
        self.retrain = arguments['retrain']
        self.scratch_model = arguments['scratch_model']
        self.atk_eps = arguments['atk_eps']
        self.atk_test_eps = arguments['atk_test_eps']
        self.scale_factor = arguments['scale']
        self.flatten_net_avg = None

        logger.info("Posion type! {}".format(self.poison_type))

        if self.attack_method == "pgd":
            self.pgd_attack = True
        elif self.attack_method == "neurotoxin":
            self.neurotoxin_attack = True
        else:
            self.pgd_attack = False
            self.neurotoxin_attack = False


        if arguments["defense_technique"] == "no-defense":
            self._defender = None
        elif arguments["defense_technique"] == "norm-clipping" or arguments["defense_technique"] == "norm-clipping-adaptive":
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "weak-dp":
            # doesn't really add noise. just clips
            self._defender = WeightDiffClippingDefense(norm_bound=arguments['norm_bound'])
        elif arguments["defense_technique"] == "krum":
            self._defender = Krum(mode='krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "multi-krum":
            self._defender = Krum(mode='multi-krum', num_workers=self.part_nets_per_round, num_adv=1)
        elif arguments["defense_technique"] == "rfa":
            self._defender = RFA()
        elif arguments["defense_technique"] == "rlr":
            pytorch_total_params = sum(p.numel() for p in self.net_avg.parameters())
            args_rlr={
                'aggr':'avg',
                'noise':0,
                'clip': 0,
                'server_lr': self.args_lr,
            }
            # theta = 4 if arguments['dataset'] == 'cifar10' else 8
            # theta = int(arguments['part_nets_per_round']*self.attacker_pool_size/self.num_nets)+1
            theta = 8
            self._defender = RLR(n_params=pytorch_total_params, device=self.device, args=args_rlr, robustLR_threshold=theta)
        elif arguments["defense_technique"] == "foolsgold":
            pytorch_total_params = sum(p.numel() for p in self.net_avg.parameters())
            self._defender = FoolsGold(num_clients=self.part_nets_per_round, num_classes=10, num_features=pytorch_total_params)
        else:
            NotImplementedError("Unsupported defense method !")

        self.__attacker_pool = np.random.choice(self.num_nets, self.attacker_pool_size, replace=False)

    def run(self, wandb_ins):
        expected_atk_threshold = 0.90
        # if self.dataset == "mnist":
        #     expected_atk_threshold = 0.90
        # elif self.dataset == "cifar10":
        #     expected_atk_threshold = 0.90
        main_task_acc = []
        raw_task_acc = []
        backdoor_task_acc = []
        fl_iter_list = []
        adv_norm_diff_list = []
        wg_norm_list = []
        print(colored('Under defense technique = {}'.format(self.defense_technique), 'red'))
        
        # variables for LIRA algorithms only
        trainlosses = []
        best_acc_clean = 0
        best_acc_poison = 0
        avoid_cls_reinit = True
        clip_image = get_clip_image(self.dataset)
        attack_train_epoch = self.lira_args['train_epoch']
        target_transform = get_target_transform(self.lira_args)
        tgt_optimizer = None
        loss_list = []

        pytorch_total_params = sum(p.numel() for p in self.net_avg.parameters())
        # The number of previous iterations to use FoolsGold on
        memory_size = 0
        delta_memory = np.zeros((self.num_nets, pytorch_total_params, memory_size))
        summed_deltas = np.zeros((self.num_nets, pytorch_total_params))
        
        basepath, checkpoint_path, bestmodel_path = create_paths(self.lira_args)
        print('========== PATHS ==========')
        print(f'Basepath: {basepath}')
        print(f'Checkpoint Model: {checkpoint_path}')
        print(f'Best Model: {bestmodel_path}')
        
        # LOAD_ATK_MODEL = False
        LOAD_ATK_MODEL = True
        checkpoint_path = bestmodel_path
        if os.path.exists(checkpoint_path) and not self.retrain:
            # Load previously saved models
            checkpoint = torch.load(checkpoint_path)
            print(colored('Load existing attack model from path {}'.format(checkpoint_path), 'red'))
            self.atkmodel.load_state_dict(checkpoint['atkmodel'], strict=True)
            atk_model_loaded = True
        else:
            # Create new model
            print(colored('Create new model from {}'.format(checkpoint_path), 'blue'))
            best_acc_clean = 0
            best_acc_poison = 0
            trainlosses = []
            start_epoch = 1
            
        subpath_trigger_saved = f"{self.dataset}baseline_{self.baseline}_atkepc_{attack_train_epoch}_eps_{self.atk_eps}_test_eps_{self.atk_test_eps}"
        if self.baseline:
            subpath_trigger_saved = ""
            
        acc_clean, backdoor_acc = test_updated(self.lira_args, self.device, self.atkmodel, self.net_avg, target_transform, 
                                self.clean_train_loader, self.vanilla_emnist_test_loader, 0, self.lira_args['train_epoch'], clip_image, 
                                testoptimizer=None, log_prefix='External', epochs_per_test=3, dataset=self.dataset, criterion=self.criterion, subpath_saved=subpath_trigger_saved)
        
        print(f"\n----------TEST FOR GLOBAL MODEL BEFORE FEDERATED TRAINING----------------")
        print(f"Main task acc: {round(acc_clean*100, 2)}%")
        print(f"Backdoor task acc: {round(backdoor_acc*100, 2)}%")
        print(f"--------------------------------------------------------------------------\n")
        local_acc_clean, local_acc_poison = 0.0, 0.0
        best_main_acc = 0.0
        # let's conduct multi-round training
        cur_training_eps = self.atk_eps
        start_decay_r = 0
        start_training_atk_model = True
        tgt_optimizer = optim.Adam(self.tgtmodel.parameters(), lr=self.atk_lr)
        atk_optimizer = optim.Adam(self.atkmodel.parameters(), lr=self.atk_lr)
        self.alternative_training = False
        self.flatten_net_avg = flatten_model(self.net_avg)
        list_ni = []
        list_tau = []
        list_ai = []
        
        for flr in range(1, self.fl_round+1):
            # randomly select participating clients
            # 在当前版本中，我们在每 FL 轮中对“part_nets_per_round”进行采样，因为我们假设攻击者始终会参与
            if local_acc_poison >= expected_atk_threshold and self.retrain:
                self.retrain = False
                print(colored('Starting sub-training phase from flr = {}'.format(flr), 'red'))
                self.lira_args["test_eps"] = self.atk_test_eps
                start_decay_r = flr
            
            if self.retrain == False:
                cur_training_eps = max(self.atk_test_eps, exponential_decay(self.atk_eps, self.eps_decay, flr-start_decay_r))
            g_user_indices = []

            selected_node_indices = np.random.choice(self.num_nets, size=self.part_nets_per_round, replace=False)
            selected_attackers = [idx for idx in selected_node_indices if idx in self.__attacker_pool]
            selected_honest_users = [idx for idx in selected_node_indices if idx not in self.__attacker_pool]
            logger.info("Selected Attackers in FL iteration-{}: {}".format(flr, selected_attackers))

            num_data_points = []
            for sni in selected_node_indices:
                if sni in selected_attackers:
                    num_data_points.append(self.num_dps_poisoned_dataset)
                else:
                    num_data_points.append(len(self.net_dataidx_map[sni]))

            total_num_dps_per_round = sum(num_data_points)
            net_freq = [num_data_points[i]/total_num_dps_per_round for i in range(self.part_nets_per_round)]
            logger.info("Net freq: {}, FL round: {} with adversary".format(net_freq, flr)) 

            # we need to reconstruct the net list at the beginning
            total_attackers = len(selected_attackers)
            net_list = [copy.deepcopy(self.net_avg) for _ in range(self.part_nets_per_round)]
            # atk_model_list = [copy.deepcopy(self.atkmodel) for _ in range(total_attackers)] # copy list of attack model to train independently
            atk_model_list = {g_id: copy.deepcopy(self.atkmodel) for g_id in selected_attackers} # copy list of attack model to train independently

            logger.info("################## Starting fl round: {}".format(flr))
            model_original = list(self.net_avg.parameters())
            # super hacky but I'm doing this for the prox-attack
            wg_clone = copy.deepcopy(self.net_avg)
            wg_hat = None
            v0 = torch.nn.utils.parameters_to_vector(model_original)
            wg_norm_list.append(torch.norm(v0).item())

            # calculate the tau for each client
            for sni in selected_node_indices:
                if sni in selected_attackers:
                    list_ni.append(self.num_dps_poisoned_dataset)
                    list_tau.append(np.ceil(self.num_dps_poisoned_dataset/ self.batch_size) * self.adversarial_local_training_period)
                else:
                    list_ni.append(len(self.net_dataidx_map[sni]))
                    list_tau.append(np.ceil(len(self.net_dataidx_map[sni])/ self.batch_size) * self.local_training_period)
                    

            #     # start the FL process
            for net_idx, global_user_idx in enumerate(selected_node_indices):
                net  = net_list[net_idx]
                if global_user_idx in selected_attackers:
                    pass
                else:
                    dataidxs = self.net_dataidx_map[global_user_idx]
         
                    # add p-percent edge-case attack here:
                    if self.attack_case == "edge-case":
                        train_dl_local, _ = get_dataloader(self.dataset, './data', self.batch_size, 
                                                       self.test_batch_size, dataidxs) # also get the data loader
                    elif self.attack_case in ("normal-case", "almost-edge-case"):
                        train_dl_local, _ = get_dataloader_normal_case(self.dataset, './data', self.batch_size, 
                                                            self.test_batch_size, dataidxs, user_id=global_user_idx,
                                                            num_total_users=self.num_nets,
                                                            poison_type=self.poison_type,
                                                            ardis_dataset=self.ardis_dataset,
                                                            attack_case=self.attack_case) # also get the data loader
                    else:
                        NotImplementedError("Unsupported attack case ...")
                
                logger.info("@@@@@@@@ Working on client (global-index): {}, which {}-th user in the current round".format(global_user_idx, net_idx))

           
                criterion = nn.CrossEntropyLoss()
                optimizer = optim.SGD(net.parameters(), lr=self.args_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4) # epoch, net, train_loader, optimizer, criterion
                adv_optimizer = optim.SGD(net.parameters(), lr=self.adv_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4) # looks like adversary needs same lr to hide with others
                prox_optimizer = optim.SGD(wg_clone.parameters(), lr=self.args_lr*self.args_gamma**(flr-1), momentum=0.9, weight_decay=1e-4)
                for param_group in optimizer.param_groups:
                    logger.info("Effective lr in FL round: {} is {}".format(flr, param_group['lr']))


                current_adv_norm_diff_list = []
                if global_user_idx in selected_attackers:
                    """ for attacker only """
                    # local_atkmodel = atk_model_list[global_user_idx]
                    atk_optimizer = optim.Adam(atk_model_list[global_user_idx].parameters(), lr=self.atk_lr)
                    attack_portion = 0.0 if self.retrain else self.lira_args['attack_portion']
                    atk_alpha = self.lira_args['attack_alpha'] if not self.retrain else 1.0
                    
                    print(colored(f"At flr {flr}, test eps is {cur_training_eps}", "green"))

                 
                    if self.prox_attack:
                        # estimate w_hat
                        for inner_epoch in range(1, self.local_training_period+1):
                            estimate_wg(wg_clone, self.device, self.clean_train_loader, prox_optimizer, inner_epoch, log_interval=self.log_interval, criterion=self.criterion)
                        wg_hat = wg_clone

                 
                    """LIRA CODE"""
                    for e in range(1, self.adversarial_local_training_period+1):
                        #if self.defense_technique in ('krum', 'multi-krum'):  
                        pdg_eps = self.eps*self.args_gamma**(flr-1) if self.defense_technique in ('krum', 'multi-krum') else self.eps
                        
                        train_lira(net, atk_model_list[global_user_idx], self.tgtmodel, optimizer, 
                                atk_optimizer, self.clean_train_loader, self.criterion, cur_training_eps, 
                                atk_alpha, attack_portion, clip_image, target_transform, 1, 
                                atkmodel_train=False, device=self.device,
                                pgd_attack=self.pgd_attack, batch_idx=e, project_frequency=self.project_frequency,
                                pgd_eps=pdg_eps, model_original=model_original, adv_optimizer=adv_optimizer)
                    if self.retrain or self.alternative_training:
                        cp_net = copy.deepcopy(net)
                        for e in range(1, attack_train_epoch+1):
                            train_lira(net, atk_model_list[global_user_idx], self.tgtmodel, None, atk_optimizer, self.clean_train_loader,
                                    criterion, cur_training_eps, atk_alpha, 1.0, clip_image, target_transform, 1, 
                                    atkmodel_train=True, 
                                    device=self.device)

                    test(net, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
                    
                    # Testing for local model and local attack model
                    local_acc_clean, local_acc_poison = test_updated(self.lira_args, self.device, atk_model_list[global_user_idx], net, target_transform, 
                        self.clean_train_loader, self.vanilla_emnist_test_loader, e, self.lira_args['train_epoch'], clip_image,
                        log_prefix='Internal', dataset=self.dataset, criterion=criterion, subpath_saved=subpath_trigger_saved, test_eps=cur_training_eps)
                    # test(net, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)

                    # if model_replacement scale models
                    if self.model_replacement:
                        v = torch.nn.utils.parameters_to_vector(net.parameters())
                        logger.info("Attacker before scaling : Norm = {}".format(torch.norm(v)))
                        # adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")
                        # logger.info("||w_bad - w_avg|| before scaling = {}".format(adv_norm_diff))

                        for idx, param in enumerate(net.parameters()):
                            param.data = (param.data - model_original[idx])*(total_num_dps_per_round/self.num_dps_poisoned_dataset) + model_original[idx]
                        v = torch.nn.utils.parameters_to_vector(net.parameters())
                        logger.info("Attacker after scaling : Norm = {}".format(torch.norm(v)))

                    # at here we can check the distance between w_bad and w_g i.e. `\|w_bad - w_g\|_2`
                    # we can print the norm diff out for debugging
                    adv_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="bad")
                    current_adv_norm_diff_list.append(adv_norm_diff)

                    if self.defense_technique == "norm-clipping-adaptive":
                        # experimental
                        norm_diff_collector.append(adv_norm_diff)
                else:
                    # for e in range(1, self.local_training_period+1):
                    #    train(net, self.device, train_dl_local, optimizer, e, log_interval=self.log_interval, criterion=self.criterion)                
                    # # at here we can check the distance between w_normal and w_g i.e. `\|w_bad - w_g\|_2`
                    # calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="normal")

                    for e in range(1, self.local_training_period+1):
                        train_lira(net, self.atkmodel, self.tgtmodel, optimizer, atk_optimizer, train_dl_local,
                                    self.criterion, cur_training_eps, 1.0, 0.0, clip_image, target_transform, 1, 
                                    atkmodel_train=False, device=self.device)
                    #    train(net, self.device, train_dl_local, optimizer, e, log_interval=self.log_interval, criterion=self.criterion)                
                       # at here we can check the distance between w_normal and w_g i.e. `\|w_bad - w_g\|_2`
                    # we can print the norm diff out for debugging
                    honest_norm_diff = calc_norm_diff(gs_model=net, vanilla_model=self.net_avg, epoch=e, fl_round=flr, mode="normal")
                    
                    if self.defense_technique == "norm-clipping-adaptive":
                        # experimental
                        norm_diff_collector.append(honest_norm_diff)
            if self.defense_technique == "foolsgold":
                delta = np.zeros((self.num_nets, pytorch_total_params))
                if memory_size > 0:
                    for net_idx, global_client_indx in enumerate(selected_node_indices):
                        flatten_local_model = flatten_model(net_list[net_idx])
                        local_update = flatten_local_model - self.flatten_net_avg
                        delta[global_client_indx,:] = local_update
                        # normalize delta
                        if np.linalg.norm(delta[global_client_indx, :]) > 1:
                            delta[global_client_indx, :] = delta[global_client_indx, :] / np.linalg.norm(delta[global_client_indx, :])

                        delta_memory[global_client_indx, :, flr % memory_size] = delta[global_client_indx, :]
                    # Track the total vector from each individual client
                    summed_deltas = np.sum(delta_memory, axis=2)      
                else:
                    for net_idx, global_client_indx in enumerate(selected_node_indices):
                        flatten_local_model = flatten_model(net_list[net_idx])
                        local_update = flatten_local_model - self.flatten_net_avg
                        local_update = local_update.detach().cpu().numpy()
                        delta[global_client_indx,:] = local_update
                        # normalize delta
                        if np.linalg.norm(delta[global_client_indx, :]) > 1:
                            delta[global_client_indx, :] = delta[global_client_indx, :] / np.linalg.norm(delta[global_client_indx, :])
                    # Track the total vector from each individual client
                    # print(f"delta={delta[selected_node_indices,:]}")
                    # print(f"summed_deltas[selected_node_indices,:].shape is: {summed_deltas[selected_node_indices,:].shape}")

                    summed_deltas[selected_node_indices,:] = summed_deltas[selected_node_indices,:] + delta[selected_node_indices,:]
                    # print(f"summed_deltas.shape is: {summed_deltas.shape}")
                    # print(f"summed_deltas={summed_deltas[selected_node_indices,:]}")
 
            ### conduct defense here:
            if self.defense_technique == "no-defense":
                pass
            elif self.defense_technique == "norm-clipping":
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "norm-clipping-adaptive":
                # we will need to adapt the norm diff first before the norm diff clipping
                logger.info("#### Let's Look at the Nom Diff Collector : {} ....; Mean: {}".format(norm_diff_collector, 
                    np.mean(norm_diff_collector)))
                self._defender.norm_bound = np.mean(norm_diff_collector)
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net, global_model=self.net_avg)
            elif self.defense_technique == "weak-dp":
                # this guy is just going to clip norm. No noise added here
                for net_idx, net in enumerate(net_list):
                    self._defender.exec(client_model=net,
                                        global_model=self.net_avg,)
            elif self.defense_technique == "krum":
                net_list, net_freq = self._defender.exec(client_models=net_list, 
                                                        num_dps=num_data_points,
                                                        g_user_indices=selected_node_indices,
                                                        device=self.device)
            elif self.defense_technique == "multi-krum":
                net_list, net_freq = self._defender.exec(client_models=net_list, 
                                                        num_dps=num_data_pointsos,
                                                        g_user_indices=selected_node_indices,
                                                        device=self.device)
            elif self.defense_technique == 'rlr':
                print(f"num_data_points: {num_data_points}")
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                        num_dps=num_data_points,
                                                        global_model=self.net_avg)
            elif self.defense_technique == "foolsgold":
                print(f"Performing foolsgold...")
                net_list, net_freq = self._defender.exec(client_models=net_list, delta = delta[selected_node_indices,:] ,summed_deltas=summed_deltas[selected_node_indices,:], net_avg=self.net_avg, r=flr, device=self.device)
            elif self.defense_technique == "rfa":
                net_list, net_freq = self._defender.exec(client_models=net_list,
                                                        net_freq=net_freq,
                                                        maxiter=500,
                                                        eps=1e-5,
                                                        ftol=1e-7,
                                                        device=self.device)
            else:
                NotImplementedError("Unsupported defense method !")

            # aggregation of the attack models distributed among participants
            atk_model_freq = [1.0/total_attackers for _ in range(total_attackers)]
            if total_attackers:
                # atk_vectorized_nets = [vectorize_net(cm).detach().cpu().numpy() for cm in list(atk_model_list.values())]
                # avg_atk_net_vec = np.average(atk_vectorized_nets, weights=atk_model_freq, axis=0).astype(np.float32)
                # load_model_weight(self.atkmodel, torch.from_numpy(avg_atk_net_vec.astype(np.float32)).to(self.device))
                self.atkmodel = fed_avg_aggregator(self.atkmodel, list(atk_model_list.values()), atk_model_freq, device=self.device, model="atkmodel")
                self.tgtmodel.load_state_dict(self.atkmodel.state_dict())
            # testing for aggregated attack model
            local_acc_clean, local_acc_poison = test_updated(self.lira_args, self.device, self.tgtmodel, net, target_transform, 
                self.clean_train_loader, self.vanilla_emnist_test_loader, e, self.lira_args['train_epoch'], clip_image,
                log_prefix='Internal', dataset=self.dataset, criterion=criterion, subpath_saved=subpath_trigger_saved, test_eps=cur_training_eps)
            # after local training periods
            
            fed_aggregator = self.aggregator
            print(f"\nStart performing fed aggregator: {fed_aggregator}...")
            fed_nova_rho = 0.1
            list_ai = [(tau - fed_nova_rho * (1 - pow(fed_nova_rho, tau)) / (1 - fed_nova_rho)) / (1 - fed_nova_rho) for tau in list_tau]
                
            if fed_aggregator == "fednova":
                self.net_avg = fed_nova_aggregator(self.net_avg, net_list, list_ni, device, list_ai)
            else:
                self.net_avg = fed_avg_aggregator(self.net_avg, net_list, net_freq, device=self.device, model=self.model)
            
            self.flatten_net_avg = flatten_model(self.net_avg)
            #overall_acc = test(self.net_avg, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
            #backdoor_acc = test(self.net_avg, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)
            # overall_acc, raw_acc = test(self.net_avg, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
            # backdoor_acc, _ = test(self.net_avg, self.device, self.targetted_task_test_loader, test_batch_size=self.test_batch_size, criterion=self.criterion, mode="targetted-task", dataset=self.dataset, poison_type=self.poison_type)
 
            """GLOBAL TESTING FOR LIRA"""
            overall_acc, raw_acc = test(self.net_avg, self.device, self.vanilla_emnist_test_loader, test_batch_size=self.test_batch_size, 
                                        criterion=self.criterion, mode="raw-task", dataset=self.dataset, poison_type=self.poison_type)
            

            acc_clean, backdoor_acc = test_updated(args=self.lira_args, device=self.device, atkmodel=self.tgtmodel, model=self.net_avg, target_transform=target_transform, 
                                    train_loader=self.clean_train_loader, test_loader=self.vanilla_emnist_test_loader,epoch=flr, trainepoch=self.lira_args['train_epoch'], 
                                    clip_image=clip_image, testoptimizer=None, log_prefix='External', epochs_per_test=3, dataset=self.dataset, criterion=criterion, 
                                    subpath_saved=subpath_trigger_saved, test_eps=cur_training_eps)

            if self.defense_technique == "weak-dp":
                # add noise to self.net_avg
                noise_adder = AddNoise(stddev=self.stddev)
                noise_adder.exec(client_model=self.net_avg,
                                                device=self.device)

            calc_norm_diff(gs_model=self.net_avg, vanilla_model=self.net_avg, epoch=0, fl_round=flr, mode="avg")
            logger.info("Measuring the accuracy of the averaged global model, FL round: {} ...".format(flr))

            if overall_acc > best_main_acc:
                best_main_acc = overall_acc
                saved_path = f"checkpoint/bestmodel"
                atk_saved_path = f"checkpoint/atkmodel"
                
                torch.save(self.net_avg.state_dict(), f"{saved_path}/{self.dataset}_{self.model}.pt")
                torch.save(self.net_avg.state_dict(), f"{atk_saved_path}/{self.dataset}_{self.model}_{self.lira_args['attack_model']}.pt")
            
            if acc_clean > best_acc_clean or (acc_clean+self.lira_args['best_threshold'] > best_acc_clean and best_acc_poison < backdoor_acc) and self.retrain:
                best_acc_poison = backdoor_acc
                best_acc_clean = acc_clean
                print(f"Saving atk model with: backdoor_acc_{backdoor_acc} and best_acc_clean_{best_acc_clean}")
                torch.save({'atkmodel': self.atkmodel.state_dict(), 'clsmodel': self.net_avg.state_dict(), 
                            'best_acc_poison': best_acc_poison, 'best_acc_clean': best_acc_clean}, 
                           bestmodel_path)
            
            wandb_logging_items = {
                'fl_iter': flr,
                'main_task_acc': acc_clean*100.0,
                'backdoor_task_acc': backdoor_acc*100.0, 
                'local_best_acc_clean': best_acc_clean,
                'local_best_acc_poison': best_acc_poison,
                'local_MA': local_acc_clean*100.0,
                'local_BA': local_acc_poison*100.0
            }
            if wandb_ins:
                print(f"start logging to wandb")
                wandb_ins.log({"General Information": wandb_logging_items})

            raw_acc = 0
            overall_acc = acc_clean                
            fl_iter_list.append(flr)
            main_task_acc.append(overall_acc)
            raw_task_acc.append(raw_acc)
            backdoor_task_acc.append(backdoor_acc)
            if len(current_adv_norm_diff_list) == 0:
                adv_norm_diff_list.append(0)
            else:
                # if you have multiple adversaries in a round, average their norm diff
                adv_norm_diff_list.append(1.0*sum(current_adv_norm_diff_list)/len(current_adv_norm_diff_list))
        
        df = pd.DataFrame({'fl_iter': fl_iter_list, 
                            'main_task_acc': main_task_acc, 
                            'backdoor_acc': backdoor_task_acc, 
                            'raw_task_acc':raw_task_acc, 
                            'adv_norm_diff': adv_norm_diff_list, 
                            'wg_norm': wg_norm_list
                            })
       
        if self.poison_type == 'ardis':
            # add a row showing initial accuracies
            df1 = pd.DataFrame({'fl_iter': [0], 'main_task_acc': [88], 'backdoor_acc': [11], 'raw_task_acc': [0], 'adv_norm_diff': [0], 'wg_norm': [0]})
            df = pd.concat([df1, df])

        results_filename = get_results_filename(self.poison_type, self.attack_method, self.model_replacement, self.project_frequency,
                self.defense_technique, self.norm_bound, self.prox_attack, fixed_pool=True, model_arch=self.model)
        df.to_csv(results_filename, index=False)
        logger.info("Wrote accuracy results to: {}".format(results_filename))
