# ================================
# 0. PATCH pour huggingface_hub (contourne l'absence de HfFolder)
# ================================
import huggingface_hub
if not hasattr(huggingface_hub, 'HfFolder'):
    class HfFolder:
        _token = None
        @staticmethod
        def get_token():
            return HfFolder._token
        @staticmethod
        def save_token(token):
            HfFolder._token = token
    huggingface_hub.HfFolder = HfFolder

# ================================
# 1. PATCH pour contourner le bug de Gradio 4.44.0
#    (TypeError: argument of type 'bool' is not iterable)
# ================================
import gradio_client.utils

original_get_type = gradio_client.utils.get_type

def patched_get_type(schema):
    if isinstance(schema, bool):
        return "boolean"
    return original_get_type(schema)

gradio_client.utils.get_type = patched_get_type

# ================================
# 2. IMPORTS STANDARDS
# ================================
import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image

# Forcer la précision float32 (évite les warnings de mixed precision sur CPU)
tf.keras.mixed_precision.set_global_policy('float32')

# ================================
# 3. CHARGEMENT DU MODÈLE (binaire)
# ================================
MODEL_PATH = "final_model.keras"
# Charger le modèle sans compiler pour éviter l'avertissement sur l'optimiseur
model = tf.keras.models.load_model(MODEL_PATH, compile=False)

# Recompiler avec une configuration simple (nécessaire pour faire des prédictions)
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])

# ================================
# 4. CONSTANTES
# ================================
IMG_SIZE = 380  # Taille d'entrée du modèle (EfficientNetB4)

# ================================
# 5. FONCTION DE PRÉDICTION
# ================================
def preprocess_image(img):
    """
    Redimensionne et normalise l'image pour le modèle.
    L'entraînement a utilisé : cv2.resize + /255.0
    """
    img = img.resize((IMG_SIZE, IMG_SIZE))
    img_array = np.array(img, dtype=np.float32) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

def predict(img):
    """
    Retourne un dictionnaire {classe: probabilité} pour les deux classes.
    La sortie du modèle est la probabilité d'être une mauvaise herbe (weed).
    """
    processed = preprocess_image(img)
    prob_weed = float(model.predict(processed, verbose=0)[0][0])
    prob_crop = 1 - prob_weed
    return {
        "crop": prob_crop,
        "weed": prob_weed
    }

# ================================
# 6. INTERFACE GRADIO
# ================================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Chargez une image de champ (culture ou mauvaise herbe)"),
    outputs=gr.Label(num_top_classes=2, label="Prédiction"),
    title="Classification Culture / Mauvaise Herbe",
    description="Ce modèle (EfficientNetB4) distingue les cultures des mauvaises herbes. \
                  Il a été entraîné sur un dataset avec bounding boxes, pour une tâche de classification binaire."
)

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