import glob
import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import imageio.v3 as iio
import einops

import torch
from torch.utils.data import DataLoader
from torchvision import transforms

from UNET_perso import UNET
from src.medicalDataLoader import MedicalImageDataset
from src.utils import getTargetSegmentation, IOU, dice_score

def disp_init(num):
    lst_patients = glob.glob('./Data/val/Img/*.png')
    patient = lst_patients[num]
    im = iio.imread(patient)
    fig =  plt.figure()
    plt.imshow(im)
    plt.axis('off')
    return fig

def disp_prediction(num, alpha=0.45):

    lst_patients = glob.glob('./Data/val/Img/*.png')
    patient = lst_patients[num]

    model = UNET(in_channels=1, out_channels=4).to('cpu')

    filepath = 'UNET_perso'
    model.load_state_dict(torch.load(filepath, map_location=torch.device('cpu')))
    model.eval()

    transform = transforms.Compose([transforms.ToTensor()])
    val_set = MedicalImageDataset('val',
                                './Data',
                                transform=transform,
                                mask_transform=transform,
                                equalize=False)
    val_loader = DataLoader(val_set,
                            batch_size=1,
                            shuffle=False)
    
    for _, (img, _, name) in enumerate(val_loader):
        if name[0] == patient:
            im = img

    pred = model(im)

    pred = pred.detach().numpy()
    pred = pred.reshape(4,256,256)
    pred_mask = np.argmax(pred, axis=0)

    im = im.detach().numpy()
    im = im.reshape(256,256)
    
    fig = plt.figure()
    try:
        plt.imshow((alpha*pred_mask + (1-alpha)*im)/2)
    except ValueError:
        pass
    plt.axis('off')

    return fig

def disp_label(num):
    lst_GT = glob.glob('./Data/val/GT/*.png')
    GT = lst_GT[num]
    label = iio.imread(GT)
    fig = plt.figure()
    try: 
        plt.imshow(label)
    except ValueError:
        pass
    plt.axis('off')

    return fig


def compute_metrics(num):

    lst_patients = glob.glob('./Data/val/Img/*.png')
    patient = lst_patients[num]

    model = UNET(in_channels=1, out_channels=4).to('cpu')

    filepath = 'UNET_perso'
    model.load_state_dict(torch.load(filepath, map_location=torch.device('cpu')))
    model.eval()

    transform = transforms.Compose([transforms.ToTensor()])
    val_set = MedicalImageDataset('val',
                                './Data',
                                transform=transform,
                                mask_transform=transform,
                                equalize=False)
    val_loader = DataLoader(val_set,
                            batch_size=1,
                            shuffle=False)
    
    for _, (img, label, name) in enumerate(val_loader):
        if name[0] == patient:
            im = img
            lab = label.reshape(256,256)

    pred = model(im)

    pred = pred.detach().numpy()
    pred = pred.reshape(4,256,256)
    pred = np.argmax(pred, axis=0)
    pred = torch.from_numpy(pred)
    pred = torch.nn.functional.one_hot(pred, num_classes=4)
    pred = einops.rearrange(pred, 'h w class -> class h w')

    lab = getTargetSegmentation(lab)
    lab = torch.nn.functional.one_hot(lab, num_classes=4)
    lab = einops.rearrange(lab, 'h w class -> class h w')
    
    dsc1 = dice_score(pred[1], lab[1])
    dsc2 = dice_score(pred[2], lab[2])
    dsc3 = dice_score(pred[3], lab[3])

    iou1 = IOU(pred[1], lab[1])
    iou2 = IOU(pred[2], lab[2])
    iou3 = IOU(pred[3], lab[3])

    df = pd.DataFrame(columns=['class 1', 'class 2', 'class 3'])
    df.loc[len(df)] = [round(dsc1,2), round(dsc2,2), round(dsc3,2)]
    df.loc[len(df)] = [round(iou1,2), round(iou2,2), round(iou3,2)]
    df = df.assign(metric=['Dice Score', 'IoU'])
    df = df[['metric','class 1','class 2','class 3']]
    
    return df

with gr.Blocks() as demo:
    with gr.Row() as row1:
        slide = gr.Slider(minimum=0, maximum=89, value=0, step=1, label='inference with validation split data (choose patient)')
    with gr.Row() as row2:
        slide.release(fn=disp_init, inputs=[slide], outputs=gr.Plot(label='initial image'))
        slide.release(fn=disp_prediction, inputs=[slide], outputs=gr.Plot(label='prediction'))
        slide.release(fn=disp_label, inputs=[slide], outputs=gr.Plot(label='groundtruth'))
    with gr.Row() as row3:
        slide.release(fn=compute_metrics, inputs=[slide], outputs=gr.DataFrame())

demo.queue()
demo.launch()