from __future__ import print_function
from typing import Dict, List, NamedTuple, Optional, Tuple
import argparse
import os
import sys
import shutil
import time
import random
import numpy as np
from tqdm import tqdm 
import pickle
import torch.nn.utils.prune as prune
# from sparselearning.models import AlexNet, VGG16, LeNet_300_100, LeNet_5_Caffe, WideResNet, MLP_CIFAR10
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms
from pathlib import Path
import operator as op
from contextlib import contextmanager
import copy
from copy import deepcopy
#from resnet import ResNet18,ResNet50,ResNet34
from resnet_width import ResNet18,  ResNet50,  ResNet34, LogitNormalizationWrapper
from functorch import make_functional_with_buffers, vmap, grad
from statistic_utils import eval_APGD_sharpness
from statistic_utils import eval_average_sharpness
from functorch import make_functional_with_buffers, vmap, grad


sys.path.append('../')
import utils
#from count_flops import print_model_param_nums, count_model_param_flops, count_full_param
# from sam.sam import NSAM
# from PyHessian.pyhessian import hessian
import augmentations





parser = argparse.ArgumentParser(
    description='Trains a CIFAR Classifier',
    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument(
    '--dataset',
    type=str,
    default='cifar10',
    choices=['cifar10', 'cifar100'],
    help='Choose between CIFAR-10, CIFAR-100.')
parser.add_argument(
    '--model',
    '-m',
    type=str,
    default='wrn',
    choices=['wrn', 'allconv', 'densenet', 'resnext','resnet18','resnet34', 'wrn-28-10'],
    help='Choose architecture.')
# Optimization options
parser.add_argument(
    '--epochs', '-e', type=int, default=100, help='Number of epochs to train.')
parser.add_argument(
    '--lr',
    '-lr',
    type=float,
    default=0.1,
    help='Initial learning rate.')
parser.add_argument(
    '--norm',
    type=str)
parser.add_argument(
    '--batch-size', '-b', type=int, default=128, help='Batch size.')
parser.add_argument('--eval_batch_size', type=int, default=5)
parser.add_argument('--momentum', type=float, default=0.9, help='Momentum.')
parser.add_argument(
    '--decay',
    '-wd',
    type=float,
    default=0.0005,
    help='Weight decay (L2 penalty).')
# WRN Architecture options
parser.add_argument(
    '--layers', default=18, type=int, help='total number of layers') # 
parser.add_argument('--width', default=64, type=int, help='Widen factor')
parser.add_argument(
    '--droprate', default=0.0, type=float, help='Dropout probability')
# AugMix options
parser.add_argument(
    '--mixture-width',
    default=3,
    type=int,
    help='Number of augmentation chains to mix per augmented example')
parser.add_argument(
    '--mixture-depth',
    default=-1,
    type=int,
    help='Depth of augmentation chains. -1 denotes stochastic depth in [1, 3]')
parser.add_argument(
    '--aug-severity',
    default=3,
    type=int,
    help='Severity of base augmentation operators')
parser.add_argument(
    '--no-jsd',
    '-nj',
    action='store_true',
    help='Turn off JSD consistency loss.')
parser.add_argument(
    '--all-ops',
    '-all',
    action='store_true',
    help='Turn on all operations (+brightness,contrast,color,sharpness).')
# Checkpointing options
parser.add_argument(
    '--pruned_model',
    default=False,
    type=bool
)
parser.add_argument('--adaptive', type=str, default='True') #action='store_true', 
parser.add_argument(
    '--save',
    '-s',
    type=str,
    default='./snapshots',
    help='Folder to save checkpoints.')
parser.add_argument(
    '--resume',
    nargs='+',
    help='Checkpoint path for resume / test.')
parser.add_argument(
    '--mask_resumes',
    nargs='+',
    help='Checkpoint path for resume / test.')
parser.add_argument('--evaluate', type=bool, default=True, help='Eval only.')
parser.add_argument(
    '--print-freq',
    type=int,
    default=50,
    help='Training loss print frequency (batches).')
# Acceleration
parser.add_argument(
    '--num-workers',
    type=int,
    default=6,
    help='Number of pre-fetching threads.')
parser.add_argument(
    '--seed',
    type=int,
    default=13,
    help='')
parser.add_argument(
    "--mode",
    type=str
)
parser.add_argument(
    '--sam_epochs',
    type=int,
    default=15,
    help=''
)
parser.add_argument(
    '--tryy',
    type=int,
    default=0,
    help=''
)
parser.add_argument(
    '--gamma',
    type=str,
    default='0.5',
    help=''
)
parser.add_argument(
    '--sam_iter',
    type=int,
    default=1
)
parser.add_argument(
    '--retrain_epochs',
    type=int,
    default=20,
    help=''
)
parser.add_argument(
    '--retrain_iter',
    type=int,
    default=1
)
parser.add_argument(
    '--augmix',
    type=bool,
    default=False
)
parser.add_argument(
    '--flat_ratio',
    type=float,
    default=0.2
)
parser.add_argument(
    '--compare_ratio',
    type=float,
    default=0.2
)
parser.add_argument(
    '--lrsche',
    type=str,
    default='constant'
)
parser.add_argument(
    '--optimizer',
    type=str,
    default='SGD'
)
parser.add_argument(
   '--only_flat',
   type=str,
   default='False'
)
parser.add_argument(
    '--flat_trial',
    default=5,
    type=int)
parser.add_argument(
    '--no_aug',
    default=True,
    type=bool)
parser.add_argument(
    '--batch-num',
    type=int,
    default=100
)
parser.add_argument(
    '--logit-norm',
    type=str,
    default='False'
)


parser.add_argument("--train_idx_path",type=str, default='./dataset/train_idx.npy')
parser.add_argument("--val_idx_path",type=str, default='./dataset/valid_idx.npy')
parser.add_argument("--data_path",type=str)

parser.add_argument('--mod',type=int,default=0)
                    
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
global fmodel
CORRUPTIONS = [
    'gaussian_noise', 'shot_noise', 'impulse_noise', 'defocus_blur',
    'glass_blur', 'motion_blur', 'zoom_blur', 'snow', 'frost', 'fog',
    'brightness', 'contrast', 'elastic_transform', 'pixelate',
    'jpeg_compression'
]
class CARD_Deck(nn.Module):

    def __init__(self, num_cards,layers,widen_factor,numclasses,droprate):
        super(CARD_Deck, self).__init__()
        self.num_cards = num_cards
        self.models = nn.ModuleList()
        for i in range(self.num_cards):
            if 'wrn' in args.model:
               self.models.append(WideResNet(layers,widen_factor ,numclasses,droprate))
          
            elif args.model =='resnet18':
                self.models.append(ResNet18(num_classes=numclasses))
           
            elif args.model=='resnet50':
                self.models.append(ResNet50(num_classes=numclasses))
           
            elif args.model=='resnet34':
               self.models.append(ResNet18(num_classes=numclasses))
                           
   
    def forward(self, x):
    
        x_cl = x.clone(
        )  # clone to make sure x is not changed by inplace methods
        out_list = []
        for i in range(self.num_cards):
            # Evaluate model i at input
            out = self.models[i](x_cl)
            # Compute softmax
            out = torch.softmax(out, dim=1)
            # Append output to list of logits
            out_list.append(out)

        return torch.mean(torch.stack(out_list), dim=0)

def get_lr(step, total_steps, lr_max, lr_min):
  """Compute learning rate according to cosine annealing schedule."""
  return lr_min + (lr_max - lr_min) * 0.5 * (1 +
                                             np.cos(step / total_steps * np.pi))


def aug(image, preprocess):
  """Perform AugMix augmentations and compute mixture.

  Args:
    image: PIL.Image input image
    preprocess: Preprocessing function which should return a torch tensor.

  Returns:
    mixed: Augmented and mixed image.
  """
  aug_list = augmentations.augmentations
  if args.all_ops:
    aug_list = augmentations.augmentations_all

  ws = np.float32(np.random.dirichlet([1] * args.mixture_width))
  m = np.float32(np.random.beta(1, 1))

  mix = torch.zeros_like(preprocess(image))
  for i in range(args.mixture_width):
    image_aug = image.copy()
    depth = args.mixture_depth if args.mixture_depth > 0 else np.random.randint(
        1, 4)
    for _ in range(depth):
      op = np.random.choice(aug_list)
      image_aug = op(image_aug, args.aug_severity)
    # Preprocessing commutes since all coefficients are convex
    mix += ws[i] * preprocess(image_aug)

  mixed = (1 - m) * preprocess(image) + m * mix
  return mixed


class AugMixDataset(torch.utils.data.Dataset):
  """Dataset wrapper to perform AugMix augmentation."""

  def __init__(self, dataset, preprocess, no_jsd=False):
    self.dataset = dataset
    self.preprocess = preprocess
    self.no_jsd = no_jsd

  def __getitem__(self, i):
    x, y = self.dataset[i]
    if self.no_jsd:
      return i,aug(x, self.preprocess), y
    else:
      im_tuple = (self.preprocess(x), aug(x, self.preprocess),
                  aug(x, self.preprocess))
      return i,im_tuple, y

  def __len__(self):
    return len(self.dataset)

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, dataset):
        self.dataset=dataset
    def __getitem__(self, i):
        x, y = self.dataset[i]
        return i,x,y
    def __len__(self):
        return len(self.dataset)

