#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Stage bench of warm forwards: the numbers OPT_BASELINE.md, OPT_REPORT.md and the card quote.

    bin/devrun -t 1800 -- python code/scripts/bench.py --iters 100 --json logs/bevdet/bench_eth_1cq.json
    bin/devrun -t 3600 -- python code/scripts/bench.py --dispatch worker --num-cqs 2 --json ...   # the D14 matrix

Each input is a short frame sequence cycled through ONE stream (every frame pushes the BEV history ring, as in
steady operation) after 6 warm frames:

- ``synthetic[=preset]``: six 1600x900 camera frames made by ``tt_bevdet.tt.warmup.synthetic_images`` (data
  generated by us, the node's native size: no resize) with a calibration preset (default ``synthetic_6cam``, the
  shipped sample's rig; OPT_BASELINE.md of 2026-10-08 used ``pandaset_019``, which now waits in the git-ignored
  ``staging_samples_pandaset/calib``), three seeds, Autoware semantics (identity ego pose: the AlignBEV grid never
  changes, so it is never re-uploaded);
- ``pandaset[=dir]``: ``pandaset_019_f{38,39,40}.json`` from ``dir``, ``samples/`` or the git-ignored
  ``staging_samples_pandaset/samples`` (1920x1080: the node's ``cv::resize`` to 1600x900 runs on the host), Autoware
  semantics;
- ``sample=<manifest>[+<manifest>...]``: any sample manifest(s) of ``load_sample`` (e.g. the shipped sample);
- ``nuscenes=<manifest>``: key-frames of a ``research/bevdet/public_data`` manifest (nuScenes v1.0-mini, CC
  BY-NC-SA 4.0: local measurements only, never shipped), 1600x900, with the manifest's ego poses
  (``stream.T_world_from_ego``: a new AlignBEV grid is uploaded every frame).

Images are decoded once before timing (Autoware hands the node decoded images), so ``host_pre`` is the node
transport + Preprocess gather + normalisation + NHWC rows + the per-calibration cache lookups + the temporal plan.

Two loops per input (``ttaw.profiling.StageBench``; p50 / p99 / mean / min / max in ms):

1. **instrumented**: one ``model(...)`` call split into synchronized stages by wrapping the model's host hooks and the
   ``TraceRunner`` I/O: ``host_pre``; ``host_in.encode`` (numpy -> ttnn host tensors) and ``h2d.encode`` (enqueue +
   synchronize); ``trace.encode`` (one replay + synchronize); ``d2h.encode`` (device -> host copy) and
   ``host_out.encode`` (TILE -> numpy conversion); ``bevpool`` (host BEVPool + history grid + plan); the same four
   for ``decode``; ``trace.push`` / ``trace.fill`` (ring update); ``host_post`` (decode + scale-NMS + node filter);
   ``e2e_instrumented``.
2. **plain**: ``model(...)`` exactly as a client calls it (no extra synchronization): ``e2e`` and the API's own
   ``timing_ms`` (``api.preprocess`` ... ``api.total``).

Device-only, once per process (input independent): ``replay.<variant>`` one blocking replay; ``b2b.<variant>``
back-to-back replays per replay; ``b2b.frame`` = encode + decode + push back to back per frame (device-bound
throughput). AICLK / power / temperature are sampled during the loops (``AiclkSampler``), and an agreement check of
PandaSet f38-40 against the stored CPU-reference boxes (``tests/goldens``, ``host.agreement``) is run first, so every
dispatch / CQ configuration's numbers come with its accuracy.
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import subprocess
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np

HERE = Path(__file__).resolve()
CODE = HERE.parents[1]
sys.path.insert(0, str(CODE))
SAMPLES = CODE / "tt_bevdet" / "samples"
GOLDENS = CODE / "tt_bevdet" / "tests" / "goldens"
ROOT = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))
DEFAULT_NUSCENES = ROOT / "research/bevdet/public_data/nuscenes/manifests/port2hz__scene-0103.json"


# ------------------------------------------------------------------------------------------------- inputs
def synthetic_frames(cfg, n: int = 3, preset: str = "") -> List[Dict[str, Any]]:
    from tt_bevdet.host.inputs import CALIB_DIR
    from tt_bevdet.tt.warmup import WARMUP_PRESET, synthetic_images
    from tt_bevdet.ttaw import io as tio

    preset = preset or WARMUP_PRESET
    calib_dir = CALIB_DIR
    if not (CALIB_DIR / f"{preset}.json").is_file() and (CODE.parent / "staging_samples_pandaset" / "calib"
                                                          / f"{preset}.json").is_file():
        calib_dir = CODE.parent / "staging_samples_pandaset" / "calib"
    calib = tio.resolve_calibration({"preset": preset}, calib_dir)
    return [{"images": synthetic_images(cfg, seed=s + 1), "calibration": calib} for s in range(n)]


