import argparse
import random
import time
import numpy as np
import os
import copy

import torch
import torch.nn as nn
import torch.optim as optim
from torch.nn.functional import normalize

from dataloaders.dataset import getDataset
from models.conv_net import Net
from models.alexnet import MiniAlexNet, AlexNet
from models.mlp import MLP, MLP4
from models.resnet import resnet18
from mi import MI



def set_data(res, total, m):
    for k,v in res.items():
        if k not in total['total']:
            total['total'][k] = []

        if m in ['noise', 'benign']:
            if k not in total[m]:
                total[m][k] = []

        if v is not None:
            total['total'][k].append(v)
            if m in ['noise', 'benign']:
                total[m][k].append(v)

def set_mean(total, keys, tensor=True):
    res = {}
    for l in keys:
        res[l] = []
        for k,v in total[l].items():
            if tensor:
                m = torch.tensor(v).mean(dim=0)
            else:
                m = torch.stack(v).mean(dim=0)
            res[l].append(m)

    return res

def train(net, trainloader, trainloader_test, optimizer, criterion, e, mi_interval=5):
    running_loss = 0.0
    clock0=time.time()
    for i, data in enumerate(trainloader, 0):
        # get the inputs; data is a list of [inputs, labels]
        inputs = data[0].to(device)
        labels = data[1].to(device)

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs, _, _ = net(inputs, {'mi': False, 'gram': False})
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

        # if i > 5:
        #     break

    clock1=time.time()
    print('| Epoch {}, time={:5.1f}ms | \n\tTrain: loss={:.3f} '.format(e+1,
                   1000*(clock1-clock0),running_loss/i),end='\n')

    return running_loss/i

def eval_separate(net, testloader, criterion, mi_interval=5, keys=[], keys_res=[], test_res=None):
    correct = {'noise': 0, 'benign': 0, 'total': 0}
    total = {'noise': 0, 'benign': 0, 'total': 0}
    mi_total = {'noise': {}, 'benign': {}, 'total': {}}
    gram_total_input = {'noise': {}, 'benign': {}, 'total': {}}
    gram_total_layer = {'noise': {}, 'benign': {}, 'total': {}}

    running_loss = 0.0
    total_data = 0
    total_cls = {'pred':{},'real':{}}
    with torch.no_grad():
        for m in keys: #['noise', 'benign']
            if testloader[m]:
                for i, data in enumerate(testloader[m]):
                    images = data[0].to(device)
                    labels = data[1].to(device)

                    get_res = {'mi': False, 'gram': False}
                    if m != 'total':
                        if test_res['TC']:
                            get_res['mi'] = True
                        if test_res['gram_matrix']:
                            get_res['gram'] = True

                    outputs, mi, gram_m = net(images, get_res)
                    _, predicted = torch.max(outputs.data, 1)

                    if test_res['accuracy']:
                        total[m] += labels.squeeze().size(0)
                        correct[m] += (predicted.squeeze() == labels.squeeze()).sum().item()

                    if m == 'total':
                        loss = criterion(outputs, labels)
                        running_loss += loss.item()
                    total_data += labels.squeeze().size(0)

                    if m in ['noise', 'benign'] and test_res is not None:
                        if test_res['TC']:
                            set_data(mi, mi_total, m)

                        if test_res['gram_matrix']:
                            set_data(gram_m['input'], gram_total_input, m)
                            set_data(gram_m['layer'], gram_total_layer, m)                      

                    # if i > 5:
                    #     break

    mi_layers = {}
    g_input = {}
    g_layer = {}
    res = {}

    if test_res['TC']:
        mi_layers = set_mean(mi_total, keys_res, True)

    if test_res['gram_matrix']:
        g_input = set_mean(gram_total_input, keys_res, False)
        g_layer = set_mean(gram_total_layer, keys_res, False)

    if test_res['accuracy']:
        for k in ['total', 'noise', 'benign']:
            if total[k] > 0:
                res[k] = correct[k] / total[k]
            else:
                res[k] = 0

    return res, mi_layers, running_loss/total_data, {'input': g_input, 'layer': g_layer}

