import os
import tensorflow as tf
import gradio as gr
import numpy as np
import cv2
from PIL import Image
from tensorflow.keras.applications.efficientnet import preprocess_input as efficientnet_preprocess

# Désactiver l'utilisation du GPU pour éviter l'erreur CUDA
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

model_dir = os.path.join('models')
custom_objects = {"mae": tf.keras.losses.MeanAbsoluteError(), "mse": tf.keras.losses.MeanSquaredError()}

try:
    model_genre = tf.keras.models.load_model(os.path.join(model_dir, "model_genre.keras"), custom_objects=custom_objects)
except Exception as e:
    print(f"Erreur lors du chargement du modèle d'âge: {e}")
    model_age = None

try:
    model_age = tf.keras.models.load_model(os.path.join(model_dir, "modelAge.keras"), custom_objects=custom_objects)
except Exception as e:
    print(f"Erreur lors du chargement du modèle d'âge: {e}")
    model_age = None

try:
    model_gender_age = tf.keras.models.load_model(os.path.join(model_dir, "age_gender_model.h5"), custom_objects=custom_objects)
except Exception as e:
    print(f"Erreur lors du chargement du modèle genre+âge: {e}")
    model_gender_age = None

try:
    model_pre_entraine = tf.keras.models.load_model(os.path.join(model_dir, "model_efficientnetB2.h5"), custom_objects=custom_objects)
except Exception as e:
    print(f"Erreur lors du chargement du modèle pré-entraîné: {e}")
    model_pre_entraine = None

img_height, img_width = 224, 224

def normalize_age(age):
    return tf.cast(age, tf.float32) / 116.0

def denormalize_age(age):
    return age * 116.0

def preprocess_image(image,model_choice):
    if image is None:
        print("Erreur: image non fournie")
        return None
    if model_choice == "Pré-entraîné":
        img = np.array(image)
        if img.shape[-1] == 4:
            img = img[:, :, :3]
        img = cv2.resize(img, (img_width, img_height))
        img_array = np.expand_dims(img, axis=0)
        return efficientnet_preprocess(img_array)

    elif model_choice == "Genre":
        image = image.resize((100, 100))
        image = np.array(image) / 255.0
    else:
        image = image.resize((128, 128))
        image = np.array(image) / 255.0
    
    image = np.expand_dims(image, axis=0).astype(np.float32)
    return image

def predict(image, model_choice):
    if image is None:
        return "Erreur : aucune image fournie."
    
    preprocessed_image = preprocess_image(image, model_choice)
    if preprocessed_image is None:
        return "Erreur : problème lors du prétraitement de l'image."
        
    if model_choice == "Âge":
        result = model_age.predict(preprocessed_image) 
        return f"Prédiction Âge : {int(denormalize_age(result[0][0]))} ans"
    
        
    elif model_choice == "Genre + Âge":
        result = model_gender_age.predict(preprocessed_image)
        age = int(denormalize_age(result[0][0].item())) 
        gender_value = result[1][0][0].item()
        gender = "Homme" if gender_value <= 0.5 else "Femme"
        return f"Prédiction Genre : {gender}, Âge : {age} ans"

            
    elif model_choice == "Genre":
        result = model_genre.predict(preprocessed_image)
        print(f"Résultat de la prédiction Genre : {result}")  # Ajout de la ligne pour afficher les résultats
        gender_value = result[0][0]
        gender = "Homme" if gender_value >= 0.5 else "Femme"
        return f"Prédiction Genre : {gender}"

        
    elif model_choice == "Pré-entraîné":
        if model_pre_entraine is None:
            return "Erreur : modèle non chargé."
        
        predictions = model_pre_entraine.predict(preprocessed_image)
        
        if isinstance(predictions, list) and len(predictions) == 2:
            pred_age = predictions[0][0][0] * 116
            pred_gender_prob = predictions[1][0][0]
            gender = "Féminin" if pred_gender_prob >= 0.5 else "Masculin"
            return f"Prédiction Âge: {pred_age:.0f} ans, Prédiction Genre: {gender}"
        else:
            return "Erreur : Format de sortie inattendu. Vérifie le modèle."
    
    return "Modèle non reconnu."


with gr.Blocks(theme='Nymbo/Nymbo_Theme') as interface:
    gr.Markdown("## 🎭 Prédiction d'Âge et de Genre")
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📷 Upload ton image")
            img_input = gr.Image(type="pil", label="Image de visage")

            gr.Markdown("### 🧠 Sélectionne un modèle")
            model_choice = gr.Radio(
                ["Genre", "Âge", "Genre + Âge", "Pré-entraîné"],
                label="Modèle",
                value="Pré-entraîné"
            )

        with gr.Column(scale=1):
            gr.Markdown("### 📊 Résultat de la prédiction")
            output_text = gr.Textbox(label="Résultat", interactive=False)
            img_preview = gr.Image(label="Prévisualisation", interactive=False, visible=False)

    # Mise à jour des résultats et de la prévisualisation
    def process_and_display(image, model_choice):
        result = predict(image, model_choice)
        return result, image

    img_input.change(
        fn=process_and_display,
        inputs=[img_input, model_choice],
        outputs=[output_text, img_preview]
    )

# Lancer l'application
if __name__ == "__main__":
    interface.launch()