import torch
import torch.nn as nn
import torchvision.transforms as transforms
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from NeuralNetwork import NeuralNetwork

# import matplotlib.pyplot as plt
# import numpy as np
# import cv2

def trainNeuralNetwork():
    neural_net = NeuralNetwork()
    train_set = ImageFolder(root='./resources/trash_dataset/train', transform=transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]))
    trainloader = DataLoader(
        train_set, batch_size=2, shuffle=True, num_workers=2)

    # potrzebne do wyświetlania loss w każdej iteracji
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(neural_net.parameters(), lr=0.001, momentum=0.9)

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('device:')
    print(device)
    neural_net.to(device)

    epoch_num = 60 # najlepiej 10, dla lepszej wiarygodności
    for epoch in range(epoch_num):
        measure_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            inputs, labels = data[0].to(device), data[1].to(device)
            # czyszczenie gradientu f-cji
            optimizer.zero_grad()
            outputs = neural_net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            measure_loss += loss.item()
            if i:
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, measure_loss))
                measure_loss = 0.0
        
        if(epoch % 5 == 0):
            torch.save(neural_net.state_dict(), './NeuralNetwork/trained_networks/trained_nn_' + str(epoch) + '.pth')

    print('Finished.')
    # PATH = './trained_nn.pth'
    # torch.save(neural_net.state_dict(), PATH)

def main():
    trainNeuralNetwork()

if __name__ == '__main__':
    main()