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

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/training/data/flower_photos_sagemaker"),
        help="ImageFolder-style dataset used for training.",
    )
    parser.add_argument(
        "--calibration-dir",
        type=Path,
        default=Path("examples/training/data/aihub_calibration"),
        help="Directory where .npy calibration samples will be written.",
    )
    parser.add_argument(
        "--input-file",
        type=Path,
        default=Path("examples/training/data/inputs.npz"),
        help="Validation .npz input file for qc-cli ai-hub validate.",
    )
    parser.add_argument("--input-name", default="input", help="ONNX input name.")
    parser.add_argument("--image-size", type=int, default=160, help="Square image size used by training.")
    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:
    image = Image.open(path).convert("RGB").resize((image_size, image_size), Image.Resampling.BILINEAR)
    array = np.asarray(image, dtype=np.float32) / 255.0
    array = np.transpose(array, (2, 0, 1))
    mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)[:, None, None]
    std = np.array([0.229, 0.224, 0.225], dtype=np.float32)[:, None, None]
    return ((array - mean) / std)[None, ...].astype("float32")


def main() -> None:
    args = parse_args()
    images = sorted(p for p in args.dataset_dir.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS)
    if not images:
        raise SystemExit(f"No images found under {args.dataset_dir}")
    if args.samples < 1:
        raise SystemExit("--samples must be at least 1")

    args.calibration_dir.mkdir(parents=True, exist_ok=True)
    args.input_file.parent.mkdir(parents=True, exist_ok=True)

    sample_count = min(args.samples, len(images))
    prepared = []
    for index, image_path in enumerate(images[:sample_count]):
        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]})
    print(f"Wrote {sample_count} calibration samples to {args.calibration_dir}")
    print(f"Wrote validation input to {args.input_file}")


if __name__ == "__main__":
    main()