"""One-shot script to capture golden mouse trajectories from the current torch
implementation. Run BEFORE the migration so we can verify the numpy/ORT rewrite
in Phase 4 produces equivalent output.

Output: tests/unit/data/golden_mouse.npz
"""
from __future__ import annotations

import random
import sys
from pathlib import Path

# Allow running as `uv run python scripts/build_golden_mouse.py` from project root.
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

import numpy as np
import torch

from ai_mouse import generate

CASES: list[tuple[tuple[int, int], tuple[int, int]]] = [
    ((100, 200), (900, 400)),    # horizontal 800px
    ((500, 500), (500, 100)),    # vertical 400px upward
    ((200, 600), (800, 200)),    # 720px diagonal
    ((100, 100), (130, 110)),    # very short 31px
    ((50,  50),  (1500, 900)),   # very long 1700px
    ((400, 300), (500, 300)),    # short horizontal 100px
    ((300, 300), (700, 700)),    # 45° diagonal
    ((600, 400), (200, 100)),    # reverse diagonal
]
SEEDS = (0, 1, 2, 3)


def main() -> None:
    out: dict[str, np.ndarray] = {}
    for case_idx, (start, end) in enumerate(CASES):
        for seed in SEEDS:
            random.seed(seed)
            np.random.seed(seed)
            torch.manual_seed(seed)
            pts = generate(start=start, end=end)
            out[f"case{case_idx}_seed{seed}"] = np.array(pts, dtype=np.int64)
    out_path = Path("tests/unit/data/golden_mouse.npz")
    np.savez_compressed(out_path, **out)
    print(f"Wrote {len(out)} golden traces to {out_path}")


if __name__ == "__main__":
    main()