import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader, TensorDataset, random_split, RandomSampler
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)
    #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)

    # Create DataLoader
    dataset = train_data
    if B > train_size:
        print("Error: The number of samples cannot be larger than the train set size!")
    else:
        sampler = RandomSampler(dataset, replacement=False, num_samples=B)
        #batch_loader = DataLoader(dataset=dataset, batch_size=bs, sampler=sampler)
        single_loader = DataLoader(dataset=dataset, batch_size=1, sampler=sampler)
        #full_loader = DataLoader(dataset=dataset, batch_size=N_train)
    '''
    # Compute the full gradient
    optimizer.zero_grad()
    FullLoss = criterion(model(inputs), labels)
    FullLoss.backward()
    FullGradient = torch.cat((model.linear.weight.grad.clone(), model.linearminus.weight.grad.clone()), 1)
    #print("The full gradient is {}".format(FullGradient))
    optimizer.zero_grad()
    '''
    # Compute the full gradient, Hessian and the Covariance Matrix
    optimizer.zero_grad()
    for idx, (image, label) in enumerate(single_loader):
    #for idx, (x, y) in enumerate(batch_loader):
        #torch.cuda.empty_cache()
        image = Variable(image)
        label = Variable(label)
        image = image.to(device)
        label = label.to(device)
        output = model(image)[0]
        output = output.to(torch.float32)
        #label = label.to(torch.float32)
        #print(type(output[0][0]))
        loss = criterion(output, label)
        loss.backward()
        FullLoss += loss.item()
        Gradient = torch.empty(0, 1)
        Gradient = FullGradient.to(device)
        for p in model.parameters():
            if p.requires_grad:
                Gradient = torch.cat((Gradient, p.grad.view(1, -1)), 0)
        Gradient = torch.transpose(Gradient)
        #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)))
    print('The full loss is {}'.format(FullLoss))
    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