def pandaset_dir(override: str = "") -> Path:
    """Where the PandaSet 019 f38-40 manifests are: ``override``, else ``samples/``, else the git-ignored
    ``staging_samples_pandaset/samples`` of the bundle (where they wait while they may not ship)."""
    cands = [Path(override)] if override else []
    cands += [SAMPLES, CODE.parent / "staging_samples_pandaset" / "samples"]
    for c in cands:
        if all((c / f"pandaset_019_f{n}.json").is_file() for n in (38, 39, 40)):
            return c
    return cands[0]


def pandaset_frames(override: str = "") -> List[Dict[str, Any]]:
    from tt_bevdet import load_sample

    d = pandaset_dir(override)
    paths = [d / f"pandaset_019_f{n}.json" for n in (38, 39, 40)]
    if not all(p.is_file() for p in paths):
        return []
    return [load_sample(p) for p in paths]


def nuscenes_frames(manifest: Path, n: int = 3) -> List[Dict[str, Any]]:
    """The first ``n`` key-frames of a public-data manifest as decoded requests with the ego pose (no reset flag:
    the cycle is a steady stream)."""
    sys.path.insert(0, str(HERE.parent))
    from run_public_tt import frame_request
    from tt_bevdet.ttaw import io as tio

    m = json.loads(manifest.read_text())
    out = []
    for f in [f for f in m["frames"] if f.get("key_frame", True)][:n]:
        req = frame_request(m, f, "bench")
        req["images"] = {k: tio.load_image(v) for k, v in req["images"].items()}
        stream = {k: v for k, v in req.pop("stream").items() if k == "T_world_from_ego"}
        req["_stream"] = stream
        out.append(req)
    return out


def describe_frames(frames: List[Dict[str, Any]]) -> Dict[str, Any]:
    shapes = set()
    for fr in frames:
        imgs = fr["images"]
        vals = imgs.values() if isinstance(imgs, dict) else imgs
        for v in vals:
            a = getattr(v, "image", v)
            shapes.add(tuple(int(s) for s in np.asarray(a).shape))
    return {"frames": len(frames), "image_shapes": sorted(shapes)}


# ------------------------------------------------------------------------------------------- instrumentation
class Probe:
    """Splits one ``model(...)`` call into synchronized stages (see the module docstring). Installed on the model
    instance and its runner (instance attributes shadow the methods); ``active`` switches the timing on."""

    def __init__(self, model, bench):
        import ttnn

        import tt_bevdet.ttaw.trace as trace_mod

        self.model, self.bench, self.active = model, bench, False
        self._seg = "encode"
        r, dev = model.runner, model.device
        sync = lambda: ttnn.synchronize_device(dev)  # noqa: E731
        orig_collect, orig_enqueue, orig_execute, orig_read = (r._collect_uploads, r._enqueue_uploads, r._execute,
                                                               r.read)
        orig_prepare, orig_post = model._prepare, model._postprocess
        orig_to_numpy = trace_mod.to_numpy
        self._restore = [(r, "_collect_uploads"), (r, "_enqueue_uploads"), (r, "_execute"), (r, "read"),
                         (model, "_prepare"), (model, "_postprocess")]
        self._trace_mod, self._orig_to_numpy = trace_mod, orig_to_numpy
        self._conv_ms = 0.0

        def timed(name, fn, *a, **k):
            if not self.active:
                return fn(*a, **k)
            t0 = time.perf_counter()
            out = fn(*a, **k)
            self.bench.add(name, (time.perf_counter() - t0) * 1e3)
            return out

        def collect(inputs, params):
            self._seg = "encode" if inputs and ("img" in inputs or "crop" in inputs) else "decode"
            return timed(f"host_in.{self._seg}", orig_collect, inputs, params)

        def enqueue(uploads):
            def run():
                orig_enqueue(uploads)
                sync()
            return timed(f"h2d.{self._seg}", run)

        def execute(trace):
            def run():
                orig_execute(trace)
                sync()
            return timed(f"trace.{trace.variant}", run)

        def to_numpy(*a, **k):
            t0 = time.perf_counter()
            out = orig_to_numpy(*a, **k)
            self._conv_ms += (time.perf_counter() - t0) * 1e3
            return out

        def read(variant=None, **kw):
            if not self.active:
                return orig_read(variant, **kw)
            self._conv_ms = 0.0
            t0 = time.perf_counter()
            out = orig_read(variant, **kw)
            total = (time.perf_counter() - t0) * 1e3
            self.bench.add(f"d2h.{variant}", total - self._conv_ms)
            self.bench.add(f"host_out.{variant}", self._conv_ms)
            return out

        r._collect_uploads, r._enqueue_uploads, r._execute, r.read = collect, enqueue, execute, read
        model._prepare = lambda *a, **k: timed("host_pre", orig_prepare, *a, **k)
        model._postprocess = lambda *a, **k: timed("host_post", orig_post, *a, **k)
        trace_mod.to_numpy = to_numpy

    def remove(self) -> None:
        for obj, name in self._restore:
            obj.__dict__.pop(name, None)
        self._trace_mod.to_numpy = self._orig_to_numpy


