import torch
import torchvision
import torchvision.transforms as transforms
from mixbatch_pytorch import MixBatch_torch
from tensorboardX import SummaryWriter
from tqdm import tqdm
import numpy as np
from models import ResNet50

SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.backends.cudnn.deterministic = True
np.random.seed(SEED)


writer = SummaryWriter(log_dir='stl10_res50_runs/baseline')
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")



from STL10.stl10_input import  LABEL_PATH,DATA_PATH,read_all_images,read_labels  # test to check if the whole dataset is read correctly

LABEL_PATH = LABEL_PATH.replace("./data","./STL10/data")
DATA_PATH = DATA_PATH.replace("./data","./STL10/data")
images = read_all_images(DATA_PATH)
print(images.shape)

labels = read_labels(LABEL_PATH)
print(labels.shape)
 
# 1/0
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

images = torch.from_numpy(images)
labels = torch.from_numpy(labels)

images,labels=images.type(torch.FloatTensor),labels.type(torch.FloatTensor)
dataset = torch.utils.data.TensorDataset(images, labels)

trainloader = torch.utils.data.DataLoader(dataset, batch_size=4,shuffle=True, num_workers=0)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

import torch.nn as nn
import torch.nn.functional as F


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)
        self.mb = MixBatch_torch(p=0.2)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.pool(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.pool(x)
        x = x.view(-1, 16 * 5 * 5)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.mb(x)
        x = self.fc2(x)
        x = F.relu(x)
        x = self.fc3(x)
        return x

net = ResNet50()
net = net.to(device)

print(torch.cuda.is_available())
import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

step_per_epoch = len(trainloader)
total_epoch = 100

for epoch in range(total_epoch):  # loop over the dataset multiple times

    net.train()
    running_loss = 0.0
    for i, data in tqdm(enumerate(trainloader, 0),total=step_per_epoch):
        # get the inputs; data is a list of [inputs, labels]
        inputs, labels = data

        # zero the parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = net(inputs.to(device))
        loss = criterion(outputs, labels.to(device))
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss.item()
        # if i % 2000 == 1999:    # print every 2000 mini-batches
        #     print('[%d, %5d] loss: %.3f' %
        #           (epoch + 1, i + 1, running_loss / 2000))
        #     running_loss = 0.0

        if i % 100 == 99:
            global_step = epoch*step_per_epoch + i
            writer.add_scalar('train_loss', loss, global_step)


    net.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in testloader:
            data, target = data.to(device), target.to(device)
            output = net(data)
            test_loss += criterion(outputs, labels.to(device))
            pred = output.argmax(dim=1, keepdim=True) # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(testloader.dataset)

    acc = 100. * correct / len(testloader.dataset)

    global_step = epoch*step_per_epoch
    writer.add_scalar('eval_loss', loss, global_step)
    writer.add_scalar('eval_acc', acc, global_step)


print('Finished Training')
writer.close()