# eval_weights(net, testloader, criterion, mi_interval=5, keys=[], keys_res=[])
# Runs model with either the benign, noise or test dataloaders
# Returns a dictionary with accuracies for each set
def eval_weights(net, testloader, criterion, mi_interval=5, keys=[], keys_res=[]):
    correct = {'noise': 0, 'benign': 0, 'total': 0}
    total = {'noise': 0, 'benign': 0, 'total': 0}
    mi_total = {'noise': {}, 'benign': {}, 'total': {}}
    running_loss = 0.0
    total_cls = {'pred':{},'real':{}}
    i_final = 1
    with torch.no_grad():
        for m in keys:
            if testloader[m]:
                print("Total Samples to Process-{}: {}".format(m.upper(), len(testloader[m].dataset)))
                for i, data in enumerate(testloader[m]):
                    images = data[0].to(device)
                    labels = data[1].to(device)

                    outputs, mi, gram_me = net(images, {'mi': False, 'gram': False})
                    _, predicted = torch.max(outputs.data, 1)

                    if m == 'total':
                        total[m] += labels.size(0)
                        correct[m] += (predicted == labels).sum().item()

                        loss = criterion(outputs, labels)
                        running_loss += loss.item()

                        i_final = i
                        
                    if m in ['noise', 'benign']:
                        total[m] += labels.size(0)
                        correct[m] += (predicted == labels.squeeze()).sum().item()

    res = {}
    for k in ['total', 'noise', 'benign']:
        if total[k] > 0:
            res[k] = correct[k] / total[k]
        else:
            res[k] = 0

    return res, None, running_loss/i_final


# weight_compare(model, weights1, weights2, interval=0.1, normalized=True, layers, method)
# Takes as input:
# - model: the model you are testing the weights on
# - weights1: state_dict() from model that generalizes well
# - weights2: state_dict() from model that starts overfitting
# - interval: float that indicates what percent of the average weight in model1 should one use
#             as a criterion for the threshold
# - normalized: whether to normalize weights to unitary L2 norm or not. We found this to be useful
# - because the last layer usually increases the value of the weights. Scale is not as important as their
#   relative values.
# - layers: takes 3 values, 0: affect only the n_layerth layer, 
#                           1: affect layers starting from n_layer, 
#                          -1: affect layers from n_layer and back
# - n_layer: depending on layers, which layer, or up to which layer, of from which layer to do the masks1
# - method: one of ['regular', 'purge_diff', 'purge_and_keep']
#   Returns model with weights modified by the appropriate mask
def weight_compare(model, 
                   weights1, 
                   weights2, 
                   interval=0.1, 
                   normalized=True, 
                   layers=0,
                   n_layer = 0,
                   method='regular',
                  net='mini-alex',
                  return_masks=False):
    masks = []
    ig = []
    dst = []
    if net == 'mini-alex':
        total_layers = ['features.0.weight', 
                               'features.4.weight'
                        , 'classifier.0.weight'
                        , 'classifier.3.weight'
                        , 'classifier.6.weight']
    elif net == 'mlp4':
        total_layers = [  'classifier.0.weight'
                        , 'classifier.2.weight'
                        , 'classifier.4.weight'
                        , 'classifier.6.weight',
                       'classifier.8.weight']
    elif net == 'resnet-18':
        total_layers = [    'conv1.weight',
                            'layer1.0.conv1.weight',
                            'layer1.0.conv2.weight',
                            'layer1.1.conv1.weight',
                            'layer1.1.conv2.weight',
                            'layer2.0.conv1.weight',
                            'layer2.0.conv2.weight',
                            'layer2.1.conv1.weight',
                            'layer2.1.conv2.weight',
                            'layer3.0.conv1.weight',
                            'layer3.0.conv2.weight',
                            'layer3.1.conv1.weight',
                            'layer3.1.conv2.weight',
                            'layer4.0.conv1.weight',
                            'layer4.0.conv2.weight',
                            'layer4.1.conv1.weight',
                            'layer4.1.conv2.weight',
                            'linear.weight'
                       ]
    if layers == 0:
        layers = [total_layers[n_layer]]
    elif layers ==1:
        layers = total_layers[:n_layer+1]
    elif layers ==-1:
        layers = total_layers[len(total_layers)-1-n_layer:]
    
    print("Layers to affect: " + str(layers)) 
    for i, layer in enumerate(layers):
        
        s = weights1[layer].shape
        
        if normalized:
            if len(s) == 2:        # Capa Lineal
                p1 = normalize(weights1[layer], p=2, dim=1)
                p2 = normalize(weights2[layer], p=2, dim=1)

            elif len(s) == 4:      # Capa Convolucional
                cin, cout, h, w = weights1[layer].shape
                p1 = weights1[layer].view(cin, cout, -1)
                p2 = weights2[layer].view(cin, cout, -1)
                p1 = normalize(weights1[layer], p=2, dim=-1)
                p2 = normalize(weights2[layer], p=2, dim=-1)
        else:
            p1 = weights1[layer]
            p2 = weights2[layer]

        diff = p2 - p1
        # Calculating means for weights in both models
        # It was not very useful
        avg1 = p1.abs().mean()
        avg2 = p2.abs().mean()

        print("Avg. Weight W1: {:.2f} - Avg. Weight W2: {:.2f}".format(avg1, avg2))

        thresh = interval*avg1
        print("Threshold: {:.3f}".format(thresh))
        res = diff.abs() > thresh
        print("N° Weights above threshold (Different): {}".format(res.sum()))
        # Conv Layers
        if len(s) == 4:
            cin, cout, h, w = weights1[layer].shape
            res = res.view(cin, cout, h, w)
        
        masks.append(res)
        # Creating mask and applying to model    
        if method == 'purge_diff':
            weights2[layer] = weights2[layer]*(~res)
        elif method == 'purge_and_keep':
            weights2[layer] = weights2[layer]*(~res) + (res)*weights1[layer]
        
        print("N° Weights Preserved (Equal): {}".format((~res).sum()))
    
    model.load_state_dict(weights2)
    
    if return_masks:
        return masks
    else:
        return model


