""" 
|Función para preprocesar una imagen antes de la predicción------------------------------------------
|                                                                                         LIBRERIAS |
|___________________________________________________________________________________________________| 
"""

# Machine Learning
import tensorflow as tf  # TensorFlow para el aprendizaje automático
from tensorflow.keras.models import load_model  # load_model de Keras para cargar modelos guardados
from tensorflow.keras.preprocessing.image import img_to_array  #  img_to_array de Keras para convertir imágenes en arrays
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input  # preprocess_input de Keras para preprocesar imágenes
 
# Procesamiento (imágenes y operaciones)
import cv2  # OpenCV para el procesamiento de imágenes
import numpy as np  # NumPy para manipulación de matrices y vectores
import os  # Importa el módulo os para interactuar con el sistema operativo
from revelado.revelado import calculate_ela, calculate_difference_image, apply_auto_contrast, equalize_histogram_color  #  funciones personalizadas del módulo revelado
 
#Interfaz y aplicación web
import gradio as gr  #  Biblioteca Gradio para crear la interfaz
from gradio import Image  #  Clase Image de Gradio para mostrar imágenes en la interfaz
from flask import Flask, redirect  # Flask y redirect de Flask para la aplicación web
import threading  #Threading para ejecutar tareas en hilos separados
from gradio import Interface
import matplotlib.pyplot as plt
from gradio import Image as GrImage, Label, Dropdown, Image as GrImageOut
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.preprocessing.image import load_img, img_to_array
from tensorflow.keras.applications.mobilenet import preprocess_input
#Mapa de calor
from PIL import Image
import numpy as np
from tf_keras_vis.gradcam import GradcamPlusPlus
from tf_keras_vis.utils.model_modifiers import ReplaceToLinear
from tf_keras_vis.utils import normalize

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
#_______________________________________________
# Importar la clase Flask de Flask para la aplicación web
app = Flask(__name__)

# Título y descripción de la interfaz
title = "D2AS: Deep Synthetic Art Detector"
description = "Una aplicación para realizar predicciones sobre imágenes artísticas generativas vs manuales"

# Cargar el modelo guardado
model_clase_unica = load_model("model.h5")
model_dos_clases = load_model("paper.h5")
 

""" 
|Función para preprocesar una imagen antes de la predicción------------------------------------------
|                                                                                         FUNCIONES |
|___________________________________________________________________________________________________| 
"""
 

  # Función para preprocesar una imagen antes de la predicción
def preprocess_image(image):
    """
    Tipo de función: Preprocesamiento de imagen
    Parámetros de entrada: image (objeto imagen)
    Parámetros de salida: imagen_preprocesada (arreglo numpy)
    Descripción: Convierte una imagen en un arreglo numpy, expande las dimensiones
                 y aplica preprocesamiento necesario para el modelo.
    """
    # Convertir la imagen en un arreglo numpy
    image_array = img_to_array(image)
    
    # Expandir las dimensiones del arreglo para que coincida con la entrada esperada por el modelo
    image_array = np.expand_dims(image_array, axis=0)
    
    # Preprocesar la imagen
    imagen_preprocesada = preprocess_input(image_array)
    
    return imagen_preprocesada


""" 
|Función para guardar la imagen original en la carpeta de entrada------------------------------------
|                                                                                         FUNCIONES |
|___________________________________________________________________________________________________| 
"""

def guardar_imagen_original(imagen, input_folder):
    """
    Tipo de función: Guardar imagen
    Parámetros de entrada: imagen (objeto imagen), input_folder (ruta de la carpeta de entrada)
    Parámetros de salida: Ninguno
    Descripción: Guarda la imagen original en la carpeta de entrada con la máxima calidad.
    """
    # Obtener la extensión de archivo original de la imagen
    filename = "original_image.jpg"
    
    # Guardar la imagen en la carpeta de entrada con la máxima calidad
    cv2.imwrite(os.path.join(input_folder, filename), imagen, [cv2.IMWRITE_JPEG_QUALITY, 100])
    

