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

    bin/devrun -t 1800 -- python code/scripts/bench.py --iters 100 --json out.json
    bin/devrun -t 1800 -- python code/scripts/bench.py --input a.npz --input b.npz --dispatch worker --cqs 2
    python code/scripts/bench.py --host-only --iters 50 --input a.npz --json host.json     # no device

For every ``--input`` cloud (default: the shipped sample ``samples/test_pcd.npz``), batch 1, after ``--warmup``
calls: ``--iters`` iterations of the served path cut into stages, each synchronised on its own (p50 / p99 / mean /
min in ms, so the rows do not add up to ``e2e`` exactly), then ``--iters`` plain ``model(points)`` calls, then
back-to-back replays:

=====================  ================================================================================================
``pre``                host pre-processing as served (``model._prepare``): decode + densification + hard voxelization
                       + VFE mean + the C13 sparse plan (4 sub-manifold + 4 strided neighbour maps, bucket choice)
``host_in``            ``TtBEVFusion.frame_inputs``: the bucket's device inputs in numpy (two-term VFE table, the eight
                       tile-ordered im2col indices, the two to-dense indices)
``h2d.convert``        numpy -> ttnn host tensors (``TraceRunner._collect_uploads``: bf16 / uint32 conversion)
``h2d.copy``           the 11 host -> device copies + device synchronise (CQ0, or CQ1 + events with 2 CQs)
``h2d``                ``h2d.convert`` + ``h2d.copy`` (the served upload)
``h2d.<input>``        one input's copy alone (outside the served path; same data)
``trace.sparse``       one blocking replay of the frame's ``sparse_<bucket>`` trace (encoder + to-dense -> ``bev``)
``trace.dense``        one blocking replay of the ``dense`` trace (SECOND -> query head -> packed output)
``trace``              ``trace.sparse`` + ``trace.dense``
``d2h``                ``TraceRunner.read("dense")`` as served: the packed fp32 output -> ``{heads, query_heat, indices}``
``d2h.copy`` /         the readback split: the device-to-host copy into a preallocated host tensor, then its host
``d2h.to_numpy``       conversion
``unpack``             ``TtBEVFusion.assemble``: ``bbox_pred`` / ``score`` / ``label_pred`` in float32
``post``               host post-processing (decode, distance-binned thresholds, circle NMS, ROS conversion, IoU-BEV
                       NMS, remapper)
``e2e``                ``model(points)`` as served (its ``timing_ms`` medians are reported too)
``b2b``                back-to-back frames (``sparse_<bucket>`` + ``dense`` replays), one synchronise per round of
                       ``--b2b-iters``: device time per frame; ``b2b.sparse`` / ``b2b.dense`` the segments alone
``b2b.sparse_<b>``     the input's own plan uploaded into every bucket ``b`` that holds it, the sparse replay b2b
                       there (the cost of the capacity: gathers and matmuls run over every capacity row)
=====================  ================================================================================================

``--host-only`` (no device, no ttnn import) splits ``pre`` and ``host_in`` instead: ``pre.load`` (``.npz`` decode),
``pre.densify``, ``pre.voxelize`` (hard voxelization + VFE mean), ``pre.plan`` (``build_sparse_plan``) and its parts
``plan.level`` (hash tables of the active sets), ``plan.subm<i>`` / ``plan.<down>`` (the neighbour maps), and
``host_in.vfe`` / ``host_in.im2col`` / ``host_in.dense``.

AICLK / power / temperature are sampled from sysfs during a device run (``ttaw.profiling.AiclkSampler``), and the host
load average is recorded (the host stages move with it). ``--check`` (default) compares the served output of each
input with a stored fp32 CPU-reference ``/predict`` body (``--reference PATH`` per input, ``-`` for none; default:
``<stem>.reference.json`` next to the input) with the container-smoke gates (``server/smoke_test.py``
``REFERENCE_GATES``). Always quote the configuration line (dispatch / CQs / grid) with the numbers.
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import statistics
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np

