"""Minimal inference example for microsoft/resnet-18 INT8 using ExecuTorch.

Loads a quantized .pte model and runs classification inference on a single image.
Saves top-5 predictions to predictions.json and an annotated sample_output.jpg
in the same directory as this script.
"""

import json
from pathlib import Path

import torch
from executorch.runtime import Runtime
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms


# ── Configuration ──────────────────────────────────────────────────────────────
MODEL_PATH = "resnet-18_raspberry_executorch_optimized.pte"
IMAGE_PATH = "sample_input.jpg"
INPUT_SIZE = (224, 224)
TOP_K = 5

# ImageNet normalization constants (mean and std per channel)
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

SCRIPT_DIR = Path(__file__).resolve().parent

# ImageNet class labels (1000 classes)
with (SCRIPT_DIR / "imagenet_classes.json").open(encoding="utf-8") as file:
    IMAGENET_CLASSES = json.load(file)


# ── Model Loading ──────────────────────────────────────────────────────────────
def load_model(pte_path: str):
    """Load ExecuTorch .pte model and return the forward method."""
    script_dir = Path(__file__).parent
    runtime = Runtime.get()
    program = runtime.load_program(str(script_dir / pte_path))
    return program.load_method("forward")


# ── Preprocessing ──────────────────────────────────────────────────────────────
def preprocess(image_path: str) -> torch.Tensor:
    """Load and preprocess image for model input.

    Pipeline: resize to 256 px (shortest edge) -> center crop to 224x224
    -> convert to [0, 1] float32 tensor -> normalize with ImageNet mean/std.
    """
    script_dir = Path(__file__).parent
    image = Image.open(script_dir / image_path).convert("RGB")

    transform = transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(INPUT_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ]
    )

    input_tensor = transform(image).unsqueeze(0)  # [1, 3, 224, 224]
    return input_tensor


# ── Inference ──────────────────────────────────────────────────────────────────
def run_inference(method, input_tensor: torch.Tensor) -> torch.Tensor:
    """Run forward pass and return the raw logits tensor."""
    outputs = method.execute([input_tensor])
    return outputs[0]


# ── Postprocessing ─────────────────────────────────────────────────────────────
def postprocess(raw_output: torch.Tensor, labels: dict[str, str]) -> list[dict]:
    """Convert raw logits to top-k class predictions with probabilities.

    Applies softmax to get class probabilities, then returns the top-5
    predictions as a list of {class_id, class_name, probability} dicts.
    """
    # Ensure output is 2D [1, 1000]
    if raw_output.dim() == 1:
        raw_output = raw_output.unsqueeze(0)

    probabilities = torch.softmax(raw_output, dim=-1)
    k = min(TOP_K, raw_output.shape[-1])
    top_scores, top_indices = probabilities.topk(k, dim=-1)

    # Quantized logits land on a discrete grid, so exact ties are common and
    # topk orders them arbitrarily. Break ties by class id for a stable result.
    ranked = sorted(
        zip(top_scores[0].tolist(), top_indices[0].tolist(), strict=False),
        key=lambda item: (-item[0], item[1]),
    )

    results = []
    for score, idx in ranked:
        results.append(
            {
                "class_id": idx,
                "class_name": labels[str(idx)],
                "probability": round(score, 6),
            }
        )
    return results


# ── Save Results ───────────────────────────────────────────────────────────────
def save_results(results: list[dict]) -> None:
    """Save top-k predictions to predictions.json in the script's directory."""
    script_dir = Path(__file__).parent
    output_path = script_dir / "predictions.json"
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Predictions saved to {output_path}")


def save_sample_output(image_path: str, predictions: list[dict]) -> None:
    """Overlay top-5 predictions on the input image and save as sample_output.jpg."""
    script_dir = Path(__file__).parent
    image = Image.open(script_dir / image_path).convert("RGB")
    draw = ImageDraw.Draw(image)

    try:
        font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size=18)
        small_font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", size=15)
    except OSError:
        font = ImageFont.load_default()
        small_font = font

    # Semi-transparent dark banner at the top
    banner_h = 30 + len(predictions) * 22
    draw.rectangle([(0, 0), (image.width, banner_h)], fill=(0, 0, 0, 180))

    draw.text(
        (10, 6), "Top-5 predictions (ResNet-18 INT8)", font=font, fill=(255, 255, 255)
    )
    for i, pred in enumerate(predictions):
        label = pred["class_name"].split(",")[0]
        pct = pred["probability"] * 100
        y = 30 + i * 22
        draw.text(
            (10, y),
            f"{i + 1}. {label}  {pct:.1f}%",
            font=small_font,
            fill=(255, 220, 50),
        )

    out_path = script_dir / "sample_output.jpg"
    image.save(out_path, quality=95)
    print(f"Sample output saved to {out_path}")


# ── Main ───────────────────────────────────────────────────────────────────────
def main() -> None:
    script_dir = Path(__file__).parent

    # Load model
    print(f"Loading model: {MODEL_PATH}")
    method = load_model(MODEL_PATH)

    # Preprocess input image
    print(f"Preprocessing image: {IMAGE_PATH}")
    input_tensor = preprocess(IMAGE_PATH)

    # Run inference
    print("Running inference...")
    raw_output = run_inference(method, input_tensor)

    # Postprocess outputs
    results = postprocess(raw_output, IMAGENET_CLASSES)

    # Print top-k predictions
    print(f"\nTop-{TOP_K} predictions:")
    for i, pred in enumerate(results, 1):
        print(
            f"  {i}. {pred['class_name']} ({pred['class_id']}): {pred['probability']:.4f}"
        )

    # Save predictions JSON and annotated output image
    save_results(results)
    save_sample_output(IMAGE_PATH, results)


if __name__ == "__main__":
    main()