import numpy as np
import torch
from numpy import asarray
from torch import nn
from torch import optim
from torch.autograd import Variable
from torch.utils.data import Dataset
import os
import pandas as pd
from PIL import Image
from torchvision.io import read_image

from torchvision.transforms import Resize, Lambda, transforms, ToTensor
from mnist import MNIST

from Constants import *
from torch.utils.data import DataLoader


class Neurons:
    def __init__(self):
        dataset = MNIST('C:\\Users\\Konrad\\PycharmProjects\\Projekt_AI-Automatyczny_saper\\Engine', gz=True, return_type='numpy')
        images, labels = dataset.load_training()
        test_images, test_labels = dataset.load_testing()
        images = images / 255
        test_images = test_images / 255

        self.train_images = [torch.tensor(image, dtype=torch.float32) for image in images]
        self.train_labels = [torch.tensor(label, dtype=torch.long) for label in labels]
        test_images = [torch.tensor(image, dtype=torch.float32) for image in test_images]
        test_labels = [torch.tensor(label, dtype=torch.long) for label in test_labels]

        input_dim = 28 * 28
        output_dim = 10

        self.model = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.LogSoftmax()
        )
        self.model.eval()
        torch.save(self.train(self.model, 100), 'model.pth')
        imsize = 28
        loader = transforms.Compose([transforms.Resize((imsize, imsize)), transforms.ToTensor()])
        image = Image.open("C:\\Users\\Konrad\\PycharmProjects\\JPG-PNG-to-MNIST-NN-Format\\test-images\\1\\claymore (1).png")
        image = image.convert("1")
        image = loader(image).float()
        image = Variable(image, requires_grad=True)
        predicted = self.model(image)
        print(predicted)

    def train(self,model, n_iter):
        criterion = nn.NLLLoss()
        optimizer = optim.SGD(model.parameters(), lr=0.001)

        for epoch in range(n_iter):
            for image, label in zip(self.train_images, self.train_labels):
                optimizer.zero_grad()

                output = model(image)
                loss = criterion(output.unsqueeze(0), label.unsqueeze(0))
                loss.backward()
                optimizer.step()

            print(f'epoch: {epoch:03}')
        return model


    def accuracy(self,expected, predicted):
        return len([int for e, p in zip(expected, predicted) if e == p])/len(expected)


    def loadImages(self):
        col_dir = 'D:\\fotyProjekt\\clowns'
        images = []
        labels = []
        imgs = os.listdir(col_dir)
        for i in imgs:
            images.append(i)
        print("Load clowns")
        print(imgs)
        for i in imgs:
            labels.append(self.getLabelsForImages(DECOY))
        col_dir = 'D:\\fotyProjekt\\chemical'
        imgs = os.listdir(col_dir)
        for i in imgs:
            images.append(i)
        print("Load chemical")
        for i in imgs:
            labels.append(self.getLabelsForImages(CHEMICAL_BOMB))
        col_dir = 'D:\\fotyProjekt\\atomic'
        imgs = os.listdir(col_dir)
        for i in imgs:
            images.append(i)
        print("Load atomic")
        for i in imgs:
            labels.append(self.getLabelsForImages(ATOMIC_BOMB))
        col_dir = 'D:\\fotyProjekt\\landmine'
        imgs = os.listdir(col_dir)
        for i in imgs:
            images.append(i)
        print("Load landmine")
        for i in imgs:
            labels.append(self.getLabelsForImages(LAND_MINE))
        col_dir = 'D:\\fotyProjekt\\claymore'
        imgs = os.listdir(col_dir)
        for i in imgs:
            images.append(i)
        print("Load claymore")
        for i in imgs:
            labels.append(self.getLabelsForImages(CLAYMORE))
        f = open('labels.csv', 'w')
        while images:
            img = images.pop()
            lb = labels.pop()
            f.write(os.path.basename(img) + ", " + lb.__str__() + "\n")
        f.close()

    def getLabelsForImages(self,imageType):
        return self.d.get(imageType)



if __name__ == '__main__':
    Neurons()