import torch
from torch.autograd import Variable
from torch.utils.data import DataLoader, TensorDataset, random_split, RandomSampler
import numpy as np
from numpy import linalg as LA

bootstrap_bs = 1
BatchNum = 40

def BootstrapLoss(model, inputs, labels, N_train, bs=bootstrap_bs):
    BatchNum = N_train
    Losses = np.zeros(BatchNum)
    criterion = torch.nn.MSELoss()
    '''
    # Converting inputs and labels to Variable
    if torch.cuda.is_available():
        inputs = Variable(torch.from_numpy(inputs).cuda())
        labels = Variable(torch.from_numpy(labels).cuda())
    else:
        inputs = Variable(torch.from_numpy(inputs))
        labels = Variable(torch.from_numpy(labels))
    '''
    # Create DataLoader
    dataset = TensorDataset(inputs, labels)
    sampler = RandomSampler(dataset, replacement=True, num_samples=bootstrap_bs * BatchNum)
    batch_loader = DataLoader(dataset=dataset, batch_size=bs, sampler=sampler)
    single_loader = DataLoader(dataset=dataset, batch_size=1)

    # Compute the loss for each batch
    for idx, (x, y) in enumerate(single_loader):
        #print("idx is {}".format(idx))
        #print(x)
        #print(y)
        y_pred = model(x)
        Losses[idx] = criterion(y_pred, y)

    MeanLoss = Losses.mean()
    LossDifferences = Losses - MeanLoss
    #print(LossDifferences)
    #print(LA.norm(LossDifferences))
    return LA.norm(LossDifferences) ** 2 / BatchNum