def test(net, test_loader):
  """Evaluate network on given dataset."""
  net.eval()
  total_loss = 0.
  total_correct = 0
  with torch.no_grad():
    for images, targets in test_loader:
      images, targets = images.cuda(), targets.cuda()
      logits = net(images)
      loss = F.cross_entropy(logits, targets)
      pred = logits.data.max(1)[1]
      total_loss += float(loss.data)
      total_correct += pred.eq(targets.data).sum().item()

  return total_loss / len(test_loader.dataset), total_correct / len(
      test_loader.dataset)

def test_accuracy(model,test_loader):
    model.eval()
    n = 0
    accuracy=torch.zeros(len(test_loader.dataset))
    with torch.no_grad():
        for idx, data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
    
            pred = output.argmax(dim=1)
            correct= pred.eq(target).float().cpu()
            accuracy[idx]=correct
    return accuracy

@contextmanager
def _perturbed_model(
  model,
  sigma: float,
  rng,
  magnitude_eps: Optional[float] = None
):
  device = next(model.parameters()).device
  if magnitude_eps is not None:
    noise = [torch.normal(0,sigma**2 * torch.abs(p) ** 2 + magnitude_eps ** 2, generator=rng) for p in model.parameters()]
  else:
    noise = [torch.normal(0,sigma**2,p.shape, generator=rng).to(device) for p in model.parameters()]
  model = deepcopy(model)
  try:
    [p.add_(n) for p,n in zip(model.parameters(), noise)]
    yield model
  finally:
    [p.sub_(n) for p,n in zip(model.parameters(), noise)]
    del model
