#!/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 900 -- python code/scripts/bench.py --iters 100 --json logs/baseline_bench.json
    bin/devrun -t 900 -- python code/scripts/bench.py --input a.npz --input b.npz --dispatch worker --cqs 2

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 + sampling +
                       deterministic voxelization (first 20 points per pillar, N_MAX 65,536) + the 11 decorated
                       features + the canvas index
``pre.load``           the file decode alone (``io.load_points``)
``pre.densify``        ``densified_points``: the (N, 5) x, y, z, intensity (uint8 semantics), time_lag cloud
``pre.voxelize``       ``voxelize``: sampling mode, strict range filter, ``assign_pillars``
``pre.decorate``       ``decorate_pillars``: the (65536, 20, 11) float32 encoder input
``pre.index``          ``canvas_gather_index``: cell -> pillar row (the C19 gather index)
``host_in``            ``stage_inputs`` (C24 ``FrameStaging.write``): the kept pillars -> bf16 in the persistent buffers
``h2d``                upload of both trace inputs + device synchronise (``TraceRunner.upload``)
``h2d.<input>``        the upload split per input (``features`` 64 MiB bf16 rows, ``canvas_index`` 2.25 MiB uint32)
``trace``              one blocking replay of the whole-graph trace
``d2h``                ``TraceRunner.read`` as served: the one packed fp32 output row -> ``{heads, query_heat, indices}``
``d2h.copy``           the readback split: the device-to-host copy into a preallocated host tensor ...
``d2h.to_numpy``       ... and its host conversion to numpy
``unpack``             ``TtTransFusion.assemble``: the graph outputs ``cls_score0`` / ``bbox_pred0`` / ``dir_cls_pred0``
``post``               host post-processing (decode, score threshold, circle NMS, ROS conversion, IoU-BEV NMS, remapper)
``e2e``                ``model(points)`` as served (its ``timing_ms`` medians are reported too)
``b2b``                back-to-back replays, one synchronise per round of ``--b2b-iters``: device time per forward
=====================  ================================================================================================

AICLK / power / temperature are sampled from sysfs during the 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_transfusion import TransFusion
from tt_transfusion.ttaw.profiling import AiclkSampler, StageBench, time_b2b

CODE = Path(__file__).resolve().parents[1]
SAMPLE = CODE / "tt_transfusion" / "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_transfusion.server import smoke_test as st
    from tt_transfusion.ttaw.server import smoke

    out = model(src).to_dict()
    res: Dict[str, Any] = {"num_detections": out["num_detections"], "counts": out["meta"].get("counts"),
                           "num_pillars": out["meta"].get("num_pillars")}
    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


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

    from tt_transfusion import io as tio
    from tt_transfusion.host.preprocess import decorate_pillars
    from tt_transfusion.tt.model import stage_inputs
    from tt_transfusion.ttaw.tensors import to_numpy
    from tt_transfusion.ttaw.voxelize import canvas_gather_index

    tr = model.runner                       # ttaw.trace.TraceRunner, variant "default"
    dev, pre = model.device, model.pre
    params = model.validate_params({})
    bench = StageBench(name)
    for _ in range(warmup):
        model(src)
    # PFN_BUCKETS: the frame's trace variant (pillar capacity) and its features input
    variant = getattr(model.staging, "variant", "default")
    feat_name = getattr(model.staging, "input_name", "features")
    out_dev = tr._trace_for(variant, 0).leaves[0][1]     # the packed output row (device)
    host_out = ttnn.allocate_tensor_on_host(out_dev.spec, dev)
    slots = {k: tr._slot(k, "input") for k in (feat_name, "canvas_index")}
    for _ in range(iters):
        t0 = time.perf_counter()
        prepared = model._prepare(points=src)
        t1 = time.perf_counter()
        inputs = stage_inputs(model.staging, prepared)
        t2 = time.perf_counter()
        tr.upload(inputs)
        ttnn.synchronize_device(dev)
        t3 = time.perf_counter()
        tr.replay(variant)
        ttnn.synchronize_device(dev)
        t4 = time.perf_counter()
        packed = tr.read(variant)
        t5 = time.perf_counter()
        raw = model.graph.assemble(packed)
        t6 = time.perf_counter()
        model._postprocess(raw, prepared, params)
        t7 = time.perf_counter()
        for stage, a, b in (("pre", t0, t1), ("host_in", t1, t2), ("h2d", t2, t3), ("trace", t3, t4),
                            ("d2h", t4, t5), ("unpack", t5, t6), ("post", t6, t7)):
            bench.add(stage, (b - a) * 1e3)
        # the split rows (outside the served path; same data)
        with bench.stage("pre.load"):
            tio.load_points(src)
        with bench.stage("pre.densify"):
            pts, _, _, _ = pre.densified_points(src)
        with bench.stage("pre.voxelize"):
            pillars, _ = pre.voxelize(pts, None)
        with bench.stage("pre.decorate"):
            decorate_pillars(pillars, pre.scale_zyx, pre.offset_zyx)
        with bench.stage("pre.index"):
            canvas_gather_index(pillars.coords, pillars.num_pillars, pre.spec, sentinel=pre.spec.max_pillars)
        for k, slot in slots.items():
            with bench.stage(f"h2d.{k}"):
                ttnn.copy_host_to_device_tensor(inputs[k], slot.buffers[0], cq_id=0)
                ttnn.synchronize_device(dev)
        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))
    rounds = [time_b2b(lambda: tr.replay(variant), lambda: ttnn.synchronize_device(dev), n=b2b_iters)
              for _ in range(b2b_rounds)]
    for r in rounds:
        bench.add("b2b", r)
    b2b = statistics.median(rounds)
    stats = prepared.stats()
    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"],
            "frame": {k: stats.get(k) for k in ("num_points", "num_points_in_range", "num_pillars",
                                                "num_pillars_total", "pillar_overflow")},
            "variant": variant,
            "bytes": {"features": int(np.prod(slots[feat_name].shape)) * 2,
                      "canvas_index": int(np.prod(slots["canvas_index"].shape)) * 4,
                      "output": int(np.prod(out_dev.shape)) * 4}}


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=10)
    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: TRANSFUSION_DISPATCH or eth")
    ap.add_argument("--cqs", type=int, default=None, choices=[1, 2],
                    help="command queues (default: TRANSFUSION_NUM_CQS or 1)")
    ap.add_argument("--weights-dir", default=None, help="holds transfusion.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("--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()
    t0 = time.perf_counter()
    with TransFusion.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("TRANSFUSION_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: Dict[str, Any] = {"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": {"cpus": os.cpu_count(), "machine": platform.processor() or platform.machine(),
                                        "loadavg_start": load0}, "inputs": {}}
        with AiclkSampler(interval_s=0.05) as clk:
            for k, (path, ref) in enumerate(zip(inputs, refs)):
                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}")
                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']]}  "
                      f"timing_ms medians: { {k2: round(v, 2) for k2, v in r['timing_ms_median'].items()} }  "
                      f"frame: {r['frame']}", flush=True)
                if a.check:
                    print("check:", json.dumps(r["check"], default=str), flush=True)
        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),
                      "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()