"""Plant Disease Assistant — Hugging Face Space (CPU, DINOv2-only).

Loads the DINOv2-L checkpoint from a HF model repo at startup, then runs
classification + template-based responses from a bundled knowledge file.

Configurable via environment variables:
    DINOV2_REPO    HF model repo containing best.pt and splits.json
                   (default: iamcode6/dinov2-l-ccmt-mi300x)
    DINOV2_CKPT    Filename of the checkpoint inside the repo
                   (default: best.pt)
"""
from __future__ import annotations

import json
import os
from pathlib import Path

import gradio as gr
import numpy as np
import timm
import torch
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.data import create_transform


HERE = Path(__file__).parent
KNOWLEDGE_PATH = HERE / "treatment_knowledge.json"
SPLITS_PATH = HERE / "splits.json"

DINOV2_REPO = os.environ.get("DINOV2_REPO", "iamcode6/dinov2-l-ccmt-mi300x")
DINOV2_CKPT = os.environ.get("DINOV2_CKPT", "best.pt")
DEVICE = "cpu"


class PlantClassifier:
    def __init__(self, checkpoint_path: Path, splits_path: Path):
        self.device = torch.device(DEVICE)

        splits = json.loads(splits_path.read_text())
        self.idx_to_class = {v: k for k, v in splits["class_to_idx"].items()}
        self.class_names = [self.idx_to_class[i] for i in range(len(self.idx_to_class))]
        self.num_classes = len(self.class_names)

        ckpt = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
        if isinstance(ckpt, dict) and "state_dict" in ckpt:
            state_dict = ckpt["state_dict"]
            cfg = ckpt.get("cfg", {})
        else:
            state_dict = ckpt
            cfg = {}
        state_dict = {k.replace("_orig_mod.", "", 1): v for k, v in state_dict.items()}

        model_name = cfg.get("model", {}).get("name", "vit_large_patch14_dinov2.lvd142m")
        img_size = cfg.get("model", {}).get("img_size", 224)

        self.model = timm.create_model(
            model_name, pretrained=False,
            num_classes=self.num_classes, img_size=img_size,
        )
        self.model.load_state_dict(state_dict)
        self.model.to(self.device).eval()

        self.transform = create_transform(
            input_size=img_size, is_training=False,
            mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225),
            interpolation="bicubic", crop_pct=0.95,
        )

    @torch.no_grad()
    def predict(self, image: Image.Image, top_k: int = 3) -> list[dict]:
        x = self.transform(image).unsqueeze(0).to(self.device)
        logits = self.model(x)
        probs = F.softmax(logits, dim=-1).squeeze(0).float().cpu().numpy()

        top_indices = np.argsort(probs)[::-1][:top_k]
        return [
            {"class": self.class_names[i], "confidence": float(probs[i]), "index": int(i)}
            for i in top_indices
        ]


class KnowledgeResponder:
    def __init__(self, path: Path):
        self.knowledge = json.loads(path.read_text())

    def format_label(self, label: str) -> str:
        return label.replace("_", " ").title()

    def respond(self, predictions: list[dict]) -> str:
        top = predictions[0]
        label = top["class"]
        confidence = top["confidence"]

        if label not in self.knowledge:
            return (
                f"**Prediction:** {self.format_label(label)} "
                f"(confidence: {confidence:.1%})\n\n"
                "No detailed information available for this condition."
            )

        k = self.knowledge[label]
        is_healthy = k["disease"] == "Healthy"

        lines = []
        if is_healthy:
            lines.append(f"## {k['crop']} — Healthy")
            lines.append(f"**Confidence:** {confidence:.1%}\n")
            lines.append(f"{k['symptoms']}")
            lines.append("\nKeep monitoring regularly and continue your current care routine.")
        else:
            lines.append(f"## {k['crop']} — {k['disease']}")
            lines.append(f"**Confidence:** {confidence:.1%}\n")
            if k.get("pathogen"):
                lines.append(f"**Pathogen:** *{k['pathogen']}*\n")
            lines.append("### Symptoms")
            lines.append(f"{k['symptoms']}\n")
            lines.append("### Severity Guide")
            for level, desc in k["severity_cues"].items():
                lines.append(f"- **{level.title()}:** {desc}")
            lines.append("")
            lines.append("### Treatment")
            lines.append(f"{k['treatment']}\n")
            lines.append("### Prevention")
            lines.append(f"{k['prevention']}")

        if len(predictions) > 1:
            lines.append("\n---\n### Other Possibilities")
            for p in predictions[1:]:
                if p["confidence"] > 0.05:
                    lines.append(f"- {self.format_label(p['class'])} ({p['confidence']:.1%})")

        return "\n".join(lines)