# to_weights(args, device)
# It just computes (args, device)accuracy results for each method and interval over many epochs
# Most important output is saving results to disk
def to_weights(args, device):
    filters = [args.filters]*args.layers
    (trainloader, trainloader_test, testloader, grad_loaders), channels, num_classes = getDataset(args.dataset, args)
    # Where we store results for all of our experiments

    mi = None 
    filename1 = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, 
                                                                     args.dataset,
                                                                     args.noise,
                                                                     args.noise_input, 
                                                                     args.seed,
                                                                     args.perct_original_data,
                                                                     args.noise_classes,
                                                                     args.max_validation_epoch-1,
                                                                              args.noise_data)
    w1 = torch.load(os.path.join('saved_models',filename1))['checkpoint']
    net, optimizer, criterion, lr_schedule = get_models(args, filters, num_classes, channels, mi, device)
    
    # We run experiments for given intervals (% of average weight)
    # For given methods:
    # purge_diff = For weights above threshold, set to 0
    # purge_and_keep = For weights above threshold, set to weight from epoch where validation acc was maximum
    # regular = No changes to weights (baseline).
    #layer_order = 1 # 0 single layer, 1 from beginning to layer, -1 backwards
    if args.net == 'mini-alex':
        n_layers = 5
    elif args.net == 'resnet-18':
        n_layers = 18
    else:
        n_layers = 5
    for n_layer in range(n_layers):
        final_accs = {i: {k: {'benign': [],'noise': [], 'test': []} for k in ['purge_diff',
                                                                          'purge_and_keep',
                                                                          'regular']} for i in [0.05, 0.1, 0.2]} 
        for interval in [0.1]:#[0.05, 0.1, 0.2]:
            for method in ['purge_diff', 'purge_and_keep', 'regular']:
                print(method.upper())
                for e in range(args.max_validation_epoch, args.epochs):
                    print("Epoch {}".format(e+1))
                    w2 = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net,
                                                                               args.dataset,
                                                                               args.noise,
                                                                               args.noise_input,
                                                                               args.seed,
                                                                               args.perct_original_data,
                                                                               args.noise_classes,
                                                                               e,
                                                                                 args.noise_data)
                    w2 = torch.load(os.path.join('saved_models',w2))['checkpoint']


                    net = loadModel(os.path.join('saved_models',filename1), net, device)
                    net = weight_compare(net,
                                         w1,
                                         w2,
                                         interval=interval,
                                         normalized=True,
                                         layers=args.layerorder,
                                         n_layer=n_layer,
                                         method=method,
                                         net=args.net)
                    # Noise and benign data
                    acc, mi, loss = eval_weights(net,
                                                 trainloader_test,
                                                 criterion,
                                                 args.mi_interval,
                                                 ['noise','benign'],
                                                 ['noise','benign'])
                    final_accs[interval][method]['benign'].append(acc['benign'])
                    final_accs[interval][method]['noise'].append(acc['noise'])
                    # Test set
                    acc, mi, loss = eval_weights(net,
                                                               testloader,
                                                               criterion,
                                                               args.mi_interval,
                                                               ['total'],
                                                               ['total'])
                    final_accs[interval][method]['test'].append(acc['total'])

        # The fruit of our work...
        if args.save_model:
            save_file = 'affected_{}_{}_{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_accs_{}_{}.pth'.format(n_layer,
                                                                                   args.layerorder,     
                                                                                   args.net, 
                                                                                   args.dataset,
                                                                                   args.noise,
                                                                                   args.noise_input, 
                                                                                   args.seed, 
                                                                                   args.perct_original_data,
                                                                                   args.noise_classes,
                                                                                   e+1,
                                                                                   args.noise_data,                                                                                                                args.max_validation_epoch)
            print("Saving results to file: {}".format(save_file))

            torch.save(final_accs, os.path.join('accs', save_file))