# ---------------------------------------------------------------------------------------------- environment
def environment() -> Dict[str, Any]:
    env: Dict[str, Any] = {"host": platform.node(), "python": platform.python_version(), "loadavg_start": os.getloadavg(),
                           "cpus": os.cpu_count(), "omp_num_threads": os.environ.get("OMP_NUM_THREADS")}
    try:
        env["cpu_model"] = next(line.split(":", 1)[1].strip() for line in open("/proc/cpuinfo")
                                if line.startswith("model name"))
    except (OSError, StopIteration):
        pass
    gov = Path("/sys/devices/system/cpu/cpu0/cpufreq/scaling_governor")
    env["cpu_governor"] = gov.read_text().strip() if gov.is_file() else None
    tm = ROOT / "tt-metal"
    try:
        env["tt_metal_commit"] = subprocess.run(["git", "-C", str(tm), "rev-parse", "HEAD"], capture_output=True,
                                                text=True, env=dict(os.environ, GIT_OPTIONAL_LOCKS="0")).stdout.strip()
    except OSError:
        pass
    topo = tm / "tt_metal/impl/dispatch/topology.cpp"
    env["eth_patch_marker"] = topo.is_file() and "single_chip_arch_1cq_no_dispatch_s" in topo.read_text()
    try:
        from tt_bevdet.ttaw import __version__ as ttaw_version

        env["ttaw"] = ttaw_version
        env["ttaw_vendored"] = json.loads((CODE / "tt_bevdet/ttaw/VENDORED.json").read_text()).get("source_commit")
    except (ImportError, OSError, ValueError):
        pass
    try:
        env["bundle_commit"] = subprocess.run(["git", "-C", str(CODE.parent), "rev-parse", "HEAD"],
                                              capture_output=True, text=True,
                                              env=dict(os.environ, GIT_OPTIONAL_LOCKS="0")).stdout.strip()
        env["bundle_dirty"] = bool(subprocess.run(["git", "-C", str(CODE.parent), "status", "--porcelain", "code"],
                                                  capture_output=True, text=True,
                                                  env=dict(os.environ, GIT_OPTIONAL_LOCKS="0")).stdout.strip())
    except OSError:
        pass
    return env


# ------------------------------------------------------------------------------------------------- measure
def agreement_check(model) -> Dict[str, Any]:
    """PandaSet f38-40 through a fresh stream vs the stored CPU-reference boxes (the e2e gate's definition)."""
    from tt_bevdet.host.agreement import agreement, vendor_boxes

    from tt_bevdet.tests import paths

    gold = paths.pandaset_golden() or GOLDENS / paths.PANDASET_GOLDEN
    frames = pandaset_frames()
    if not frames or not gold.is_file():
        return {"skipped": "PandaSet sample or its goldens not in the bundle"}
    rec = json.loads(gold.read_text())
    res = {}
    for n, req in zip((38, 39, 40), frames):
        out = model(**req, stream={"id": "bench-agreement"}, score_threshold=0.0)
        a = agreement(vendor_boxes(out), rec["frames"][f"f{n}"]["boxes"])
        res[f"f{n}"] = {k: (round(v, 5) if isinstance(v, float) else v) for k, v in a.items() if k != "score_errs"}
    return res