print(f"[app] Downloading DINOv2-L checkpoint from {DINOV2_REPO}...")
checkpoint_path = Path(hf_hub_download(repo_id=DINOV2_REPO, filename=DINOV2_CKPT))

print("[app] Loading classifier on CPU (~30s)...")
classifier = PlantClassifier(checkpoint_path, SPLITS_PATH)
print(f"[app] Loaded {classifier.num_classes} classes")

knowledge = KnowledgeResponder(KNOWLEDGE_PATH)


def diagnose(image: Image.Image | None):
    if image is None:
        return "Please upload an image.", ""

    image = image.convert("RGB")
    predictions = classifier.predict(image, top_k=3)

    table = "**DINOv2-L Classification (97% accuracy)**\n\n"
    table += "| Rank | Disease | Confidence |\n"
    table += "|------|---------|------------|\n"
    for i, p in enumerate(predictions):
        marker = " ←" if i == 0 else ""
        table += (
            f"| {i+1} | {knowledge.format_label(p['class'])} | "
            f"{p['confidence']:.1%}{marker} |\n"
        )

    return table, knowledge.respond(predictions)


CUSTOM_CSS = """
.prose, .prose *, [class*="markdown"], [class*="markdown"] * {
    color: #1a1a1a !important;
    opacity: 1 !important;
}
.prose strong, .prose h1, .prose h2, .prose h3 {
    color: #000 !important;
    font-weight: 700 !important;
}
.dark .prose, .dark .prose *,
.dark [class*="markdown"], .dark [class*="markdown"] * {
    color: #f5f5f5 !important;
}
.dark .prose strong, .dark .prose h1, .dark .prose h2, .dark .prose h3 {
    color: #ffffff !important;
}
.prose table { border-collapse: collapse; }
.prose th, .prose td { padding: 6px 10px; border: 1px solid #888; }
"""

with gr.Blocks(title="Plant Disease Assistant", css=CUSTOM_CSS) as app:
    gr.Markdown(
        "# 🌱 Plant Disease Assistant\n"
        "Upload a photo of a plant leaf to get an instant diagnosis, "
        "severity assessment, and treatment recommendations.\n\n"
        "*DINOv2-Large fine-tuned on AMD Instinct MI300X (ROCm) — "
        "97.06% accuracy on the CCMT crop disease dataset.*"
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(type="pil", label="Upload a plant leaf photo")
            diagnose_btn = gr.Button("Diagnose", variant="primary", size="lg")

            example_paths = sorted(str(p) for p in (HERE / "examples").glob("*.jpg"))
            if example_paths:
                gr.Examples(
                    examples=[[p] for p in example_paths],
                    inputs=image_input,
                    label="Or try one of these (click a thumbnail)",
                    examples_per_page=11,
                )
        with gr.Column(scale=2):
            classification_output = gr.Markdown()
            response_output = gr.Markdown()

    diagnose_btn.click(
        fn=diagnose, inputs=image_input,
        outputs=[classification_output, response_output],
    )
    image_input.change(
        fn=diagnose, inputs=image_input,
        outputs=[classification_output, response_output],
    )

    gr.Markdown(
        "---\n"
        "**Model:** DINOv2-Large (304M params) — 97.06% accuracy, 0.9713 macro F1\n\n"
        "**Hardware:** Fine-tuned on AMD Instinct MI300X (192 GB HBM3) via AMD Developer Cloud\n\n"
        "**Dataset:** CCMT Crop Pest and Disease Detection — 22 classes across cashew, cassava, maize, and tomato\n\n"
        "*Built for the lablab.ai AMD Developer Hackathon*"
    )


if __name__ == "__main__":
    app.launch(server_name="0.0.0.0", server_port=7860, show_api=False)