#!/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 --variant qt128 --input a.npz --input b.npz --dispatch worker --cqs 2

For every ``--input`` cloud (default: the shipped synthetic sample of the variant), 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, calibration=...)`` calls, then
back-to-back replays:

=====================  ================================================================================================
``pre``                host pre-processing as served (``model._prepare``): decode, ego crop, float32 projection, fp16
                       interpolation, unique cells, profile checks, the device plan (sort, CSR offsets, gather tables,
                       8-d features with ``pre_norm`` folded)
``pre.load``           the file decode alone (``ttaw.io.load_points``)
``pre.crop``           the ego crop box (``crop_keep_mask``: float32 ``R p + t``)
``pre.interp``         the fp16 interpolation grid and the interpolated points (``interpolation_grid`` +
                       ``interpolate_points``)
``pre.project``        float32 projection of every point to the frustum grid (``project2d``)
``pre.unique``         unique cells and the inverse map (``unique_coors``)
``pre.plan``           ``build_plan``: the sort, CSR offsets, ``point2pix_s`` / ``pix2vox_s`` and the 8-d features
``host_in``            ``FrameStaging.write``: range checks + the plan written into the persistent host buffers
``h2d``                upload of the 9 trace inputs + device synchronise (``TraceRunner.upload``)
``h2d.<input>``        the upload split per input (``features`` fp32, ``offsets``, ``p2p_<s>``, ``v2p_<s>`` uint32)
``trace``              one blocking replay of the whole-graph trace
``d2h``                ``TraceRunner.read`` as served: the fp32 ``pred_probs`` ``[N_cap, 32]`` -> numpy
``d2h.copy``           the readback split: the device-to-host copy into a preallocated host tensor ...
``d2h.to_numpy``       ... and its host conversion to numpy (host untilize of the TILE rows)
``unsort``             the first N rows / 27 columns back in the network-input order (``plan.unsort``)
``post``               host post-processing (first-max argmax, probability, class filter, output assembly)
``e2e``                ``model(points, calibration=...)`` 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>.<variant>.reference.json`` next to the input) with the container-smoke gates (``server/smoke_test.py``
``REFERENCE_GATES``). The calibration of an input (``--calib`` per input): ``preset:<name>`` (``calib/``), a JSON
file holding the calibration dict, ``none`` (crop disabled); default: the preset named like the input's stem, else
``<stem>.calib.json`` next to the input, else the crop disabled. 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_frnet import FRNet
from tt_frnet.ttaw.profiling import AiclkSampler, StageBench, time_b2b

CODE = Path(__file__).resolve().parents[1]
SAMPLES = CODE / "tt_frnet" / "samples"           # the shipped synthetic_<variant>.npz
CALIB_DIR = CODE / "tt_frnet" / "calib"


