import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import gradio as gr

# ==========================
# 1) Device
# ==========================
device = torch.device("cpu")  # Hugging Face = CPU par défaut


# ==========================
# 2) Définition du modèle (PlantCNN) EXACTEMENT comme à l'entraînement
# ==========================
class PlantCNN(nn.Module):
    def __init__(self, num_classes):
        super(PlantCNN, self).__init__()
        # 🧩 Bloc 1
        self.conv_block1 = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, padding=0),
            nn.ELU(),
            nn.Conv2d(32, 32, kernel_size=3, padding=0),
            nn.ELU(),
            nn.MaxPool2d(2)
        )
        # 🧩 Bloc 2
        self.conv_block2 = nn.Sequential(
            nn.Conv2d(32, 64, kernel_size=3, padding=0),
            nn.ELU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=0),
            nn.ELU(),
            nn.MaxPool2d(2)
        )
        # 🧩 Bloc 3
        self.conv_block3 = nn.Sequential(
            nn.Conv2d(64, 128, kernel_size=3, padding=0),
            nn.ELU(),
            nn.Conv2d(128, 128, kernel_size=3, padding=0),
            nn.ELU(),
            nn.MaxPool2d(2)
        )
        # 🔁 Global Average Pooling
        self.gap = nn.AdaptiveAvgPool2d((1, 1))

        # 🧱 Dense layers
        self.fc1 = nn.Sequential(
            nn.Linear(128, 256),
            nn.ELU(),
            nn.BatchNorm1d(256),
            nn.Dropout(0.2)
        )
        self.fc2 = nn.Sequential(
            nn.Linear(256, 128),
            nn.ELU(),
            nn.BatchNorm1d(128),
            nn.Dropout(0.2)
        )
        self.fc3 = nn.Sequential(
            nn.Linear(128, 64),
            nn.ELU(),
            nn.BatchNorm1d(64),
            nn.Dropout(0.2)
        )
        self.out = nn.Linear(64, num_classes)  # CrossEntropyLoss fait softmax

    def forward(self, x):
        x = self.conv_block1(x)
        x = self.conv_block2(x)
        x = self.conv_block3(x)
        x = self.gap(x)           # [B, 128, 1, 1]
        x = torch.flatten(x, 1)   # [B, 128]
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        x = self.out(x)
        return x


# ==========================
# 3) Classes (dans le même ordre que train_dataset.classes)
# ==========================
class_names = [
    "Apple___Apple_scab",
    "Apple___Black_rot",
    "Apple___Cedar_apple_rust",
    "Apple___healthy",
    "Blueberry___healthy",
    "Cherry_(including_sour)___Powdery_mildew",
    "Cherry_(including_sour)___healthy",
    "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot",
    "Corn_(maize)___Common_rust_",
    "Corn_(maize)___Northern_Leaf_Blight",
    "Corn_(maize)___healthy",
    "Grape___Black_rot",
    "Grape___Esca_(Black_Measles)",
    "Grape___Leaf_blight_(Isariopsis_Leaf_Spot)",
    "Grape___healthy",
    "Orange___Haunglongbing_(Citrus_greening)",
    "Peach___Bacterial_spot",
    "Peach___healthy",
    "Pepper,_bell___Bacterial_spot",
    "Pepper,_bell___healthy",
    "Potato___Early_blight",
    "Potato___Late_blight",
    "Potato___healthy",
    "Raspberry___healthy",
    "Soybean___healthy",
    "Squash___Powdery_mildew",
    "Strawberry___Leaf_scorch",
    "Strawberry___healthy",
    "Tomato___Bacterial_spot",
    "Tomato___Early_blight",
    "Tomato___Late_blight",
    "Tomato___Leaf_Mold",
    "Tomato___Septoria_leaf_spot",
    "Tomato___Spider_mites Two-spotted_spider_mite",
    "Tomato___Target_Spot",
    "Tomato___Tomato_Yellow_Leaf_Curl_Virus",
    "Tomato___Tomato_mosaic_virus",
    "Tomato___healthy",
]

num_classes = len(class_names)


# ==========================
# 4) Chargement du modèle
# ==========================
MODEL_PATH = "best_model.pth"  # ⚠️ mets exactement le même nom que le fichier uploadé

