#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Round-6 probe (OPT_PLAN status after round 5, item 2): ttnn's DRAM-sliced decoder convs under explicit slice /
Conv2dConfig choices. Each decoder kxk conv that runs the DRAM-sliced path in the served graph (bf16 DRAM input, one
bf16 weight term, fp32 output, the decoder fidelity) is timed in a trace (10 calls, 5 replays) per configuration and
its output compared bit for bit with the automatic configuration.

    bin/devrun -t 1500 -- python -u code/scripts/diag/slice_probe.py [--only dl5,z4] [--cfg auto,h2,...]
"""
from __future__ import annotations

import argparse
import sys
import time

import numpy as np

# name -> (H, W, Cin, Cout, k): the served graph's DRAM-sliced convs (input map size)
SHAPES = {"dl5": (80, 160, 512, 256, 3), "z4": (160, 320, 128, 512, 3), "dl7": (160, 320, 256, 128, 3),
          "dl9": (320, 640, 128, 64, 3), "mrg": (320, 640, 64, 32, 3),
          # L1-path convs (CONV_L1 = 3200 px): second 3x3s of blocks 0-1, S1 Z convs (k3 forms) of blocks 0-2,
          # block 3's Z phases (80x160, DRAM-sliced)
          "dl1": (20, 40, 768, 768, 3), "dl3": (40, 80, 512, 512, 3), "z0": (10, 20, 1280, 3072, 2),
          "z1": (20, 40, 768, 2048, 2), "z2": (40, 80, 512, 2048, 2), "z3p": (80, 160, 256, 256, 2),
          # backbone 1x1 convs (matmul path; stacked [W_hi ; W_lo] outputs): expand of stages 2 / 3 / 4 / 6,
          # project of stages 2 / 3
          "ex2": (160, 320, 16, 192, 1), "ex3": (80, 160, 24, 288, 1), "ex4": (40, 80, 40, 480, 1),
          "ex6": (20, 40, 112, 1344, 1), "pr2": (80, 160, 96, 64, 1), "pr3": (40, 80, 144, 64, 1)}


def _configs():
    import ttnn

    hs, bs = ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.TensorMemoryLayout.BLOCK_SHARDED
    db = {"enable_act_double_buffer": True, "enable_weights_double_buffer": True}
    c = {"auto": ("auto", {}), "db": ("auto", db)}
    for n in (2, 3, 4, 6, 8):
        c[f"h{n}"] = (("height", n), {})
        c[f"h{n}db"] = (("height", n), db)
    for n in (2, 4, 8):
        c[f"w{n}"] = (("width", n), {})
    c["l1"] = ("l1", {})
    c["l1db"] = ("l1", db)
    c["l1adb"] = ("l1", {"enable_act_double_buffer": True})
    c["l1wdb"] = ("l1", {"enable_weights_double_buffer": True})
    c["adb"] = ("auto", {"enable_act_double_buffer": True})
    c["wdb"] = ("auto", {"enable_weights_double_buffer": True})
    for n in (2, 4):
        c[f"h{n}hs"] = (("height", n), {"shard_layout": hs})
        c[f"h{n}bs"] = (("height", n), {"shard_layout": bs})
        for a in (64, 128, 256):
            c[f"h{n}a{a}"] = (("height", n), {"act_block_h_override": a})
    c["hs"] = ("auto", {"shard_layout": hs})
    c["bs"] = ("auto", {"shard_layout": bs})
    c["fid"] = ("auto", {"full_inner_dim": True})
    return c


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", default="")
    ap.add_argument("--cfg", default="")
    ap.add_argument("--fid", default="HiFi3")
    a = ap.parse_args()
    import torch
    import ttnn

    from tt_sceneseg.device import close_device, open_device
    from tt_sceneseg.tt.layers import _BiasFreeConv2d
    from tt_sceneseg.ttaw.ops.conv import CONV_PRECISION, FeatureMap
    from tt_sceneseg.ttaw.precision import Precision
    from tt_sceneseg.ttaw.tensors import round_to_bf16, to_numpy

    prec = Precision.parse(CONV_PRECISION).with_(fidelity=a.fid)
    dev = open_device()
    cfgs = _configs()
    want = [c for c in a.cfg.split(",") if c] or list(cfgs)
    base = "l1" if want[0].startswith("l1") else "auto"
    if base not in want:
        want = [base] + want
    rng = np.random.default_rng(6)
    for name, (H, W, ci, co, k) in SHAPES.items():
        if a.only and name not in a.only.split(","):
            continue
        w = round_to_bf16((rng.standard_normal((co, ci, k, k)) / np.sqrt(ci * k * k)).astype(np.float32))
        x = round_to_bf16(rng.standard_normal((H * W, ci)).astype(np.float32)).reshape(1, 1, H * W, ci)
        xt = ttnn.from_torch(torch.from_numpy(x), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
                             memory_config=ttnn.DRAM_MEMORY_CONFIG)
        fm = FeatureMap(xt, 1, H, W, ci)
        ref = None
        for cname in want:
            slicing, cc = cfgs[cname]
            conv = _BiasFreeConv2d(w, None, padding=k // 2, precision=prec, output_dtype="float32",
                                   conv_config=cc, name=f"{name}.{cname}")
            conv.slicing = slicing
            try:
                y = conv(fm).tensor
                v = to_numpy(y).reshape(-1, y.shape[-1])[:H * W, :co].copy()
                ttnn.deallocate(y)
                if ref is None:
                    ref = v
                same = "IDENTICAL" if np.array_equal(v.view(np.uint32), ref.view(np.uint32)) else \
                    f"DIFF max {float(np.abs(v - ref).max() / np.abs(ref).max()):.2e}"
                tid = ttnn.begin_trace_capture(dev, cq_id=0)
                for _ in range(10):
                    ttnn.deallocate(conv(fm).tensor)
                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)
                print(f"{name} {H}x{W} {ci}->{co} {cname:8s} {us:9.1f} us  {same}", flush=True)
            except Exception as e:  # noqa: BLE001 (a rejected configuration is a result)
                msg = str(e).strip().splitlines()
                print(f"{name} {H}x{W} {ci}->{co} {cname:8s}   FAILED  {msg[0][:160] if msg else type(e).__name__}",
                      flush=True)
            conv.release()
        ttnn.deallocate(xt)
    close_device(dev)
    return 0


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