#!/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/bevformer/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 (every frame reads and writes the device's
``prev_bev`` history, as in steady operation; no timestamps, so no history reset) after ``--warm`` frames:

- ``synthetic[=preset]``: six 1600x900 camera images made by ``tt_bevformer.tt.warmup.synthetic_images`` (data
  generated by us, the node's native size: no host resize), seeds 1-3, a calibration preset from ``calib/`` or the
  git-ignored ``staging_samples_pandaset/calib`` (default ``synthetic_6cam``, the shipped rig; OPT_BASELINE.md
  measured ``synthetic=pandaset_019``), identity ego pose;
- ``pandaset[=dir]``: ``pandaset_019_f{30,35,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; each
  frame has its own time-compensated calibration, so the per-calibration SCA tables are re-uploaded every frame);
- ``pandaset_f30``: frame 30 alone (one static calibration: the tables are uploaded once);
- ``sample=<manifest>[+<manifest>...]``: any sample manifest(s) of ``load_sample``, e.g. the shipped
  ``code/tt_bevformer/samples/synthetic_6cam.json`` (its stored reference is checked).

Images are decoded once before timing (Autoware hands the node decoded images); ``decode`` is measured separately
(PIL, the server's path). 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`` (all of ``_prepare``) with ``pre.request`` (request parsing, RGB -> BGR,
   lidar2img), ``pre.node_resize`` (the node's uint8 ``cv::resize`` emulation, 6 cameras), ``pre.c14`` (C14
   ``bevformer_preprocess``: normalise + x0.8 float resize + pad), ``pre.geometry`` (per-calibration projection and
   SCA tables, cached), ``pre.rotation_map``, ``pre.frame_inputs`` (NCHW -> NHWC rows); ``host_in`` (numpy -> ttnn
   host tensors) and ``h2d`` (enqueue + synchronize); ``trace.image`` / ``trace.bev`` (one replay + synchronize);
   ``d2h`` (device -> host copy) and ``host_out`` (TILE -> numpy); ``unpack``; ``host_post``; ``e2e_instrumented``.
2. **plain**: ``model(...)`` exactly as a client calls it: ``e2e`` and the API's own ``timing_ms`` (``api.*``).

Device-only, once per process (input independent): ``replay.<variant>`` one blocking replay; ``b2b.<variant>``
back-to-back replays of one variant; ``b2b.frame`` = ``image`` + ``bev`` back to back (the device time per frame).
AICLK / power / temperature are sampled during the loops (``AiclkSampler``). An agreement check of every input
manifest that has a stored CPU reference next to it (``<stem>.reference.json``: first frame, fresh stream) runs
first with the container-smoke gates, so every 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, Optional, Tuple

import numpy as np

HERE = Path(__file__).resolve()
CODE = HERE.parents[1]
sys.path.insert(0, str(CODE))
SAMPLES = CODE / "tt_bevformer" / "samples"
STAGING = CODE.parent / "staging_samples_pandaset"
ROOT = Path(os.environ.get("TT_MODELS_ROOT", "/home/ubuntu/experiments/tt-models"))


# ------------------------------------------------------------------------------------------------- inputs
def calib_dir_for(preset: str) -> Path:
    from tt_bevformer.host.inputs import CALIB_DIR

    for d in (CALIB_DIR, STAGING / "calib"):
        if (d / f"{preset}.json").is_file():
            return d
    return CALIB_DIR


def synthetic_frames(preset: str = "") -> List[Dict[str, Any]]:
    from tt_bevformer.tt.warmup import synthetic_images
    from tt_bevformer.ttaw import io as tio

    preset = preset or "synthetic_6cam"
    calib = tio.resolve_calibration({"preset": preset}, calib_dir_for(preset))
    return [{"images": synthetic_images(seed=s + 1), "calibration": calib} for s in range(3)]


def pandaset_dir(override: str = "") -> Path:
    cands = [Path(override)] if override else []
    cands += [SAMPLES, STAGING / "samples", STAGING]
    for c in cands:
        if all((c / f"pandaset_019_f{n}.json").is_file() for n in (30, 35, 40)):
            return c
    return cands[0]


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

    out = []
    for p in paths:
        req = load_sample(p)
        st = req.pop("stream", {}) or {}
        req["_stream"] = {k: v for k, v in st.items() if k == "T_world_from_ego"}
        req["_manifest"] = str(p)
        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"]
        for v in (imgs.values() if isinstance(imgs, dict) else imgs):
            shapes.add(tuple(int(s) for s in np.asarray(getattr(v, "image", v)).shape))
    return {"frames": len(frames), "image_shapes": sorted(shapes),
            "manifests": [fr.get("_manifest") for fr in frames if fr.get("_manifest")]}


def decode_ms(manifest: Path, reps: int = 5) -> Dict[str, float]:
    """JPEG decode of the six camera files of a manifest (``ttaw.io.load_image``: PIL, the server's path)."""
    from tt_bevformer.ttaw import io as tio

    spec = json.loads(Path(manifest).read_text())
    files = [Path(manifest).parent / rel for rel in spec["images"].values()]
    ts = []
    for _ in range(reps):
        t0 = time.perf_counter()
        for f in files:
            tio.load_image(str(f))
        ts.append((time.perf_counter() - t0) * 1e3)
    return {"p50": float(np.median(ts)), "min": float(min(ts)), "n": reps,
            "bytes": int(sum(f.stat().st_size for f in files))}


# ------------------------------------------------------------------------------------------- instrumentation
class Probe:
    """Splits one ``model(...)`` call into synchronized stages (module docstring). Installed on the model instance,
    its runner and the host modules ``_prepare`` reads at call time; ``active`` switches the timing on."""

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

        import tt_bevformer.host.geometry as hgeo
        import tt_bevformer.host.inputs as hin
        import tt_bevformer.host.preprocess as hpre
        import tt_bevformer.tt.model as tmodel
        import tt_bevformer.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
        self._acc: Dict[str, float] = {}
        self._inst: List[Tuple[Any, str]] = []        # instance attributes to drop
        self._mods: List[Tuple[Any, str, Any]] = []   # module attributes to restore

        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 summed(name, fn):
            def wrapper(*a, **k):
                if not self.active:
                    return fn(*a, **k)
                t0 = time.perf_counter()
                out = fn(*a, **k)
                self._acc[name] = self._acc.get(name, 0.0) + (time.perf_counter() - t0) * 1e3
                return out
            return wrapper

        def patch_inst(obj, name, fn):
            setattr(obj, name, fn)
            self._inst.append((obj, name))

        def patch_mod(mod, name, fn):
            self._mods.append((mod, name, getattr(mod, name)))
            setattr(mod, name, fn)

        orig_collect, orig_enqueue, orig_execute, orig_read = r._collect_uploads, r._enqueue_uploads, r._execute, r.read
        orig_prepare, orig_post, orig_geo, orig_unpack = model._prepare, model._postprocess, model._geometry_for, \
            model.tt.unpack

        def prepare(*a, **k):
            if not self.active:
                return orig_prepare(*a, **k)
            self._acc = {}
            out = timed("host_pre", orig_prepare, *a, **k)
            for name in ("pre.request", "pre.node_resize", "pre.c14", "pre.geometry", "pre.rotation_map",
                         "pre.frame_inputs"):
                self.bench.add(name, self._acc.get(name, 0.0))
            return out

        def enqueue(uploads):
            def run():
                orig_enqueue(uploads)
                sync()
            if self.active:
                self.bench.add("h2d.tensors", float(len(uploads)))
            return timed("h2d", run)

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

        orig_to_numpy = trace_mod.to_numpy

        def to_numpy(*a, **k):
            t0 = time.perf_counter()
            out = orig_to_numpy(*a, **k)
            self._acc["host_out"] = self._acc.get("host_out", 0.0) + (time.perf_counter() - t0) * 1e3
            return out

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

        patch_inst(r, "_collect_uploads", lambda i, p: timed("host_in", orig_collect, i, p))
        patch_inst(r, "_enqueue_uploads", enqueue)
        patch_inst(r, "_execute", execute)
        patch_inst(r, "read", read)
        patch_inst(model, "_prepare", prepare)
        patch_inst(model, "_postprocess", lambda *a, **k: timed("host_post", orig_post, *a, **k))
        patch_inst(model, "_geometry_for", summed("pre.geometry", orig_geo))
        patch_inst(model.tt, "unpack", lambda *a, **k: timed("unpack", orig_unpack, *a, **k))
        patch_mod(trace_mod, "to_numpy", to_numpy)
        patch_mod(hin, "prepare_request", summed("pre.request", hin.prepare_request))
        patch_mod(hpre, "node_transport", summed("pre.node_resize", hpre.node_transport))
        patch_mod(hpre, "bevformer_preprocess", summed("pre.c14", hpre.bevformer_preprocess))
        patch_mod(hgeo, "rotation_index_map", summed("pre.rotation_map", hgeo.rotation_index_map))
        patch_mod(tmodel, "frame_inputs", summed("pre.frame_inputs", tmodel.frame_inputs))

    def remove(self) -> None:
        for obj, name in self._inst:
            obj.__dict__.pop(name, None)
        for mod, name, orig in reversed(self._mods):
            setattr(mod, name, orig)


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


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"),
                           "env": {k: v for k, v in os.environ.items() if k.startswith("BEVFORMER_")}}
    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
    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:
        vend = json.loads((CODE / "tt_bevformer/ttaw/VENDORED.json").read_text())
        env["ttaw"] = vend.get("version")
        env["ttaw_vendored"] = vend.get("source_commit")
    except (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 find_reference(manifest: Path) -> Optional[Path]:
    for d in (manifest.parent, STAGING / "samples", STAGING):
        p = d / f"{manifest.stem}.reference.json"
        if p.is_file():
            return p
    return None


def agreement_check(model, manifests: List[Path]) -> Dict[str, Any]:
    """Each manifest with a stored CPU reference: its frame as a fresh stream through ``model(...)``, compared with
    the container-smoke gates (``server/smoke_test.py`` ``REFERENCE_GATES``)."""
    from tt_bevformer import load_sample
    from tt_bevformer.server.smoke_test import REFERENCE_GATES
    from tt_bevformer.ttaw.server import smoke as ttaw_smoke

    res: Dict[str, Any] = {}
    for m in manifests:
        ref = find_reference(m)
        if ref is None:
            continue
        req = load_sample(m)
        req["stream"] = dict(req.get("stream") or {}, id=f"bench-agreement-{m.stem}", reset=True)
        out = model(**req)
        metrics, fails = ttaw_smoke.compare_with_reference(out.to_dict(), ttaw_smoke.load_json(ref),
                                                           gates=REFERENCE_GATES)
        res[m.stem] = {"reference": str(ref), "metrics": metrics, "fails": fails,
                       "summary": ttaw_smoke.describe_metrics(metrics), "num_detections": len(out.scores)}
        model.streams.forget(f"bench-agreement-{m.stem}")
    return res


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

    sid = f"bench-{name}"

    def call(i: int):
        fr = {k: v for k, v in frames[i % len(frames)].items() if not k.startswith("_")}
        stream = dict(frames[i % len(frames)].get("_stream", {}), id=sid)
        return model(**fr, stream=stream)

    for i in range(warm):
        call(i)
    res: Dict[str, Any] = {"input": name, **describe_frames(frames), "iters": iters, "warm": warm}
    if "instrumented" in loops:
        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()
        res.update({"instrumented": inst.summary(), "aiclk_instrumented": clk_i.summary()})
    if "plain" in loops:
        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))
        res.update({"plain": plain.summary(), "aiclk_plain": clk_p.summary(),
                    "detections_per_frame": {"min": int(min(ndet)), "max": int(max(ndet)),
                                             "mean": float(np.mean(ndet))}})
    res["loadavg_end"] = os.getloadavg()
    model.streams.forget(sid)
    return res