def bench_input(model, name: str, frames: List[Dict[str, Any]], iters: int, warm: int) -> Dict[str, Any]:
    from tt_bevdet.ttaw.profiling import AiclkSampler, StageBench

    sid = f"bench-{name}"

    def call(i: int):
        fr = dict(frames[i % len(frames)])
        stream = dict(fr.pop("_stream", {}), id=sid)
        return model(**fr, stream=stream)

    for i in range(warm):
        call(i)
    inst = StageBench(f"{name} instrumented")
    probe = Probe(model, inst)
    try:
        probe.active = True
        with AiclkSampler(interval_s=0.05) as clk_i:
            for i in range(iters):
                t0 = time.perf_counter()
                call(warm + i)
                inst.add("e2e_instrumented", (time.perf_counter() - t0) * 1e3)
        probe.active = False
    finally:
        probe.remove()
    plain = StageBench(f"{name} plain")
    ndet = []
    with AiclkSampler(interval_s=0.05) as clk_p:
        for i in range(iters):
            t0 = time.perf_counter()
            out = call(warm + iters + i)
            plain.add("e2e", (time.perf_counter() - t0) * 1e3)
            for k, v in out.timing_ms.items():
                plain.add(f"api.{k}", v)
            ndet.append(len(out.scores))
    return {"input": name, **describe_frames(frames), "iters": iters, "warm": warm,
            "instrumented": inst.summary(), "plain": plain.summary(),
            "detections_per_frame": {"min": int(min(ndet)), "max": int(max(ndet)), "mean": float(np.mean(ndet))},
            "aiclk_instrumented": clk_i.summary(), "aiclk_plain": clk_p.summary(), "loadavg": os.getloadavg()}