def add_results(results, train, acc, mi, loss, gram_m, test_res):
    if train:
        t = 'train'
    else:
        t = 'test'

    results['{}_loss'.format(t)].append(loss)  

    if test_res['accuracy']:
        results['acc_{}_total'.format(t)].append(acc['total'])
        if train:
            results['acc_{}_benign'.format(t)].append(acc['benign'])
            results['acc_{}_noise'.format(t)].append(acc['noise'])

    if test_res['TC']:
        results['mi_{}_total'.format(t)].append(mi['total'])
        if train:
            results['mi_{}_benign'.format(t)].append(mi['benign'])
            results['mi_{}_noise'.format(t)].append(mi['noise'])

    if test_res['gram_matrix']:
        results['gram_{}_input_total'.format(t)].append(gram_m['input']['total'])
        if train:
            results['gram_{}_input_benign'.format(t)].append(gram_m['input']['benign'])
            results['gram_{}_input_noise'.format(t)].append(gram_m['input']['noise'])

        results['gram_{}_layer_total'.format(t)].append(gram_m['layer']['total'])
        if train:
            results['gram_{}_layer_benign'.format(t)].append(gram_m['layer']['benign'])
            results['gram_{}_layer_noise'.format(t)].append(gram_m['layer']['noise'])

def reset_results(results, test_res):
    results['train_loss'] = []
    results['test_loss'] = []

    if test_res['TC']:
        results['mi_train_total'] = []
        results['mi_train_benign'] = []
        results['mi_train_noise'] = []
        results['mi_test_total'] = []

    if test_res['gram_matrix']:
        results['gram_train_input_total'] = []
        results['gram_train_input_benign'] = []
        results['gram_train_input_noise'] = []
        results['gram_test_input_total'] = []

        results['gram_train_layer_total'] = []
        results['gram_train_layer_benign'] = []
        results['gram_train_layer_noise'] = []
        results['gram_test_layer_total'] = []

    if test_res['accuracy']:
        results['acc_train_total'] = []
        results['acc_train_benign'] = []
        results['acc_train_noise'] = []
        results['acc_test_total'] = []

def saveValues(name_file, results, model, args, save_model = True):
    res = {
            'results': results,
            'args': args,
            }
    if save_model:
        res['checkpoint'] = model.state_dict()
    torch.save(res, name_file)

def get_models(args, filters, num_classes, channels, mi, device):
    lr_schedule = None
    if args.net == 'conv':
        net = Net(filters, num_classes, channels, mi, args.add_batch_norm, args.add_dropout, args.noise, args.activation).to(device)

    elif args.net == 'mini-alex':
        net = MiniAlexNet(num_classes, channels, mi).to(device)
        optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9)
        if args.scheduler:
            lr_schedule = optim.lr_scheduler.StepLR(optimizer, 15, gamma=0.5)
    
    elif args.net =='mini-inception':
        net = Inception3Mod(num_classes)
    
    elif args.net == 'mlp':
        net = MLP(num_classes, mi).to(device)
        optimizer = optim.SGD(net.parameters(), lr=args.lr)
        if args.scheduler:
            lr_schedule = optim.lr_scheduler.StepLR(optimizer, 50, gamma=0.5)

    elif args.net == 'mlp4':
        net = MLP4(num_classes, mi).to(device)
        optimizer = optim.SGD(net.parameters(), lr=args.lr)
        if args.scheduler:
            lr_schedule = optim.lr_scheduler.StepLR(optimizer, 50, gamma=0.5)

    elif args.net == 'resnet-18':  
        net = resnet18(num_channels=channels)
        #num_ftrs = net.fc.in_features
        #net.fc = nn.Linear(num_ftrs, num_classes)
        net.to(device)
        optimizer = optim.SGD(net.parameters(), lr=args.lr, momentum=0.9)
        #optimizer = optim.Adam(net.parameters(), lr=0.01)
        if args.scheduler:
            lr_schedule = optim.lr_scheduler.StepLR(optimizer, 100, gamma=0.1)

    criterion = nn.CrossEntropyLoss() #reduction='sum'

    return net, optimizer, criterion, lr_schedule

def loadModel(name_file, net, device):
    checkpoint = torch.load(name_file, map_location=device)
    net.load_state_dict(checkpoint['checkpoint']) 
    return net

def get_mask(args,trainloader_test,net,criterion, optimizer,results):
    for m in ['noise', 'benign']:
        #masks = { n: torch.zeros_like(p) for n,p in net.named_parameters() }
        masks = {}
        if trainloader_test[m]:
            for i, data in enumerate(trainloader_test[m], 0):
                # get the inputs; data is a list of [inputs, labels]
                inputs = data[0].to(device)
                labels = data[1].to(device)

                # zero the parameter gradients
                optimizer.zero_grad()
                outputs, _, _ = net(inputs, {'mi': False, 'gram': False})
                loss = criterion(outputs, labels.squeeze())
                loss.backward()

                for n,p in net.named_parameters():
                    #masks[n] += (p.grad.abs()/2 > args.grads_threshold).int()
                    masks[n] = (p.grad.abs() > args.grads_threshold)

        results[m].append(masks)