def compute_grad(model,sample, target):
    model.eval()
    model.zero_grad()
    sample = sample.unsqueeze(0)  # prepend batch dimension for processing
    target = target.unsqueeze(0)
    prediction = model(sample)
    loss_fn=nn.CrossEntropyLoss()
    
    loss = loss_fn(prediction, target)
    return torch.autograd.grad(loss, list(model.parameters()))


def compute_sample_grads(model,data, targets):
    """ manually process each sample with per sample gradient """
    sample_grads = [compute_grad(model,data[i], targets[i]) for i in range(data.shape[0])]
    sample_grads = zip(*sample_grads)
    sample_grads = [torch.stack(shards) for shards in sample_grads]
    return sample_grads
@torch.no_grad()
def pac_bayes(net,train_data,accuracy,seed,  
    magnitude_eps: Optional[float] = None,
    search_depth: int = 15,
    montecarlo_samples: int = 10,
    accuracy_displacement: float = 0.1,
    displacement_tolerance: float = 1e-2):

    train_loader = torch.utils.data.DataLoader(
                train_data,
                batch_size=1,
                shuffle=False,
                num_workers=args.num_workers,
                pin_memory=True)

    BIG_NUMBER = 10348628753
    device = next(net.parameters()).device
    rng = torch.Generator(device=device) if magnitude_eps is not None else torch.Generator()
    rng.manual_seed(BIG_NUMBER + seed)

    flatness=torch.zeros(len(train_loader.dataset))
    for i, dt in enumerate(tqdm(train_loader)):
        idx,data,target=dt
        lower, upper = 0, 2
        for _ in range(search_depth):
            sigma=(lower+upper)/2
            accuracy_samples = []
            for _ in range(montecarlo_samples):
                with _perturbed_model(net, sigma, rng, magnitude_eps) as p_model:
                    data, target = data.to(device, dtype=torch.float), target.to(device)
                    logits = p_model(data)
                    pred = logits.argmax(dim=1)
                    correct= pred.eq(target).float().cpu()
                    loss_estimate = correct.sum()
                    accuracy_samples.append(loss_estimate)
            displacement=abs(np.mean(accuracy_samples) - accuracy[idx])
            if abs(displacement - accuracy_displacement) < displacement_tolerance:
                break
            elif displacement > accuracy_displacement:
                # Too much perturbation
                upper = sigma
            else:
                # Not perturbed enough to reach target displacement
                lower = sigma
        flatness[idx]=sigma
    return flatness.numpy()

