#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Stem kernel probe (OPT_PLAN K-IN kernel form / MK0, ``tt/stem.py``), per the hang protocol.

``--smallest``: a 32 x 64 input (16 x 32 output, 16 tiles) with random uint8 values and the real stem weights:
eager x2 identical, a 2-call trace, and the error vs a float64 reference of the same math (conv of the uint8 values
with W_hi + W_lo, + bias map, SiLU). Default: the real 320 x 640 shape on the shipped sample, the kernel vs the stock
stem (``TtSceneSeg.stem`` with the knob off: bf16 cast, reshape, one conv per weight term, EPI), bit for bit, the
float64 error of both, and the time of each in a trace.

    bin/devrun -t 600 -- env TT_METAL_WATCHER=2 python -u code/scripts/diag/stem_probe.py --smallest
    bin/devrun -t 900 -- python -u code/scripts/diag/stem_probe.py
"""
from __future__ import annotations

import argparse
import sys
import time

import numpy as np


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--smallest", action="store_true")
    ap.add_argument("--debug", action="store_true", help="the kernel's intermediate stages vs the host (smallest)")
    a = ap.parse_args()

    import torch
    import torch.nn.functional as F
    import ttnn

    from tt_sceneseg.device import close_device, open_device
    from tt_sceneseg.reference import config
    from tt_sceneseg.reference.weights import load_weights
    from tt_sceneseg.reference import assets
    from tt_sceneseg.tt.layers import split_terms, stem_fold
    from tt_sceneseg.tt.model import STEM_KEY
    from tt_sceneseg.tt.stem import StemKernel
    from tt_sceneseg.ttaw.tensors import to_numpy

    weights = load_weights(assets.onnx_path())
    lw = weights[STEM_KEY]
    dev = open_device()
    rc = 0
    hw = (32, 64) if a.smallest else tuple(config.INPUT_HW)
    w_u8, bias_rows = stem_fold(lw.weight, lw.bias, "bgr", hw)
    rng = np.random.default_rng(1)
    if a.smallest:
        u8 = rng.integers(0, 256, size=(hw[0], hw[1], 3), dtype=np.uint8)
    else:
        from tt_sceneseg.host.preprocess import prepare_image

        prep = prepare_image(str(config_sample()), channel_order="bgr")
        u8 = np.asarray(prep.device_input(), np.uint8).reshape(hw[0], hw[1], 3)
    img = ttnn.from_torch(torch.from_numpy(u8.reshape(1, 1, hw[0], hw[1] * 3).copy()), dtype=ttnn.uint8,
                          layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    bias = ttnn.from_torch(torch.from_numpy(bias_rows.reshape(1, 1, *bias_rows.shape).copy()), dtype=ttnn.float32,
                           layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
    sk = StemKernel(w_u8, input_hw=hw)
    # float64 reference of the same math
    (wh, _), (wl, _) = split_terms(w_u8, np.zeros(32, np.float32), 2)
    x = torch.from_numpy(u8.transpose(2, 0, 1)[None].astype(np.float64))
    conv = F.conv2d(x, torch.from_numpy(wh.astype(np.float64) + wl.astype(np.float64)), stride=2, padding=1)
    ref = F.silu(conv[0].permute(1, 2, 0).reshape(-1, 32) + torch.from_numpy(bias_rows.astype(np.float64))).numpy()
    if a.debug:
        A = np.zeros((hw[0] // 2 * hw[1] // 2, 32))
        xp = np.pad(u8.astype(np.float64), ((1, 1), (1, 1), (0, 0)))
        for oy in range(hw[0] // 2):
            for ox in range(hw[1] // 2):
                A[oy * (hw[1] // 2) + ox, :27] = xp[2 * oy:2 * oy + 3, 2 * ox:2 * ox + 3, :].reshape(27)
        wt = sk.w_tiles.astype(np.float64)
        exp = {1: A @ wt[:32], 2: bias_rows.astype(np.float64), 3: A @ wt[:32] + A @ wt[32:]}
        for d in (1, 2, 3):
            o = sk(img, bias, debug=d)
            got = to_numpy(o).reshape(-1, 32).astype(np.float64)
            ttnn.deallocate(o)
            e = np.abs(got - exp[d]).max() / max(np.abs(exp[d]).max(), 1e-30)
            bad = np.argwhere(np.abs(got - exp[d]) > 1e-3 * np.abs(exp[d]).max())
            print(f"debug {d}: rel err {e:.2e}, {len(bad)} bad values, first {bad[:4].tolist()}", flush=True)
            if len(bad):
                r, c = bad[0]
                print("   got", got[r, :8], "\n   exp", exp[d][r, :8], flush=True)
    outs = []
    for _ in range(2):
        o = sk(img, bias)
        outs.append(to_numpy(o).reshape(-1, 32))
        ttnn.deallocate(o)
    same = np.array_equal(outs[0].view(np.uint32), outs[1].view(np.uint32))
    err = float(np.abs(outs[0] - ref).max() / np.abs(ref).max())
    tid = ttnn.begin_trace_capture(dev, cq_id=0)
    o1 = sk(img, bias)
    o2 = sk(img, bias)
    ttnn.end_trace_capture(dev, tid, cq_id=0)
    ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
    tr = np.array_equal(to_numpy(o2).reshape(-1, 32).view(np.uint32), outs[0].view(np.uint32))
    ttnn.release_trace(dev, tid)
    for o in (o1, o2):
        ttnn.deallocate(o)
    # the FPU's in-tile accumulation of the 27 products is not exact fp32 (rel. err ~2e-4 of the output range on
    # the smallest case, the A @ W_hi stage alone; the full run compares with the stock conv's own error)
    ok = same and tr and err < 1e-3
    line = f"stem kernel {hw}: eager x2 identical {same}, trace identical {tr}, rel err vs float64 {err:.2e}"
    if not a.smallest:
        from tt_sceneseg.tt.model import TtSceneSeg

        # the served model's stem on the pixel-stick canvas (converted to row pages first) == the kernel on rows
        mk = TtSceneSeg(weights, "bgr", stem_kernel=True)
        pix = ttnn.from_torch(torch.from_numpy(u8.reshape(1, 1, hw[0] * hw[1], 3).copy()), dtype=ttnn.uint8,
                              layout=ttnn.ROW_MAJOR_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
        sp = mk.stem(pix)
        pix_same = np.array_equal(to_numpy(sp.tensor).reshape(-1, 32).view(np.uint32), outs[0].view(np.uint32))
        ttnn.deallocate(sp.tensor)
        mk.release()
        ok = ok and pix_same
        line += f" | pixel-stick canvas identical {pix_same}"
        model = TtSceneSeg(weights, "bgr", stem_kernel=False)
        st = model.stem(img)
        sn = to_numpy(st.tensor).reshape(-1, 32)
        nd = int(np.sum(sn.view(np.uint32) != outs[0].view(np.uint32)))
        serr = float(np.abs(sn - ref).max() / np.abs(ref).max())
        ttnn.deallocate(st.tensor)

        def tm(fn):
            r = fn()
            ttnn.deallocate(r)
            tid = ttnn.begin_trace_capture(dev, cq_id=0)
            rs = [fn() for _ in range(5)]
            ttnn.end_trace_capture(dev, tid, cq_id=0)
            ttnn.execute_trace(dev, tid, cq_id=0, blocking=True)
            t1 = time.perf_counter()
            for _ in range(5):
                ttnn.execute_trace(dev, tid, cq_id=0, blocking=False)
            ttnn.synchronize_device(dev)
            us = (time.perf_counter() - t1) / 25 * 1e6
            ttnn.release_trace(dev, tid)
            for r in rs:
                ttnn.deallocate(r)
            return us

        t_k = tm(lambda: sk(img, bias))
        t_s = tm(lambda: model.stem(img).tensor)
        ok = ok and err <= 1.5 * serr + 1e-7
        line += (f" | vs stock stem: {nd} / {sn.size} values differ, stock rel err vs float64 {serr:.2e} | "
                 f"kernel {t_k:.1f} us, stock {t_s:.1f} us")
        model.release()
    print(line + ("  OK" if ok else "  FAIL"), flush=True)
    rc |= 0 if ok else 1
    sk.release()
    close_device(dev)
    print("ALL OK" if rc == 0 else "FAIL")
    return rc


def config_sample():
    from pathlib import Path

    import tt_sceneseg

    return Path(tt_sceneseg.__file__).parent / "samples" / "highway_normal_1.png"


if __name__ == "__main__":
    sys.exit(main())