#!/usr/bin/env python3
"""Prepare Qualcomm AI Hub calibration and validation inputs for the meter detector."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
from PIL import Image

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        default=Path("examples/meter-detection/data/electric-meter-detection"),
        help="Root of the extracted Roboflow dataset.",
    )
    parser.add_argument(
        "--calibration-dir",
        type=Path,
        default=Path("examples/meter-detection/data/aihub_calibration"),
        help="Directory where .npy calibration samples will be written.",
    )
    parser.add_argument(
        "--input-file",
        type=Path,
        default=Path("examples/meter-detection/data/inputs.npz"),
        help="Validation .npz input file for qc-cli ai-hub validate.",
    )
    parser.add_argument("--input-name", default="images", help="ONNX input name.")
    parser.add_argument("--image-size", type=int, default=640, help="Square image size used for ONNX export.")
    parser.add_argument("--samples", type=int, default=16, help="Number of calibration samples to write.")
    return parser.parse_args()


def preprocess_image(path: Path, image_size: int) -> np.ndarray:
    """Apply Ultralytics-style letterboxing and produce an NCHW float32 tensor."""
    with Image.open(path) as source:
        image = source.convert("RGB")

    scale = min(image_size / image.width, image_size / image.height)
    resized_width = round(image.width * scale)
    resized_height = round(image.height * scale)
    image = image.resize((resized_width, resized_height), Image.Resampling.BILINEAR)

    canvas = Image.new("RGB", (image_size, image_size), (114, 114, 114))
    left = round((image_size - resized_width) / 2 - 0.1)
    top = round((image_size - resized_height) / 2 - 0.1)
    canvas.paste(image, (left, top))

    array = np.asarray(canvas, dtype=np.float32) / 255.0
    return np.transpose(array, (2, 0, 1))[None, ...].astype(np.float32)


def main() -> None:
    args = parse_args()
    if args.image_size < 1:
        raise SystemExit("--image-size must be at least 1")
    if args.samples < 1:
        raise SystemExit("--samples must be at least 1")

    images = sorted(
        path
        for path in args.dataset_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS and path.parent.name == "images"
    )
    if not images:
        raise SystemExit(f"No images found under {args.dataset_dir}")

    args.calibration_dir.mkdir(parents=True, exist_ok=True)
    args.input_file.parent.mkdir(parents=True, exist_ok=True)
    for stale_sample in args.calibration_dir.glob("sample_*.npy"):
        stale_sample.unlink()

    prepared: list[np.ndarray] = []
    for index, image_path in enumerate(images[: args.samples]):
        sample = preprocess_image(image_path, args.image_size)
        np.save(args.calibration_dir / f"sample_{index:03d}.npy", sample)
        prepared.append(sample)

    np.savez(args.input_file, **{args.input_name: prepared[0]})  # pyright: ignore[reportArgumentType]
    print(f"Wrote {len(prepared)} calibration samples to {args.calibration_dir}")
    print(f"Wrote validation input to {args.input_file}")


if __name__ == "__main__":
    main()