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'
    }
)

with st.sidebar:
    st.write("contact@patacotron.tech")

st.title("Entorno de ejecución")
cnn, autoencoder, svm, iforest, gan, docc = st.tabs(["CNN", "Autoencoder", "OC-SVM", 'iForest', 'GAN', 'DOCC'])

with cnn:

    col_a, col_b, = st.columns(2)
    
    with col_a:
        st.title("Redes neuronales convolucionales")
        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 = []
    
        @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 if model not in selected_models]
                        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.')

    with autoencoder:
        st.write('Próximamente')
    with svm:
        st.write('Próximamente')
    with iforest:
        st.write('Próximamente')
    with gan:
        st.write('Próximamente')
    with docc:
        st.write('Próximamente')