model = PlantCNN(num_classes=num_classes)
state_dict = torch.load(MODEL_PATH, map_location=device)
model.load_state_dict(state_dict)
model.to(device)
model.eval()

print("✅ Modèle PyTorch chargé sur CPU, prêt pour la prédiction.")


# ==========================
# 5) Transformations des images
# ==========================
image_size = 224

transform = transforms.Compose([
    transforms.Resize((image_size, image_size)),
    transforms.ToTensor(),
])


def prettify_class_name(raw_name: str) -> str:
    # Transforme "Tomato___Early_blight" -> "Tomato – Early blight"
    name = raw_name.replace("___", " – ")
    name = name.replace("_", " ")
    return name


# ==========================
# 6) Fonction de prédiction
# ==========================
def predict(image: Image.Image):
    if image is None:
        return "Merci d'uploader une image 🌸"

    img = transform(image).unsqueeze(0).to(device)  # [1, 3, H, W]

    with torch.no_grad():
        outputs = model(img)
        probs = torch.softmax(outputs, dim=1)[0]

    top_prob, top_idx = torch.max(probs, dim=0)
    top_prob = float(top_prob.item())
    top_idx = int(top_idx.item())

    pred_class_raw = class_names[top_idx]
    pred_class = prettify_class_name(pred_class_raw)

    # Top 3 prédictions
    topk = 3
    top_probs, top_indices = torch.topk(probs, k=topk)
    top_probs = top_probs.cpu().numpy()
    top_indices = top_indices.cpu().numpy()

    md = f"### 🌿 Résultat de la prédiction\n"
    md += f"**Classe prédite :** `{pred_class}`\n\n"
    md += f"**Confiance :** `{top_prob*100:.2f}%` 💖\n\n"
    md += "---\n"
    md += "### 🌈 Top 3 des classes probables\n"

    for i in range(topk):
        cls_raw = class_names[int(top_indices[i])]
        cls_name = prettify_class_name(cls_raw)
        cls_prob = float(top_probs[i])
        md += f"- `{cls_name}` → **{cls_prob*100:.2f}%**\n"

    md += "\n> 🍃 *Modèle entraîné sur PlantVillage (38 classes). À utiliser pour l’exploration, pas comme outil médical.*\n"

    return md


# ==========================
# 7) CSS girly (injecté via HTML, compatible ancienne version Gradio)
# ==========================
custom_css = """
body {
    background: #ffeef7;
}
#root, .gradio-container {
    background: linear-gradient(135deg, #ffeef7 0%, #e4f9ff 100%) !important;
}
.gradio-container {
    font-family: 'Segoe UI', system-ui, -apple-system, BlinkMacSystemFont, sans-serif;
}
h1, h2, h3 {
    color: #d6478f !important;
}
button {
    border-radius: 999px !important;
}
"""


# ==========================
# 8) Interface Gradio
# ==========================
with gr.Blocks() as demo:
    # Injecter le CSS custom
    gr.HTML(f"<style>{custom_css}</style>")

    gr.Markdown(
        """
        <div style="text-align:center; margin-bottom: 10px;">
            <h1>🌸 Plant Disease Detector by Hanen 🌸</h1>
            <p style="font-size:16px; color:#555;">
                Upload une feuille de plante 🍃 et laisse ton modèle PyTorch super entraîné
                deviner la maladie (ou si elle est healthy 💚).
            </p>
        </div>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="🌿 Upload une image de feuille",
                height=320
            )
            gr.Markdown(
                """
                <span style="font-size:14px; color:#666;">
                👉 Utilise une image du dataset PlantVillage ou une photo claire d'une feuille sur fond neutre.<br>
                Formats supportés : JPG, PNG.
                </span>
                """
            )
        with gr.Column(scale=1):
            output_md = gr.Markdown(
                value="Le résultat apparaîtra ici 💖",
            )

    btn = gr.Button("✨ Analyser la feuille ✨")
    btn.click(fn=predict, inputs=image_input, outputs=output_md)

    gr.Markdown(
        """
        <div style="text-align:center; font-size:13px; color:#777; margin-top:20px;">
            Modèle : <b>PlantCNN (PyTorch)</b> – Accuracy test ≈ <b>98%</b> 🌟<br>
            Déployé avec 💕 sur Hugging Face & Gradio.
        </div>
        """
    )

if __name__ == "__main__":
    demo.launch()