import streamlit as st
import tensorflow as tf
from tensorflow.keras.models import load_model
import os

st.set_page_config(
    page_title = 'Patacotrón',
    initial_sidebar_state = 'collapsed',
    menu_items = {
        "About" : 'Proyecto ideado para la investigación de "Clasificación de imágenes de una sola clase con algortimos de Inteligencia Artificial".',
        "Report a Bug" : 'https://docs.google.com/forms/d/e/1FAIpQLScH0ZxAV8aSqs7TPYi86u0nkxvQG3iuHCStWNB-BoQnSW2V0g/viewform?usp=sf_link'
    }
)

col_a, col_b, = st.columns(2)

with col_a:
    st.title("Entorno de ejecución")
    st.caption("Los modelos no están en orden de eficacia, sino en orden de creación.")
    
    # Get the absolute path to the current directory
    current_dir = os.path.abspath(os.path.dirname(__file__))
    
    # Get the absolute path to the parent directory of the current directory
    root_dir = os.path.abspath(os.path.join(current_dir, os.pardir))
    
    # Join the path to the models folder
    DIR = os.path.join(root_dir, "models")
    
    threshold = .75
    
    models = os.listdir(DIR)
    
    model_dict = dict()
    for model in models:
        model_name = model.split(DIR)
        model_name = str(model.split('.h5')[0])
        model_dir = os.path.join(DIR, model)
        model_dict[model_name] = model_dir

    ultraversions = ['ptctrn_v1.4', 'ptctrn_v1.5', 'ptctrn_v1.6', 'ptctrn_v1.12']
    ultra_button = st.checkbox('Ultra-Patacotrón (mejores resultados)')
    ultra_flag = False
    if ultra_button:
        ultra_flag = True
    
    # Create a dropdown menu to select the model
    model_choice = st.multiselect("Seleccione uno o varios modelos de clasificación", model_dict.keys())

    selected_models = []
    
    def ensemble_model(model_list, img):
        y_gorrito = np.zeros((1, 1))
        for model in model_list:
            instance_model = load_model(model_dict[model])
            y_gorrito += float(instance_model.predict(np.expand_dims(img/255., 0)))
            #clear_session()
        return y_gorrito/len(model_list)

    @tf.function
    def predict(model_list, img):
        y_gorrito = 0
        for model in model_list:
            y_gorrito += tf.cast(model(tf.expand_dims(img/255., 0)), dtype=tf.float32)
        return y_gorrito / len(model_list)
    
    # Set the image dimensions
    IMAGE_WIDTH = IMAGE_HEIGHT = 224

    uploaded_file = st.file_uploader(label = '',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic'])
    executed = False
    
with col_b:
    if st.button('¿Hay un patacón en la imagen?'):
        if len(selected_models) > 0 and ultra_flag:
            st.write('Debe elegir un solo método: Ultra-Patacotrón o selección múltiple.')
        elif uploaded_file is not None:

            raw_img = tf.image.decode_image(uploaded_file.read(), channels=3)
            img = tf.image.resize(raw_img,(IMAGE_WIDTH, IMAGE_HEIGHT))
          
            # Pass the image to the model and get the prediction
            if ultra_flag:
                with st.spinner('Cargando ultra-predicción...'):
                    if not executed: 
                        ultraptctrn = [load_model(model_dict[model]) for model in ultraversions]
                        executed = True
                    y_gorrito = predict(ultraptctrn, img)
            else:
                with st.spinner('Cargando predicción...'):
                    selected_models = [load_model(model_dict[model]) for model in model_choice]
                    y_gorrito = predict(selected_models, img)
    
            if y_gorrito >= threshold:
                st.success("¡Patacón Detectado!")
            else:
                st.error("No se encontró rastro de patacón.")
            st.caption(f'La probabilidad de que la imagen tenga un patacón es del: {round(float(y_gorrito), 2)*100}%')
            st.image(raw_img.numpy())
            st.caption('Si los resultados no fueron los esperados, por favor, [haz click aquí](https://docs.google.com/forms/d/e/1FAIpQLScH0ZxAV8aSqs7TPYi86u0nkxvQG3iuHCStWNB-BoQnSW2V0g/viewform?usp=sf_link)')
        else:
            st.write('Revisa haber seleccionado los modelos y la imagen correctamente.')