#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Stage bench of warm frames: 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/streampetr/baseline/bench_eth-1cq.json
    bin/devrun -t 1800 -- 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 (the temporal memory is carried from frame to frame,
as in steady operation) after ``--warm`` warm frames; the time stamps advance monotonically across cycles:

- ``synthetic``: five 1920x1080 camera frames generated by us (seeded smooth colour fields + noise, three seeds) with
  the calibration preset ``--synthetic-preset`` (default ``pandaset_019``: 1920x1080, from the git-ignored
  ``staging_samples_pandaset/calib`` while it does not ship; the OPT_BASELINE input), 10 Hz
  stamps, no ego pose (identity: a static ego);
- ``pandaset[=dir]``: ``pandaset_019_f{40..43}.json`` (``samples/``, else the git-ignored
  ``staging_samples_pandaset/samples``, else ``dir``): PandaSet 019 (CC BY 4.0), 1920x1080 JPEGs, real stamps and
  ego poses (each cycle shifts the stamps by 0.4 s; the pose jumps back once per cycle);
- ``sample=<manifest>[+<manifest>...]``: any sample manifests of ``load_sample``.

Images are decoded before timing (Autoware hands the node decoded images), so ``host_pre`` is the request parsing +
the C14 triangle resize / normalisation of the five cameras + the per-calibration cache lookup + the stream 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``; ``img_rows`` (NCHW float32 -> NHWC rows); ``host_in`` (numpy -> ttnn host
   tensors: the bf16 image rows, the ages, the changed params); ``h2d`` (enqueue + synchronize); ``trace`` (one
   replay + synchronize); ``d2h`` (device -> host copy) and ``host_out`` (TILE -> numpy); ``host_post`` (decode of
   the 6 layers + IoU-BEV NMS + Detections3D); ``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.frame`` one blocking replay; ``b2b.frame`` back-to-back