""" 
|Función para realizar predicciones sobre una imagenes-----------------------------------------------
|                                                                                           FUNCION |
|___________________________________________________________________________________________________| 
"""
#Función para crear el mapa de calor para una clase
def generate_heatmap(original_image, processed_image, model):
    # Convert the original image to the required format
    original_image_color = Image.fromarray(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGBA)).resize((300, 300))
    original_image_gray = original_image_color.convert("L").convert("RGBA")

    # Prepare the processed image
    processed_image_resized = cv2.resize(processed_image, (300, 300))
    processed_image_preprocessed = preprocess_image(processed_image_resized)
      # Define the scoring function for the detected class
    def score_for_class(output):
        return output[:, 0]  # Score for the detected class  
    # Get the predicted probabilities
    probabilities = model.predict(processed_image_preprocessed)[0]
    percentage_class_1 = int(probabilities[0] * 100)
    percentage_class_0 = 100 - percentage_class_1

    # Create an instance of GradcamPlusPlus
    
    gradcam = GradcamPlusPlus(model, model_modifier=ReplaceToLinear(), clone=True)

 
    # Get the Grad-CAM heatmap for the detected class
 
    cam_class = gradcam(score_for_class, processed_image_preprocessed) # Use preprocessed image
    cam_class = normalize(cam_class)

    # Reduce the heatmap dimensionality to 2D and scale to 8-bit valueçs
    cam_class = np.uint8(255 * np.squeeze(cam_class))

    # Save the heatmap and load as an RGBA image
    plt.imsave('temporal.png', cam_class)
    cam_image = Image.open('temporal.png').convert("RGBA")
   
    # Superimpose the heatmap on the original grayscale image
    superimposed_image_class = Image.blend(original_image_gray, cam_image, alpha=0.5)
    # Tamaño final deseado
    final_height, final_width = 300, 300
    # Concatenate the original image, processed image, and heatmaps
    # Redimensionar todas las imágenes a la misma altura
    # Convertir la imagen PIL a una matriz NumPy
    original_image_color_array = np.array(original_image_color)

    if original_image_color_array.shape[2] == 4:
        original_image_color_array = cv2.cvtColor(original_image_color_array, cv2.COLOR_RGBA2RGB)

 
    processed_image_array = np.array(processed_image)
 
    superimposed_image_class_array = np.array(superimposed_image_class)
    if superimposed_image_class_array.shape[2] == 4:
        superimposed_image_class_array = cv2.cvtColor(superimposed_image_class_array, cv2.COLOR_RGBA2RGB)

    # Redimensionar todas las imágenes a las mismas dimensiones
    original_image_color_resized = cv2.resize(original_image_color_array, (final_width, final_height))
    processed_image_resized = cv2.resize(processed_image_array, (final_width, final_height))
    superimposed_image_class_resized = cv2.resize(superimposed_image_class_array, (final_width, final_height))

    # Concatenar las imágenes redimensionadas
    final_image = np.hstack([original_image_color_resized, processed_image_resized, superimposed_image_class_resized])
    # Verificar la imagen procesada
    print("Processed Image Resized Shape:", processed_image_resized.shape)
    # Verificar el modelo
    print("Model Summary:", model.summary())

    def score_for_class(output):
        score = output[:, 0]  # Score for the detected class
        print("Score Shape:", score.shape)
        return score

    # Return the final concatenated image
    return final_image

def cargar_y_preprocesar_imagen(imagen_path, target_size=(300, 300)):
    # Crear una instancia de ImageDataGenerator con las mismas configuracion para el conjunto de entrenamiento
    datagen =  ImageDataGenerator(rescale=1.0/255.)  
    
    # Cargar la imagen
    img = load_img(imagen_path, target_size=target_size)

    # Convertir la imagen a un array y expandimos las dimensiones para que pueda ser procesada por el ImageDataGenerator
    img = img_to_array(img)
    img = np.expand_dims(img, axis=0)

    # Usar el ImageDataGenerator para procesar la imagen
    img = datagen.flow(img, batch_size=1)

    # El ImageDataGenerator retorna un generador, así que usamos next() para obtener la imagen procesada
    img = next(img)

    # Devolver la imagen procesada
    return img[0] 
#Generamos el mapa de calor para dos clases
def generate_heatmap_2_classes(original_image, processed_image, model):
    # Convertir la imagen original al formato requerido y redimensionar
    # original_image_color = Image.fromarray(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGBA)).resize((300, 300))
    original = Image.open("img.png").resize((300, 300)).convert("RGBA")
    original_image_color = Image.open("img.png").resize((300, 300)).convert("RGBA")
 
    # Leer la imagen procesada desde el archivo "img.png"
    processed_image = cargar_y_preprocesar_imagen("img.png")
    image = np.expand_dims(processed_image, axis=0)  # Añade una dimensión extra para el batch

    # Obtén las probabilidades de pertenencia a cada clase
    probabilities = model.predict(image)[0]
 
    # Crear una instancia de GradcamPlusPlus
    gradcam = GradcamPlusPlus(model, model_modifier=ReplaceToLinear(), clone=True)
     # Obtén las probabilidades de pertenencia a cada clase
    probabilities = model.predict(image)[0]
    percentage_human = int(probabilities[0] * 100)
    percentage_synthetic = int(probabilities[1] * 100)
  
    # Definir funciones de puntuación para ambas clases
    def score_for_human(output):
        return output[:, 0]  # Puntuación para la clase "humana"

    def score_for_synthetic(output):
        return output[:, 1]  # Puntuación para la clase "sintética"

    # Obtén los mapas de calor de Grad-CAM para cada clase
    cam_human = gradcam(score_for_human, image)
    cam_human = normalize(cam_human)
    cam_synthetic = gradcam(score_for_synthetic, image)
    cam_synthetic = normalize(cam_synthetic)

   # Reduce la dimensionalidad de los mapas de calor a 2D y escálalos a valores de 8 bits
    cam_human = np.uint8(255 * np.squeeze(cam_human))
    cam_synthetic = np.uint8(255 * np.squeeze(cam_synthetic))

    images_combined = [original_image_color]
    for cam in [cam_human, cam_synthetic]:
        # Guardar el mapa de calor
        plt.imsave('temporal.png', cam)

        # Cargar el mapa de calor y convertirlo a RGBA
        cam_image = Image.open('temporal.png').convert("RGBA")

        # Superponer el mapa de calor a la imagen original
        superimposed_image = Image.blend(original, cam_image, alpha=0.5)

        # Guardar la imagen en la lista de imágenes a combinarse
        images_combined.append(superimposed_image)

    # Concatenar las imágenes y guardar la imagen final
    final_image = Image.fromarray(np.hstack([np.array(im) for im in images_combined]))

    return final_image

