# %%
import torch as th
import torch.nn as nn
import pandas as pd
from os import mkdir
from shutil import copyfile, copytree
from torchvision import datasets, transforms, models
import matplotlib.pyplot as plt
from collections import OrderedDict
from tqdm import tqdm
import numpy as np

# %%

train_transformation = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize(224),
    transforms.RandomRotation(10),
])

train_dataset = datasets.ImageFolder (
    root="C:\\Users\\preda\\Desktop\\Github\\Facultate\\Anul 2\\IA\\train_folder",
    transform=train_transformation,
)

train_loader = th.utils.data.DataLoader(
    train_dataset, batch_size=32, shuffle=True
)

validation_transformation = transforms.Compose([
    transforms.ToTensor(),
    transforms.Resize(224),
])

validation_dataset = datasets.ImageFolder (
    root="C:\\Users\\preda\\Desktop\\Github\\Facultate\\Anul 2\\IA\\validation_folder",
    transform=validation_transformation,
)

validation_loader = th.utils.data.DataLoader(
    validation_dataset, batch_size=32, shuffle=True
)

# %%
class Network(nn.Module):
    def __init__(self):
        super().__init__()

        self.net = nn.Sequential(
            nn.Conv2d(3, 16, 3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(16),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, 3, padding=1),
            nn.ReLU(),
            nn.BatchNorm2d(32),
            nn.MaxPool2d(2),

            nn.Flatten(),

            nn.Linear(32 * 50 * 50, 1000),
            nn.ReLU(),
            nn.Dropout(0.7),
            nn.Linear(1000, 3),
            # nn.ReLU(),
            # nn.Dropout(0.5),
            # nn.Linear(200, 3),
        )

    def forward(self, x):
        result = self.net(x)
        return result

# %%

def train_one_epoch():
    net.train()

    hit = 0
    total = 0
    avg_loss = 0.
    for images, labels in tqdm(train_loader):
        images = images.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()
        predictions = net.forward(images)
        loss = criterion(predictions, labels)
        avg_loss += loss.item()

        loss.backward()
        optimizer.step()

        hit += th.sum(labels == predictions.argmax(dim=1))
        total += len(images)

    return hit / total, avg_loss / len(train_loader)

def test_data(data_loader):
    net.eval()

    with th.no_grad():
        hit = 0
        total = 0
        avg_loss = 0.
        for images, labels in data_loader:
            images = images.cuda()
            labels = labels.cuda()

            predictions = net.forward(images)

            loss = criterion(predictions, labels)
            avg_loss += loss.item()
            hit += th.sum(labels == predictions.argmax(dim=1))
            total += len(images)

        return (hit / total, avg_loss / len(data_loader))

def print_confussion_matrix(data_loader):
    ans = [[0 for i in range(3)] for j in range(3)]
    with th.no_grad():
        for images, labels in tqdm(data_loader):
            images = images.cuda()
            labels = labels.cuda()

            predictions = net(images)
            prediction = predictions.argmax(dim=1)
            for i in range(len(labels)):
                ans[labels[i].item()][prediction[i].item()] += 1

    plt.imshow(ans, cmap='gray')
    print(np.array(ans))


# %%

# definim datele
vgg = models.alexnet(pretrained=True).cuda()
# for param in vgg.parameters():
#     param.requires_grad = False
vgg.classifier = nn.Sequential(
                            nn.Dropout(p=0.5, inplace=False),
                            nn.Linear(in_features=9216, out_features=1700, bias=True),
                            nn.ReLU(inplace=True),
                            # nn.Dropout(p=0.5, inplace=False),
                            # nn.Linear(in_features=4096, out_features=4096, bias=True),
                            # nn.ReLU(inplace=True),
                            nn.Linear(in_features=1700, out_features=3, bias=True)
                          ).cuda()

net = vgg
criterion = nn.CrossEntropyLoss().cuda()
loss_hist = []
validation_loss_hist = []

# %%

optimizer = th.optim.Adam(net.parameters(), lr=1e-5)

for epoch in range(100000):
    accuracy, loss = train_one_epoch()
    print(f"Loss: {loss}")
    validation_accuracy, validation_loss = test_data(validation_loader)
    loss_hist.append(loss)
    validation_loss_hist.append(validation_loss)
    epoch = len(loss_hist)
    print(f"Epoch: {epoch}")
    print(f"Train: loss: {loss}, accuracy: {accuracy}")
    print(f"Validation: loss: {validation_loss}, accuracy: {validation_accuracy}")

    # min_epoch = 2
    # if(epoch > min_epoch):
    #     plt.plot(loss_hist[min_epoch:], color='blue')
    #     plt.plot(validation_loss_hist[min_epoch:], color='red')
    #     plt.show()

# %%

test_dataset = datasets.ImageFolder (
    root="C:\\Users\\preda\\Desktop\\Github\\Facultate\\Anul 2\\IA\\test_folder",
    transform=validation_transformation,
)

test_loader = th.utils.data.DataLoader(
    test_dataset, batch_size=1, shuffle=False
)

# %%

def print_test():
    tests = []
    for file in os.listdir("C:\\Users\\preda\\Desktop\\Github\\Facultate\\Anul 2\\IA\\test_folder\\test"):
        tests.append(file)
    lastFile = 0
    net.eval()

    f = open("C:\\Users\\preda\\Desktop\\Github\\Facultate\\Anul 2\\IA\\solution.txt", "w")
    f.write("id,label\n")

    with th.no_grad():
        for images, labels in tqdm(test_loader):
            images = images.cuda()
            labels = labels.cuda()

            predictions = net(images)
            prediction = predictions.argmax(dim=1)
            f.write(tests[lastFile] + "," + str(prediction.item()) + "\n")

            lastFile += 1

    f.close()

# %%