import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader, TensorDataset, random_split, RandomSampler
#from torch.func import functional_call, vmap, grad
from torch.nn.utils._per_sample_grad import call_for_per_sample_grads
import numpy as np
import copy
import random

#bootstrap_bs = 5
#BatchNum = 100
#NumSample = 20
#B = 50
seed = 2
npseed = 6
torchseed = 5

def LossScaledTrace(test_model, train_data, d, train_size, B=3200):
    #torch.manual_seed(seed)
    model = copy.deepcopy(test_model)
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    # Extract the images and labels from the train_data set
    images = Variable(train_data.dataset.data.to(torch.float))
    labels = Variable(train_data.dataset.targets.to(torch.float))
    images = images.to(device)
    labels = labels.to(device)
    print('Images size is')
    print(images.size())
    #model = nn.DataParallel(model)
    FullGradient = torch.zeros(d, 1)
    FullGradient = FullGradient.to(device)
    #Gradients = torch.zeros(BatchNum, d * 2)
    CovarianceMatrix = torch.zeros(d, d)
    CovarianceMatrix = CovarianceMatrix.to(device)
    Hessian = torch.zeros(d, d)
    Hessian = Hessian.to(device)
    FullLoss = 0
    #FullLoss = FullLoss.to(device)
    criterion = torch.nn.CrossEntropyLoss()
    #optimizer = torch.optim.SGD(model.parameters(), lr=1)
    # Compute the per-sample gradients with call_for_per_sample_grads()
    res = call_for_per_sample_grads(model)(images)[0]
    criterion(res, labels).backward()
    for param in model.parameters():
        param.grad_sample  # instead of the grad, per sample gradients are in the grad_sample field
        print(param.grad_sample)
    '''
    for i in range(ft_per_sample_grads.size()[0]):
        FullLoss += 1#loss.item()
        Gradient = torch.transpose(torch.cat((model.module.conv1[0].weight.grad.view(1, -1), model.module.conv1[0].bias.grad.view(1, -1), model.module.conv2[0].weight.grad.view(1, -1)
                              , model.module.conv2[0].bias.grad.view(1, -1), model.module.out.weight.grad.view(1, -1), model.module.out.bias.grad.view(1, -1)),1), 0, 1)
        #print("The gradient is {}".format(Gradient))
        #print("The size of the gradient is {}".format(Gradient.size()))
        FullGradient += Gradient
        Hessian += torch.mm(torch.reshape(Gradient, (d, 1)), torch.reshape(Gradient, (1, d))) / max([(loss.item() * 2), 10e-15])
        #print("The size of the Hessian is {}".format(Hessian.size()))
        CovarianceMatrix += torch.mm(torch.reshape(Gradient, (d, 1)), torch.reshape(Gradient, (1, d)))
        optimizer.zero_grad()
        #print(torch.matmul(Gradients[idx], torch.transpose(Gradients[idx])).size())
    FullGradient = FullGradient / B
    Hessian = Hessian / B
    CovarianceMatrix = CovarianceMatrix / B
    CovarianceMatrix -= torch.mm(torch.reshape(FullGradient, (d, 1)), torch.reshape(FullGradient, (1, d)))
    FullLoss =  FullLoss / B
    print(type(torch.trace(Hessian)))
    '''
    return torch.trace(torch.mm(Hessian, CovarianceMatrix)).cpu().numpy() / max([(FullLoss * 2), 10e-15]), \
           torch.norm(Hessian, p='fro').cpu().numpy() / 1, torch.trace(Hessian).cpu().numpy() / 1
