from dataset import my_dataset
from dataset import matching_dataset
import numpy as np
import torch
import torchvision as tv
import torchvision.transforms as transforms
import torch.nn as nn
import models
import time
import resnet

lr = 0.02
device = torch.device("cuda")
transform_train = transforms.Compose([
    # transforms.RandomCrop(32, padding=4),
    # transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    # transforms.Normalize((0.5, 0.5, 0.5, 0.5), (0.5, 0.5, 0.5, 0.5)),
])


def train(trainloader, testloader, model, criterion, optimizer, epoch, use_cuda, size=None, save=True):
    # switch to train mode
    loss_list = []
    acc_list = []
    model = nn.DataParallel(model).cuda()
    model.train()
    if use_cuda:
        model = model.to(device)
    best_acc = 0.0
    for i in range(epoch):
        for param_group in optimizer.param_groups:
            if i in [200, 400]:
                param_group['lr'] *= 0.1
        loss = 0.0
        print("%d | %d: " % (i, epoch))
        for data in trainloader:
            model.train(True)
            image, label = data['image'], data['label']
            if use_cuda:
                image = torch.autograd.Variable(image.cuda())
                label = torch.autograd.Variable(label.cuda())

            optimizer.zero_grad()

            output = criterion(model(image), label)
            output.backward()
            optimizer.step()
            loss += output.item()
        print("loss: %.04f" % (loss))
        loss_list.append(loss)
        # test_loss, test_acc = test(testloader, model, criterion, 1, use_cuda = True)
        # if test_acc > best_acc:
        #    best_acc = test_acc
        # print('the accuracy is %.03f, the best accuracy is %.03f'%(test_acc, best_acc))
        acc = test1(model, testloader)
        acc_list.append(acc)

        if save == True and best_acc < acc:
            torch.save(model.state_dict(), "task1resnet18.pkl")
        best_acc = max(best_acc, acc)
        print("accuracy : %.03f, best : %.03f" % (acc, best_acc))
    # np.save("acc_50.npy", acc_list)
    # np.save("loss_50.npy", loss_list)


def test1(model, loader):
    model.eval()
    correct = 0
    total = 0
    for data in loader:
        image, label = data['image'], data['label']
        image, label = image.cuda(), label.cuda()
        image, label = torch.autograd.Variable(image), torch.autograd.Variable(label)
        output = model(image)
        pred = torch.argmax(output, 1)
        correct += (pred == label).sum().float()
        total += len(label)
    acc = (100 * correct * 1.0 / total)
    print("accuracy : %.03f" % (acc))
    return acc


resnet = models.ResNet(block=models.BasicBlock, num_blocks=[3,3,3])
resnet = nn.DataParallel(resnet).cuda()
resnet.load_state_dict(torch.load("task1resnet18.pkl"))
# resnet = resnet.resnet50()
mlp = models.MLP()
criterion = nn.CrossEntropyLoss().cuda()
# optimizer = torch.optim.Adam(resnet.parameters(), lr = lr)
optimizer = torch.optim.SGD(resnet.parameters(), lr=lr, momentum=0.9, weight_decay=1e-4)
ds = my_dataset("train", transform_train)

train_ds, val_ds = torch.utils.data.random_split(ds, [2000, 249])
train_loader = torch.utils.data.DataLoader(train_ds, 32, False, num_workers=20)
val_loader = torch.utils.data.DataLoader(val_ds, 16, False, num_workers=20)

# print(len(task2_ds))

# dataset2 = matching_dataset(mode="train")
# x = dataset2.__getitem__(5, 6)
# print(x['raw'][1].shape)

test1(resnet, val_loader)
# train(train_loader, val_loader, resnet, criterion, optimizer, 500, use_cuda=True, save = False)
