#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device BEVPool kernel timing probe (tt.bevpool_k, OPT_REPORT round 7): the model shape (shipped sample's
calibration) for table variants, timed as a trace of 10 calls (best of 3), plus the error vs the float64 model."""
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from tt_bevdet.tests import test_bevpool_k_device as T  # noqa: E402
from tt_bevdet.tt.bevpool_k import BEVPoolK  # noqa: E402


def main():
    import ttnn

    from tt_bevdet.device import close_device, open_device
    from tt_bevdet.host.calib import bevpool_tables
    from tt_bevdet.host.inputs import load_sample, prepare_frame

    cfg = T.CFG
    req = load_sample(Path(__file__).resolve().parents[1] / "tt_bevdet/samples/synthetic_6cam.json")
    tables = bevpool_tables(cfg, prepare_frame(req["images"], dict(req["calibration"]), cfg).calib)
    hw = 704
    rd = tables.ranks_depth.astype(np.int64)
    rb, pix, dbin = tables.ranks_bev, tables.ranks_feat, (rd // hw) % cfg.n_depth
    src = T._src(4224, 3)
    ref = BEVPoolK.reference(src, rb, pix, dbin, 16384, T.DW, T.C)
    # argv: variants "nb=4", "nb=4,inexact", "nb=4,BPK_NODEPTH", ... (default: the served form)
    variants = []
    for a in sys.argv[1:] or ["nb=4"]:
        v = dict(distinct=True, nb=4, defines=())
        for part in a.split(","):
            if part.startswith("nb="):
                v["nb"] = int(part[3:])
            elif part == "inexact":
                v["distinct"] = False
            elif part.startswith("BPK_"):
                v["defines"] = v["defines"] + (part,)
        variants.append(v)
    dev = open_device(None, allow_fallback=False)
    try:
        gx, gy = T._grid(dev)
        for v in variants:
            k = BEVPoolK(n_cells=16384, n_pix=4224, depth_w=T.DW, channels=T.C, n_cores=gx * gy, grid_x=gx, batch=v.get("nb", 4),
                         block_cap=40000, defines=v["defines"])
            tab = k.table(rb, pix, dbin, exact=v["distinct"])
            s_d, t_d = T._to_dev(dev, src, True), T._table_dev(dev, tab)
            o = T._out(dev, 16384)
            k(s_d, t_d, o)
            out = T._read(o)
            tid = ttnn.begin_trace_capture(dev, cq_id=0)
            for _ in range(10):
                k(s_d, t_d, o)
            ttnn.end_trace_capture(dev, tid, cq_id=0)
            best = 1e9
            for _ in range(3):
                ttnn.synchronize_device(dev)
                t0 = time.perf_counter()
                ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
                best = min(best, (time.perf_counter() - t0) / 10 * 1e6)
            ttnn.release_trace(dev, tid)
            u = T._ulps(out, ref)
            print(f"{v}: {k.last_stats} | {best:.1f} us per call | != model {int((u > 0).sum())}, max ulp "
                  f"{int(u.max())}", flush=True)
            for t in (s_d, t_d, o):
                ttnn.deallocate(t)
    finally:
        close_device(dev)


if __name__ == "__main__":
    main()