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

For every ``--input`` cloud (default: the shipped sample), 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 (``e2e``; with the variant
``seg_det`` also ``e2e_seg`` = ``model(points, run_detection=False)``, the device path of the ``seg`` serve
profile), then back-to-back replays:

=================  ================================================================================================
``pre``            host pre-processing as served (``model._prepare``): decode + ``host.preprocess.prepare_frame``
``pre.load``       the file decode alone (``io.load_points``)
``pre.voxelize``   features + voxelization (``Frame.timing_ms``, the same call)
``pre.pooling``    the 4 serialized-pooling levels (clusters, Morton orders, CSR)
``pre.neighbours`` the 5 sub-manifold neighbour maps (numpy, alias-first by default)
``tables``         ``tt.inputs.frame_inputs``: the per-frame device tables of the bucket (feature table, 5 im2col
                   indices, 10 attention + 10 inverse gathers, 4 K1 offsets, 4 unpool rows, 9 BEV rows)
``host_in``        the tables -> ttnn host tensors (``TraceRunner._collect_uploads``)
``h2d``            their upload into the persistent trace inputs + device synchronise
``trace_a``        one blocking replay of trace A ``encseg_<b>`` (encoder + seg head)
``d2h_a``          the read of the ``seg`` output (label, probability, entropy per voxel; fp32)
``unpack_a``       ``tt.model.unpack_seg``
``trace_b``        one blocking replay of trace B ``det_<b>`` (det head; variant ``seg_det``)
``d2h_b``          the read of the packed det output
``unpack_b``       ``TtPTv3.assemble`` (the det ONNX outputs)
``post``           host post-processing (segmentation reconstruction to the input points, classification /
                   entropy / obstacle filter; decode + IoU-BEV NMS + remapper for the boxes)
``e2e``            ``model(points)`` as served (its ``timing_ms`` medians are reported too)
``b2b_a/b/ab``     back-to-back replays of A, of B, and of A then B, one synchronise per round of ``--b2b-iters``
=================  ================================================================================================

