import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
from huggingface_hub import hf_hub_download
import time
import os

# Registrar las funciones personalizadas
from tensorflow.keras.saving import register_keras_serializable

@register_keras_serializable()
def fourier_transform(x):
    fourier = tf.signal.fft2d(tf.cast(x, tf.complex64))
    fourier = tf.complex(tf.math.real(fourier), tf.math.imag(fourier))
    fourier = tf.abs(fourier)
    return tf.concat([tf.math.real(fourier), tf.math.imag(fourier)], axis=-1)

@register_keras_serializable()
def inverse_fourier_transform(x):
    real_part, imag_part = tf.split(x, num_or_size_splits=2, axis=-1)
    complex_fourier = tf.complex(real_part, imag_part)
    return tf.abs(tf.signal.ifft2d(complex_fourier))

# Configuración de GPU para TensorFlow
physical_devices = tf.config.list_physical_devices('GPU')
if physical_devices:
    print("GPU disponible. Configurando...")
    try:
        # Permitir crecimiento de memoria según sea necesario
        for gpu in physical_devices:
            tf.config.experimental.set_memory_growth(gpu, True)
        print("Configuración de GPU completada")
    except Exception as e:
        print(f"Error en configuración de GPU: {e}")
else:
    print("No se detectó GPU. El procesamiento será más lento.")

# Descargar modelo desde Hugging Face (con caché)
cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
model_path = hf_hub_download(
    repo_id="Bmo411/DenoisingAutoencoder", 
    filename="autoencoder_complete_model_Fourier.keras",
    cache_dir=cache_dir
)
print(f"Modelo cargado desde: {model_path}")

# Cargar el modelo una sola vez
model = tf.keras.models.load_model(model_path)

# Crear versión optimizada para inferencia
@tf.function(jit_compile=True)
def predict_optimized(input_tensor):
    return model(input_tensor, training=False)

# Funciones de preprocesamiento optimizadas
def degrade_image(image, downscale_factor=4):
    """Reduce la calidad de la imagen reduciendo su tamaño y volviéndola a escalar."""
    h, w = image.shape[:2]
    
    # Verificar tamaño mínimo
    if h < downscale_factor*4 or w < downscale_factor*4:
        return image  # Evitar downscaling excesivo en imágenes pequeñas
    
    # Reducir tamaño (forzando pérdida de calidad)
    small_img = cv2.resize(image, (w // downscale_factor, h // downscale_factor), 
                          interpolation=cv2.INTER_AREA)

    # Volver a escalarla al tamaño original
    degraded_img = cv2.resize(small_img, (w, h), interpolation=cv2.INTER_NEAREST)

    return degraded_img

def preprocess_image(image, std_dev=0.1, downscale_factor=4, target_size=(256, 256)):
    """Recibe una imagen en numpy array, la degrada en calidad, le agrega ruido y la normaliza."""
    
    # Verificar si la imagen es a color o en escala de grises
    if len(image.shape) == 2:
        # Convertir a RGB si es escala de grises
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 4:
        # Convertir de RGBA a RGB si tiene canal alpha
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
        
    # Reducir calidad
    degraded_img = degrade_image(image, downscale_factor)

    # Redimensionar a tamaño esperado por el modelo
    resized_img = cv2.resize(degraded_img, target_size, interpolation=cv2.INTER_AREA)
    
    # Normalizar
    resized_img = resized_img.astype(np.float32) / 255.0

    # Agregar ruido gaussiano
    noise = np.random.normal(0, std_dev, resized_img.shape)
    noisy_img = resized_img + noise
    noisy_img = np.clip(noisy_img, 0, 1)

    return noisy_img

# Variable para medir el tiempo total de la primera ejecución
first_run = True

def Denoiser(imagen):
    """Aplica el modelo autoencoder para eliminar ruido de la imagen."""
    global first_run
    
    # Verificar que la imagen no sea None
    if imagen is None:
        return None, None
    
    # Convertir imagen de entrada a array NumPy
    imagen = np.array(imagen)
    
    # Verificar dimensiones de la imagen
    if len(imagen.shape) < 2:
        return None, None
    
    try:
        # Preprocesar imagen (degradarla y agregar ruido)
        noisy_image = preprocess_image(imagen)
        
        # Expandir dimensiones para el formato del modelo
        noisy_image_input = np.expand_dims(noisy_image, axis=0)
        
        # Medir el tiempo de la predicción
        start_time = time.time()
        
        # Predecir con el autoencoder utilizando la función optimizada
        reconstructed = predict_optimized(noisy_image_input).numpy()[0]
        
        prediction_time = time.time() - start_time
        
        if first_run:
            print(f"Primera ejecución: {prediction_time:.2f} segundos")
            first_run = False
        else:
            print(f"Tiempo de predicción: {prediction_time:.2f} segundos")
        
        # Asegurarse de que las imágenes estén en el rango [0, 255]
        noisy_image = np.uint8(noisy_image * 255)
        reconstructed = np.uint8(reconstructed * 255)
        
        return noisy_image, reconstructed
        
    except Exception as e:
        print(f"Error en el procesamiento: {e}")
        return None, None

# Ejemplo para precalentamiento (warm-up) del modelo
try:
    dummy_input = np.zeros((1, 256, 256, 3), dtype=np.float32)
    _ = predict_optimized(dummy_input)
    print("Modelo precalentado con éxito")
except Exception as e:
    print(f"Error en precalentamiento: {e}")

# Crear interfaz en Gradio
demo = gr.Interface(
    fn=Denoiser, 
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Image(type="numpy", label="Imagen con Ruido"), 
        gr.Image(type="numpy", label="Imagen Restaurada")
    ],
    title="Autoencoder para Denoising",
    description="Este modelo de autoencoder reduce el ruido en imágenes. Sube una imagen y el modelo generará una versión restaurada.",
    examples=[
        "https://raw.githubusercontent.com/gradio-app/gradio/main/demo/english_htr/images/Create%20a%20free%20Gradio%20account%20to%20access%20our%20most%20powerful%20features.jpeg"
    ],
    cache_examples=True
)

# Lanzar la aplicación en Gradio
if __name__ == "__main__":
    demo.launch(share=False)