def get_grads_sets(args, grad_loaders, net, criterion, optimizer, device, results):
    noise = []
    benign = []
    total = []
    for loader in grad_loaders:
        mask_t = {}
        mask_b = {}
        mask_lt = {}
        for data in loader:
            # get the inputs; data is a list of [inputs, labels]
            inputs = data[0].to(device)
            labels = data[1].to(device)
                    
            outputs, _, _ = net(inputs, {'mi': False, 'gram': False})
            optimizer.zero_grad()
            loss = criterion(outputs, labels)
            loss.backward()
            for n,p in net.named_parameters():
                mask_t[n] = (p.grad.abs() > args.grads_threshold).int()

            outputs, _, _ = net(inputs[data[2] == 0], {'mi': False, 'gram': False})
            optimizer.zero_grad()
            loss = criterion(outputs, labels[data[2] == 0])
            loss.backward()
            for n,p in net.named_parameters():
                mask_b[n] = (p.grad.abs() > args.grads_threshold).int()

            outputs, _, _ = net(inputs[data[2] == 1], {'mi': False, 'gram': False})
            optimizer.zero_grad()
            loss = criterion(outputs, labels[data[2] == 1])
            loss.backward()
            for n,p in net.named_parameters():
                mask_lt[n] = (p.grad.abs() > args.grads_threshold).int()

        t_noise = 0.0
        t_benign = 0.0
        t_total = 0.0
        num_total = 0.0
        for n,p in net.named_parameters():
            t_benign += (mask_t[n]*((mask_b[n] - mask_lt[n]) > 0 ).int()).sum()
            t_noise += (mask_t[n]*((mask_lt[n] - mask_b[n]) > 0 ).int()).sum()
            t_total += mask_t[n].sum()
            num_total += p.numel()

        noise.append((t_noise/num_total).item())
        benign.append((t_benign/num_total).item())
        total.append((t_total/num_total).item())

    results['noise'].append(np.mean(noise))
    results['benign'].append(np.mean(benign))
    results['total'].append(np.mean(total))

def to_train(args, device):
    filters = [args.filters]*args.layers
    (trainloader, trainloader_test, testloader), channels, num_classes = getDataset(args.dataset, args)

    mi = None
    if args.use_mi:
        mi = MI()

    net, optimizer, criterion, lr_schedule = get_models(args, filters, num_classes, channels, mi, device)
    test_res = {}
    test_res['TC'] = args.get_tc
    test_res['gram_matrix'] = args.get_gram
    test_res['accuracy'] = args.get_acc

    for e in range(args.epochs):
        loss = train(net, trainloader, trainloader_test, optimizer, criterion, e, args.mi_interval)
        acc, _, _, _ = eval_separate(net, trainloader_test, criterion, args.mi_interval, ['total','noise','benign'], ['total','noise','benign'], test_res)
        print('| Epoch {}, |\tTrain Acc: Total={:.3f} Benign={:.3f} Noise={:.3f} '.format(e+1,acc['total'],acc['benign'],acc['noise']),end='\n')
        print('-'*100)

        test_acc, _, _, _ = eval_separate(net, testloader, criterion, args.mi_interval*5, ['total'], ['total'], test_res)
        print('| Epoch {} | Test acc={:5.1f}% |'.format(e+1,100*test_acc['total']),end='\n')

        if lr_schedule:
            lr_schedule.step(loss)

        if args.save_model:
            save_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                                args.seed, args.perct_original_data, args.noise_classes, e, args.noise_data)
            saveValues(os.path.join('saved_models',save_file), {}, net, args)

