#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Run a public-dataset sequence through the Python API (the device graph) and score it against the CPU goldens.

    bin/devrun -t 1800 -- python code/scripts/run_public_tt.py --scene ps019 [--frames 0:80] [--out DIR]

Sequences (``research/streampetr/public_data/inputs/<scene>``, built by the research agents; never shipped):
``ps019`` (PandaSet 019 frames 0-79, CC BY 4.0), ``ns0103`` / ``ns0916`` (nuScenes v1.0-mini scenes 0103 / 0916,
every CAM_FRONT image; CC BY-NC-SA 4.0: internal validation only). One stream per run, reset at the first frame:
images (paths), inline calibration (K / P of the image as sent, ``T_ego_cam``), the per-camera stamps (the node's
0.15 s sync check) and the anchor stamp + ego pose (``ego_motion.npz`` ``T_map_base``).

Writes ``<out>/<scene>/detections.json`` in the research format (``{"meta", "results": [{"frame", "dets": [...]}]}``,
the input of ``research/streampetr/public_data/scripts/compare_runs.py det`` and ``eval_nuscenes_ap.py
--det-root``) and ``<out>/<scene>/f1.json``: per-frame detection F1 against the research ORT goldens (1 m, same label;
the free-running gate of PLAN 2.9 is the mean over >= 30 frames).
"""
from __future__ import annotations

import argparse
import json
import os
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" / "streampetr" / "public_data"
SLOTS = ("CAM_FRONT", "CAM_FRONT_LEFT", "CAM_BACK_LEFT", "CAM_FRONT_RIGHT", "CAM_BACK_RIGHT")


def requests(scene: str, frames=None):
    """Yield ``(frame, request kwargs)`` for the API (module docstring)."""
    src = PUBLIC / "inputs" / scene
    man = json.loads((src / "manifest.json").read_text())
    em = np.load(src / "ego_motion.npz")
    slots = man["slots"]
    h, w = man["img_hw"]
    cams = {}
    for s in SLOTS:
        c = man["cams"][slots[s]]
        cams[s] = {"intrinsics": c["K"], "P": c["P"], "T_ref_from_camera": c["T_ego_cam"]}
    calib = {"frame_id": "base_link", "image_size": [w, h], "cameras": cams}
    a, b = frames or (0, len(man["frames"]))
    for k in range(a, b):
        fr = man["frames"][k]
        images = [{"camera": s, "path": str((src / fr["imgs"][slots[s]]).resolve()),
                   "timestamp_s": float(fr["stamps"][s])} for s in SLOTS]
        stream = {"id": f"public-{scene}", "timestamp_s": float(fr["anchor_stamp"]),
                  "T_world_from_ego": np.asarray(em["T_map_base"][k], np.float64).tolist()}
        if k == a:
            stream["reset"] = True
        yield int(fr["frame"]), {"images": images, "calibration": calib, "stream": stream}


def research_dets(out) -> list:
    """API output -> the research ``dets`` rows (``index`` / ``layer`` / ``query`` from ``meta``)."""
    body = out.to_dict()
    meta = body.get("meta") or {}
    layers, queries = meta.get("decoder_layer") or [], meta.get("query") or []
    rows = []
    for i, d in enumerate(body["detections"]):
        layer = int(layers[i]) if i < len(layers) else -1
        query = int(queries[i]) if i < len(queries) else -1
        cx, cy, cz = d["center"]
        sx, sy, sz = d["size"]
        vx, vy = (d.get("velocity") or [0.0, 0.0])[:2]
        rows.append({"index": layer * 900 + query if layer >= 0 else -1, "layer": layer, "query": query,
                     "label": int(d["label_id"]), "score": float(d["score"]), "x": cx, "y": cy, "z": cz,
                     "dim_x": sx, "dim_y": sy, "dim_z": sz, "yaw": float(d["yaw"]), "vx": vx, "vy": vy})
    return rows


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", default="ps019", choices=["ps019", "ns0103", "ns0916"])
    ap.add_argument("--frames", default=None, help="a:b (default: all)")
    ap.add_argument("--out", default=str(ROOT / "logs" / "streampetr" / "public"))
    ap.add_argument("--weights-dir", default=None)
    a = ap.parse_args()

    from tt_streampetr import StreamPETR
    from tt_streampetr.host.agreement import detection_f1

    frames = tuple(int(x) for x in a.frames.split(":")) if a.frames else None
    out_dir = Path(a.out) / a.scene
    out_dir.mkdir(parents=True, exist_ok=True)
    gold_path = PUBLIC / "golden" / a.scene / "detections.json"
    gold = {int(r["frame"]): r["dets"] for r in json.loads(gold_path.read_text())["results"]} \
        if gold_path.is_file() else {}
    results, f1 = [], []
    with StreamPETR.from_pretrained(weights_dir=a.weights_dir) as model:
        info = model.info
        for f, req in requests(a.scene, frames):
            t0 = time.perf_counter()
            out = model(**req)
            dets = research_dets(out)
            results.append({"frame": f, "num_final": len(dets), "timing_ms": dict(out.timing_ms),
                            "wall_ms": (time.perf_counter() - t0) * 1e3, "dets": dets})
            if f in gold:
                m = detection_f1(dets, gold[f], dist=1.0)
                f1.append({"frame": f, **m})
                print(f"{a.scene} frame {f}: {len(dets)} objects, F1 {m['f1']:.3f} (tp {m['tp']}/{m['n_ref']}), "
                      f"trace {out.timing_ms.get('trace', float('nan')):.1f} ms, total {out.timing_ms['total']:.0f} ms",
                      flush=True)
    meta = {"scene": a.scene, "frames": frames, "model": info.get("model"), "device": info.get("device"),
            "weights": info.get("weights"), "knobs": info.get("knobs"),
            "pipeline": "tt_streampetr API (device graph: backbone + head + memory update in one trace)"}
    (out_dir / "detections.json").write_text(json.dumps({"meta": meta, "results": results}, indent=1, default=float))
    summary = {"frames": len(f1), "mean_f1": float(np.mean([x["f1"] for x in f1])) if f1 else None,
               "min_f1": float(np.min([x["f1"] for x in f1])) if f1 else None,
               "mean_f1_from_frame1": float(np.mean([x["f1"] for x in f1[1:]])) if len(f1) > 1 else None,
               "trace_ms_p50": float(np.median([r["timing_ms"].get("trace", np.nan) for r in results])),
               "total_ms_p50": float(np.median([r["timing_ms"]["total"] for r in results]))}
    (out_dir / "f1.json").write_text(json.dumps({"summary": summary, "frames": f1}, indent=1, default=float))
    print(json.dumps(summary, indent=1))


if __name__ == "__main__":
    main()