'''
def forgot(model_lst,train_loader,train_data):
    for iter in range(args.retrain_iter):
'''
def sam_idx(save_path,flats,i,ascend=True):
    print(i) 
    indice=np.arange(flats[0].shape[0])
    if(args.compare_ratio!=0):
        total=int(indice.shape[0]*args.compare_ratio)
        for j in range(len(flats)):

            flat=flats[j]
            indice1=np.argsort(flat)
            if(ascend):
                indice=np.intersect1d(indice,indice1[:total])
            else:
                indice=np.intersect1d(indice,indice1[-total:])
        print(indice.shape)
        flat_intection=np.zeros((len(flats),indice.shape[0]))
        for k in range(len(flats)):
            flat_intection[k]=flats[k][indice]           
        #print(np.mean(np.min(flat_intection,axis=0)))
        #print(np.median(np.min(flat_intection,axis=0)))
        static=[]
        this_flats=np.array(flats)
        np.save(save_path,this_flats)        
        this_indice=indice[i::len(flats)]
        return np.array(list(set(indice)-set(this_indice)))
    else:
        return np.array([-1])



def mask_gradient(net,mask):
    for name, tensor in net.named_parameters():
        #print(name)
        if name in mask:
            #print(name)
            tensor.grad = tensor.grad*mask[name]
            
            
            
def test_worstcase(net,train_loader,mask):
    net.eval()
    total_loss=0
    base_optimizer = torch.optim.SGD
    optimizer = NSAM(net.parameters(),base_optimizer,rho=args.gamma,lr=0.1,momentum=0.9,weight_decay=5e-4)     
    flatness=torch.zeros(len(train_loader.dataset))
    for i,dt in enumerate(tqdm(train_loader)):
        optimizer.zero_grad()
        idx,data,target=dt
        data, target = data.cuda(), target.cuda()
        logits = net(data)
        
        loss = F.cross_entropy(logits, target)
        total_loss+=loss
        loss.backward()
        mask_gradient(net,mask)
        optimizer.first_step()
        
        output=net(data)
        rloss = F.cross_entropy(output, target)
        sharpness=rloss-loss
        flatness[idx]=sharpness.detach()
        
        optimizer.back_step()
    return flatness.numpy()
        

