import utils
import gradio as gr
import tensorflow as tf
import matplotlib.pyplot as plt
from ttictoc import tic,toc

# '''--------------------------- Preprocesamiento ----------------------------'''
# tic()
# 3D U-Net
path_3d_unet = 'F:/Desktop/Universidad/Semestres/NovenoSemestre/Proyecto_de_Grado/Codigo/3D_U-Net/outputs/checkpoints/model.49-0.97.h5'

with tf.device("cpu:0"):
    model_unet = utils.import_3d_unet(path_3d_unet)

# # Cargar imagen
# img = utils.load_img('F:/Downloads/ADNI_002_S_0295_MR_MP-RAGE__br_raw_20070525135721811_1_S32678_I55275.nii')

# # Extraer cerebro
# with tf.device("cpu:0"):
#     brain = utils.brain_stripping(img, model_unet)
# print(toc())

# '''---------------------------- Procesamiento ------------------------------'''
# # Med net
# weight_path = 'F:/Desktop/Universidad/Semestres/NovenoSemestre/Proyecto_de_Grado/Codigo/Procesamiento/mednet_weights/pretrain/resnet_50_23dataset.pth'
# device_ids = [0]
# mednet = utils.create_mednet(weight_path, device_ids)

# # Extraer características
# features = utils.get_features(brain, mednet)

def load_img(file):
    sitk, array = utils.load_img(file.name)     
    return sitk, array

def show_img(img, mri_slice):
    fig = plt.figure()
    plt.imshow(img[:,:,mri_slice], cmap='gray')
    
    return fig, gr.update(visible=True)

def show_brain(brain, brain_slice):
    fig = plt.figure()
    plt.imshow(brain[brain_slice,:,:], cmap='gray')
    
    return fig, gr.update(visible=True)

def process_img(img, brain_slice):
    with tf.device("cpu:0"):
        brain = utils.brain_stripping(img, model_unet)
        
        fig, update = show_brain(brain, brain_slice)
        
    return brain, fig, update

def clear():
    return gr.File.update(value=None), gr.Plot.update(value=None), gr.update(visible=False)

# gr.Textbox.update(placeholder='Ingrese nombre del paciente'), gr.Number.update(value=0), 

# demo = gr.Interface(fn=load_img,
#                     inputs=gr.File(file_count="single", file_type=[".nii"]),
#                     outputs=gr.Plot()
#                     # outputs='text'
#                     )

with gr.Blocks() as demo:
    with gr.Row():
        gr.Markdown("""
                    # SIMCI
                    Interfaz de SIMCI
                    """)
    
    # Inputs
    with gr.Row():
        with gr.Column(scale=1):  
            # Objeto para subir archivo nifti
            input_name = gr.Textbox(placeholder='Ingrese nombre del paciente', label='Nombre')
            input_age = gr.Number(label='Edad')
            
            
            input_file = gr.File(file_count="single", file_type=[".nii"], label="Archivo Imagen MRI")
                
            with gr.Row():
                # Botón para cargar imagen
                load_img_button = gr.Button(value="Load")
                
                # Botón para borrar
                clear_button = gr.Button(value="Clear")
            
            # Botón para procesar imagen
            process_button = gr.Button(value="Procesar")
    
        # Outputs 
        with gr.Column(scale=1):
            # Plot para imágen original
            plot_img_original = gr.Plot(label="Imagen MRI original")
            
            # Slider para imágen original
            mri_slider = gr.Slider(minimum=0,
                                    maximum=166,
                                    value=100,
                                    step=1,
                                    label="MRI Slice",
                                    visible=False) 
            
            # Plot para imágen procesada
            plot_brain = gr.Plot(label="Imagen MRI procesada")
            
            # Slider para imágen procesada
            brain_slider = gr.Slider(minimum=0,
                                     maximum=192,
                                     value=100,
                                     step=1,
                                     label="MRI Slice",
                                     visible=False)
    
    # componentes = 
    
    # Variables
    original_input_sitk = gr.State()
    original_input_img = gr.State()
    brain_img = gr.State()
    
    # Cambios
    # Cargar imagen nueva
    input_file.change(load_img, 
                      input_file, 
                      [original_input_sitk, original_input_img])
    
    # Mostrar imagen nueva
    load_img_button.click(show_img, 
                          [original_input_img, mri_slider], 
                          [plot_img_original, mri_slider])
    
    # Limpiar campos
    clear_button.click(fn=clear, 
                       outputs=[input_file, plot_img_original, mri_slider])
    
    # Actualizar imagen original
    mri_slider.change(show_img, 
                      [original_input_img, mri_slider], 
                      [plot_img_original,mri_slider])
    
    # Procesar imagen
    process_button.click(fn=process_img, 
                         inputs=[original_input_sitk, brain_slider], 
                         outputs=[brain_img,plot_brain,brain_slider])
    
    # Actualizar imagen procesada
    brain_slider.change(show_brain,
                        [brain_img, brain_slider], 
                        [plot_brain,brain_slider])


if __name__ == "__main__":
    demo.launch()

# # Visualización resultados 
# mri_slice = 100

# # Plot Comparación máscaras
# fig, axs = plt.subplots(1,2)
# fig.subplots_adjust(bottom=0.15)
# fig.suptitle('Comparación Máscaras Obtenidas')

# axs[0].set_title('MRI original')
# axs[0].imshow(img[mri_slice,:,:],cmap='gray')

# axs[1].set_title('Cerebro extraido con 3D U-Net')
# axs[1].imshow(brain[mri_slice,:,:],cmap='gray')


# # Slider para cambiar slice
# ax_slider = plt.axes([0.15, 0.05, 0.75, 0.03])
# mri_slice_slider = Slider(ax_slider, 'Slice', 0, 192, 100, valstep=1)

# def update(val):
#     mri_slice = mri_slice_slider.val
    
#     axs[0].imshow(img[:,:,mri_slice],cmap='gray')
#     axs[1].imshow(brain[mri_slice,:,:],cmap='gray')
    
# # Actualizar plot comparación máscaras
# mri_slice_slider.on_changed(update)