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

# Carga los modelos previamente entrenados
autoencoder = load_model("models/autoencoder.h5", compile=False)
decoder = load_model("models/decoder.h5")

latent_dim = 128 


def add_gaussian_noise(image, noise_factor=0.2):
    noisy_image = image + noise_factor * np.random.normal(size=image.shape)
    noisy_image = np.clip(noisy_image, 0., 1.)
    return noisy_image

def denoise_and_generate(image, num_images):
    image = np.array(image) / 255.0
    image = tf.image.resize(image, (224, 224))  # Redimensionar la imagen
    noisy_image = add_gaussian_noise(np.expand_dims(image, axis=0))
    
    denoised_image = autoencoder.predict(noisy_image).squeeze()
    denoised_image = (denoised_image * 255).astype(np.uint8)

    noisy_image_display = (noisy_image.squeeze() * 255).astype(np.uint8)
    
    # Genera imágenes con el VAE
    random_latent_vectors = np.random.normal(size=(num_images, latent_dim))
    generated_images = decoder.predict(random_latent_vectors)
    
    # Prepara las imágenes para devolverlas
    outputs = [noisy_image_display, denoised_image] + [generated_images[i].squeeze() for i in range(num_images)]
    
    return outputs


# Define la interfaz
inputs = [
    gr.Image(label="Imagen de Entrada"),
    gr.Slider(1, 2, step=1, label="Número de Imágenes Generadas")
]

outputs = [
    gr.Image(label="Imagen con Ruido"),
    gr.Image(label="Imagen Reconstruida (DAE)"),
] + [
    gr.Image(label=f"Imagen Generada {i+1} (VAE)") for i in range(2)
]

# Crea la aplicación Gradio
interface = gr.Interface(
    fn=denoise_and_generate,
    inputs=inputs,
    outputs=outputs,
    title="Interfaz Interactiva para DAE y VAE",
    description="Sube una imagen para agregar ruido, reconstruir con DAE y generar imágenes nuevas con VAE."
)

# Lanza la aplicación
if __name__ == "__main__":
    interface.launch(share=True)