import streamlit as st
import tensorflow as tf
import numpy as np
from keras.models import load_model
from tensorflow.keras.backend import clear_session
import cv2
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 = .8
    
    ultra_button = st.checkbox('Usar ensamblaje de los modelos con mayor eficacia hasta la fecha (mejores resultados)')
    ultra_flag = False
    if ultra_button:
        ultra_flag = True
    
    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
        
    ultraptctrn = ['ptctrn_v1.6', 'ptctrn_v1.8', 'ptctrn_v1.9.1', 'ptctrn_v1.12']
    
    # 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, 0)))
            clear_session()
        return y_gorrito/len(model_list)
    
    for model in model_choice:
        selected_models.append(model)
    
    # Set the image dimensions
    IMAGE_WIDTH = IMAGE_HEIGHT = 224

    uploaded_file = st.file_uploader(label = '',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic'])

with col_b:
    if st.button('¿Hay un patacón en la imagen?'):
        if uploaded_file is not None and len(selected_models) > 0 or ultra_flag:
            # Load the image and resize it to the required dimensions
            img = np.frombuffer(uploaded_file.read(), np.uint8)
            img = cv2.imdecode(img, cv2.IMREAD_COLOR)
            raw_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (IMAGE_WIDTH, IMAGE_HEIGHT))
        
            # Convert the image to RGB and preprocess it for the model
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = img / 255.
        
            # Pass the image to the model and get the prediction
            if ultra_flag:
                with st.spinner('Cargando ultra-predicción...'):
                    y_gorrito = ensemble_model(ultraptctrn, img)
            else:
                with st.spinner('Cargando predicción...'):
                    y_gorrito = ensemble_model(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)
            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.')