import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.optimizers import Adam

# Configurar TensorFlow en modo compatibilidad (para evitar algunos warnings)
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)

# Cargar el modelo entrenado
model = load_model('colchones/modelo_colchones.h5')  # Asegúrate de tener el archivo del modelo correcto y la ruta adecuada

# Compilación del modelo
optimizer = Adam(learning_rate=0.001)
model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Forzar la construcción de métricas
model.built = True

# Función para realizar la predicción
def predict(firmeza, tamaño, posición, presupuesto):
    # Convertir las opciones seleccionadas a números
    firmeza_dict = {'Suave': 0, 'Media': 1, 'Firme': 2}
    tamaño_dict = {'Individual': 0, 'Matrimonial': 1, 'Queen': 2, 'King': 3}
    posición_dict = {'Lateral': 0, 'Espalda': 1, 'Estómago': 2}
    presupuesto_dict = {'Bajo': 0, 'Medio': 1, 'Alto': 2}

    colchon_data = [
        firmeza_dict.get(firmeza, 0),
        tamaño_dict.get(tamaño, 0),
        posición_dict.get(posición, 0),
        presupuesto_dict.get(presupuesto, 0),
        firmeza_dict.get(firmeza, 0),
        presupuesto_dict.get(presupuesto, 0)
    ]

    # Realizar la predicción
    X_pred = np.array([colchon_data])  # Convertir a un arreglo numpy
    prediction = model.predict(X_pred)
    predicted_class = np.argmax(prediction)

    tipos_colchones = ['Muelles Bicónicos', 'Látex', 'Muelles Ensacados', 'Espumación HR', 'Híbrido']
    imagenes_colchones = {
        0: 'colchones/static/images/muelles_biconicos.jpeg',
        1: 'colchones/static/images/latex.jpg',
        2: 'colchones/static/images/muelle_ensacado.jpg',
        3: 'colchones/static/images/espumacion_hr.jpg',
        4: 'colchones/static/images/hibrido.jpg'
    }

    return tipos_colchones[predicted_class], imagenes_colchones.get(predicted_class)

# Definir las opciones para cada característica
firmeza_options = ['Suave', 'Media', 'Firme']
tamaño_options = ['Individual', 'Matrimonial', 'Queen', 'King']
posición_options = ['Lateral', 'Espalda', 'Estómago']
presupuesto_options = ['Bajo', 'Medio', 'Alto']

# Crear la interfaz de Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Predicción de Tipo de Colchón")
    firmeza = gr.Dropdown(choices=firmeza_options, label='Firmeza')
    tamaño = gr.Dropdown(choices=tamaño_options, label='Tamaño')
    posición = gr.Dropdown(choices=posición_options, label='Posición de dormir')
    presupuesto = gr.Dropdown(choices=presupuesto_options, label='Presupuesto')
    recomendacion = gr.Textbox(label='Tipo de Colchón Recomendado')
    imagen = gr.Image(label='Imagen del Colchón', type='filepath', visible=False)

    def make_prediction(firmeza, tamaño, posición, presupuesto):
        tipo_colchon, img_colchon = predict(firmeza, tamaño, posición, presupuesto)
        return tipo_colchon, gr.update(value=img_colchon, visible=True)

    def clear_image(*args):
        return '', gr.update(value=None, visible=False)

    firmeza.change(clear_image, inputs=[firmeza, tamaño, posición, presupuesto], outputs=[recomendacion, imagen])
    tamaño.change(clear_image, inputs=[firmeza, tamaño, posición, presupuesto], outputs=[recomendacion, imagen])
    posición.change(clear_image, inputs=[firmeza, tamaño, posición, presupuesto], outputs=[recomendacion, imagen])
    presupuesto.change(clear_image, inputs=[firmeza, tamaño, posición, presupuesto], outputs=[recomendacion, imagen])

    btn = gr.Button("Recomendar")
    btn.click(fn=make_prediction, inputs=[firmeza, tamaño, posición, presupuesto], outputs=[recomendacion, imagen])

# Ejecutar la aplicación
demo.launch()