#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""SE-K probe (OPT_PLAN SE-K, ``tt/sek.py``), per the hang protocol: the one-program SE MLP vs the stock two
``layers.Linear`` (fc1 + SiLU, fc2 + sigmoid; two weight terms, two activation terms, EPI) on the 7 distinct SE shapes
of SceneSeg (C, squeeze), random fp32 rows: eager x2 identical, a 2-call trace, the error of both vs float64, and the
time of each in a trace.

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

import argparse
import sys
import time

import numpy as np

SHAPES = [(32, 8), (96, 4), (144, 6), (240, 10), (480, 20), (672, 28), (1152, 48)]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--smallest", action="store_true")
    a = ap.parse_args()

    import torch
    import ttnn

    from tt_sceneseg.device import close_device, open_device
    from tt_sceneseg.tt.layers import Linear
    from tt_sceneseg.tt.sek import SEKernel
    from tt_sceneseg.ttaw.tensors import to_numpy

    dev = open_device()
    rng = np.random.default_rng(2)
    rc = 0

    def tm(fn):
        r = fn()
        ttnn.deallocate(r)
        tid = ttnn.begin_trace_capture(dev, cq_id=0)
        rs = [fn() for _ in range(10)]
        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) / 50 * 1e6
        ttnn.release_trace(dev, tid)
        for r in rs:
            ttnn.deallocate(r)
        return us

    for c, s in (SHAPES[:2] if a.smallest else SHAPES):
        w1 = (rng.standard_normal((s, c)) / np.sqrt(c) / 50).astype(np.float32)      # 1 / N folded: small weights
        b1 = (rng.standard_normal(s) * 0.5).astype(np.float32)
        w2 = (rng.standard_normal((c, s)) / np.sqrt(s)).astype(np.float32)
        b2 = (rng.standard_normal(c) * 0.5).astype(np.float32)
        x = (rng.standard_normal((1, 1, 1, c)) * 30).astype(np.float32)
        fc1 = Linear(w1, b1, activation="silu", output_dtype="float32", weight_terms=2, act_terms=2, name="fc1",
                     epi=True)
        fc2 = Linear(w2, b2, activation="sigmoid", output_dtype="float32", weight_terms=2, act_terms=2, name="fc2",
                     epi=True)
        sek = SEKernel(fc1, fc2)
        xt = ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=dev,
                             memory_config=ttnn.DRAM_MEMORY_CONFIG)
        x64 = x.reshape(-1).astype(np.float64)
        h = x64 @ w1.T.astype(np.float64) + b1
        h = h / (1 + np.exp(-h))
        ref = 1 / (1 + np.exp(-(h @ w2.T.astype(np.float64) + b2)))
        st = fc2(fc1(xt))
        sn = to_numpy(st).reshape(-1)[:c].astype(np.float64)
        ttnn.deallocate(st)
        res = []
        for _ in range(2):
            o = sek(xt)
            res.append(to_numpy(o).reshape(-1)[:c])
            ttnn.deallocate(o)
        same = np.array_equal(res[0].view(np.uint32), res[1].view(np.uint32))
        tid = ttnn.begin_trace_capture(dev, cq_id=0)
        o1 = sek(xt)
        o2 = sek(xt)
        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)[:c].view(np.uint32), res[0].view(np.uint32))
        ttnn.release_trace(dev, tid)
        ttnn.deallocate(o1)
        ttnn.deallocate(o2)
        ek = float(np.abs(res[0] - ref).max())
        es = float(np.abs(sn - ref).max())
        t_k = tm(lambda: sek(xt))
        t_s = tm(lambda: fc2(fc1(xt)))
        ok = same and tr and ek <= max(2 * es, 1e-6)
        rc |= 0 if ok else 1
        print(f"SE C {c:5d} S {s:3d}: eager x2 identical {same}, trace identical {tr}, max abs err vs float64: "
              f"kernel {ek:.2e}, stock {es:.2e} | kernel {t_k:7.1f} us, stock {t_s:7.1f} us  {'OK' if ok else 'FAIL'}",
              flush=True)
        sek.release()
        fc1.release()
        fc2.release()
        ttnn.deallocate(xt)
    close_device(dev)
    print("ALL OK" if rc == 0 else "FAIL")
    return rc


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