def hessian_ensemble(deck,masks,train_loader,savepaths):

    deck.eval() 
    flatness=torch.zeros(len(train_loader.dataset))
    criterion = nn.CrossEntropyLoss()
    for i, dt in enumerate(tqdm(train_loader)):

        idx,data,target=dt
        if(args.eval_batch_size==1 and idx%10!=args.mod):
            continue      
        hessian_dataloader = (data, target)
        
        hessian_comp = hessian(deck,
                        criterion,
                        data=hessian_dataloader,
                        cuda=True)
        trace = hessian_comp.trace()
        print('\n***Trace: ', np.mean(trace))
        flatness[idx]=(np.mean(trace))
    Path(savepaths).mkdir(parents=True, exist_ok=True)
    path=os.path.join(savepaths,'hessian_ensemble_largebatch.npy')
    np.save(path,flatness)
    return


def apply_mask_graient_deck(deck,ft_per_sample_grads,masks):
    i=0
    for name,param in deck.named_parameters():
        index=int(name[7])
        mask=masks[index]
        if name[9:] in mask.masks:
            ft_per_sample_grads[i]=ft_per_sample_grads[i]*torch.unsqueeze(mask.masks[name[9:]],dim=0)
        i=i+1
    return ft_per_sample_grads
            

def test_hessian(net, train_loader, mask, mode='trace'):
    net.eval() 
    flatness=[]
    hessian_dataloader = []
    criterion = nn.CrossEntropyLoss()
    for i, dt in enumerate(tqdm(train_loader)):
        if(i == args.batch_num):
            break
        idx,data,target=dt
        hessian_dataloader.append((data, target))
    
    hessian_comp = hessian(net,
                    criterion,
                    mask,
                    dataloader=hessian_dataloader,
                    cuda=True)
        
    if mode == 'trace':
        trace = hessian_comp.trace()        
        flatness.append((np.mean(trace)))
        print('\n***Trace: ', np.mean(trace))
    elif mode == 'top_e':
        top_eigenvalues, _ = hessian_comp.eigenvalues()
        flatness.append(top_eigenvalues)
        print('\n***top_eigenvalues: ', np.mean(top_eigenvalues))
    
    flatness=np.array(flatness)
    return flatness

def hessians(model_lst, savepaths, masks, mode='trace'):
    nets_flatness=[]
    for i in range(len(model_lst)):
        train_loader,_ = get_trainloader()
        net=model_lst[i]
        if(len(masks)>0):
            mask=masks[i].masks
        else:
            mask=None
        flatness=test_hessian(net, train_loader, mask, mode)
        nets_flatness.append(flatness)
        
    savepaths = os.path.join(savepaths, f'bs{args.eval_batch_size}_bn{args.batch_num}')
    Path(savepaths).mkdir(parents=True, exist_ok=True)
    path=os.path.join(savepaths,f'hessian_{mode}_indiv_largebatch.npy')
    np.save(path,nets_flatness)
    
    
