#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Run the served model (``METEOR.from_pretrained``: the ``frame`` trace, ETH 12x10) on the public-dataset frames of
``research/meteor/public_data`` and write the TT outputs in the layout of the CPU goldens (``fNNNN.npz`` compact +
``fNNNN.json`` decoded, ``run_public_reference.save_frame``), so ``compare_tt.py`` scores them against the CPU
goldens and ``make_demo.py`` renders them. Development tool of the workspace (needs the research scripts and the
converted public inputs, which never ship; nuScenes is CC BY-NC-SA: local validation only).

    bin/devrun -t 1800 -- python -u code/scripts/run_public_frames.py --out logs/public_tt [--ids ps019,ps090,ns0103]

Every frame's graph feed is checked against the golden's ``input_sha256`` first (the same pixels, K, T and v0 as the
CPU reference). Frames are processed per sequence (one rig each: the lift tables are written once per sequence).
Also writes ``<out>/run.json`` (versions, device, per-frame seconds) and, with ``--full`` frames, the 19 outputs.
"""
from __future__ import annotations

import argparse
import glob
import json
import os
import sys
import time
from pathlib import Path

import numpy as np

ROOT = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))
PUBLIC = ROOT / "research" / "meteor" / "public_data"
sys.path.insert(0, str(PUBLIC / "scripts"))
sys.path.insert(0, str(ROOT / "research" / "meteor" / "scripts"))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True, type=Path)
    ap.add_argument("--ids", default="ps019,ps090,ns0103")
    ap.add_argument("--full", default="ps019:20,ps090:20,ns0103:9", help="id:frame pairs that also get full_fNNNN.npz")
    a = ap.parse_args()
    import me_public_common as C  # noqa: F401  (research decode, metrics)
    import meteor_io
    from run_public_reference import save_frame

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

    full = {tuple(p.split(":")) for p in a.full.split(",") if p}
    jobs = []
    for gid in a.ids.split(","):
        gdir = PUBLIC / "golden" / gid
        frames = sorted(int(Path(f).stem[1:]) for f in glob.glob(str(gdir / "f[0-9][0-9][0-9][0-9].json")))
        jobs.append((gid, frames))
    run = {"bundle_version": __version__, "frames": {}, "started": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    t_load = time.time()
    with METEOR.from_pretrained() as model:
        run["load_s"] = round(time.time() - t_load, 1)
        run["device"] = {k: model.device_info.get(k) for k in ("dispatch", "grid", "num_command_queues", "cores")}
        print("device", run["device"], "load", run["load_s"], "s", flush=True)
        for gid, frames in jobs:
            root = PUBLIC / "inputs" / gid
            sd = meteor_io.scene_dir_of(str(root)) if hasattr(meteor_io, "scene_dir_of") else \
                str(root / (root / "scenes.txt").read_text().split()[0])
            man = meteor_io.load_manifest(sd)
            ego_gt = dict(np.load(os.path.join(sd, "ego_motion.npz")))
            present = np.asarray(man.get("present_mask", [1] * 8), bool)
            out_dir = a.out / gid
            out_dir.mkdir(parents=True, exist_ok=True)
            for i in frames:
                feed = meteor_io.load_frame(sd, i, man)
                gold = json.loads((PUBLIC / "golden" / gid / f"f{i:04d}.json").read_text())
                sha = {k: C.sha256_array(v) for k, v in feed.items()}
                if sha != gold["input_sha256"]:
                    raise SystemExit(f"{gid} f{i}: the feed differs from the golden's input_sha256")
                fr = MeteorFrame(feed["imgs"], feed["K"], feed["T_cam_ego"], feed["v0"], present)
                t0 = time.time()
                o = {k: np.asarray(v) for k, v in model._forward({"frame": fr}).items()}
                dt = time.time() - t0
                save_frame(str(out_dir), i, o, feed, man, ego_gt, full=(gid, str(i)) in full)
                run["frames"][f"{gid}/f{i:04d}"] = round(dt, 3)
                print(f"{gid} f{i:04d}: {dt * 1e3:.0f} ms", flush=True)
    run["finished"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    (a.out / "run.json").write_text(json.dumps(run, indent=1) + "\n")


if __name__ == "__main__":
    main()