# ================================
# 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 (72 classes)
# ================================
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)
# On utilise une perte factice car on ne fera pas d'entraînement
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# ================================
# 4. NOMS DES CLASSES (issus de l'entraînement)
# ================================
class_names = [
    "Apple___alternaria_leaf_spot",
    "Apple___black_rot",
    "Apple___brown_spot",
    "Apple___gray_spot",
    "Apple___healthy",
    "Apple___rust",
    "Apple___scab",
    "Bell_pepper___bacterial_spot",
    "Bell_pepper___healthy",
    "Blueberry___healthy",
    "Cassava___bacterial_blight",
    "Cassava___brown_streak_disease",
    "Cassava___green_mottle",
    "Cassava___healthy",
    "Cassava___mosaic_disease",
    "Cherry___healthy",
    "Cherry___powdery_mildew",
    "Coffee___healthy",
    "Coffee___red_spider_mite",
    "Coffee___rust",
    "Corn___common_rust",
    "Corn___gray_leaf_spot",
    "Corn___healthy",
    "Corn___northern_leaf_blight",
    "Grape___Leaf_blight",
    "Grape___black_measles",
    "Grape___black_rot",
    "Grape___healthy",
    "Orange___citrus_greening",
    "Peach___bacterial_spot",
    "Peach___healthy",
    "Potato___bacterial_wilt",
    "Potato___early_blight",
    "Potato___healthy",
    "Potato___late_blight",
    "Potato___leafroll_virus",
    "Potato___mosaic_virus",
    "Potato___nematode",
    "Potato___pests",
    "Potato___phytophthora",
    "Raspberry___healthy",
    "Rice___bacterial_blight",
    "Rice___blast",
    "Rice___brown_spot",
    "Rice___tungro",
    "Rose___healthy",
    "Rose___rust",
    "Rose___slug_sawfly",
    "Soybean___healthy",
    "Squash___powdery_mildew",
    "Strawberry___healthy",
    "Strawberry___leaf_scorch",
    "Sugercane___healthy",
    "Sugercane___mosaic",
    "Sugercane___red_rot",
    "Sugercane___rust",
    "Sugercane___yellow_leaf",
    "Tomato___bacterial_spot",
    "Tomato___early_blight",
    "Tomato___healthy",
    "Tomato___late_blight",
    "Tomato___leaf_curl",
    "Tomato___leaf_mold",
    "Tomato___mosaic_virus",
    "Tomato___septoria_leaf_spot",
    "Tomato___spider_mites",
    "Tomato___target_spot",
    "Watermelon___anthracnose",
    "Watermelon___downy_mildew",
    "Watermelon___healthy",
    "Watermelon___mosa"
]

# ================================
# 5. FONCTION DE PRÉDICTION
# ================================
def preprocess_image(img):
    """Redimensionne et normalise l'image pour le modèle EfficientNet."""
    img = img.resize((224, 224))
    img_array = np.array(img)
    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):
    """
    img : image PIL fournie par gr.Image(type="pil")
    Retourne un dictionnaire {classe: probabilité} pour le composant gr.Label
    """
    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

# ================================
# 6. INTERFACE GRADIO (sans allow_flagging)
# ================================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Chargez une image de feuille"),
    outputs=gr.Label(num_top_classes=3, label="Maladie prédite (top 3)"),
    title="Classification des maladies des plantes (72 classes)",
    description="Chargez une photo de feuille et le modèle prédira la maladie parmi 72 classes. Modèle basé sur EfficientNetB0 avec fine-tuning.",
    examples=None
)

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