from tt_bevfusion import BEVFusion
from tt_bevfusion.ttaw.profiling import AiclkSampler, StageBench, time_b2b

CODE = Path(__file__).resolve().parents[1]
SAMPLE = CODE / "tt_bevfusion" / "samples" / "test_pcd.npz"


def check_output(model, src: str, reference: Optional[Path]) -> Dict[str, Any]:
    """Served output vs the stored fp32 CPU-reference ``/predict`` body (the container-smoke comparison)."""
    from tt_bevfusion.server import smoke_test as st
    from tt_bevfusion.ttaw.server import smoke

    out = model(src).to_dict()
    meta = out.get("meta", {})
    res: Dict[str, Any] = {"num_detections": out["num_detections"], "counts": meta.get("counts"),
                           "num_voxels": meta.get("num_voxels"), "sparse": meta.get("sparse")}
    if reference is not None and reference.is_file():
        metrics, fails = smoke.compare_with_reference(out, smoke.load_json(reference), gates=st.REFERENCE_GATES)
        res["reference"] = {"file": str(reference), "metrics": metrics, "failures": fails}
    return res


# ----------------------------------------------------------------------------------------------------- host only

def bench_host(model, src: str, *, iters: int, warmup: int, name: str) -> Dict[str, Any]:
    """The host stages split into their parts (no device)."""
    from tt_bevfusion import io as tio
    from tt_bevfusion.host.preprocess import STRIDED_LAYERS, build_sparse_plan, select_bucket
    from tt_bevfusion.reference.config import SPARSE_SHAPE
    from tt_bevfusion.tt import sparse as tsp
    from tt_bevfusion.ttaw import sparse as sp
    from tt_bevfusion.ttaw.models.sparse_encoder import pack_table
    from tt_bevfusion.tt.model import KNOBS

    pre = model.pre
    bench = StageBench(name)
    for _ in range(warmup):
        pre(src)
    prepared = None
    for _ in range(iters):
        with bench.stage("pre"):
            prepared = pre(src)
        with bench.stage("pre.load"):
            tio.load_points(src)
        with bench.stage("pre.densify"):
            pts, _, _ = pre.densified_points(src)
        with bench.stage("pre.voxelize"):
            vs, vfe = pre.voxelize(pts)
        with bench.stage("pre.plan"):
            build_sparse_plan(vs.coors, rule=pre.rule, backend=pre.plan_backend)
        if pre.plan_backend == "numba":
            from tt_bevfusion.host import fast
            from tt_bevfusion.reference.config import SPARSE_LEVELS

            c = np.ascontiguousarray(np.asarray(vs.coors, dtype=np.int64)[:, ::-1])
            shape = SPARSE_SHAPE
            for i, (down, nxt) in enumerate(zip(("down1", "down2", "down3", "conv_out"),
                                                ("l2", "l3", "l4", "out")), 1):
                with bench.stage(f"plan.subm{i}"):
                    fast.subm_map(c, shape)
                k, s_, p_ = STRIDED_LAYERS[down]
                with bench.stage(f"plan.{down}"):
                    c, shape, _, _ = fast.strided_map(c, shape, k, s_, p_)
            skip_numpy_parts = True
        else:
            skip_numpy_parts = False
        # the plan's parts (build_sparse_plan, step by step; numpy oracle)
        t = time.perf_counter()
        c = np.asarray(vs.coors, dtype=np.int64)
        lvl = sp.SparseLevel(c[:, ::-1].copy(), SPARSE_SHAPE)
        level_ms = (time.perf_counter() - t) * 1e3
        levels = {"l1": lvl}
        for i, (down, nxt) in enumerate(zip(("down1", "down2", "down3", "conv_out"), ("l2", "l3", "l4", "out")), 1):
            if skip_numpy_parts:
                break
            with bench.stage(f"plan.subm{i}"):
                sp.subm_neighbor_map(lvl, 3, rule=pre.rule)
            k, s, p = STRIDED_LAYERS[down]
            t = time.perf_counter()
            lvl, _ = sp.strided_neighbor_map(lvl, k, s, p)
            bench.add(f"plan.{down}", (time.perf_counter() - t) * 1e3)
            levels[nxt] = lvl
        if not skip_numpy_parts:
            bench.add("plan.level_l1", level_ms)
            select_bucket({k: v.num_active for k, v in levels.items()})
        plan = prepared.plan
        g_vfe_width, terms = 32, 2
        index = str(KNOBS.read().HOST_INDEX)
        with bench.stage("host_in"):
            tsp.frame_inputs(prepared.vfe, plan, plan.bucket, vfe_terms=terms, vfe_width=g_vfe_width, index=index)
        caps = tsp.bucket_capacities(plan.bucket)
        with bench.stage("host_in.vfe"):
            pack_table(np.asarray(prepared.vfe, np.float32), caps["l1"], terms=terms, width=g_vfe_width)
        with bench.stage("host_in.im2col"):
            for n in tsp.INDEX_NAMES:
                lin, lout = tsp.LEVEL_OF_MAP[n]
                if index == "numba":
                    from tt_bevfusion.host import fast

                    fast.im2col_index_into(plan.maps[n], caps[lin], caps[lout])
                else:
                    sp.im2col_index(plan.maps[n], sentinel=caps[lin], rows=caps[lout])
        with bench.stage("host_in.dense"):
            sp.dense_gather_index(plan.levels["out"], sentinel=caps["out"])
    stats = prepared.stats()
    return {"stages_ms": bench.summary(), "table": bench.table(),
            "frame": {"num_points": int(len(prepared.points)), "num_voxels": stats.get("num_voxels"),
                      "sparse": stats.get("sparse")}}