def device_only(model, iters: int, b2b: int) -> Dict[str, Any]:
    """Replays only (no host work between them, no uploads): one blocking replay, back-to-back replays, and the
    frame sequence encode + decode + push back to back. The AlignBEV grid is set to the identity of a full history
    (Autoware mode, every slot a bit-exact copy); when the last input uploaded another grid (the nuScenes ego-pose
    frames), ``decode`` and the frame are measured again with that grid (``*_moving``): grid_sample's cost depends
    on the sampling positions."""
    import ttnn

    from tt_bevdet.tt.bev_fusion import RING_SLOTS, history_grid
    from tt_bevdet.ttaw.profiling import AiclkSampler, StageBench, time_b2b

    r, dev = model.runner, model.device
    sync = lambda: ttnn.synchronize_device(dev)  # noqa: E731
    slot = r._slot("grid", "param")
    last = None if slot.last_value is None else np.frombuffer(slot.last_value, np.float32).reshape(slot.shape).copy()
    identity = history_grid(np.tile(np.array([1, 0, 0, 0, 1, 0], np.float32), (RING_SLOTS, 1)),
                            list(range(RING_SLOTS)), 1, model.cfg.bev_hw)
    moving = None if last is None or np.array_equal(last, identity) else last
    sb = StageBench("device only")
    mv = "decode_move" if getattr(model.tt, "tile_ring", None) is not None else "decode"   # BEVDET_RING=tile

    after = getattr(model.tt, "after_variant", lambda v, d: v)       # push_x after an identity decode (R4.3)

    def frame(dec="decode"):
        r.replay("encode")
        r.replay(dec)
        r.replay(after("push", dec))

    with AiclkSampler(interval_s=0.05) as clk:
        r.upload(params={"grid": identity})
        sync()
        for v in ("encode", "decode", "push", "fill") + (("push_x",) if getattr(model.tt, "push_x", False) else ()):
            r.replay(v)
            sync()
            for _ in range(iters):
                t0 = time.perf_counter()
                r.replay(v)
                sync()
                sb.add(f"replay.{v}", (time.perf_counter() - t0) * 1e3)
            for _ in range(3):
                sb.add(f"b2b.{v}", time_b2b(lambda: r.replay(v), sync, n=b2b, warmup=2))
        for _ in range(3):
            sb.add("b2b.frame", time_b2b(frame, sync, n=b2b, warmup=2))
        if moving is not None:
            r.upload(params={"grid": moving})
            sync()
            for _ in range(iters):
                t0 = time.perf_counter()
                r.replay(mv)
                sync()
                sb.add("replay.decode_moving", (time.perf_counter() - t0) * 1e3)
            for _ in range(3):
                sb.add("b2b.decode_moving", time_b2b(lambda: r.replay(mv), sync, n=b2b, warmup=2))
                sb.add("b2b.frame_moving", time_b2b(lambda: frame(mv), sync, n=b2b, warmup=2))
    return {"stages_ms": sb.summary(), "aiclk": clk.summary(), "moving_grid": moving is not None}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--iters", type=int, default=100, help="timed frames per loop and input (>= 50)")
    ap.add_argument("--warm", type=int, default=6)
    ap.add_argument("--b2b", type=int, default=50, help="back-to-back replays per measurement")
    ap.add_argument("--inputs", default="synthetic,pandaset",
                    help="comma list of synthetic, pandaset[=dir], nuscenes[=manifest], sample=<manifest>[+<manifest>...] "
                         f"(default nuScenes manifest: {DEFAULT_NUSCENES})")
    ap.add_argument("--dispatch", default=None, choices=[None, "eth", "worker"])
    ap.add_argument("--num-cqs", type=int, default=None)
    ap.add_argument("--weights-dir", default=None)
    ap.add_argument("--no-agreement", action="store_true")
    ap.add_argument("--json")
    a = ap.parse_args()
    env = environment()
    from tt_bevdet import BEVDet

    t0 = time.perf_counter()
    model = BEVDet.from_pretrained(dispatch=a.dispatch, num_command_queues=a.num_cqs, weights_dir=a.weights_dir)
    res: Dict[str, Any] = {"env": env}
    try:
        res.update({"config": model.device_info, "knobs": model.knobs.as_dict(),
                    "load_s": round(time.perf_counter() - t0, 2),
                    "warmup_ms": {k: round(v, 1) for k, v in model.warmup_ms.items()},
                    "program_cache_entries": getattr(model.device, "num_program_cache_entries", lambda: None)()})
        print("config:", json.dumps(res["config"]), "| knobs:", json.dumps(res["knobs"]), flush=True)
        res["conv_placement"] = model.tt.encoder.placements()
        res["bottleneck_tails"] = model.tt.encoder.tails()
        if res["conv_placement"]:
            from collections import Counter
            print("conv placement:", dict(Counter(res["conv_placement"].values())),
                  {k: v for k, v in res["conv_placement"].items() if v != "l1"}, flush=True)
        if not a.no_agreement:
            res["agreement_pandaset_f38_40"] = agreement_check(model)
            print("agreement:", json.dumps(res["agreement_pandaset_f38_40"]), flush=True)
        inputs: List[Tuple[str, List[Dict[str, Any]]]] = []
        for spec in [s.strip() for s in a.inputs.split(",") if s.strip()]:
            if spec.split("=", 1)[0] == "synthetic":
                inputs.append(("synthetic", synthetic_frames(model.cfg, preset=spec.partition("=")[2])))
            elif spec.split("=", 1)[0] == "pandaset":
                fr = pandaset_frames(spec.split("=", 1)[1] if "=" in spec else "")
                if fr:
                    inputs.append(("pandaset", fr))
                else:
                    print("skip pandaset: pandaset_019_f38-40.json not found", flush=True)
            elif spec.startswith("sample="):
                from tt_bevdet import load_sample

                paths = [Path(x) for x in spec.split("=", 1)[1].split("+")]
                inputs.append((paths[0].stem, [load_sample(x) for x in paths]))
            elif spec.startswith("nuscenes"):
                man = Path(spec.split("=", 1)[1]) if "=" in spec else DEFAULT_NUSCENES
                if man.is_file():
                    inputs.append(("nuscenes", nuscenes_frames(man)))
                else:
                    print(f"skip nuscenes: {man} not found", flush=True)
            else:
                raise SystemExit(f"unknown input {spec!r}")
        res["inputs"] = {}
        for name, frames in inputs:
            r = bench_input(model, name, frames, a.iters, a.warm)
            res["inputs"][name] = r
            p, i = r["plain"], r["instrumented"]
            print(f"[{name}] e2e p50 {p['e2e']['p50']:.1f} / p99 {p['e2e']['p99']:.1f} ms | "
                  + " | ".join(f"{k} {v['p50']:.2f}" for k, v in i.items()), flush=True)
        res["device_only"] = device_only(model, a.iters, a.b2b)
        print("device only:", " | ".join(f"{k} p50 {v['p50']:.3f}" for k, v in
                                          res["device_only"]["stages_ms"].items()), flush=True)
        res["trace"] = model.runner.describe()
        res["env"]["loadavg_end"] = os.getloadavg()
    finally:
        model.close()
    if a.json:
        Path(a.json).parent.mkdir(parents=True, exist_ok=True)
        Path(a.json).write_text(json.dumps(res, indent=1, default=str) + "\n")
        print("wrote", a.json)
    return 0


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