#!/usr/bin/env python
"""
Quick-start inference script for kavin-aravindhan/vit-oct-wamd.

Usage:
    pip install -r requirements.txt
    python infer.py path/to/scan.png [more_images.png ...]
    python infer.py path/to/scan.png --device cpu

Downloads the checkpoint + model code from the Hub on first run (cached
afterwards), runs the classifier, and prints normal/wet_amd + confidence
for each image.

NOTE: requirements.txt pins transformers/torch/huggingface_hub to the exact
versions used to produce best_model.pt. An unpinned `pip install
transformers` can resolve to a version whose SiglipVisionModel uses
different internal parameter names, which will fail to load this
checkpoint's state dict.
"""
import argparse
import os
import sys

import cv2
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

REPO_ID = "kavin-aravindhan/vit-oct-wamd"
IMAGE_SIZE = 384
LABELS = {0: "normal", 1: "wet_amd"}


def fetch_repo_files(local_dir=None):
    """Download modeling.py + its dependencies + the checkpoint, return local dir."""
    modeling_path = hf_hub_download(repo_id=REPO_ID, filename="modeling.py", local_dir=local_dir)
    local_dir = os.path.dirname(modeling_path)
    for fname in ["alignment.py", "embedder.py", "best_model.pt"]:
        hf_hub_download(repo_id=REPO_ID, filename=fname, local_dir=local_dir)
    return local_dir


def preprocess_image(image_path):
    image = Image.open(image_path).convert("RGB")
    img_array = np.array(image, dtype=np.float32) / 255.0
    img_resized = cv2.resize(img_array, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_LINEAR)
    img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1)  # (H,W,C) -> (C,H,W)
    img_tensor = (img_tensor - 0.5) / 0.5  # [-1, 1], matches SigLIP normalization
    return img_tensor


def main():
    parser = argparse.ArgumentParser(description="Run vit-oct-wamd on one or more OCT images.")
    parser.add_argument("images", nargs="+", help="Path(s) to OCT B-scan image(s).")
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--local-dir", default=None, help="Where to cache the downloaded model files.")
    args = parser.parse_args()

    local_dir = fetch_repo_files(args.local_dir)
    sys.path.insert(0, local_dir)
    from modeling import load_model  # noqa: E402

    print(f"Loading model on {args.device} ...")
    model = load_model(os.path.join(local_dir, "best_model.pt"), device=args.device)

    batch = torch.stack([preprocess_image(p) for p in args.images]).to(args.device)
    with torch.no_grad():
        logits = model(batch)
        probs = torch.softmax(logits, dim=-1)
        preds = probs.argmax(-1)

    for path, pred, prob in zip(args.images, preds, probs):
        label = LABELS[pred.item()]
        confidence = prob[pred].item()
        print(f"{path}: {label} (confidence={confidence:.3f})")


if __name__ == "__main__":
    main()