def to_test(args, device):
    filters = [args.filters]*args.layers
    (trainloader, trainloader_test, testloader, grad_loaders), channels, num_classes = getDataset(args.dataset, args)

    save_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                                args.seed, args.perct_original_data, args.noise_classes)

    if os.path.exists(os.path.join('results',save_file)):
        checkpoint = torch.load(os.path.join('results',save_file), map_location=device)
        results = checkpoint['results']
        print("Results from {} loaded".format(save_file))
    else:
        results = {
            'train_loss': [],
            'test_loss': [],

            'acc_train_total': [],
            'acc_train_benign': [],
            'acc_train_noise': [],
            'acc_test_total': [],

            'mi_train_total': [],
            'mi_train_benign': [],
            'mi_train_noise': [],
            'mi_test_total': [],

            'gram_train_input_total': [],
            'gram_train_input_benign': [],
            'gram_train_input_noise': [],
            'gram_test_input_total': [],

            'gram_train_layer_total': [],
            'gram_train_layer_benign': [],
            'gram_train_layer_noise': [],
            'gram_test_layer_total': [],

        }

    mi = None
    if args.use_mi:
        mi = MI()

    net, optimizer, criterion, lr_schedule = get_models(args, filters, num_classes, channels, mi, device)
    
    test_res = {}
    test_res['TC'] = args.get_tc
    test_res['gram_matrix'] = args.get_gram
    test_res['accuracy'] = args.get_acc
    reset_results(results, test_res)

    print('*'*100)
    #print(results)
    acc, mi, loss, gram_m = eval_separate(net, trainloader_test, criterion, args.mi_interval, ['noise','benign'], ['noise','benign'], test_res)
    add_results(results, True, acc, mi, loss, gram_m, test_res)
    #print(results)
    for e in range(args.epochs):
        if args.noise_data == 'gaussian' or args.net == 'resnet-18':
            load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                                args.seed, args.perct_original_data, args.noise_classes, e, args.noise_data)
        else:
            load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                               args.seed, args.perct_original_data, args.noise_classes, e)

        net = loadModel(os.path.join('saved_models',load_file), net, device)
        acc, mi, loss, gram_m = eval_separate(net, trainloader_test, criterion, args.mi_interval, ['noise','benign'], ['noise','benign'], test_res)
        add_results(results, True, acc, mi, loss, gram_m, test_res)

        print('-'*100)
        test_acc, mi_test, loss, gram_m = eval_separate(net, testloader, criterion, args.mi_interval*5, ['total'], ['total'], test_res)
        add_results(results, False, test_acc, mi_test, loss, gram_m, test_res)
        print('| Epoch {} | Test acc={:5.1f}% |'.format(e+1,100*test_acc['total']),end='\n')
        print('-'*100)
        #print(results)

        if args.save_model:
            save_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                                args.seed, args.perct_original_data, args.noise_classes, args.noise_data)
            saveValues(os.path.join('results',save_file), results, net, args, False)

    print('Finished Training')
    # print(results)

def to_score(args, device):
    filters = [args.filters]*args.layers
    mi = MI()
    
    if args.dataset == 'cifar10':
        chnls = 3
    else:
        chnls = 1

    net, _, _, _ = get_models(args, filters, 10, chnls, mi, device)
    #print(net)
    
    if args.net == 'mini-alex':
        layers = ['features.0.weight', 'features.4.weight', 'classifier.0.weight', 'classifier.3.weight', 'classifier.6.weight']
    elif args.net == 'mlp4':
        layers = ['classifier.0.weight', 'classifier.2.weight', 'classifier.4.weight', 'classifier.6.weight', 'classifier.8.weight']
    elif args.net == 'resnet-18':
        layers = [ n for n,p in net.named_parameters() if 'conv' in n or 'linear' in n ][:-1]

    print(len(layers))
    params_init = {}
    for name, param in net.named_parameters():
        if name in layers:
            params_init[name] = param.clone()

    res = {}
    #for noise in [0.0, 0.01, 0.05, 0.1, 0.2, 0.4, 0.6, 0.8]:
    for noise in [0.0, 0.1,  0.4, 0.8]:

        interval = 0.1

        params_gen = {}
        params_mean = {}
        #keep = {'init': {}, 'gen': {}}
        #std = {'init': {}, 'gen': {}}
        total_keep = {'init': [], 'gen': []}

        # for name, param in net.named_parameters():
        #     #params_mean[name] = [param.clone()]
        #     keep['init'][name] = []
        #     keep['gen'][name] = []
        #     std['init'][name] = []
        #     std['gen'][name] = []

        if args.noise_data == 'gaussian' or args.net == 'resnet-18':
            load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, args.dataset, args.noise, noise, 
                                                                    args.seed, args.perct_original_data, args.noise_classes, 7, args.noise_data)
        else:
            load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}.pth'.format(args.net, args.dataset, args.noise, noise, 
                                                                    args.seed, args.perct_original_data, args.noise_classes, 5)
        net_gen = loadModel(os.path.join('saved_models',load_file), net, device)

        for name, param in net_gen.named_parameters():
            if name in layers:
                params_gen[name] = param.clone()

        for e in range(0,args.epochs):
            if args.noise_data == 'gaussian' or args.net == 'resnet-18':
                load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, args.dataset, args.noise, noise, 
                                                                    args.seed, args.perct_original_data, args.noise_classes, e, args.noise_data)
            else:
                load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}.pth'.format(args.net, args.dataset, args.noise, noise, 
                                                                    args.seed, args.perct_original_data, args.noise_classes, e)

            net = loadModel(os.path.join('saved_models',load_file), net, device)

            total_param = 0
            used_init = 0
            used_gen = 0
            for name, param in net.named_parameters():
                #params_mean[name].append(param.clone())
                if name in layers:
                    p1 = normalize(param.clone(), p=2, dim=1)
                    p2 = normalize(params_init[name], p=2, dim=1)
                    p3 = normalize(params_gen[name], p=2, dim=1)
                    diff_init = p2 - p1
                    diff_gen = p3 - p1

                    avg_p1 = p1.abs().mean()
                    thresh = interval*avg_p1

                    res_init = diff_init.abs() > thresh 
                    res_gen = diff_gen.abs() > thresh 

                    #keep['init'][name].append(res_init.sum().item()/param.numel())
                    #std['init'][name].append((diff_init*(res_init)).std(dim=1).mean().item())

                    used_init += res_init.sum().item()
                    total_param += param.numel()

                    if (e >= 5 and args.net == 'mini-alex') or (e >= 8 and args.net == 'resnet-18'):
                        #keep['gen'][name].append(res_gen.sum().item()/param.numel())
                        #std['gen'][name].append((diff_gen*(res_gen)).std(dim=1).mean().item())

                        used_gen += res_gen.sum().item()

            total_keep['init'].append(used_init/total_param)
            if used_gen > 0:
                total_keep['gen'].append(used_gen/total_param)

        res[noise] = {
                #'init_keep': keep['init'],
                #'gen_keep': keep['gen'],
                #'init_std': std['init'],
                #'gen_std': std['gen'],
                'total_init': total_keep['init'],
                'total_gen': total_keep['gen']
                }

    torch.save(res, './scores/check_score_results_{}_{}_{}_{}.pth'.format(args.noise, args.noise_data, args.net, args.dataset))

