import numpy as np
import pandas as pd
#import cv2
import os
import pydicom

import streamlit as st
from streamlit_cropper import st_cropper

#import tensorflow as tf
#import tensorflow_io as tfio 

#from tensorflow import keras

import matplotlib.pyplot as plt
from PIL import Image

from fastai.vision.all import *

import pathlib
import platform

plt = platform.system()
if plt == 'Linux':
    pathlib.WindowsPath = pathlib.PosixPath

# Get the directory of the current script
script_dir = os.path.dirname(os.path.realpath(__file__))

st.set_page_config(page_title="Aplicación cancer de mama - IXCHEN")

os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'

class Modelo:

    def __init__(self, model_name):
        self.model_name = model_name
        self.modelo = self.cargar_modelo(self.model_name)
    
    def cargar_modelo(self, modelo):
        if modelo == "alexnet":
            return load_learner(os.path.join(script_dir, 'clasificacion_alexnet.pkl'))
        elif modelo == 'restnet50':
            return load_learner(os.path.join(script_dir, 'clasificacion.pkl'))
        
    def predecir(self, imagen):
        _, _, probs = self.modelo.predict(imagen)
        return probs
    
modelo_alexnet = Modelo('alexnet')
modelo_resnet50 = Modelo('restnet50')

realtime_update = st.sidebar.checkbox(label="Actualizar en tiempo real", value=True)
box_color = st.sidebar.color_picker(label="Color caja", value='#0000FF')
aspect_choice = st.sidebar.radio(label="Relación aspecto", options=["1:1", "16:9", "4:3", "2:3", "Free"])
aspect_dict = {
    "1:1": (1, 1),
    "16:9": (16, 9),
    "4:3": (4, 3),
    "2:3": (2, 3),
    "Free": None
}

aspect_ratio = aspect_dict[aspect_choice]
class_names = ['Benigno', 'Maligno', 'Normal']

def cargar_dicom_file(image_bytes, tipo):
    if tipo != 'storage':
        dcm = pydicom.dcmread(image_bytes)
    else:
        dcm = pydicom.dcmread(image_bytes)
    numpy_lossy_image = dcm.pixel_array
    return np.squeeze(numpy_lossy_image)

def mostrar_resultados_prediccion(crop_image):
    resultado_alexnet = modelo_alexnet.predecir(crop_image)
    resultado_resnet50 = modelo_resnet50.predecir(crop_image)
    return resultado_alexnet, resultado_resnet50

def obtener_valor_tag(image_byte, tag: int) -> str:
    raw_tag_value = tfio.image.decode_dicom_data(image_byte, tag).numpy()
    return raw_tag_value.decode('utf-8')


st.write("# Modelo de Machine Learning para el diagnostico del cancer de mama")
uploaded_file = st.file_uploader("Subir archivo DICOM", type=[".dcm"])
if uploaded_file:
    dicom_file = cargar_dicom_file(uploaded_file, tipo='')
    img = Image.fromarray(dicom_file)
    rect = st_cropper(img,realtime_update=realtime_update, box_color=box_color,
                    aspect_ratio=aspect_ratio, return_type='box')
    realizar_prediccion = st.button("Realizar predicción")
    st.session_state['uploaded_file'] = uploaded_file.getvalue()
    if realizar_prediccion:
            st.empty()
            crop = dicom_file[
                rect['top']:rect['height'] + rect['top'], 
                rect['left']:rect['width'] + rect['left']]
            resultado_alexnet, resultado_restnet50 = mostrar_resultados_prediccion( crop)
            resultado_alexnet = np.squeeze(resultado_alexnet) * 100
            resultado_resnet50 = np.squeeze(resultado_restnet50) * 100
            datos_alexnet = {
                "Probabilidad": resultado_alexnet,
                "Clase": class_names
            }
            df_alexnet = pd.DataFrame(data=datos_alexnet).set_index('Clase')
            datos_resnet50 = {
                "Probabilidad": resultado_resnet50,
                "Clase": class_names
            }
            df_resnet50 = pd.DataFrame(data=datos_resnet50).set_index('Clase')
            st.write("### Probabilidades según tipo de tumor - AlexNet")
            st.bar_chart(df_alexnet)
            st.write("### Probabilidades según tipo de tumor - ResNet50")
            st.bar_chart(df_resnet50)