#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Run golden frames' graph inputs through the served model and save the 19 raw device outputs (npz) for offline
analysis against the goldens (e.g. which boxes a set-based gate misses and why).

    bin/devrun -t 900 -- python code/scripts/dump_device_outputs.py --out DIR meteor_valday_f040 ...
"""
from __future__ import annotations

import argparse
import os
from pathlib import Path

import numpy as np

from tt_meteor import METEOR
from tt_meteor.host.preprocess import MeteorFrame

GOLDENS = Path(os.environ.get("METEOR_GOLDENS", "/home/ubuntu/experiments/tt-models/research/meteor/goldens"))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("frames", nargs="+")
    ap.add_argument("--out", required=True)
    ap.add_argument("--repeat", type=int, default=2, help="runs per frame (outputs must be bit-identical)")
    a = ap.parse_args()
    out = Path(a.out)
    out.mkdir(parents=True, exist_ok=True)
    with METEOR.from_pretrained() as model:
        for name in a.frames:
            with np.load(GOLDENS / name / "taps.npz") as z:
                g = {k: z[k] for k in ("input.imgs", "input.K", "input.T_cam_ego", "input.v0", "input.present")}
            fr = MeteorFrame(g["input.imgs"], g["input.K"], g["input.T_cam_ego"], g["input.v0"], g["input.present"])
            runs = [{k: np.array(v) for k, v in model._forward({"frame": fr}).items()} for _ in range(a.repeat)]
            for r in runs[1:]:
                for k in r:
                    assert np.array_equal(r[k], runs[0][k]), f"{name}: {k} differs between runs"
            np.savez(out / f"{name}.npz", **runs[0])
            print(f"{name}: saved {len(runs[0])} outputs, {a.repeat} runs bit-identical")


if __name__ == "__main__":
    main()