def get_ntk(networks, savepaths, masks):
    #, batch_num=-1
    # if recalbn > 0:
    #     network = recal_bn(network, xloader, recalbn, device)
    #     if network_2 is not None:
    #         network_2 = recal_bn(network_2, xloader, recalbn, device)
    # ntks = []
        # for i, dt in enumerate(train_loader):
    #     _, inputs, targets = dt
    #     inputs = inputs.cuda(non_blocking=True)
    train_loader,_ = get_trainloader()
    log_det_ntk_over_data = []
    for i, dt in enumerate(train_loader):
        if i == args.batch_num:
            break
        _, inputs, targets = dt
        inputs = inputs.cuda(non_blocking=True)
        log_det_ntk_lst = []
        for network in networks:
            network.eval()
            grads = []

            if len(masks)>0:
                sparse_model=True
            else:
                sparse_model=False

            network.zero_grad()
            inputs_ = inputs.clone().cuda(non_blocking=True)
            logit = network(inputs_)
            if isinstance(logit, tuple):
                logit = logit[1]  # 201 networks: return features and logits

            for _idx in range(len(inputs_)):
                logit[_idx:_idx+1].backward(torch.ones_like(logit[_idx:_idx+1]), retain_graph=True)
                grad = []
                for name, W in network.named_parameters():
                    if 'weight' in name and W.grad is not None:
                        # the model is a sparse model
                        if sparse_model:
                            grad.append(W.grad[W != 0].view(-1).detach())
                        else:
                            #rint('Dense Model')
                            grad.append(W.grad.view(-1).detach())
                grads.append(torch.cat(grad, -1))
                network.zero_grad()
                torch.cuda.empty_cache()

            grads = torch.stack(grads, 0) 
            # print("stacking matrices")
            print("grads.shape", len(grads), grads.shape)
            # print("computing the matrix multiplication")
            ntk = torch.einsum('nc,mc->nm', [grads, grads])
            # print("matrix multiplication completed")
            print(len(ntk), ntk.shape)

            eigenvalue = torch.linalg.eigvalsh(ntk)
            log_det_ntk_lst.append(np.log(eigenvalue.detach().cpu().numpy()).sum())

            print(len(eigenvalue), log_det_ntk_lst)
        log_det_ntk_over_data.append(log_det_ntk_lst)

    log_det_ntk_over_data = np.array(log_det_ntk_over_data)
    savepaths = os.path.join(savepaths, f'bs{args.eval_batch_size}_bn{args.batch_num}')
    Path(savepaths).mkdir(parents=True, exist_ok=True)
    path=os.path.join(savepaths,f'log_det_ntk.npy')
    print(f'--------------------> save to {path}')
    np.save(path, log_det_ntk_over_data)
    


