import os
import json
import torch
import timm
import gradio as gr
import numpy as np
from PIL import Image
from torchvision import transforms
from huggingface_hub import hf_hub_download

MODEL_REPO = os.environ.get("HF_MODEL_REPO", "alikh02/weed-classifier")

print(f"Loading model from {MODEL_REPO}...")

class_names_path = hf_hub_download(repo_id=MODEL_REPO, filename="class_names.json")
with open(class_names_path) as f:

    raw = json.load(f)
    idx_to_species = {int(k): v for k, v in raw.items()}

NUM_CLASSES = len(idx_to_species)
species_list = [idx_to_species[i] for i in range(NUM_CLASSES)]

model_path = hf_hub_download(repo_id=MODEL_REPO, filename="model.pt")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = timm.create_model("efficientnet_b0", pretrained=False, num_classes=NUM_CLASSES)
model.load_state_dict(torch.load(model_path, map_location=device))
model.eval().to(device)

print(f"Model loaded! {NUM_CLASSES} classes: {species_list}")

# ── Preprocessing ─────────────────────────────────────────────────────────────
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406],
                         [0.229, 0.224, 0.225]),
])

# ── Inference ─────────────────────────────────────────────────────────────────
def classify_weed(image: Image.Image):
    """
    Takes a PIL image, returns a dict of {species: confidence} for Gradio Label.
    """
    if image is None:
        return {}

    tensor = preprocess(image.convert("RGB")).unsqueeze(0).to(device)

    with torch.no_grad():
        logits = model(tensor)
        probs = torch.softmax(logits, dim=1)[0].cpu().numpy()

    # Return top-5 as a dict (Gradio Label component expects {label: confidence})
    top5_idx = np.argsort(probs)[::-1][:5]
    results = {idx_to_species[i]: float(probs[i]) for i in top5_idx}
    return results


# ── Gradio UI ─────────────────────────────────────────────────────────────────
DESCRIPTION = """
## 🌿 Weed Species Classifier

Upload a photo of a plant and the model will identify whether it's a weed species
or crop, and tell you which species it is.

**Supported species:** """ + " · ".join(species_list) + """

*Model: EfficientNet-B0 fine-tuned on the [DeepWeeds](https://github.com/AlexOlsen/DeepWeeds) dataset*
"""

EXAMPLES = [
    # ["examples/chinee_apple.jpg"],
    # ["examples/snakeweed.jpg"],
]

with gr.Blocks(title="Weed Classifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown(DESCRIPTION)

    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                type="pil",
                label="Upload plant image",
                sources=["upload", "webcam"],
            )
            classify_btn = gr.Button("🔍 Classify", variant="primary")

        with gr.Column(scale=1):
            label_output = gr.Label(
                num_top_classes=5,
                label="Predicted Species (top 5)",
            )
            info_box = gr.Markdown(visible=False)

    # Wire up
    def run_and_annotate(image):
        results = classify_weed(image)
        if not results:
            return {}, gr.update(visible=False, value="")
        top_species, top_conf = list(results.items())[0]
        note = f"**Top prediction:** `{top_species}` with **{top_conf*100:.1f}%** confidence"
        return results, gr.update(visible=True, value=note)

    classify_btn.click(
        fn=run_and_annotate,
        inputs=image_input,
        outputs=[label_output, info_box],
    )
    image_input.change(          # also run when image changes
        fn=run_and_annotate,
        inputs=image_input,
        outputs=[label_output, info_box],
    )

    if EXAMPLES:
        gr.Examples(examples=EXAMPLES, inputs=image_input)

    gr.Markdown("""
---
### 📖 How to use
1. Upload or drag-and-drop a plant photo
2. Hit **Classify** — results appear on the right
3. The bar chart shows the top 5 most likely species with confidence scores

### ℹ️ About the model
- Architecture: EfficientNet-B0 (pretrained on ImageNet, fine-tuned on DeepWeeds)
- Training: Two-phase — head-only warm-up, then full fine-tuning
- Dataset: [DeepWeeds](https://github.com/AlexOlsen/DeepWeeds) — 17,509 labelled images
""")


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