# SPDX-License-Identifier: Apache-2.0
"""Probe (OPT round 3, R3.1): per-conv ``Conv2dConfig`` sweep of the model's kxk convs at their real shapes.

The encoder and the decoder are built at the defaults and run once eagerly (random crop / gates / ``bev_in``) with
``Conv2d.__call__`` hooked to record each kxk conv (its settled slicing and config overrides, its input spec). Every
distinct conv (by geometry + settled config) is then rebuilt with the same weights, bias, activation, precision and
output dtype and called on a random DRAM TILE input of that spec with candidate configs: activation / weight double
buffering, other shard layouts, activation block heights, ``full_inner_dim``, split reader / activation reuse. Per
candidate: the bf16 / fp32 output bit for bit against the settled config on the same input, and the device time per
call (one trace of ``reps`` calls, best of 3 blocking replays: the conv's whole call incl. its I2S / halo / S2I).

    bin/devrun -t 1500 -- python -u code/scripts/probe_conv.py --part enc --json logs/bevdet/opt_r3/probe_conv_enc.json
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np


def bits(t):
    import torch
    import ttnn

    x = ttnn.to_torch(ttnn.from_device(t))
    if x.dtype == torch.float32:
        return x.view(torch.int32).numpy()
    return x.to(torch.bfloat16).view(torch.int16).numpy()


def timed(device, fn, reps):
    import ttnn

    outs = []
    fn().deallocate()
    ttnn.synchronize_device(device)
    tid = ttnn.begin_trace_capture(device, cq_id=0)
    for _ in range(reps):
        outs.append(fn())
    ttnn.end_trace_capture(device, tid, cq_id=0)
    best = 1e9
    for _ in range(3):
        ttnn.synchronize_device(device)
        t0 = time.perf_counter()
        ttnn.execute_trace(device, tid, cq_id=0, blocking=True)
        best = min(best, (time.perf_counter() - t0) * 1e3 / reps)
    ttnn.release_trace(device, tid)
    for o in outs:
        o.deallocate()
    return best


def candidates(base_slicing, base_ov, cin, rows):
    """(label, slicing, overrides) list; the first is the settled config."""
    import ttnn

    HS, BS, WS = (ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.TensorMemoryLayout.BLOCK_SHARDED,
                  ttnn.TensorMemoryLayout.WIDTH_SHARDED)
    b = dict(base_ov)
    out = [("base", base_slicing, b)]
    db = {"enable_act_double_buffer": True, "enable_weights_double_buffer": True}
    out += [("adb", base_slicing, {**b, "enable_act_double_buffer": True}),
            ("wdb", base_slicing, {**b, "enable_weights_double_buffer": True}),
            ("db", base_slicing, {**b, **db})]
    if base_slicing == "l1":
        out += [("db+fid", "l1", {**b, **db, "full_inner_dim": True})]
    else:
        out += [(f"db+abh{h}", base_slicing, {**b, **db, "act_block_h_override": h}) for h in (256, 128, 64)]
        out += [(f"hs+db{h or ''}", base_slicing, {**b, **db, "shard_layout": HS, **({"act_block_h_override": h} if h else {})})
                for h in (0, 128, 64)]
    l1 = {k: v for k, v in b.items() if k != "act_block_h_override"}
    for name, lay in (("hs", HS), ("bs", BS), ("ws", WS)):
        for abh in (0, 128, 64):
            if name == "ws" and abh:
                continue
            ov = {**l1, "shard_layout": lay}
            if abh:
                ov["act_block_h_override"] = abh
            out.append((f"{name}{abh or ''}", "l1", ov))
            out.append((f"{name}{abh or ''}+db", "l1", {**ov, **db}))
        if name == "hs":
            out.append(("hs+db+reuse", "l1", {**l1, "shard_layout": HS, **db, "enable_activation_reuse": True}))
            out.append(("hs+db+split", "l1", {**l1, "shard_layout": HS, **db, "force_split_reader": True}))
        if name == "bs":
            out.append(("bs+db+fid", "l1", {**l1, "shard_layout": BS, **db, "full_inner_dim": True}))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--part", choices=("enc", "dec"), required=True)
    ap.add_argument("--json")
    ap.add_argument("--only", default="", help="comma list of conv-name substrings")
    ap.add_argument("--reps", type=int, default=8)
    ap.add_argument("--quick", action="store_true", help="only base / double-buffer candidates")
    a = ap.parse_args()
    import torch
    import ttnn

    from tt_bevdet.api import KNOBS
    from tt_bevdet.device import close_device, open_device
    from tt_bevdet.reference.config import default_config
    from tt_bevdet.reference.weights import BEVDetWeights
    from tt_bevdet.tests import paths
    from tt_bevdet.tt.params import BEVDetParams
    from tt_bevdet.ttaw.ops import conv as convmod
    from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap

    cfg = default_config()
    params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg)
    k = KNOBS.read()
    prec = str(k.PRECISION)
    rec = {}
    orig = Conv2d.__call__

    def hook(self, x, **kw):
        out = orig(self, x, **kw)
        if not self.is_matmul and isinstance(x, FeatureMap):
            t = x.tensor
            rec[self.name] = {"conv": self, "spec": (x.batch, x.height, x.width, x.channels),
                              "slicing": self.slicing, "ov": dict(self.conv_config_overrides),
                              "dram_in": bool(convmod.is_dram(t)), "dtype": str(t.dtype)}
        return out

    dev = open_device(trace_region_size=1024 << 20)
    rng = np.random.default_rng(0)
    res = []
    try:
        Conv2d.__call__ = hook
        if a.part == "enc":
            from tt_bevdet.tt.image_encoder import TtImageEncoder

            mod = TtImageEncoder(params, precision=prec, stem_in="kin", conv_l1=str(k.CONV_L1),
                                 backbone_res=str(k.BACKBONE_RES))
            n = cfg.n_cams
            ih, iw = cfg.input_hw
            crop = ttnn.from_torch(torch.from_numpy(rng.integers(0, 256, (1, 1, n * ih, iw * 3), dtype=np.uint8)),
                                   dtype=ttnn.uint8, layout=ttnn.ROW_MAJOR_LAYOUT, device=dev,
                                   memory_config=ttnn.DRAM_MEMORY_CONFIG)
            g = ttnn.from_torch(torch.full((n, 1, 1, 512), 0.5), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT,
                                device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            d, f = mod(crop, g, g)
            d.deallocate()
            f.deallocate()
        else:
            from tt_bevdet.tt.bev_decoder import TtBEVDecoder
            from tt_bevdet.tt.placement import parse_groups

            dg = parse_groups(str(k.CONV_L1))
            mod = TtBEVDecoder(params, precision=prec,
                               conv_l1="all" if "decall" in dg else ("exact" if "dec" in dg else False),
                               merge=str(k.DEC_MERGE) not in ("0", "false", "False"))
            ny, nx = cfg.bev_hw
            x = ttnn.from_torch(torch.randn(1, 1, ny * nx, 864).relu(), dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT,
                                device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            mod(FeatureMap(x, 1, ny, nx, 864)).deallocate()
            ttnn.deallocate(x)
        Conv2d.__call__ = orig
        mod.release()
        seen = {}
        for name, r in rec.items():
            if not r["dram_in"]:          # e.g. the K-IN stem conv on its own L1 shards
                continue
            c = r["conv"]
            key = (r["spec"], c.out_channels, c.kernel, c.stride, c.padding, c.dilation, c.activation,
                   c.output_dtype, r["slicing"], tuple(sorted(r["ov"].items())))
            seen.setdefault(key, []).append(name)
        print(f"{len(rec)} kxk convs, {len(seen)} distinct", flush=True)
        only = [s for s in a.only.split(",") if s]
        for key, names in seen.items():
            if only and not any(o in nm for o in only for nm in names):
                continue
            r = rec[names[0]]
            c = r["conv"]
            bsz, h, w, cin = r["spec"]
            rows = bsz * h * w
            xin = ttnn.from_torch(torch.randn(1, 1, rows, cin).relu(), dtype=ttnn.bfloat16,
                                  layout=ttnn.TILE_LAYOUT, device=dev, memory_config=ttnn.DRAM_MEMORY_CONFIG)
            fm = FeatureMap(xin, bsz, h, w, cin)
            row = {"names": names, "spec": list(r["spec"]), "cout": c.out_channels, "kernel": list(c.kernel),
                   "stride": list(c.stride), "dilation": list(c.dilation), "slicing": r["slicing"],
                   "ov": {kk: str(v) for kk, v in r["ov"].items()}, "cands": []}
            print(f"== {names} {r['spec']} -> {c.out_channels} k{c.kernel} s{c.stride} d{c.dilation} "
                  f"slicing={r['slicing']} ov={r['ov']}", flush=True)
            ref = None
            cands = candidates(r["slicing"], r["ov"], cin, rows)
            if a.quick:
                cands = cands[:4]
            for label, slicing, ov in cands:
                cc = Conv2d(c.weight, c.bias, stride=c.stride, padding=c.padding, dilation=c.dilation,
                            groups=c.groups, activation=c.activation, precision=c.precision,
                            output_dtype=c.output_dtype, slicing=slicing, conv_config=ov, name=c.name + "." + label)
                try:
                    o = cc(fm)
                    got = bits(o.tensor)
                    o.deallocate()
                    if ref is None:
                        ref = got
                    nbad = int(np.count_nonzero(got != ref)) if got.shape == ref.shape else -1
                    t = timed(dev, lambda: cc(fm), a.reps)
                    row["cands"].append({"label": label, "ms": round(t, 4), "bad": nbad})
                    print(f"   {label:14s} {t * 1e3:8.1f} us  {nbad} differ", flush=True)
                except Exception as e:  # noqa: BLE001 - rejected config / L1 clash
                    msg = str(e).splitlines()[0][:150]
                    row["cands"].append({"label": label, "error": msg})
                    print(f"   {label:14s} error: {msg}", flush=True)
                    if label == "base":
                        break
                finally:
                    cc.release()
                    dev.clear_program_cache()      # frees the cached programs' L1_SMALL config tensors
            ok = [x for x in row["cands"] if "ms" in x and x["bad"] == 0]
            if ok:
                bst = min(ok, key=lambda x: x["ms"])
                row["best_exact"] = bst
                print(f"   best exact: {bst['label']} {bst['ms'] * 1e3:.1f} us (base "
                      f"{row['cands'][0].get('ms', 0) * 1e3:.1f}) x{len(names)}", flush=True)
            res.append(row)
            ttnn.deallocate(xin)
            if a.json:
                Path(a.json).write_text(json.dumps(res, indent=1))
    finally:
        Conv2d.__call__ = orig
        close_device(dev)
    return 0


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