def get_trainloader():
    np.random.seed(args.seed)
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed(args.seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    os.environ["PYTHONHASHSEED"] = str(args.seed)
    print(f"Random seed set as {args.seed}")
    if args.dataset=='cifar10':
        mean=[0.4914,0.4822,0.4465]
        std=[0.2470, 0.2435, 0.2616]
    else:
        mean=[0.5071, 0.4865, 0.4409]
        std=[0.2673, 0.2564, 0.2762]

    preprocess = transforms.Compose(
        [transforms.ToTensor(),
        transforms.Normalize(mean, std)])
    test_transform = preprocess
    train_idx=np.load(args.train_idx_path)
    val_idx= np.load(args.val_idx_path)
    if args.dataset == 'cifar10':
        train_data = torch.utils.data.Subset(datasets.CIFAR10(
                    f'{args.data_path}/cifar10', train=True, 
                    transform=test_transform, download=True),indices=train_idx)
        test_data = datasets.CIFAR10(
                    f'{args.data_path}/cifar10', train=False, 
                    transform=test_transform, download=True)
        base_c_path = f'{args.data_path}/CIFAR-10-C/'
        num_classes = 10
    else:
        train_data = torch.utils.data.Subset(datasets.CIFAR100(
                    f'{args.data_path}/cifar100', train=True, 
                    transform=test_transform, download=True),indices=train_idx)

        test_data = datasets.CIFAR100(
                    f'{args.data_path}/cifar100', train=False, transform=test_transform, download=True)
        base_c_path = f'{args.data_path}/CIFAR-100-C/'
        num_classes = 100  

    train_data=CustomDataset(train_data)
    print('train_data without augmix')

    train_loader = torch.utils.data.DataLoader(
        train_data,
        batch_size=args.eval_batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=True)

    test_loader = torch.utils.data.DataLoader(
        test_data,
        batch_size=args.eval_batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True)
    return train_loader, test_loader


def APGD(model_lst, savepaths, masks, mode):
    nets_flatness=[]
    loss_f = lambda logits, y: F.cross_entropy(logits, y, reduction='mean')
    for i in range(len(model_lst)):
        net=model_lst[i]
        net.eval()
        net = LogitNormalizationWrapper(net, normalize_logits=args.logit_norm == 'True')
        if(len(masks)>0):
            mask=masks[i].masks
        else:
            mask=None
        
        sharpness_ablation = {}
        #, 
        for norm in ['l2', 'linf']:
            for gamma in args.gamma:
                train_loader,_ = get_trainloader()
                args.norm = norm 
                if mask is None:
                    print("--------------------> Evaluate a dense model")
                    print(f'{mode}, args.logit_norm={args.logit_norm} rho={gamma}, adaptive={args.adaptive == "True"}, {args.norm}, num_batch={args.batch_num}, train_loader batch={train_loader.batch_size}')
                    if mode == 'APGD_worst':
                        flatness=eval_APGD_sharpness(net,
                                                    mask,
                                                    train_loader,
                                                    loss_f=loss_f,
                                                    train_err=None, 
                                                    train_loss=None,
                                                    n_iters=20,
                                                    rho=gamma, 
                                                    adaptive= args.adaptive == 'True',
                                                    norm=args.norm,
                                                    T=args.batch_num)
                    elif mode == 'APGD_average':
                        flatness=eval_average_sharpness(net,
                                                    mask,
                                                    train_loader,
                                                    loss_f=loss_f,
                                                    n_iters=100,
                                                    T=args.batch_num,
                                                    rho=gamma, 
                                                    adaptive= args.adaptive == 'True',
                                                    norm=args.norm,
                                                    )
                        
                else:
                    print("--------------------> Evaluate a sparse model")
                    print(f'{mode}, args.logit_norm={args.logit_norm} rho={gamma}, adaptive={args.adaptive == "True"}, {args.norm}, num_batch={args.batch_num}, {train_loader.batch_size}')
                    if mode == 'APGD_worst':
                        flatness=eval_APGD_sharpness_sparse(net,
                                                            mask,
                                                            train_loader,
                                                            loss_f=loss_f,
                                                            train_err=None,
                                                            train_loss=None,
                                                            n_iters=20,
                                                            rho=gamma, 
                                                            adaptive=args.adaptive == 'True',
                                                            norm=args.norm,
                                                            T=args.batch_num)
                    elif mode == 'APGD_average':
                        flatness=eval_average_sharpness_sparse(net,
                                                            mask,
                                                            train_loader,
                                                            loss_f=loss_f,
                                                            n_iters=100,
                                                            T=args.batch_num,
                                                            rho=gamma, 
                                                            adaptive= args.adaptive == 'True',
                                                            norm=args.norm,
                                                            )
                sharpness_ablation[f'{args.norm}_{gamma}'] = flatness
                print(sharpness_ablation)
        
        nets_flatness.append(sharpness_ablation)
        
    savepaths = os.path.join(savepaths, f'bs{args.eval_batch_size}_bn{args.batch_num}')
    Path(savepaths).mkdir(parents=True, exist_ok=True)
    path=os.path.join(savepaths,   f'{mode}_adaptive{args.adaptive}_logitnorm{args.logit_norm}_indiv.pkl')
    f=open(path,"wb")
    pickle.dump(nets_flatness,f)  
    f.close()

def compute_loss_stateless_model (params, buffers, sample, target):
    batch = sample.unsqueeze(0)
    targets = target.unsqueeze(0)
    loss_fn=nn.CrossEntropyLoss()
    global fmodel
    predictions = fmodel(params, buffers, batch) 
    loss = loss_fn(predictions, targets)
    return loss


def apply_mask_graient(net,ft_per_sample_grads,mask):
    i=0
    for name, param in net.named_parameters():
        if name in mask:
            ft_per_sample_grads[i]=ft_per_sample_grads[i]*torch.unsqueeze(mask[name],dim=0)
        i=i+1
    return ft_per_sample_grads

def test_fisher(model_lst, savepaths, masks):
    flatness_model_lst = []

    for model_i in range(len(model_lst)):
        net = model_lst[model_i]
        net.eval()
        train_loader, _ = get_trainloader()
        if(len(masks)>0):
            mask=masks[model_i].masks
        else:
            mask=None
        flatness=[]   #torch.zeros(len(train_loader.dataset))
        global fmodel
        fmodel, params, buffers = make_functional_with_buffers(net)
        ft_compute_grad = grad(compute_loss_stateless_model)
        ft_compute_sample_grad = vmap(ft_compute_grad, in_dims=(None, None, 0, 0))

        for i, dt in enumerate(tqdm(train_loader)):
            if i == args.batch_num:
                break
            idx,data,target=dt
            
            #print(idx)
            data,target=data.to(device, dtype=torch.float), target.to(device)
            ft_per_sample_grads = ft_compute_sample_grad(params, buffers, data, target)
            ft_per_sample_grads=list(ft_per_sample_grads)
            if mask is not None:
                ft_per_sample_grads=apply_mask_graient(net,ft_per_sample_grads,mask)
                
            for j in range(len(ft_per_sample_grads)):
                ft_per_sample_grads[j]=ft_per_sample_grads[j].view(data.shape[0],-1)
            
            grads=torch.cat(ft_per_sample_grads,dim=1)
            fisher=grads.pow(2).sum(dim=1)
            
            flatness.append(fisher.cpu().detach())
        flatness = torch.cat(flatness).numpy()
        print(flatness.shape)
        flatness_model_lst.append(np.mean(flatness))
    
    flatness_model_lst = np.array(flatness_model_lst)
    savepaths = os.path.join(savepaths, f'bs{args.eval_batch_size}_bn{args.batch_num}')
    Path(savepaths).mkdir(parents=True, exist_ok=True)
    path=os.path.join(savepaths,  'fisher.npy')
    np.save(path,flatness_model_lst)

def main():

    if args.gamma:
        args.gamma = [float(item) for item in args.gamma.split(',')]
        print('args.gamma', args.gamma)
    else:
        print('No numbers provided')
    
    if_sparse = False
    if args.dataset == 'cifar10':

        num_classes = 10
    else:
        num_classes = 100  



    # Distribute model across all visible GPUs
    cudnn.benchmark = True

    net_ensemble_lst = []
    masks=[]
    savepaths=[]

    i=0
    if args.resume:
        if len(args.resume) > 0:
            for file in args.resume:
                print(file)
                if args.model == 'densenet':
                    print("hh")
                elif 'wrn' in args.model:
                    print(args.layers, args.width)
                    net = WideResNet(args.layers, args.width , num_classes, args.droprate)  
                elif args.model=='resnet18':
                    net=ResNet18(width=args.width,num_classes=num_classes)
                elif args.model=='resnet34':
                    net=ResNet34(width=args.width,num_classes=num_classes)
                elif args.model=='resnet50':
                    net=ResNet50(width=args.width,num_classes=num_classes)

                net=net.cuda()
                indice=file.rindex('/')
                savepath=file[:indice]
                savepaths.append(savepath)
                print(savepath)
                
                checkpoint=torch.load(file,map_location='cuda:0')     
                if 'mask' in checkpoint:              
                    mask=checkpoint['mask']
                    masks.append(mask)
                net.load_state_dict(checkpoint['state_dict'])
                if ('mask' in checkpoint and  args.mode=='hessian'):
                    for name,module in net.named_modules():
                        if(name+'.weight') in mask.masks:
                            print(name)
                            prune.custom_from_mask(module,name='weight',mask=mask.masks[name+'.weight'])
                print(f'Model restored from: {file} ')
                i=i+1
                net.eval()
                net_ensemble_lst.append(net)


    if args.evaluate:
        if(args.mode=='hessian_t'):
            print("--------------------> enter hessian trace")
            hessians(net_ensemble_lst       ,args.save,  masks, mode='trace')
        elif (args.mode=='hessian_e'):
            print("--------------------> enter hessian Eigen")
            hessians(net_ensemble_lst       ,args.save,masks , mode='top_e')
        elif (args.mode=='fisher'):
            print("--------------------> enter Fisher")
            test_fisher(net_ensemble_lst       ,args.save,masks)
        elif('APGD' in args.mode):
            APGD(net_ensemble_lst   , args.save, masks, mode=args.mode)
        elif(args.mode=='ntk'):
            get_ntk(net_ensemble_lst,     args.save, masks)
            

if __name__ == '__main__':
    main()
