

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

from PIL import Image
from typing import List

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  -------------------------------

# ----

GrowthDict = {
            1: "early growth", 
            2: "mature growth", 
            3: "overgrown"}

ClonesDict = {
            0: "0 clones", 
            1: "1 clone", 
            2: "2 clones", 
            3: "3 or more clones"}

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

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


class MultiPrediction():
    def __init__(self, growth_classes : List[int], growth_prop : List[float], clone_classes : List[int], clone_prop : List[float]):
        self.pred = {
           "growth_classes" : growth_classes,
           "growth_prop" : growth_prop,
           "clone_classes" : clone_classes,
           "clone_prop" : clone_prop,
           "text" : (', '.join(['%s (%0.1f%%)' % (GrowthDict[growth_classes[i]], growth_prop[i] * 100.0) for i in range(len(growth_classes))]) +
                    "\n" + ', '.join(['%s (%0.1f%%)' % (ClonesDict[clone_classes[i]], clone_prop[i] * 100.0) for i in range(len(clone_classes))]))
        }

    def getDict(self) -> dict:
        return self.pred


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...
        
        k_highest = 2 # find de 2 højest scored classes og probabilities 
        (probabilities, output_classes) = torch.topk(torch.reshape(prediction[0].detach(), MHE), k_highest)

        def toIntList(iterable) -> List[int]:
            return [int(x) for x in iterable]

        def toFloatList(iterable) -> List[float]:
            return [float(x) for x in iterable]

        growth_classes = toIntList(output_classes[0,:].numpy()) # rangeret efter højeste probability, feks [2 1]
        print(growth_classes)

        growth_probabilities = toFloatList(probabilities[0,:].numpy()) # de relaterede sandsynligheder, feks [9.9999847e+01 5.7163057e-03], 
        print(growth_probabilities)

        n_cells_classes = toIntList(output_classes[1,:].numpy()) # rangeret efter højeste probability, feks [1 2]
        print(n_cells_classes)

        n_cells_probabilities = toFloatList(probabilities[1,:].numpy()) # de relaterede sandsynligheder, feks [1.0000000e+02 2.8741499e-02]
        print(n_cells_probabilities)

        return MultiPrediction(growth_classes, growth_probabilities, n_cells_classes, n_cells_probabilities)

