#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Run public-data manifests (research/bevdet/public_data, schema ``bevdet_public_manifest/1``) through the TT model's
Python API and write the results in the CPU reference's format, for ``eval_nuscenes.py`` (paper-style mAP / NDS,
indicative only) and for agreement with the CPU results of the same manifests.

    bin/devrun -t 1800 -- python code/scripts/run_public_tt.py \\
        research/bevdet/public_data/nuscenes/manifests/port2hz__scene-0103.json ... --tag tt
    tools/research-venv/bin/python -I research/bevdet/public_data/scripts/eval_nuscenes.py --modes port2hz \\
        --results-dir research/bevdet/public_data/nuscenes/results/tt --backend-tag tt --skip-gt

Every frame of a manifest goes through ``BEVDet.__call__`` exactly like a client request: the six JPEGs of the frame
(``image_root`` + relative paths), its calibration (K of the image as sent, camera -> ego R / t), and the stream
(one fresh stream per manifest; ``reset`` on a flag-0 frame; the ego pose ``ego2global`` for the ego-motion mode
``port2hz``, none for the Autoware modes, whose AlignBEV is the identity; ``single`` / ``single_quirk`` are one reset
frame, i.e. no history, with the scaled K or the node's raw-K quirk). ``score_threshold=0`` returns every vendor box
(score > 0.1), as the CPU results hold. Output: ``<dataset>/results/<tag>/<manifest>.json`` (or ``--out-dir``; boxes
in the vendor convention: bottom-centre z) and a per-manifest agreement summary against ``<dataset>/results/
<manifest>.json`` (the CPU run) at score >= 0.2 (``tests/test_e2e_device.py`` definition).

nuScenes data are CC BY-NC-SA 4.0: local validation only, never shipped (PLAN.md D17 / D18).
"""
from __future__ import annotations

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

import numpy as np

HERE = Path(__file__).resolve()
sys.path.insert(0, str(HERE.parents[1]))

EGO_MOTION_MODES = ("port2hz", "train")
RUN_MODES = ("port2hz", "autoware", "autoware_kf", "single", "single_quirk")


def frame_request(m, f, stream_id):
    from tt_bevdet.reference.config import CAMERA_ORDER

    root = Path(m["image_root"])
    cams = {}
    for name, k, rot, t in zip(CAMERA_ORDER, f["K"], f["cam2ego_R"], f["cam2ego_t"]):
        T = np.eye(4)
        T[:3, :3] = np.asarray(rot, np.float64)
        T[:3, 3] = np.asarray(t, np.float64)
        cams[name] = {"intrinsics": k, "T_ref_from_camera": T.tolist()}
    stream = {"id": stream_id, "reset": int(f["flag"]) == 0}
    if m["mode"] in EGO_MOTION_MODES:
        stream["T_world_from_ego"] = f["ego2global"]
    images = {name: str(root / rel) for name, rel in zip(CAMERA_ORDER, f["images"])}
    # the manifests' K is the intrinsic of the image the network sees (1600x900: README of public_data, section 2),
    # so the API must not scale it again when it resizes PandaSet's 1920x1080 images like the node
    return {"images": images, "calibration": {"cameras": cams, "scale_intrinsics": False}, "stream": stream}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("manifests", nargs="+", type=Path)
    ap.add_argument("--tag", default="tt")
    ap.add_argument("--out-dir", type=Path, default=None,
                    help="write <out-dir>/<dataset>/<manifest>.json instead of <dataset>/results/<tag>/")
    ap.add_argument("--weights-dir", default=None)
    ap.add_argument("--threads", type=int, default=4)
    a = ap.parse_args()
    os.environ.setdefault("OMP_NUM_THREADS", str(a.threads))
    import torch

    torch.set_num_threads(a.threads)
    from tt_bevdet import BEVDet
    from tt_bevdet.host.agreement import agreement, vendor_boxes

    model = BEVDet.from_pretrained(weights_dir=a.weights_dir)
    summary = {}
    try:
        for man in a.manifests:
            m = json.loads(man.read_text())
            if m["mode"] not in RUN_MODES:
                print(f"skip {man.name}: mode {m['mode']} (per-frame time-compensated calibration and BEVDet's "
                      "history padding are not the port's semantics)")
                continue
            ds_dir = man.parents[1]
            out_dir = (a.out_dir / ds_dir.name) if a.out_dir else ds_dir / "results" / a.tag
            out_dir.mkdir(parents=True, exist_ok=True)
            cpu = ds_dir / "results" / man.name
            ref = {fr["name"]: fr for fr in json.loads(cpu.read_text())["frames"]} if cpu.is_file() else {}
            frames, totals, t0 = [], {"ref": 0, "dev": 0, "recalled": 0, "precise": 0}, time.time()
            sid = f"{a.tag}-{man.stem}"
            for f in m["frames"]:
                out = model(**frame_request(m, f, sid), score_threshold=0.0)
                boxes = vendor_boxes(out)
                rec = {k: f[k] for k in ("i", "name", "key_frame", "timestamp_us", "ego2global", "flag") if k in f}
                rec.update({"sample_token": f.get("sample_token"), "boxes": boxes, "n_boxes": len(boxes),
                            "n_boxes_ge_0.2": sum(b["score"] >= 0.2 for b in boxes),
                            "timing_ms": {k: round(v, 2) for k, v in out.timing_ms.items()}})
                frames.append(rec)
                if "boxes" in ref.get(f["name"], {}):           # the CPU run decodes every frame, or key-frames only
                    ag = agreement(boxes, ref[f["name"]]["boxes"])
                    for k in totals:
                        totals[k] += ag[k]
            res = {"manifest": str(man), "mode": m["mode"], "sequence": m["sequence"], "dataset": m["dataset"],
                   "backend": f"tt-{model.MODEL_NAME}", "n_frames": len(frames), "complete": True,
                   "device": model.device_info, "precision": model.knobs.as_dict()["PRECISION"],
                   "frames": frames}
            (out_dir / man.name).write_text(json.dumps(res))
            s = dict(totals)
            if totals["ref"]:
                s["recall"] = totals["recalled"] / totals["ref"]
                s["precision"] = totals["precise"] / max(totals["dev"], 1)
            s["frames"] = len(frames)
            s["s_per_frame"] = round((time.time() - t0) / max(len(frames), 1), 3)
            summary[man.name] = s
            print(f"{man.name}: {json.dumps(s)}", flush=True)
    finally:
        model.close()
    print(json.dumps(summary, indent=1))
    return 0


if __name__ == "__main__":
    sys.exit(main())