# ---------------------------------------------------------------------------------------------------- device

def bench_input(model, src: str, *, iters: int, warmup: int, b2b_iters: int, b2b_rounds: int,
                name: str, bucket_sweep: bool = True) -> Dict[str, Any]:
    import ttnn

    from tt_bevfusion.tt.model import input_name, sparse_variant
    from tt_bevfusion.ttaw.tensors import to_numpy

    tr, dev, g = model.runner, model.device, model.graph
    params = model.validate_params({})
    bench = StageBench(name)
    for _ in range(warmup):
        model(src)
    out_dev = tr._trace_for("dense", 0).leaves[0][1]      # the packed output (device)
    host_out = ttnn.allocate_tensor_on_host(out_dev.spec, dev)
    sync = lambda: ttnn.synchronize_device(dev)  # noqa: E731
    prepared = None
    nbytes: Dict[str, int] = {}
    for _ in range(iters):
        t0 = time.perf_counter()
        prepared = model._prepare(points=src)
        t1 = time.perf_counter()
        plan = prepared.plan
        sv = sparse_variant(plan.bucket)
        inputs = g.frame_inputs(prepared.vfe, plan, plan.bucket, staged=True)
        t2 = time.perf_counter()
        uploads = tr._collect_uploads(inputs, None)
        t3 = time.perf_counter()
        tr._enqueue_uploads(uploads)
        sync()
        t4 = time.perf_counter()
        tr.replay(sv)
        sync()
        t5 = time.perf_counter()
        tr.replay("dense")
        sync()
        t6 = time.perf_counter()
        packed = tr.read("dense")
        t7 = time.perf_counter()
        raw = g.assemble(packed)
        t8 = time.perf_counter()
        model._postprocess(raw, prepared, params)
        t9 = time.perf_counter()
        for stage, a, b in (("pre", t0, t1), ("host_in", t1, t2), ("h2d.convert", t2, t3), ("h2d.copy", t3, t4),
                            ("h2d", t2, t4), ("trace.sparse", t4, t5), ("trace.dense", t5, t6), ("trace", t4, t6),
                            ("d2h", t6, t7), ("unpack", t7, t8), ("post", t8, t9)):
            bench.add(stage, (b - a) * 1e3)
        # the split rows (outside the served path; same data, CQ0)
        for slot, host, _ in uploads:
            short = slot.name.split(".", 1)[-1]
            nbytes[short] = int(np.prod(slot.shape)) * (2 if "bfloat16" in str(slot.dtype).lower() else 4)
            with bench.stage(f"h2d.{short}"):
                ttnn.copy_host_to_device_tensor(host, slot.buffers[0], cq_id=0)
                sync()
        with bench.stage("d2h.copy"):
            ttnn.copy_device_to_host_tensor(out_dev, host_out, blocking=True)
        with bench.stage("d2h.to_numpy"):
            to_numpy(host_out)
    timing: List[Dict[str, float]] = []
    for _ in range(iters):
        t0 = time.perf_counter()
        out = model(src)
        bench.add("e2e", (time.perf_counter() - t0) * 1e3)
        timing.append(dict(out.timing_ms))
    plan = prepared.plan
    sv = sparse_variant(plan.bucket)
    tr.upload(g.frame_inputs(prepared.vfe, plan, plan.bucket))
    sync()

    def frame():
        tr.replay(sv)
        tr.replay("dense")

    rounds = [time_b2b(frame, sync, n=b2b_iters) for _ in range(b2b_rounds)]
    for r in rounds:
        bench.add("b2b", r)
    for _ in range(b2b_rounds):
        bench.add("b2b.sparse", time_b2b(lambda: tr.replay(sv), sync, n=b2b_iters))
        bench.add("b2b.dense", time_b2b(lambda: tr.replay("dense"), sync, n=b2b_iters))
    sweep: Dict[str, Any] = {}
    if bucket_sweep:
        counts = plan.counts()
        for b, caps in g.buckets.items():
            if any(counts[lv] > int(c) for lv, c in zip(("l1", "l2", "l3", "l4", "out"), caps)):
                sweep[b] = None
                continue
            tr.upload(g.frame_inputs(prepared.vfe, plan, b))
            sync()
            ms = [time_b2b(lambda b=b: tr.replay(sparse_variant(b)), sync, n=max(5, b2b_iters // 5))
                  for _ in range(3)]
            sweep[b] = statistics.median(ms)
            bench.add(f"b2b.sparse_{b}", sweep[b])
        tr.upload(g.frame_inputs(prepared.vfe, plan, plan.bucket))
        sync()
    b2b = statistics.median(rounds)
    stats = prepared.stats()
    input_bytes = {k: v for k, v in sorted(nbytes.items())}
    return {"stages_ms": bench.summary(), "table": bench.table(),
            "timing_ms_median": {k: statistics.median(t[k] for t in timing) for k in timing[0]},
            "b2b_rounds_ms": rounds, "b2b_iters": b2b_iters, "device_fps_b2b": 1e3 / b2b,
            "e2e_fps_sync": 1e3 / bench.summary()["e2e"]["p50"], "bucket_sweep_sparse_b2b_ms": sweep,
            "frame": {"num_points": int(len(prepared.points)), "num_voxels": stats.get("num_voxels"),
                      "sparse": stats.get("sparse")},
            "bytes": {"inputs": input_bytes, "inputs_total": sum(input_bytes.values()),
                      "output": int(np.prod(out_dev.shape)) * 4}}


def _body(out) -> str:
    d = out.to_dict()
    d.pop("timing_ms", None)
    return json.dumps(d, sort_keys=True, default=str)


def bench_pipe(model, srcs: List[str], *, frames: int, clients: int) -> Dict[str, Any]:
    """Throughput of concurrent callers (OPT_PLAN T1): ``clients`` threads share ``frames`` calls cycling over
    ``srcs``; frames/s = frames / wall time. Every output is compared with a serial call's body (``timing_ms``
    dropped): ``mismatches`` must be 0. With ``BEVFUSION_PIPELINE=0`` the calls serialize on the model lock."""
    import threading

    ref = {s: _body(model(s)) for s in srcs}
    for s in srcs:                                         # warm the serial path of every input once more
        model(s)
    jobs = [srcs[i % len(srcs)] for i in range(frames)]
    lock = threading.Lock()
    lat: List[float] = []
    waits: List[float] = []
    stages: Dict[str, List[float]] = {"preprocess": [], "stage": [], "device": [], "postprocess": []}
    bad: List[str] = []
    nxt = [0]

    def worker():
        while True:
            with lock:
                i = nxt[0]
                nxt[0] += 1
            if i >= len(jobs):
                return
            t = time.perf_counter()
            out = model(jobs[i])
            ms = (time.perf_counter() - t) * 1e3
            body = _body(out)
            with lock:
                lat.append(ms)
                waits.append(float(out.timing_ms.get("wait", 0.0)))
                for k in stages:
                    stages[k].append(float(out.timing_ms.get(k, 0.0)))
                if body != ref[jobs[i]]:
                    bad.append(jobs[i])

    threads = [threading.Thread(target=worker) for _ in range(clients)]
    t0 = time.perf_counter()
    for th in threads:
        th.start()
    for th in threads:
        th.join()
    wall = time.perf_counter() - t0
    return {"frames": frames, "clients": clients, "inputs": [Path(s).name for s in srcs], "wall_s": wall,
            "fps": frames / wall, "ms_per_frame": wall * 1e3 / frames, "latency_p50_ms": statistics.median(lat),
            "latency_max_ms": max(lat), "wait_p50_ms": statistics.median(waits),
            **{f"{k}_p50_ms": statistics.median(v) for k, v in stages.items()}, "mismatches": len(bad)}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--input", action="append", default=None,
                    help="cloud .npz / .pcd (repeatable; default: the shipped sample)")
    ap.add_argument("--reference", action="append", default=None,
                    help="fp32 CPU-reference /predict body per --input ('-' for none; default <stem>.reference.json"
                         " next to the input)")
    ap.add_argument("--iters", type=int, default=100)
    ap.add_argument("--warmup", type=int, default=5)
    ap.add_argument("--b2b-iters", type=int, default=50)
    ap.add_argument("--b2b-rounds", type=int, default=5)
    ap.add_argument("--dispatch", default=None, choices=["eth", "worker"], help="default: BEVFUSION_DISPATCH or eth")
    ap.add_argument("--cqs", type=int, default=None, choices=[1, 2],
                    help="command queues (default: BEVFUSION_NUM_CQS or 1)")
    ap.add_argument("--weights-dir", default=None, help="holds bevfusion_lidar.onnx + yaml (default: the pinned HF "
                    "snapshot)")
    ap.add_argument("--allow-fallback", action="store_true", help="accept a WORKER fallback of an ETH open")
    ap.add_argument("--no-check", dest="check", action="store_false")
    ap.add_argument("--no-bucket-sweep", dest="bucket_sweep", action="store_false")
    ap.add_argument("--host-only", action="store_true", help="split the host stages; no device")
    ap.add_argument("--pipe-frames", type=int, default=0, help="T1: frames of the concurrent-caller throughput run "
                    "per input and over all inputs mixed (0: skip)")
    ap.add_argument("--pipe-clients", type=int, default=3, help="concurrent callers of the throughput run")
    ap.add_argument("--pipe-only", action="store_true", help="skip the stage / e2e / b2b rows (throughput run only)")
    ap.add_argument("--tag", default="", help="label of this configuration in the output")
    ap.add_argument("--json")
    a = ap.parse_args()
    inputs = [Path(p) for p in (a.input or [str(SAMPLE)])]
    refs = [None if r in (None, "-") else Path(r) for r in (a.reference or [])] + [None] * len(inputs)
    load0 = os.getloadavg()
    host = {"cpus": os.cpu_count(), "machine": platform.processor() or platform.machine(), "loadavg_start": load0,
            "omp_num_threads": os.environ.get("OMP_NUM_THREADS")}
    if a.host_only:
        model = BEVFusion.__new__(BEVFusion)
        model.variant = os.environ.get("BEVFUSION_VARIANT") or BEVFusion.DEFAULT_VARIANT
        model.weights_path = BEVFusion.resolve_weights(weights_dir=a.weights_dir)
        model._build_host()
        res: Dict[str, Any] = {"tag": a.tag or "host-only", "host": host, "iters": a.iters, "inputs": {}}
        for k, path in enumerate(inputs):
            r = bench_host(model, str(path), iters=a.iters, warmup=a.warmup, name=f"{path.name} | host-only")
            r["input"] = str(path)
            res["inputs"][f"{k}:{path.name}"] = r
            print(r["table"], flush=True)
            print("frame:", json.dumps(r["frame"], default=str), flush=True)
        res["host"]["loadavg_end"] = os.getloadavg()
        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
    t0 = time.perf_counter()
    with BEVFusion.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs,
                                   weights_dir=a.weights_dir) as model:
        load_s = time.perf_counter() - t0
        info = dict(model.device_info)
        want = a.dispatch or os.environ.get("BEVFUSION_DISPATCH", "eth")
        if info.get("dispatch") != want and not a.allow_fallback:
            raise SystemExit(f"device opened with dispatch {info.get('dispatch')!r}, asked for {want!r}")
        tag = a.tag or f"{info.get('dispatch')}-{info.get('num_command_queues')}cq {info.get('grid')}"
        res = {"tag": tag, "config": info, "variant": model.variant,
               "precision": model.graph.policy.describe(), "knobs": model.graph.knobs.as_dict(),
               "trace": model.runner.describe(), "load_s": load_s, "warmup_ms": dict(model.warmup_ms),
               "iters": a.iters, "host": host, "inputs": {}}
        with AiclkSampler(interval_s=0.05) as clk:
            for k, (path, ref) in enumerate(zip(inputs, refs)):
                if a.pipe_only:
                    break
                src = str(path)
                if ref is None and a.reference is None:
                    cand = path.parent / f"{path.stem}.reference.json"
                    ref = cand if cand.is_file() else None
                r = bench_input(model, src, iters=a.iters, warmup=a.warmup, b2b_iters=a.b2b_iters,
                                b2b_rounds=a.b2b_rounds, name=f"{path.name} | {tag}", bucket_sweep=a.bucket_sweep)
                r["input"] = src
                if a.check:
                    r["check"] = check_output(model, src, ref)
                res["inputs"][f"{k}:{path.name}"] = r
                print(r["table"], flush=True)
                print(f"b2b rounds ms: {[round(x, 3) for x in r['b2b_rounds_ms']]}  bucket sweep: "
                      f"{r['bucket_sweep_sparse_b2b_ms']}  timing_ms medians: "
                      f"{ {k2: round(v, 2) for k2, v in r['timing_ms_median'].items()} }  frame: "
                      f"{json.dumps(r['frame'], default=str)}  bytes: {r['bytes']}", flush=True)
                if a.check:
                    print("check:", json.dumps(r["check"], default=str), flush=True)
            if a.pipe_frames:
                srcs = [str(p) for p in inputs]
                pipe = {}
                for name, group in [(Path(s).name, [s]) for s in srcs] + ([("mixed", srcs)] if len(srcs) > 1 else []):
                    pipe[name] = bench_pipe(model, group, frames=a.pipe_frames, clients=a.pipe_clients)
                    print(f"pipe {name}: {json.dumps({k: (round(v, 3) if isinstance(v, float) else v) for k, v in pipe[name].items()})}",
                          flush=True)
                res["pipe"] = pipe
        res["aiclk"] = clk.summary()
        res["trace_after"] = model.runner.describe()
        res["host"]["loadavg_end"] = os.getloadavg()
    print(json.dumps({"tag": tag, "config": info, "aiclk": res["aiclk"], "load_s": round(load_s, 2),
                      "warmup_ms": res["warmup_ms"],
                      "program_cache_entries": res["trace_after"].get("program_cache_entries"),
                      "loadavg": [load0, res["host"]["loadavg_end"]]}, default=str), flush=True)
    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")


if __name__ == "__main__":
    main()