"""
app.py — Plant Disease Classifier — Gradio Web UI
=================================================
Loads the trained EfficientNet-B0 weights (best_plant_model.pth)
and serves a drag-and-drop image classifier via Gradio.
"""

import json
import os
from pathlib import Path

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

# ──────────────────────────────────────────────
# 0. CONFIGURATION
# ──────────────────────────────────────────────
MODEL_PATH = "best_plant_model.pth"
CLASS_MAP_FILE = "class_names.json"
IMAGE_SIZE = 224
CONFIDENCE_THRESHOLD = 60.0
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("[INFO] App startup initiated", flush=True)

# ──────────────────────────────────────────────
# 1. LOAD CLASS NAMES
# ──────────────────────────────────────────────
def load_class_names():
    print("[INFO] Loading class names...", flush=True)

    if not Path(CLASS_MAP_FILE).exists():
        raise FileNotFoundError(f"Class map file not found: {CLASS_MAP_FILE}")

    with open(CLASS_MAP_FILE, "r", encoding="utf-8") as f:
        class_names = json.load(f)

    if not isinstance(class_names, list) or len(class_names) == 0:
        raise ValueError("class_names.json is empty or invalid.")

    print(f"[INFO] Loaded {len(class_names)} class names", flush=True)
    return class_names


# ──────────────────────────────────────────────
# 2. LOAD MODEL
# ──────────────────────────────────────────────
def load_model(num_classes: int):
    print("[INFO] Creating model...", flush=True)
    model = timm.create_model(
        "efficientnet_b0",
        pretrained=False,
        num_classes=num_classes
    )

    if not Path(MODEL_PATH).exists():
        raise FileNotFoundError(f"Model file not found: {MODEL_PATH}")

    print("[INFO] Loading model weights...", flush=True)
    state_dict = torch.load(MODEL_PATH, map_location="cpu")

    print("[INFO] Applying state dict...", flush=True)
    model.load_state_dict(state_dict)

    print(f"[INFO] Moving model to {DEVICE}...", flush=True)
    model = model.to(DEVICE)
    model.eval()

    print("[INFO] Model is ready", flush=True)
    return model


# ──────────────────────────────────────────────
# 3. PARSE CLASS LABELS
# ──────────────────────────────────────────────
def parse_class_name(raw_name: str):
    parts = raw_name.split("___")
    plant = parts[0].replace("_", " ").strip()
    disease = parts[1].replace("_", " ").strip() if len(parts) > 1 else "Unknown"

    if disease.lower() == "healthy":
        disease = "Healthy"

    return plant, disease


# ──────────────────────────────────────────────
# 4. IMAGE TRANSFORMS
# ──────────────────────────────────────────────
inference_transforms = transforms.Compose([
    transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])


# ──────────────────────────────────────────────
# 5. INIT GLOBALS
# ──────────────────────────────────────────────
class_names = load_class_names()
model = load_model(num_classes=len(class_names))
print(f"[INFO] Model loaded successfully on {DEVICE}", flush=True)


# ──────────────────────────────────────────────
# 6. PREDICTION FUNCTION
# ──────────────────────────────────────────────
@torch.no_grad()
def predict(image: Image.Image):
    if image is None:
        return "Please upload an image."

    try:
        image = image.convert("RGB")
        tensor = inference_transforms(image).unsqueeze(0).to(DEVICE)

        logits = model(tensor)
        probs = F.softmax(logits, dim=1)

        confidence, pred_idx = probs.max(dim=1)
        confidence = confidence.item() * 100.0
        pred_idx = pred_idx.item()

        if confidence < CONFIDENCE_THRESHOLD:
            return (
                f"⚠️ **Unknown or unsupported plant**\n\n"
                f"The model is only **{confidence:.1f}%** confident, which is below "
                f"the **{CONFIDENCE_THRESHOLD:.0f}%** threshold.\n\n"
                f"This likely means the uploaded image is not one of the trained classes.\n\n"
                f"**Supported plants:** Apple, Blueberry, Cherry, Corn, Grape, Orange, "
                f"Peach, Bell Pepper, Potato, Raspberry, Soybean, Squash, Strawberry, Tomato."
            )

        raw_label = class_names[pred_idx]
        plant, disease = parse_class_name(raw_label)

        top5_probs, top5_idxs = probs.topk(5, dim=1)
        top5_lines = []

        for p, i in zip(top5_probs[0], top5_idxs[0]):
            lbl = class_names[i.item()]
            pl, dis = parse_class_name(lbl)
            top5_lines.append(f"• {pl} — {dis}: {p.item() * 100:.1f}%")

        return (
            f"🌿 **Plant:** {plant}\n"
            f"🦠 **Disease:** {disease}\n"
            f"📊 **Confidence:** {confidence:.1f}%\n\n"
            f"---\n"
            f"**Top-5 Predictions:**\n" + "\n".join(top5_lines)
        )

    except Exception as e:
        return f"❌ Error during prediction: {str(e)}"


# ──────────────────────────────────────────────
# 7. UI
# ──────────────────────────────────────────────
title = "🌱 Plant Disease Classifier"
description = (
    "Upload a photo of a plant leaf and this AI model will identify the plant species "
    "and diagnose any disease it detects.\n\n"
    f"**Model:** EfficientNet-B0 fine-tuned plant disease classifier ({len(class_names)} classes).\n\n"
    "Unsupported plants may be flagged as unknown."
)

examples_dir = Path("examples")
example_images = sorted(examples_dir.glob("*")) if examples_dir.exists() else []

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload a leaf image"),
    outputs=gr.Markdown(label="Diagnosis"),
    title=title,
    description=description,
    examples=[str(p) for p in example_images] if example_images else None,
)

# ──────────────────────────────────────────────
# 8. LAUNCH
# ──────────────────────────────────────────────
if __name__ == "__main__":
    print("[INFO] Starting Gradio app...", flush=True)

    port = int(os.getenv("PORT", "7860"))

    demo.launch(
        server_name="0.0.0.0",
        server_port=port,
        share=False,
        theme=gr.themes.Soft(),
    )