# ================================
# 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
# ================================
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

# ================================
# 3. CHARGEMENT DU MODÈLE (9 classes)
# ================================
MODEL_PATH = "final_model.keras"
model = tf.keras.models.load_model(MODEL_PATH, compile=False)
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# ================================
# 4. NOMS DES CLASSES (issus de l'entraînement)
# ================================
class_names = [
    "Chinee apple",
    "Lantana",
    "Negative",
    "Parkinsonia",
    "Parthenium",
    "Prickly acacia",
    "Rubber vine",
    "Siam weed",
    "Snake weed"
]

# ================================
# 5. PARAMÈTRES
# ================================
IMG_SIZE = 380  # taille utilisée lors de l'entraînement

# ================================
# 6. PRÉTRAITEMENT (identique à l'entraînement)
# ================================
def preprocess_image(img):
    img = img.resize((IMG_SIZE, IMG_SIZE))
    img_array = np.array(img)
    # Appliquer le même preprocess_input que EfficientNetB4
    img_array = tf.keras.applications.efficientnet.preprocess_input(img_array)
    img_array = np.expand_dims(img_array, axis=0)
    return img_array

def predict(img):
    processed = preprocess_image(img)
    preds = model.predict(processed, verbose=0)[0]
    results = {class_names[i]: float(preds[i]) for i in range(len(class_names))}
    return results

# ================================
# 7. INTERFACE GRADIO
# ================================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Chargez une image de mauvaise herbe"),
    outputs=gr.Label(num_top_classes=3, label="Espèce prédite (top 3)"),
    title="🌿 Classification des mauvaises herbes",
    description="Modèle EfficientNetB4 entraîné sur 9 espèces de mauvaises herbes (Chinee apple, Lantana, Negative, Parkinsonia, Parthenium, Prickly acacia, Rubber vine, Siam weed, Snake weed)."
)

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