def device_only(model, iters: int, b2b: int) -> Dict[str, Any]:
    """Replays only (no uploads, no host work between them) on the device state the last frame left."""
    import ttnn

    from tt_bevformer.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")

    def frame():
        r.replay("image")
        r.replay("bev")

    with AiclkSampler(interval_s=0.05) as clk:
        for v in ("image", "bev"):
            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))
    return {"stages_ms": sb.summary(), "aiclk": clk.summary()}


def transfers(model) -> Dict[str, Any]:
    """Bytes and page sizes of every persistent input / param / state of the runner (what one frame uploads: the
    ``img`` rows, ``rot_index`` and the changed params every frame, the SCA tables when the calibration changes)."""
    desc = model.runner.describe()
    size = {"bfloat16": 2, "float32": 4, "uint32": 4, "int32": 4, "uint16": 2}
    out: Dict[str, Any] = {}
    for kind in ("inputs", "params", "states"):
        for name, d in (desc.get(kind) or {}).items():
            shape = [int(s) for s in d["shape"]]
            el = size.get(str(d.get("dtype")), 4)
            if d.get("layout") == "TILE":
                padded = shape[:-2] + [-(-shape[-2] // 32) * 32, -(-shape[-1] // 32) * 32]
                nbytes, page = int(np.prod(padded)) * el, 32 * 32 * el
            else:
                nbytes, page = int(np.prod(shape)) * el, shape[-1] * el
            out[f"{kind}.{name}"] = {**d, "bytes": nbytes, "page_bytes": page, "pages": nbytes // page}
    return out


def cache_entries(device) -> Optional[int]:
    fn = getattr(device, "num_program_cache_entries", None)
    return int(fn()) if fn else 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 for the record)")
    ap.add_argument("--warm", type=int, default=4)
    ap.add_argument("--b2b", type=int, default=20, help="back-to-back replays per measurement (3 rounds)")
    ap.add_argument("--inputs", default="synthetic,pandaset_f30,pandaset",
                    help="comma list of synthetic[=preset], pandaset[=dir], pandaset_f30, "
                         "sample=<manifest>[+<manifest>...]")
    ap.add_argument("--loops", default="instrumented,plain", help="instrumented and / or plain")
    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("--device-iters", type=int, default=20)
    ap.add_argument("--json")
    a = ap.parse_args()
    env = environment()
    from tt_bevformer import BEVFormer

    loops = tuple(s.strip() for s in a.loops.split(",") if s.strip())
    t0 = time.perf_counter()
    model = BEVFormer.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": cache_entries(model.device),
                    "trace": model.runner.describe(), "transfers": transfers(model)})
        print("config:", json.dumps(res["config"]), "| knobs:", json.dumps(res["knobs"]),
              "| programs:", res["program_cache_entries"], 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)))
            elif kind in ("pandaset", "pandaset_f30"):
                d = pandaset_dir(arg)
                nums = (30,) if kind == "pandaset_f30" else (30, 35, 40)
                paths = [d / f"pandaset_019_f{n}.json" for n in nums]
                if all(p.is_file() for p in paths):
                    inputs.append((kind, manifest_frames(paths)))
                else:
                    print(f"skip {kind}: manifests not found in {d}", flush=True)
            elif kind == "sample":
                paths = [Path(p) for p in arg.split("+")]
                inputs.append((f"sample:{paths[0].stem}", manifest_frames(paths)))
            else:
                raise SystemExit(f"unknown input {spec!r}")
        manifests = sorted({Path(m) for _, fr in inputs for m in describe_frames(fr)["manifests"]})
        res["decode_ms"] = {m.stem: decode_ms(m) for m in manifests}
        if not a.no_agreement:
            res["agreement"] = agreement_check(model, manifests)
            print("agreement:", json.dumps({k: v["summary"] + (" FAIL " + str(v["fails"]) if v["fails"] else "")
                                            for k, v in res["agreement"].items()}), flush=True)
        res["inputs"] = []
        for name, frames in inputs:
            r = bench_input(model, name, frames, a.iters, a.warm, loops)
            res["inputs"].append(r)
            line = {k: round(v["p50"], 2) for k, v in {**r.get("instrumented", {}), **r.get("plain", {})}.items()}
            print(f"[{name}]", json.dumps(line), flush=True)
        res["device_only"] = device_only(model, a.device_iters, a.b2b)
        print("device only:", json.dumps({k: round(v["p50"], 3) for k, v in res["device_only"]["stages_ms"].items()}),
              "| aiclk:", json.dumps(res["device_only"]["aiclk"]), flush=True)
        res["program_cache_entries_end"] = cache_entries(model.device)
    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")
    return 0


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