"""Gradio demo for the TrashNet robust classifier.

Before running:
1. Train the model with train_eval_final.py.
2. Make sure the checkpoint exists, for example: out/best_augmented.pth.
3. Run: python app.py --model-path out/best_augmented.pth
"""

import argparse
from pathlib import Path

import gradio as gr
import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import ResNet18_Weights, resnet18

CLASS_NAMES = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]


def build_model(model_path: str, device: torch.device):
    model = resnet18(weights=ResNet18_Weights.DEFAULT)
    model.fc = nn.Linear(model.fc.in_features, len(CLASS_NAMES))
    state_dict = torch.load(model_path, map_location=device)
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model


def build_transform():
    return transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])


def make_predict_fn(model, transform, device):
    def predict(image):
        if image is None:
            return {}
        image = image.convert("RGB")
        x = transform(image).unsqueeze(0).to(device)
        with torch.no_grad():
            logits = model(x)
            probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
        return {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}
    return predict


def main(args):
    model_path = Path(args.model_path)
    if not model_path.exists():
        raise FileNotFoundError(
            f"Model checkpoint not found: {model_path}. Train first or pass --model-path."
        )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = build_model(str(model_path), device)
    transform = build_transform()
    predict_fn = make_predict_fn(model, transform, device)

    description = (
        "Upload an image of waste. The model predicts one of six TrashNet classes: "
        "cardboard, glass, metal, paper, plastic, or trash."
    )

    demo = gr.Interface(
        fn=predict_fn,
        inputs=gr.Image(type="pil", label="Upload waste image"),
        outputs=gr.Label(num_top_classes=6, label="Prediction confidence"),
        title="Robust Trash Classifier",
        description=description,
        examples=args.examples if args.examples else None,
    )
    demo.launch()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Launch Gradio TrashNet classifier demo")
    parser.add_argument("--model-path", type=str, default="best_augmented.pth")
    parser.add_argument("--examples", nargs="*", default=[])
    main(parser.parse_args())