AICLK / power / temperature are sampled from sysfs (``ttaw.profiling.AiclkSampler``) and the host load average is
recorded (the host stages move with it). ``--check`` (default) compares the served output with the stored fp32
CPU-reference ``/predict`` bodies (``--reference PATH`` per input; default ``<stem>.reference.json`` and, for
``run_detection=False``, ``<stem>.seg.reference.json`` next to the input) under the container-smoke gates
(``server/smoke_test.py`` ``REFERENCE_GATES``). Always quote the configuration line (dispatch / CQs / grid).
"""
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_ptv3 import PTv3
from tt_ptv3.ttaw.profiling import AiclkSampler, StageBench, time_b2b

CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_ptv3" / "samples"


def default_sample() -> Path:
    for name in ("test_pcd.npz", "pandaset_019_f40.npz"):
        if (SAMPLES / name).is_file():
            return SAMPLES / name
    raise SystemExit(f"no sample in {SAMPLES}; pass --input")


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

    res: Dict[str, Any] = {}
    runs = [("default", {}, ref)]
    if model.with_det:
        runs.append(("run_detection=False", {"run_detection": False}, ref_seg))
    elif ref_seg is not None:
        runs = [("default", {}, ref_seg)]
    for tag, kw, path in runs:
        out = model(src, **kw).to_dict()
        r: Dict[str, Any] = {"num_detections": out.get("num_detections"), "num_points": out.get("num_points"),
                             "frame": out.get("meta", {}).get("frame")}
        if path is not None and path.is_file():
            metrics, fails = smoke.compare_with_reference(out, smoke.load_json(path), gates=st.REFERENCE_GATES)
            r["reference"] = {"file": str(path), "metrics": metrics, "failures": fails}
        res[tag] = r
    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_ptv3 import io as tio
    from tt_ptv3.tt import inputs as tin
    from tt_ptv3.tt.model import unpack_seg

    from tt_ptv3.host import preprocess as pre

    tr, dev = model.runner, model.device
    b, caps = model.bucket_for(pre.prepare_frame(np.asarray(tio.load_points(src).points, np.float32)).counts)
    va, vb = f"encseg_{b}", f"det_{b}"
    params = model.validate_params({})
    bench = StageBench(name)
    for _ in range(warmup):
        model(src)
    for _ in range(iters):
        t0 = time.perf_counter()
        prepared = model._prepare(points=src)
        t1 = time.perf_counter()
        frame = prepared["frame"]
        assert prepared["bucket"][0] == b, (prepared["bucket"], b)
        tables = {f"{b}.{k}": v for k, v in tin.frame_inputs(frame, caps, det=model.with_det,
                                                                    impl=getattr(model, "host_pre", "numpy"),
                                                                    threads=getattr(model, "host_threads", 1)).items()}
        t2 = time.perf_counter()
        uploads = tr._collect_uploads(tables, None)
        t3 = time.perf_counter()
        tr._enqueue_uploads(uploads)
        ttnn.synchronize_device(dev)
        t4 = time.perf_counter()
        tr.replay(va)
        ttnn.synchronize_device(dev)
        t5 = time.perf_counter()
        seg_raw = tr.read(va)["seg"]
        t6 = time.perf_counter()
        seg = unpack_seg(seg_raw, frame.counts[0], caps[0])
        t7 = time.perf_counter()
        raw: Dict[str, Any] = {"pred_labels": seg["labels"], "voxel_prob": seg["prob"],
                               "voxel_entropy": seg["entropy"], "det": None}
        rows = [("pre", t0, t1), ("tables", t1, t2), ("host_in", t2, t3), ("h2d", t3, t4), ("trace_a", t4, t5),
                ("d2h_a", t5, t6), ("unpack_a", t6, t7)]
        if model.with_det:
            tr.replay(vb)
            ttnn.synchronize_device(dev)
            t8 = time.perf_counter()
            packed = tr.read(vb)
            t9 = time.perf_counter()
            raw["det"] = model.graph.assemble(packed)
            t10 = time.perf_counter()
            rows += [("trace_b", t7, t8), ("d2h_b", t8, t9), ("unpack_b", t9, t10)]
            t7 = t10
        model._postprocess(raw, prepared, params)
        t11 = time.perf_counter()
        rows.append(("post", t7, t11))
        for stage, x, y in rows:
            bench.add(stage, (y - x) * 1e3)
        for k, v in frame.timing_ms.items():
            bench.add(f"pre.{k}", v)
        with bench.stage("pre.load"):
            tio.load_points(src)
    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))
    timing_seg: List[Dict[str, float]] = []
    if model.with_det:
        for _ in range(iters):
            t0 = time.perf_counter()
            out = model(src, run_detection=False)
            bench.add("e2e_seg", (time.perf_counter() - t0) * 1e3)
            timing_seg.append(dict(out.timing_ms))
    sync = lambda: ttnn.synchronize_device(dev)  # noqa: E731
    rounds: Dict[str, List[float]] = {"b2b_a": [time_b2b(lambda: tr.replay(va), sync, n=b2b_iters)
                                                for _ in range(b2b_rounds)]}
    if model.with_det:
        rounds["b2b_b"] = [time_b2b(lambda: tr.replay(vb), sync, n=b2b_iters) for _ in range(b2b_rounds)]
        rounds["b2b_ab"] = [time_b2b(lambda: (tr.replay(va), tr.replay(vb)), sync, n=b2b_iters)
                            for _ in range(b2b_rounds)]
    for k, rs in rounds.items():
        for r in rs:
            bench.add(k, r)
    summ = bench.summary()
    full = statistics.median(rounds.get("b2b_ab", rounds["b2b_a"]))
    in_bytes = {k: int(np.asarray(v).nbytes) for k, v in tables.items()}
    slot_bytes = {}
    for k in tables:
        s = tr._slot(k, "input")
        slot_bytes[k] = int(np.prod(s.shape)) * (2 if "bfloat16" in str(s.dtype).lower() else 4)
    out_shapes = {}
    for v in ([va, vb] if model.with_det else [va]):
        out_shapes[v] = [list(t.shape) for _, t, _ in tr._trace_for(v, 0).leaves]
    med = lambda rows: {k: statistics.median(t[k] for t in rows) for k in rows[0]} if rows else {}  # noqa: E731
    return {"stages_ms": summ, "table": bench.table(), "timing_ms_median": med(timing),
            "timing_ms_median_seg": med(timing_seg), "b2b_rounds_ms": rounds, "b2b_iters": b2b_iters,
            "device_fps_b2b": 1e3 / full, "e2e_fps_sync": 1e3 / summ["e2e"]["p50"],
            "frame": {"points": frame.num_points, "counts": list(frame.counts),
                      "unique_voxels": int(frame.voxels.num_unique_voxels), "bucket": b, "caps": list(caps)},
            "input_bytes_host": in_bytes, "input_bytes_device_slots": slot_bytes,
            "input_bytes_total": sum(slot_bytes.values()), "output_shapes": out_shapes}


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--input", action="append", default=None, help="cloud .npz / .pcd (repeatable)")
    ap.add_argument("--reference", action="append", default=None,
                    help="fp32 CPU-reference /predict body per --input ('-' for none; default <stem>.reference.json)")
    ap.add_argument("--iters", type=int, default=50)
    ap.add_argument("--warmup", type=int, default=5)
    ap.add_argument("--b2b-iters", type=int, default=20)
    ap.add_argument("--b2b-rounds", type=int, default=5)
    ap.add_argument("--dispatch", default=None, choices=["eth", "worker"], help="default: PTV3_DISPATCH or eth")
    ap.add_argument("--cqs", type=int, default=None, choices=[1, 2], help="default: PTV3_NUM_CQS or 1")
    ap.add_argument("--variant", default=None, choices=["seg_det", "seg"], help="default: PTV3_VARIANT or seg_det")
    ap.add_argument("--weights-dir", default=None, help="default: PTV3_WEIGHTS_DIR or 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="")
    ap.add_argument("--json")
    a = ap.parse_args()
    inputs = [Path(p) for p in (a.input or [str(default_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 PTv3.from_pretrained(dispatch=a.dispatch, num_command_queues=a.cqs, variant=a.variant,
                              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("PTV3_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')} {model.variant}"
        res: Dict[str, Any] = {"tag": tag, "config": info, "variant": model.variant,
                               "buckets": {k: list(v) for k, v in model.buckets.items()}, "knobs": model.knobs.as_dict(),
                               "precision": model.graph.policy.describe(), "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,
                                        "omp_num_threads": os.environ.get("OMP_NUM_THREADS")}, "inputs": {}}
        print("loaded in %.1f s:" % load_s, json.dumps(info, default=str), json.dumps(res["warmup_ms"]), flush=True)
        with AiclkSampler(interval_s=0.05) as clk:
            for k, (path, ref) in enumerate(zip(inputs, refs)):
                src = str(path)
                ref_seg = None
                if ref is None and a.reference is None:
                    stem = path.parent / path.stem
                    ref = Path(f"{stem}.reference.json")
                    ref = ref if ref.is_file() else None
                    rs = Path(f"{stem}.seg.reference.json")
                    ref_seg = rs if rs.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, ref_seg)
                res["inputs"][f"{k}:{path.name}"] = r
                print(r["table"], flush=True)
                print(f"b2b rounds ms: { {k2: [round(x, 3) for x in v] for k2, v in r['b2b_rounds_ms'].items()} }  "
                      f"timing_ms medians: { {k2: round(v, 2) for k2, v in r['timing_ms_median'].items()} }  "
                      f"frame: {r['frame']}  input bytes {r['input_bytes_total']}", 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()