import numpy as np
import gradio as gr

from utils import get_index_to_class_mapping
from load_data import generate_video_and_gradcam
from load_disease_info import disease_info
from PIL import Image


def display_name(pathology_str:str) -> str:
    return disease_info(pathology_str)


def show_logo(path):
    numpy_logo = np.asarray(Image.open(path))
    return numpy_logo


with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown("# Latent diffusion model (LDM) for synthetic X-ray generation")
        with gr.Row():
            gr.HTML("<img width='200' height='200' src='/file/logos/FAU_TechFak_H_RGB_white.png' alt='FAU TechFak logo'>")
            gr.HTML("<img width='200' height='200' src='/file/logos/LME_logo_english.png' alt='LME logo'>")
            gr.HTML("<img width='200' height='200' src='/file/logos/MDD-logo-weiss.png' alt='Medical data donors logo'>")
        with gr.Column():
            gr.Markdown("To create a synthetic image select the pathology from the set below and click **Generate**.")
            radio = gr.Radio(list(get_index_to_class_mapping().values()), label="Pathology")
        btn = gr.Button("Generate")
        with gr.Row():
            with gr.Tab(label="Diffusion model"):
                video = gr.Video(label="Diffusion steps & Generated X-ray")
            with gr.Tab(label="GradCAM Overlay"):                
                gr_image = gr.Image()
            with gr.Column():
                gr.Markdown("## Description of the diseases")
                disease_description = gr.HTML(label="Disease description")


    def submit(radio_selection):
        generated_video, generated_gradcam = generate_video_and_gradcam(radio_selection)
        return {
            video: generated_video,
            gr_image: generated_gradcam,
            disease_description: display_name(radio_selection)
        }

    btn.click(fn=submit, inputs=radio, outputs=[video, gr_image, disease_description])

demo.launch()