def hacer_predicciones(modelo_seleccionado, imagen):
    """
    Tipo de función: Predicción
    Parámetros de entrada: imagen (objeto imagen)
    Parámetros de salida: diccionario de predicciones, imagen procesada (objeto imagen)
    Descripción: Realiza predicciones sobre la imagen utilizando el modelo cargado.
                 Devuelve un diccionario con las probabilidades de cada clase y la imagen procesada.
    """
    if modelo_seleccionado == "Modelo 1 Clase":
        model = model_clase_unica
    else:
        model = model_dos_clases

    print(modelo_seleccionado)
 
    original_image = cv2.cvtColor(imagen, cv2.COLOR_RGB2BGR)

    num = 1
    difference_image = calculate_difference_image(imagen, num)
    difference_image2 = apply_auto_contrast(difference_image)
    cv2.imwrite("img.png", difference_image2)
    processed_image = cv2.imread("img.png")
    imagen_resized = cv2.resize(processed_image, (300, 300))
    imagen_preprocesada = preprocess_image(imagen_resized)
    
    if modelo_seleccionado == "Modelo 1 Clase":
        predicciones = model.predict(imagen_preprocesada).tolist()[0]
        probabilidad_clase_1 = predicciones[0]
        probabilidad_clase_0 = 1 - probabilidad_clase_1
        final_image = generate_heatmap(original_image, processed_image, model)
    else:
        # Usar la función cargar_y_preprocesar_imagen para el preprocesamiento correcto
        #processed_image_preprocessed = cargar_y_preprocesar_imagen("img.jpg")
        gradcam = GradcamPlusPlus(model, model_modifier=ReplaceToLinear(), clone=True)
        imagen_path = "img.png"
        # Convertir la imagen PIL en una matriz NumPy
        imagen = cargar_y_preprocesar_imagen(imagen_path)
        image = np.expand_dims(imagen, axis=0)  # Añade una dimensión extra para el batch

        probabilities = model.predict(image)[0]
        print (probabilities)
        probabilidad_clase_0 = int(probabilities[0] * 100)
        probabilidad_clase_1 = int(probabilities[1] * 100)
        predicted_label = np.argmax(probabilities)
        print(predicted_label)
        
        def score_for_human(output):
            return output[:, 0]

        def score_for_synthetic(output):
            return output[:, 1]
         
        gradcam = GradcamPlusPlus(model, model_modifier=ReplaceToLinear(), clone=True)
        cams = [gradcam(score_func, image) for score_func in [score_for_human, score_for_synthetic]]
        cams = [normalize(cam) for cam in cams]
        cams = [np.uint8(255 * np.squeeze(cam)) for cam in cams]
        # Convertir la imagen PIL a una matriz NumPy
     
 
        final_image = generate_heatmap_2_classes(original_image, image, model)

    return {"Arte Manual": probabilidad_clase_0/100, "Arte Sintetico": probabilidad_clase_1/100}, final_image
"""
|Función para generar gradio en un hilo separado----------------------------------------------------
|                                                                                         FUNCIONES |
|___________________________________________________________________________________________________| 
"""

def run_gradio():
    """
    Tipo de función: Ejecución Gradio
    Parámetros de entrada: Ninguno
    Parámetros de salida: Ninguno
    Descripción: Ejecuta la interfaz Gradio en un hilo separado.
    """
    output = iface.process([imagen])  # Ejecutar el procesamiento de la imagen
    predicciones = output[0]
    final_image = output[1]

    # Guardar la imagen en una ubicación accesible
    cv2.imwrite("output_img.jpg", final_image)
#___________________________________________________________________________________________________

# Crear una entrada de imagen en Gradio
#im = GrImage(shape=(None, None), image_mode='RGB', invert_colors=False, source="upload")
im = GrImage(image_mode='RGB', invert_colors=False, source="upload")
model_selector = Dropdown(choices=["Modelo 1 Clase", "Modelo 2 Clases"], label="Selecciona un modelo")

# Crear una interfaz de Gradio con una entrada de imagen
iface = Interface(
    fn=hacer_predicciones,
    inputs=[model_selector, im],
    outputs=[Label(), GrImageOut(type="numpy")],
    title=title,
    description=description,
    allow_flagging="never"
)

# Lanzar la aplicación Gradio en un hilo separado
threading.Thread(target=run_gradio).start()

# Lanzar la interfaz Gradio
iface.launch(share=False)

# Definir la función de ruta en Flask
@app.route('/')
def index():
    return redirect(iface.launch(share=True))

# Ejecutar la aplicación Flask
if __name__ == '__main__':
    app.run()