#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Device BEVPool kernel probe (tt.bevpool_k, OPT_REPORT round 7): error anatomy per row tile for a few chunk sizes
and debug defines, against the float64 model. ``python code/scripts/probe_bpk.py`` under bin/devrun."""
import sys
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():
    from tt_bevdet.device import close_device, open_device

    dev = open_device(None, allow_fallback=False)
    try:
        rng = np.random.default_rng(2)
        n_cells, n_pix = 128, 256
        rb = np.r_[rng.integers(0, 32, 32 * 40 + 7), rng.integers(32, 128, 300)]
        pix = rng.integers(0, n_pix, rb.size)
        dbin = rng.integers(0, T.CFG.n_depth, rb.size)
        src = T._src(n_pix, 0)
        ref = BEVPoolK.reference(src, rb, pix, dbin, n_cells, T.DW, T.C)
        mag = T._abs_sum(src, rb, pix, dbin, n_cells)
        gx, _ = T._grid(dev)
        for chunk, fid, distinct in ((16, "HiFi4", False), (16, "HiFi4", True), (1, "HiFi4", True)):
            k = BEVPoolK(n_cells=n_cells, n_pix=n_pix, depth_w=T.DW, channels=T.C, n_cores=4, grid_x=gx,
                         chunk=chunk, block_cap=512, fidelity=fid)
            tab = k.table(rb, pix, dbin, exact=distinct)
            s_d, t_d = T._to_dev(dev, src, True), T._table_dev(dev, tab)
            o = T._out(dev, n_cells)
            k(s_d, t_d, o)
            out = T._read(o)
            err = np.abs(out.astype(np.float64) - ref)
            ulp = np.maximum(np.abs(ref), 1e-38) * 2.0 ** -7
            ex = (err - ulp) / np.maximum(mag, 1e-30)
            per_rt = ex.reshape(4, 32, -1).max(axis=(1, 2))
            rel = err / np.maximum(np.abs(ref), 1e-30)
            print(f"chunk {chunk:3d} {fid} distinct {distinct}: max excess per row tile {np.array2string(per_rt, precision=3)}; "
                  f"ulps>1: {int((T._ulps(out, ref) > 1).sum())}; rt0 max rel err {rel[:32].max():.3g}", flush=True)
    finally:
        close_device(dev)


if __name__ == "__main__":
    main()