"""Minimal inference example for DoB24/fundus-9model-benchmark.

Loads one of the 9 trained checkpoints, applies the same preprocessing used
during evaluation, and predicts a class for a single fundus image.

Usage:
    python inference_example.py path/to/image.jpg --model densenet121
"""
from __future__ import annotations
import argparse, json
from pathlib import Path
import torch
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image
import timm
from huggingface_hub import hf_hub_download

REPO = "DoB24/fundus-9model-benchmark"
CLASSES = [
    "Central Serous Chorioretinopathy",
    "Diabetic Retinopathy", "Disc Edema", "Glaucoma", "Healthy",
    "Macular Scar", "Myopia", "Pterygium",
    "Retinal Detachment", "Retinitis Pigmentosa",
]
# Map repo-name to timm model id
TIMM_ID = {
    "vgg19": "vgg19", "resnet50": "resnet50", "resnet101": "resnet101",
    "densenet121": "densenet121", "inception_v3": "inception_v3",
    "swin_b": "swin_base_patch4_window7_224",
}

def load_model(name: str, num_classes: int = 10):
    weights = hf_hub_download(REPO, f"weights/{name}_best.pth")
    if name in TIMM_ID:
        model = timm.create_model(TIMM_ID[name], pretrained=False, num_classes=num_classes)
        if name == "inception_v3":
            model = timm.create_model("inception_v3", pretrained=False,
                                      num_classes=num_classes, aux_logits=False)
    elif name == "clip_openai":
        import open_clip
        model, _, _ = open_clip.create_model_and_transforms("ViT-B-16", pretrained=None)
        model.visual.proj = None
        model = torch.nn.Sequential(model.visual, torch.nn.Linear(768, num_classes))
    elif name == "dinov2_l":
        model = torch.hub.load("facebookresearch/dinov2", "dinov2_vitl14")
        model = torch.nn.Sequential(model, torch.nn.Linear(1024, num_classes))
    elif name == "retfound":
        model = timm.create_model("vit_large_patch16_224", pretrained=False,
                                  num_classes=num_classes)
    else:
        raise ValueError(f"Unknown model: {name}")
    sd = torch.load(weights, map_location="cpu", weights_only=False)
    if isinstance(sd, dict) and "state_dict" in sd:
        sd = sd["state_dict"]
    model.load_state_dict(sd, strict=False)
    model.eval()
    return model

def preprocess(img_path: str) -> torch.Tensor:
    tf = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    return tf(Image.open(img_path).convert("RGB")).unsqueeze(0)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image")
    ap.add_argument("--model", default="densenet121",
                    choices=list(TIMM_ID) + ["clip_openai", "dinov2_l", "retfound"])
    args = ap.parse_args()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = load_model(args.model).to(device)
    x = preprocess(args.image).to(device)
    with torch.no_grad():
        probs = F.softmax(model(x), dim=1)[0].cpu().numpy()
    order = probs.argsort()[::-1]
    print(f"\nTop-3 predictions for {args.image} using {args.model}:")
    for i in order[:3]:
        print(f"  {CLASSES[i]:40s}  {probs[i]*100:5.2f} %")

if __name__ == "__main__":
    main()