replays per replay (device-bound throughput). AICLK / power / temperature are sampled during the loops
(``AiclkSampler``), and the agreement checks (the shipped synthetic sample and the staged PandaSet frame 40, cold
starts, against their stored CPU-reference bodies with the container smoke's ``REFERENCE_GATES``) run first, so every dispatch / CQ configuration's numbers come
with its accuracy.
"""
from __future__ import annotations

import argparse
import copy
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_streampetr" / "samples"
ROOT = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))
CAMS = ("CAM_FRONT", "CAM_FRONT_LEFT", "CAM_BACK_LEFT", "CAM_FRONT_RIGHT", "CAM_BACK_RIGHT")


# ------------------------------------------------------------------------------------------------- inputs
def synthetic_image(seed: int, cam: int, hw: Tuple[int, int] = (1080, 1920)) -> np.ndarray:
    """A deterministic 8-bit RGB camera frame made by us: a smooth colour field (low-frequency sines) + noise."""
    rng = np.random.default_rng(1000 * seed + cam)
    h, w = hw
    y = np.linspace(0, 1, h, dtype=np.float32)[:, None]
    x = np.linspace(0, 1, w, dtype=np.float32)[None, :]
    planes = []
    for c in range(3):
        fx, fy, ph = rng.uniform(1, 6), rng.uniform(1, 6), rng.uniform(0, 6.28)
        planes.append(127 + 90 * np.sin(6.28 * fx * x + ph) * np.cos(6.28 * fy * y + ph / 2))
    img = np.stack(planes, -1) + rng.normal(0, 12, (h, w, 3)).astype(np.float32)
    return np.clip(img, 0, 255).astype(np.uint8)


def synthetic_frames(preset: str, n: int = 3) -> List[Dict[str, Any]]:
    """Seeded colour-field frames at the image size of the calibration preset ``preset`` (the package's ``calib/``,
    else the git-ignored ``staging_samples_pandaset/calib``; sent inline)."""
    from tt_streampetr.host.inputs import CALIB_DIR

    path = CALIB_DIR / f"{preset}.json"
    if not path.is_file():
        path = CODE.parent / "staging_samples_pandaset" / "calib" / f"{preset}.json"
    cal = json.loads(path.read_text())
    calib = {k: v for k, v in cal.items() if k != "_doc"}
    w, h = cal.get("image_size", [1920, 1080])
    frames = []
    for s in range(n):
        frames.append({"images": [{"camera": c, "image": synthetic_image(s + 1, i, (h, w))}
                                  for i, c in enumerate(CAMS)],
                       "calibration": calib, "_dt": 0.1})
    return frames


def pandaset_dir(override: str = "") -> Path:
    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 (40, 41, 42, 43)):
            return c
    return cands[0]


def decoded(req: Dict[str, Any]) -> Dict[str, Any]:
    from tt_streampetr.ttaw import io as tio

    out = dict(req)
    out["images"] = [{**{k: v for k, v in e.items() if k != "path"}, "image": tio.load_image(e["path"])}
                     for e in req["images"]]
    return out


def manifest_frames(paths: List[Path]) -> List[Dict[str, Any]]:
    from tt_streampetr import load_sample

    return [decoded(load_sample(p)) for p in paths]


def describe_frames(frames: List[Dict[str, Any]]) -> Dict[str, Any]:
    shapes = {tuple(int(s) for s in np.asarray(e["image"]).shape) for fr in frames for e in fr["images"]}
    return {"frames": len(frames), "image_shapes": sorted(shapes)}


class Sequencer:
    """Turns a short frame list into an endless stream: stamps shifted per cycle so time never goes backwards."""

    def __init__(self, frames: List[Dict[str, Any]], sid: str, t0: float = 1.0e9):
        self.frames, self.sid, self.t0 = frames, sid, t0
        st = [f.get("stream", {}).get("timestamp_s") for f in frames]
        self.real = all(s is not None for s in st)
        self.period = (max(st) - min(st) + 0.1) if self.real else 0.1 * len(frames)
        self.base = min(st) if self.real else 0.0

    def __call__(self, i: int) -> Dict[str, Any]:
        fr = self.frames[i % len(self.frames)]
        cyc = i // len(self.frames)
        req = {k: v for k, v in fr.items() if not k.startswith("_")}
        stream = dict(fr.get("stream") or {})
        stream.pop("reset", None)
        if self.real:
            shift = cyc * self.period
            stream["timestamp_s"] = float(stream["timestamp_s"]) + shift
            req["images"] = [dict(e, timestamp_s=float(e["timestamp_s"]) + shift) if "timestamp_s" in e else e
                             for e in fr["images"]]
        else:
            stream["timestamp_s"] = self.t0 + 0.1 * i
        stream.setdefault("T_world_from_ego", np.eye(4).tolist())   # a static ego (a pose is required per frame)
        stream["id"] = self.sid
        if i == 0:
            stream["reset"] = True
        req["stream"] = stream
        return req


# ------------------------------------------------------------------------------------------- instrumentation
class Probe:
    """Splits one ``model(...)`` call into synchronized stages (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_streampetr.tt.model as tt_model
        import tt_streampetr.ttaw.trace as trace_mod

        self.model, self.bench, self.active = model, bench, False
        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
        self._trace_mod, self._orig_to_numpy = trace_mod, trace_mod.to_numpy
        self._tt_model, self._orig_img_rows = tt_model, tt_model.img_rows
        self._restore = [(r, "_collect_uploads"), (r, "_enqueue_uploads"), (r, "_execute"), (r, "read"),
                         (model, "_prepare"), (model, "_postprocess")]
        self._conv_ms = 0.0
        orig_to_numpy, orig_img_rows = self._orig_to_numpy, self._orig_img_rows

        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 enqueue(uploads):
            def run():
                orig_enqueue(uploads)
                sync()
            return timed("h2d", run)

        def execute(trace):
            def run():
                orig_execute(trace)
                sync()
            return timed("trace", 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("d2h", total - self._conv_ms)
            self.bench.add("host_out", self._conv_ms)
            return out

        r._collect_uploads = lambda inputs, params: timed("host_in", orig_collect, inputs, params)
        r._enqueue_uploads, r._execute, r.read = 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
        tt_model.img_rows = lambda *a, **k: timed("img_rows", orig_img_rows, *a, **k)

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


# ---------------------------------------------------------------------------------------------- environment
def _git(*args: str) -> str:
    return subprocess.run(["git", *args], capture_output=True, text=True,
                          env=dict(os.environ, GIT_OPTIONAL_LOCKS="0")).stdout.strip()


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"),
                           "trace_alloc_tracking": os.environ.get("TT_METAL_TRACE_ALLOC_TRACKING")}
    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"
    env["tt_metal_commit"] = _git("-C", str(tm), "rev-parse", "HEAD")
    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_streampetr.ttaw import __version__ as ttaw_version

        env["ttaw"] = ttaw_version
        env["ttaw_vendored"] = json.loads((CODE / "tt_streampetr/ttaw/VENDORED.json").read_text()).get("source_commit")
    except (ImportError, OSError, ValueError):
        pass
    env["bundle_commit"] = _git("-C", str(CODE.parent), "rev-parse", "HEAD")
    env["bundle_dirty"] = bool(_git("-C", str(CODE.parent), "status", "--porcelain", "code"))
    return env


# ------------------------------------------------------------------------------------------------- measure
def agreement_check(model, sample: Path = None) -> Dict[str, Any]:
    """A cold-start sample frame vs its stored CPU-reference body (``tests/test_e2e_device.py``): the shipped synthetic
    sample by default, or the staged PandaSet frame 40 (``pandaset_dir() / pandaset_019_f40.json``)."""
    from tt_streampetr import load_sample
    from tt_streampetr.server.smoke_test import REFERENCE_GATES
    from tt_streampetr.ttaw.server import smoke

    sample = Path(sample) if sample is not None else SAMPLES / "synthetic_5cam.json"
    ref_path = sample.with_suffix(".reference.json")
    if not ref_path.is_file():
        return {"skipped": f"{ref_path} not found"}
    kw = load_sample(sample)
    kw["stream"] = dict(kw.get("stream") or {}, id="bench-agreement", reset=True)
    out = model(**kw)
    body = json.loads(json.dumps(out.to_dict()))
    metrics, fails = smoke.compare_with_reference(body, json.loads(ref_path.read_text()), gates=REFERENCE_GATES)
    d = metrics["detections"]
    return {"recall": d.get("recall"), "precision": d.get("precision"), "score_err_max": d.get("score_err_max"),
            "num": body.get("num_detections"), "failures": fails}


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

    seq = Sequencer(frames, f"bench-{name}")
    for i in range(warm):
        model(**seq(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):
                req = seq(warm + i)
                t0 = time.perf_counter()
                model(**req)
                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):
            req = seq(warm + iters + i)
            t0 = time.perf_counter()
            out = model(**req)
            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 and back-to-back replays of the
    ``frame`` trace on the inputs / state the last frame left (the shapes are fixed: the time does not depend on
    the values)."""
    import ttnn

    from tt_streampetr.ttaw.profiling import AiclkSampler, StageBench, time_b2b

    r, dev = model.runner, model.device
    sync = lambda: ttnn.synchronize_device(dev)  # noqa: E731
    sb = StageBench("device only")
    with AiclkSampler(interval_s=0.05) as clk:
        r.replay("frame")
        sync()
        for _ in range(iters):
            t0 = time.perf_counter()
            r.replay("frame")
            sync()
            sb.add("replay.frame", (time.perf_counter() - t0) * 1e3)
        for _ in range(3):
            sb.add("b2b.frame", time_b2b(lambda: r.replay("frame"), sync, n=b2b, warmup=2))
    return {"stages_ms": sb.summary(), "aiclk": clk.summary()}


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=30, help="back-to-back replays per measurement")
    ap.add_argument("--inputs", default="synthetic,pandaset",
                    help="comma list of synthetic, pandaset[=dir], sample=<manifest>[+<manifest>...]")
    ap.add_argument("--synthetic-preset", default="pandaset_019")
    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_streampetr import StreamPETR

    t0 = time.perf_counter()
    model = StreamPETR.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)
        if not a.no_agreement:
            res["agreement_shipped"] = agreement_check(model)
            print("agreement (shipped synthetic sample):", json.dumps(res["agreement_shipped"]), flush=True)
            res["agreement_sample_f40"] = agreement_check(model, pandaset_dir() / "pandaset_019_f40.json")
            print("agreement (staged PandaSet f40):", json.dumps(res["agreement_sample_f40"]), flush=True)
        inputs: List[Tuple[str, List[Dict[str, Any]]]] = []
        for spec in [s.strip() for s in a.inputs.split(",") if s.strip()]:
            kind, _, arg = spec.partition("=")
            if kind == "synthetic":
                inputs.append(("synthetic", synthetic_frames(arg or a.synthetic_preset)))
            elif kind == "pandaset":
                d = pandaset_dir(arg)
                paths = [d / f"pandaset_019_f{n}.json" for n in (40, 41, 42, 43)]
                if all(p.is_file() for p in paths):
                    inputs.append(("pandaset", manifest_frames(paths)))
                else:
                    print("skip pandaset: pandaset_019_f40-43.json not found", flush=True)
            elif kind == "sample":
                paths = [Path(x) for x in arg.split("+")]
                inputs.append((paths[0].stem, manifest_frames(paths)))
            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()
        bb = getattr(getattr(model, "tt", None), "backbone", None)
        if bb is not None and hasattr(bb, "placements"):
            pl = bb.placements()
            res["conv_placements"] = pl
            if pl:
                from collections import Counter
                print("conv placements (CONV_L1):", dict(Counter(pl.values())),
                      "| not on L1:", [k for k, v in pl.items() if not v.startswith("l1")], flush=True)
        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())