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

Loads a quantized .pte model and runs image classification on a single image.
Outputs top-5 predictions with class labels and probabilities.
"""

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-50_raspberry_executorch_optimized.pte"
IMAGE_PATH = "sample_input.jpg"
INPUT_SIZE = (224, 224)
TOP_K = 5

# ImageNet normalization constants
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."""
    runtime = Runtime.get()
    program = runtime.load_program(pte_path)
    return program.load_method("forward")


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

    Resizes to 256, center-crops to 224x224, converts to float32 tensor in
    [0, 1], then applies ImageNet mean/std normalization.
    """
    image = Image.open(image_path).convert("RGB")
    transform = transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(INPUT_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ]
    )
    return transform(image).unsqueeze(0)  # [1, 3, 224, 224]


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


# ── Postprocessing ─────────────────────────────────────────────────────────────
def postprocess(raw_output: torch.Tensor, labels: dict[str, str], top_k: int = TOP_K) -> list:
    """Decode raw logits into top-k class predictions with probabilities.

    Applies softmax to convert logits to probabilities, then returns the
    top-k classes sorted by descending probability.
    """
    # Ensure 2D: [1, 1000] → [1000]
    logits = raw_output
    if logits.dim() > 1:
        logits = logits.squeeze(0)

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

    # 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.tolist(), top_indices.tolist(), strict=False),
        key=lambda item: (-item[0], item[1]),
    )

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


# ── Save Results ───────────────────────────────────────────────────────────────
def save_results(results: list, script_dir: Path) -> None:
    """Save top-k predictions as predictions.json in the script directory."""
    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_output_image(image_path: str, top_prediction: dict, script_dir: Path) -> None:
    """Overlay the top prediction label on the original image and save as sample_output.jpg."""
    image = Image.open(image_path).convert("RGB")
    draw = ImageDraw.Draw(image)

    label = f"{top_prediction['class']}: {top_prediction['probability'] * 100:.1f}%"

    font_size = max(16, image.width // 20)
    font: ImageFont.ImageFont | ImageFont.FreeTypeFont
    try:
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size
        )
    except OSError:
        try:
            font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", font_size)
        except OSError:
            font = ImageFont.load_default()

    bbox = draw.textbbox((0, 0), label, font=font)
    text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
    padding = 8
    x, y = padding, image.height - text_h - padding * 3

    draw.rectangle(
        [x - padding, y - padding, x + text_w + padding, y + text_h + padding],
        fill=(0, 0, 0, 180),
    )
    draw.text((x, y), label, fill=(255, 255, 255), font=font)

    output_path = script_dir / "sample_output.jpg"
    image.save(output_path, quality=95)
    print(f"Output image saved to: {output_path}")


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

    # Resolve paths relative to script location
    model_path = script_dir / MODEL_PATH
    image_path = script_dir / IMAGE_PATH

    # Load model
    method = load_model(str(model_path))

    # Preprocess input image
    input_tensor = preprocess(str(image_path))

    # Run inference
    raw_output = run_inference(method, input_tensor)

    # Decode predictions
    results = postprocess(raw_output, IMAGENET_CLASSES)

    # Print top-k predictions
    print(f"\nTop-{TOP_K} predictions for: {image_path.name}")
    print("-" * 50)
    for rank, pred in enumerate(results, 1):
        print(f"  {rank}. {pred['class']:<40} {pred['probability']*100:.2f}%")

    # Save predictions.json
    save_results(results, script_dir)

    # Save labeled output image
    save_output_image(str(image_path), results[0], script_dir)


if __name__ == "__main__":
    main()