import gradio as gr
import numpy as np
import tensorflow
from PIL import Image

vae = tensorflow.keras.models.load_model("dae.h5")
dae = tensorflow.keras.models.load_model("dae.h5")

def preprocess_image(image):
    """Redimensiona y normaliza la imagen."""
    # Convert to PIL Image if it's not already
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # Resize image to 128x128
    image = image.resize((128, 128))
    
    # Convert to numpy array, normalize and add batch dimension
    image_array = np.array(image).astype("float32") / 255.0
    image_array = np.expand_dims(image_array, axis=0)
    
    return image_array

def reconstruct_image(image):
    """Reconstruye la imagen con el modelo seleccionado."""
    image = preprocess_image(image)

    reconstructed = dae.predict(image)[0]


    return (reconstructed * 255).astype("uint8")  # Convertir a imagen

def generate_image(z_dim_values):
    """Genera una imagen a partir de vectores latentes específicos."""

    z = np.array([z_dim_values]).astype('float32')
    decoder = vae.layers[-1]
    
    # Generar la imagen
    generated = decoder.predict(z)[0]
    
    return (generated * 255).astype("uint8")

# Interfaz con Gradio usando tabs
with gr.Blocks(title="Demo de VAE y DAE") as demo:
    gr.Markdown("# Proyecto de VAE y DAE")
    
    with gr.Tab("Reconstrucción de Imágenes"):
        gr.Markdown("## Reconstruye una imagen usando DAE")
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(label="Imagen Original")
                reconstruct_btn = gr.Button("Reconstruir")
            with gr.Column():
                output_image = gr.Image(label="Imagen Reconstruida")
        
        reconstruct_btn.click(
            fn=reconstruct_image,
            inputs=[input_image],
            outputs=output_image
        )
    
    with gr.Tab("Generación de Imágenes (VAE)"):
        gr.Markdown("## Genera nuevas imágenes manipulando el espacio latente")
        with gr.Row():
            with gr.Column():
                # Crear 2 controles deslizantes para las dimensiones latentes
                # Ajusta el número según la dimensión de tu espacio latente
                sliders = []
                for i in range(2):  # Asumiendo un espacio latente de dimensión 2
                    slider = gr.Slider(-5.0, 5.0, value=0.0, step=0.1, 
                                      label=f"Dimensión Latente {i+1}")
                    sliders.append(slider)
                generate_btn = gr.Button("Generar")
            with gr.Column():
                generated_image = gr.Image(label="Imagen Generada")
        
        generate_btn.click(
            fn=generate_image,
            inputs=sliders,
            outputs=generated_image
        )

if __name__ == "__main__":
    demo.launch(share=True)