def to_grads(args, device):
    filters = [args.filters]*args.layers
    (trainloader, trainloader_test, _), channels, num_classes = getDataset(args.dataset, args)

    mi = None
    if args.use_mi:
        mi = MI()

    net, optimizer, criterion, lr_schedule = get_models(args, filters, num_classes, channels, mi, device)

    results = { 'noise': [], 'benign': [], 'acc_total': [], 'acc_noise': [], 'acc_benign': []}

    print('*'*100)
    for e in range(args.epochs):
        #for e in [0, 2, 5, 8, 10, 20]:
        #print(e)

        for i, data in enumerate(trainloader, 0):
            # get the inputs; data is a list of [inputs, labels]
            inputs = data[0].to(device)
            labels = data[1].to(device)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs, _, _ = net(inputs, {'mi': False, 'gram': False})
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            if i % int(len(trainloader)/args.grads_interval) == 0:
                get_mask(args,trainloader_test,net,criterion, optimizer,results)
                acc, _, _, _ = eval_separate(net, trainloader_test, criterion, args.mi_interval, ['total','noise','benign'], ['total','noise','benign'], {'TC': False, 'gram_matrix': False, 'accuracy': True})
                results['acc_total'].append(acc['total'])
                results['acc_noise'].append(acc['noise'])
                results['acc_benign'].append(acc['benign'])

        print('| Epoch {}, |\tTrain Acc: Total={:.3f} Benign={:.3f} Noise={:.3f} '.format(e+1,acc['total'],acc['benign'],acc['noise']),end='\n')

        #if args.noise_data == 'gaussian':
        load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                                args.seed, args.perct_original_data, args.noise_classes, e, args.noise_data)
        #else:
        #    load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
        #                                                        args.seed, args.perct_original_data, args.noise_classes, e)

        net = loadModel(os.path.join('saved_models',load_file), net, device)

    #torch.save(results, './grads/results_{}_{}_{}_{}_{}_{}_count.pth'.format(args.noise, args.noise_data, args.net, args.dataset, args.noise_input, args.grads_threshold))
    torch.save(results, './grads/results_{}_{}_{}_{}_{}_{}.pth'.format(args.noise, args.noise_data, args.net, args.dataset, args.noise_input, args.grads_threshold))

