import cv2
import os
import torch
import torchvision
from torchvision import transforms, utils
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(3, 32, 3, 1, 1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2))
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(32, 64, 3, 1, 1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.conv3 = torch.nn.Sequential(
            torch.nn.Conv2d(64, 64, 3, 1, 1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2)
        )
        self.dense = torch.nn.Sequential(
            #torch.nn.Linear(4096, 1024),
            #torch.nn.ReLU(),
            torch.nn.Linear(1024, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 2)
            )
    
    def forward(self, x):
        conv1_out = self.conv1(x)
        conv2_out = self.conv2(conv1_out)
        conv3_out = self.conv3(conv2_out)
        res = conv3_out.view(conv3_out.size(0), -1)
        out = self.dense(res)
        return out
        
        

    
if __name__ == "__main__":
    train_data_path = "data/nn_data/trainset/"
    test_data_path = "data/nn_data/testset/"
    train_data = torchvision.datasets.ImageFolder(train_data_path,
                                                    transform=transforms.Compose([
                                                        transforms.Resize((32, 32)),
                                                        transforms.ToTensor()]))
    test_data = torchvision.datasets.ImageFolder(test_data_path,
                                                    transform=transforms.Compose([
                                                        transforms.Resize((32, 32)),
                                                        transforms.ToTensor()]))
    print(len(train_data))
    print(len(test_data))
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=4,shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_data, batch_size=4,shuffle=True)

    model = Net()
    print(model)
    optimizer = torch.optim.Adam(model.parameters())
    loss_func = torch.nn.CrossEntropyLoss()
    
    for epoch in range(50):
        print('epoch {}'.format(epoch + 1))
        # training-----------------------------
        train_loss = 0.
        train_acc = 0.
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = Variable(batch_x), Variable(batch_y)
            out = model(batch_x)
            loss = loss_func(out, batch_y)
            train_loss += loss.item()
            pred = torch.max(out, 1)[1]
            train_correct = (pred == batch_y).sum()
            train_acc += train_correct.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        print('Train Loss: {:.6f}, Acc: {:.6f}'.format(train_loss / (len(
            train_data)), train_acc / (len(train_data))))

        # evaluation--------------------------------
        model.eval()
        eval_loss = 0.
        eval_acc = 0.
        num_false = 0
        for batch_x, batch_y in test_loader:
            print(batch_x.shape)
            batch_x, batch_y = Variable(batch_x, volatile=True), Variable(batch_y, volatile=True)
            out = model(batch_x)
            loss = loss_func(out, batch_y)
            eval_loss += loss.item()
            pred = torch.max(out, 1)[1]
            if epoch == 49:
                for k in range(len(pred)):
                    if pred[k] != batch_y[k]:
                        torchvision.utils.save_image(
                            batch_x[k], os.path.join("data/k/"+str(num_false) +"-" + str(batch_y[k])+".png"))
                        num_false += 1
            num_correct = (pred == batch_y).sum()
            eval_acc += num_correct.item()
        print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len(
            test_data)), eval_acc / (len(test_data))))
    
    #torch.save(model, "CNN_model")