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

st.set_page_config(
    page_title = 'Patacotrón',
    layout = 'wide',
    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'
    }
)

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

cnn, autoencoder, svm, iforest, gan, vit = st.tabs(["CNN", "Autoencoder", "OC-SVM", 'iForest', 'GAN', 'ViT'])

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.")
        
        current_dir = os.getcwd()
        root_dir = os.path.dirname(current_dir)
        
        # Join the path to the models folder
        DIR = os.path.join(current_dir, "models")
        models = os.listdir(DIR)
        
        model_dict = dict()
        for model in models: #preprocessing of strings so the name is readable in the multiselect bar
            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())

        threshold = st.slider('¿Cuál va a ser el límite donde se considere patacón? (el valor recomendado es 75%)', 0, 100, 75)
        st.caption(f"Límite de {threshold}% seleccionado.")
        
        selected_models = []
    
        @tf.function
        def tf_predict(model_list, img): #faster, but for few formats
            y_gorrito = 0
            raw_img = tf.image.decode_image(img, channels=3)
            img = tf.image.resize(raw_img,(IMAGE_WIDTH, IMAGE_HEIGHT))
            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), raw_img.numpy()]

        def basic_predict(model_list, img): #for non-supported formats
            y_gorrito = 0
            raw_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (IMAGE_WIDTH, IMAGE_HEIGHT))
            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), raw_img]

        def preprocess(file_uploader): #makes the uploaded image readable
            img = np.frombuffer(uploaded_file.read(), np.uint8)
            img = cv2.imdecode(img, cv2.IMREAD_COLOR)
            return img
                
        # Set the image dimensions
        IMAGE_WIDTH = IMAGE_HEIGHT = 224
    
        executed = False
        
    with col_b:
        
        uploaded_file = st.file_uploader(label = '',type= ['jpg','png', 'jpeg', 'jfif', 'webp', 'heic'])
        
        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:
                img = preprocess(uploaded_file)
                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
                        try:
                            y_gorrito, raw_img = tf_predict(ultraptctrn, img)
                        except:
                            y_gorrito, raw_img = basic_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]
                        try:             
                            y_gorrito, raw_img = tf_predict(selected_models, img)
                        except:
                            y_gorrito, raw_img = basic_predict(selected_models, img)                           
        
                if round(float(y_gorrito*100)) >= threshold:
                    st.success("¡Patacón Detectado!")
                
                else:
                    st.error("No se considera que haya un patacón en la imagen")
                    
                st.caption(f'La probabilidad de que la imagen tenga un patacón es del: {round(float(y_gorrito * 100), 2)}%')
                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)')
                
                st.image(raw_img)
            
            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 vit:
    st.write('Próximamente')