def to_grad_new(args, device):
    filters = [args.filters]*args.layers
    (trainloader, trainloader_test, _, grad_loaders), channels, num_classes = getDataset(args.dataset, args)

    net, optimizer, criterion, lr_schedule = get_models(args, filters, num_classes, channels, None, device)

    results = { 'noise': [], 'benign': [], 'total': []}

    print('*'*100)
    for e in range(args.epochs):
        get_grads_sets(args, grad_loaders, net, criterion, optimizer, device, results)
        for i, data_train in enumerate(trainloader, 0):
            if i % int(len(trainloader)/args.grads_interval) == 0:
                get_grads_sets(args, grad_loaders, net, criterion, optimizer, device, results)

            inputs_train = data_train[0].to(device)
            labels_train = data_train[1].to(device)

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs_train, _, _ = net(inputs_train, {'mi': False, 'gram': False})
            loss = criterion(outputs_train, labels_train)
            loss.backward()
            optimizer.step()

        print('| Epoch {}, |\tTrain Acc: Total={:.3f} Benign={:.3f} Noise={:.3f} '.format(e+1,results['total'][-1],results['benign'][-1],results['noise'][-1]),end='\n')

        if args.noise_data == 'gaussian' or args.net == 'resnet-18':
            load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                                args.seed, args.perct_original_data, args.noise_classes, e, args.noise_data)
        else:
            load_file = '{}_{}_noise_{}_{}_seed_{}_{}_{}_epoch_{}.pth'.format(args.net, args.dataset, args.noise, args.noise_input, 
                                                               args.seed, args.perct_original_data, args.noise_classes, e)

        net = loadModel(os.path.join('saved_models',load_file), net, device)

    torch.save(results, './grads_new/results_{}_{}_{}_{}_{}_{}_{}.pth'.format(args.noise, args.noise_data, args.net, args.dataset, args.noise_input, args.grads_threshold, args.grads_interval))



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='MI Experiments')
    
    #------------------ Model --------------------------------------------#
    parser.add_argument('--filters', type=int, default=32)
    parser.add_argument('--layers', type=int, default=4)
    parser.add_argument('--add-batch-norm', action='store_true', default=False)
    parser.add_argument('--add-dropout', action='store_true', default=False)
    parser.add_argument('--activation',default='',type=str, required=False,
                        choices=['relu','tanh'])
    parser.add_argument('--net',default='',type=str, required=True,
                        choices=['conv','mini-alex','mini-inception', 'mlp', 'mlp4', 'resnet-18'])

    #------------------ MI -----------------------------------------------#
    parser.add_argument('--use-mi', action='store_true', default=False)
    parser.add_argument('--mi-interval', type=int, default=5, metavar='N',
                        help='Interval MI (default: 5)')

    #------------------ Training -----------------------------------------#
    parser.add_argument('--save-model', action='store_true', default=False,
                        help='For Saving the current Model')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--seed', type=int, default=42, metavar='S',
                        help='random seed (default: 42)')
    parser.add_argument('--batch-size-train', type=int, default=128, metavar='N',
                        help='input batch size for training (default: 128)')
    parser.add_argument('--batch-size-test', type=int, default=128, metavar='N',
                        help='input batch size for training (default: 128)')
    parser.add_argument('--epochs', type=int, default=14, metavar='N',
                        help='number of epochs to train (default: 14)')
    parser.add_argument('--scheduler', action='store_false', default=True,
                        help='For Saving the current Model')
    parser.add_argument('--train', action='store_true', default=False,
                        help='For Saving the current Model')

    #------------------ Data ---------------------------------------------#
    parser.add_argument('--dataset',default='',type=str,required=True,
                        choices=['mnist','cifar10','cifar100'])
    parser.add_argument('--noise-input', type=float, default=0.0)
    parser.add_argument('--perct-original-data', type=float, default=1.0)
    parser.add_argument('--noise-classes', type=int, default=1)
    parser.add_argument('--noise', type=str, required=True,
                        choices=['input','labels'])
    parser.add_argument('--weights', action='store_true', default=False,
                        help='For running the model for weight comparison purposes')
    parser.add_argument('--score', action='store_true', default=False,
                        help='To get Score of complexity')
    parser.add_argument('--grad-new', action='store_true', default=False)

    parser.add_argument('--grads', action='store_true', default=False,
                        help='To get grads')
    parser.add_argument('--grads-threshold', type=float, default=0.001)
    parser.add_argument('--grads-interval', type=int, default=3)
    
    parser.add_argument('--max-validation-epoch', type=int, default=1)
    parser.add_argument('--noise-data', type=str, default='cifar100',
                        choices=['cifar100','gaussian'])

    #------------------ Results ------------------------------------------#
    parser.add_argument('--get-tc', action='store_true', default=False)
    parser.add_argument('--get-acc', action='store_true', default=False)
    parser.add_argument('--get-gram', action='store_true', default=False)
    parser.add_argument('--layerorder', type=int, default=0, choices=[0, -1, 1])
    args = parser.parse_args()

    use_cuda = torch.cuda.is_available()
    print(use_cuda)
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if use_cuda:
        torch.cuda.manual_seed(args.seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False

    device = torch.device("cuda" if use_cuda else "cpu")

    print("Running on: {}".format(device))
    print(args)

    if args.weights:
        to_weights(args, device)
    elif args.score:
        to_score(args, device)
    elif args.grads:
        to_grads(args, device)
    elif args.grad_new:
        to_grad_new(args, device)
    else:
        if args.train:
            to_train(args, device)
        else:
            to_test(args, device)
