import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import Compose, Lambda, ToTensor
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from PIL import Image
import random

imageSize = (128, 128)
labels = ['carrot','corn', 'potato', 'tomato'] # musi być w kolejności alfabetycznej 
fertilizer = {labels[0]: 'kompost', labels[1]: 'saletra amonowa', labels[2]: 'superfosfat', labels[3]:'obornik kurzy'}
#labels = ['corn','tomato'] #uncomment this two lines for 2 crops only
#fertilizer = {labels[0]: 'kompost', labels[1]: 'saletra amonowa'}
torch.manual_seed(42)

#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cpu")
# device = torch.device("mps") if torch.backends.mps.is_available() else torch.device('cpu')
# print(device)

def getTransformation():
    transform=transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        transforms.Resize(imageSize),
        Lambda(lambda x: x.flatten())])
    return transform

def getDataset(train=True):
    transform = getTransformation()
    if (train):
        trainset = datasets.ImageFolder(root='dataset/train', transform=transform)
        return trainset
    else:
        testset = datasets.ImageFolder(root='dataset/test', transform=transform)
        return testset


def train(model, dataset, n_iter=100, batch_size=256):
    optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
    criterion = nn.NLLLoss()
    dl = DataLoader(dataset, batch_size=batch_size)
    model.train()
    for epoch in range(n_iter):
        for images, targets in dl:
            optimizer.zero_grad()
            out = model(images.to(device))
            loss = criterion(out, targets.to(device))
            loss.backward()
            optimizer.step()
        if epoch % 10 == 0:
            print('epoch: %3d loss: %.4f' % (epoch, loss))
    return model

def accuracy(model, dataset):
    model.eval()
    correct = sum([(model(images.to(device)).argmax(dim=1) == targets.to(device)).sum()
    for images, targets in DataLoader(dataset, batch_size=256)])
    return correct.float() / len(dataset)

def getModel():
    hidden_size = 500
    model = nn.Sequential(
        nn.Linear(imageSize[0] * imageSize[1] * 3, hidden_size),
        nn.ReLU(),
        nn.Linear(hidden_size, len(labels)),
        nn.LogSoftmax(dim=-1)
        ).to(device)
    return model
    
def saveModel(model, path):
    print("Saving model")
    torch.save(model.state_dict(), path)

def loadModel(path):
    print("Loading model")
    model = getModel()
    model.load_state_dict(torch.load(path))
    return model

def trainNewModel(n_iter=100, batch_size=256):
    trainset = getDataset(True)
    model = getModel()
    model = train(model, trainset)
    return model

def predictLabel(imagePath, model):
    image = Image.open(imagePath).convert("RGB")
    image = preprocess_image(image)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    with torch.no_grad():
        model.eval()  # Ustawienie modelu w tryb ewaluacji
        output = model(image)

    # Znalezienie indeksu klasy o największej wartości prawdopodobieństwa
    predicted_class = torch.argmax(output).item()
    return labels[predicted_class]

    # Znalezienie indeksu klasy o największej wartości prawdopodobieństwa
    predicted_class = torch.argmax(output).item()
    return labels[predicted_class]


def preprocess_image(image):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    transform = getTransformation()
    image = transform(image).unsqueeze(0)  # Add batch dimension
    image = image.to(device)  # Move the image tensor to the same device as the model
    return image