

import torch
import torchvision.transforms as transforms
import torch.nn.init

from PIL import Image

keep_prob = 0.9 
n_classes = 8  # multi hot encoded
MHE = (2,4)  # dvs 2 output values med  4 classes i hver

# Convolutional network model
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.layer1 = torch.nn.Sequential(
            torch.nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1),  # input image size 44,  1,32,3,1,1
            torch.nn.ReLU(),        
            torch.nn.MaxPool2d(kernel_size=2, stride=2),   
            torch.nn.Dropout(p=1 - keep_prob)) 

        self.layer2 = torch.nn.Sequential(
            torch.nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1), # 32,64,3,1,1
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2),
            torch.nn.Dropout(p=1 - keep_prob))

        self.layer3 = torch.nn.Sequential(
            torch.nn.Conv2d(64, 100, kernel_size=3, stride=1, padding=1), # 64,100,3,1,1
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=1),
            torch.nn.Dropout(p=1 - keep_prob))

        self.fc1 = torch.nn.Linear(6 * 6 * 100, 400, bias=True)  # 6 * 6 * 100, 400, fully connected layer 1
        torch.nn.init.xavier_uniform_(self.fc1.weight)  # _ ???
        
        self.layer4 = torch.nn.Sequential(   
            self.fc1,
            torch.nn.ReLU(),
            torch.nn.Dropout(p=1 - keep_prob))
        
        self.fc2 = torch.nn.Linear(400, n_classes, bias=True)   # 400 - > 8 classes, , fully connected layer 2
        torch.nn.init.xavier_uniform_(self.fc2.weight) # initialize weigts parameters

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = self.layer3(out)
        out = out.view(out.size(0), -1)   # Flatten them for fully connected layers
        out = self.layer4(out)  # self.fc1(out)
        out = self.fc2(out)
        return torch.nn.Sigmoid()(out)  # out sigmoid range 0..1  -------------------------------

# ----

class Prediction():
    def __init__(self, growth:int, cells:int):
        self.growth = growth
        self.cells = cells

    def __str__(self) -> str:
        growth_level = {1: "early growth", 2: "mature growth", 3: "overgrown"}.get(self.growth, "?")
        cells = {0: "0 cells", 1: "1 cell", 2: "2 cells", 3: "3 or more cells"}.get(self.cells, "?")
        return f"{growth_level}, {cells}"
    
    def getDict(self) -> dict:
        return {"growth" : self.growth, "cells" : self.cells, "text" : str(self)}

class Predictor():
    def __init__(self, pth_filename:str):
        self.cnn = CNN()  
        self.modelName = pth_filename.split('.')[0]

        # load trained model from file
        print("Loading model from: " + pth_filename)
        self.cnn.load_state_dict(torch.load(pth_filename))
        
        # set model in evaluation mode
        self.cnn.eval()

    def predict(self, input_img:Image) -> Prediction:
        
        # Resize to 44 x 44 pixels, and convert to grayscale (if not already). Model is trained on 44 x 44 pixel images
        img = input_img.resize((44, 44)).convert("L")

        im_tensor = transforms.ToTensor()(img).unsqueeze(0)  # unsqueeze(0) giver [1,44,44] -> [1,1,44,44]

        # print(np.shape(im_tensor)) # = [1,1,44,44] det er det format som modellen tager som input.  [n_batch, channel, imx, imy]
        # Hvis batch, n af flere billeder på en gang skal det pakkes i formatet [n,1,44,44]

        # run the model prediction on one image
        prediction = self.cnn(im_tensor)

        # output is something like : [2.7335e-13, 1.0000e+00, 1.6826e-11, 1.3905e-11, 9.9984e-01, 2.6207e-05, 4.9733e-08, 1.2173e-10]
        # hvilket indeholder de 8 neuroner. Skal decodes til [A, B], ud fra [A1, A2, A3, A4, B1, B2, B3, B4]
        # bruge argmax og reshape til at få final [A,B] 

        output_classes = list(torch.argmax(torch.reshape(prediction[0].detach(), MHE), dim = -1).numpy())
        # .detach() bruges for at undgå "grad_fn=<SigmoidBackward0>" dvs undgå at den aktivt tracker gradienten for hver operation...

        return Prediction(int(output_classes[0]), int(output_classes[1]))
