#!/usr/bin/env python3
"""Create Horizon hb_mapper calibration feature maps from the YOLO dataset."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
from PIL import Image


IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def resize_rgb(image: Image.Image, size: int) -> np.ndarray:
    image = image.convert("RGB")
    resized = image.resize((size, size), Image.Resampling.BILINEAR)
    array = np.asarray(resized, dtype=np.float32)
    return np.transpose(array, (2, 0, 1))[None, ...]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--samples", type=int, default=128)
    parser.add_argument("--size", type=int, default=640)
    args = parser.parse_args()

    roots = [args.dataset / "images" / "train", args.dataset / "images" / "vel"]
    images = sorted(
        p for root in roots if root.is_dir() for p in root.iterdir() if p.suffix.lower() in IMAGE_EXTS
    )
    if not images:
        raise SystemExit(f"no images found below {args.dataset}")
    count = min(args.samples, len(images))
    # Evenly spread samples over the sorted set so calibration is not dominated by one split.
    indices = np.linspace(0, len(images) - 1, count, dtype=np.int64)
    selected = [images[int(i)] for i in indices]

    args.output.mkdir(parents=True, exist_ok=True)
    for old in args.output.iterdir():
        if old.is_file() and old.suffix in {".bin", ".rgbchw", ".txt"}:
            old.unlink()
    manifest = []
    for index, path in enumerate(selected):
        data = resize_rgb(Image.open(path), args.size)
        data.tofile(args.output / f"{index:05d}.rgbchw")
        manifest.append(str(path.relative_to(args.dataset)))
    (args.output.parent / "calibration_manifest.txt").write_text(
        "\n".join(manifest) + "\n", encoding="utf-8"
    )
    print(f"wrote {count} samples ({args.size}x{args.size}) to {args.output}")


if __name__ == "__main__":
    main()