def calibration_for(path: Path, spec: Optional[str] = None) -> Dict[str, Any]:
    """The ``calibration`` argument for ``path`` (module docstring)."""
    if spec:
        if spec == "none":
            return {"ego_crop_box": None}
        if spec.startswith("preset:"):
            return {"preset": spec[len("preset:"):]}
        return json.loads(Path(spec).read_text())
    stem = path.name.split(".")[0]
    if (CALIB_DIR / f"{stem}.json").is_file():
        return {"preset": stem}
    side = path.parent / f"{stem}.calib.json"
    if side.is_file():
        return json.loads(side.read_text())
    return {"ego_crop_box": None}


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

    out = model(str(path), calibration=calib).to_dict()
    res: Dict[str, Any] = {"num_points": out.get("num_points"), "num_points_raw": out.get("num_points_raw"),
                           "skipped": (out.get("meta") or {}).get("skipped")}
    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, path: Path, calib: Dict[str, Any], *, iters: int, warmup: int, b2b_iters: int,
                b2b_rounds: int, name: str) -> Dict[str, Any]:
    import ttnn

    from tt_frnet.host.preprocess import (build_plan, crop_keep_mask, interpolate_points, interpolation_grid,
                                          project2d, unique_coors)
    from tt_frnet.host.frame import POINT_FIELDS
    from tt_frnet.ttaw import io as tio
    from tt_frnet.ttaw.tensors import to_numpy

    tr = model.runner                       # ttaw.trace.TraceRunner, variant "default"
    dev, pre, cfg = model.device, model.pre, model.config
    params = model.validate_params({})
    src = str(path)
    bench = StageBench(name)
    for _ in range(warmup):
        model(src, calibration=calib)
    out_dev = tr.outputs("default")         # the persistent trace output (device)
    host_out = ttnn.allocate_tensor_on_host(out_dev.spec, dev)
    names = list(model.graph.input_specs())
    slots = {k: tr._slot(k, "input") for k in names}
    T12, bounds, _ = pre.crop_setup(calib)
    prepared = None
    for _ in range(iters):
        t0 = time.perf_counter()
        prepared = model._prepare(points=src, calibration=calib)
        t1 = time.perf_counter()
        plan = prepared.plan
        inputs = model.staging.write(plan)
        t2 = time.perf_counter()
        tr.upload(inputs)
        ttnn.synchronize_device(dev)
        t3 = time.perf_counter()
        tr.replay("default")
        ttnn.synchronize_device(dev)
        t4 = time.perf_counter()
        rows = tr.read("default")
        t5 = time.perf_counter()
        probs = np.asarray(rows, np.float32).reshape(-1, rows.shape[-1])
        raw = plan.unsort(probs[:, : cfg.num_classes])
        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), ("unsort", t5, t6), ("post", t6, t7)):
            bench.add(stage, (b - a) * 1e3)
        # the split rows (outside the served path; the same data and functions as autoware_preprocess)
        with bench.stage("pre.load"):
            xyzi = np.ascontiguousarray(tio.load_points(src).select(POINT_FIELDS), np.float32)
        with bench.stage("pre.crop"):
            keep = (crop_keep_mask(xyzi[:, :3], T12, np.asarray(bounds, np.float64).astype(np.float32))
                    if T12 is not None else np.ones(len(xyzi), bool))
            rawp = xyzi[keep]
        with bench.stage("pre.interp"):
            grid, filled = interpolation_grid(rawp, cfg)
            interp = interpolate_points(grid, filled, cfg)
            allpts = np.concatenate([rawp, interp], axis=0) if len(interp) else rawp
        with bench.stage("pre.project"):
            fx, fy = project2d(allpts[:, :3], cfg.frustum_w, cfg.frustum_h, cfg.fov)
            coors = np.stack([np.zeros_like(fy), fy, fx], axis=1).astype(np.int64)
        with bench.stage("pre.unique"):
            unique_coors(coors, cfg.frustum_w, cfg.frustum_h)
        with bench.stage("pre.plan"):
            build_plan(prepared.inputs, cfg, pre_norm=pre.pre_norm)
        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, calibration=calib)
        bench.add("e2e", (time.perf_counter() - t0) * 1e3)
        timing.append(dict(out.timing_ms))
    rounds = [time_b2b(lambda: tr.replay("default"), 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)
    elem = {"float32": 4, "uint32": 4, "bfloat16": 2}
    in_bytes = {k: int(np.prod(shape)) * elem[dtype] for k, (shape, dtype) in model.graph.input_specs().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"],
            "frame": {k: v for k, v in prepared.stats().items() if k != "ego_crop"},
            "bytes": {**{f"h2d.{k}": v for k, v in in_bytes.items()}, "h2d": sum(in_bytes.values()),
                      "d2h": 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 / .bin; repeatable; "
                    "default: the shipped synthetic sample of the variant")
    ap.add_argument("--calib", action="append", default=None,
                    help="per --input: preset:<name> | <calibration .json> | none (default: see the docstring)")
    ap.add_argument("--reference", action="append", default=None,
                    help="fp32 CPU-reference /predict body per --input ('-' for none; default <stem>.<variant>"
                         ".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("--variant", default=None, help="ot128 | qt128 (default: FRNET_VARIANT or ot128)")
    ap.add_argument("--dispatch", default=None, choices=["eth", "worker"], help="default: FRNET_DISPATCH or eth")
    ap.add_argument("--cqs", type=int, default=None, choices=[1, 2], help="command queues (default: FRNET_NUM_CQS or 1)")
    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(SAMPLES / f"synthetic_{a.variant or os.environ.get('FRNET_VARIANT', 'ot128')}.npz")])]
    calibs = list(a.calib or []) + [None] * len(inputs)
    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 FRNet.from_pretrained(variant=a.variant, dispatch=a.dispatch, num_command_queues=a.cqs) as model:
        load_s = time.perf_counter() - t0
        info = dict(model.device_info)
        want = a.dispatch or os.environ.get("FRNET_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,
                               "graph": model.graph.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}, "inputs": {}}
        with AiclkSampler(interval_s=0.05) as clk:
            for k, (path, cspec, ref) in enumerate(zip(inputs, calibs, refs)):
                calib = calibration_for(path, cspec)
                if ref is None and a.reference is None:
                    cand = path.parent / f"{path.name.split('.')[0]}.{model.variant}.reference.json"
                    ref = cand if cand.is_file() else None
                r = bench_input(model, path, calib, iters=a.iters, warmup=a.warmup, b2b_iters=a.b2b_iters,
                                b2b_rounds=a.b2b_rounds, name=f"{path.name} | {tag} | {model.variant}")
                r["input"], r["calibration"] = str(path), calib
                if a.check:
                    r["check"] = check_output(model, path, calib, 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']}  bytes: {r['bytes']}", 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, "variant": res["variant"], "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()