# SPDX-License-Identifier: Apache-2.0
"""Probe (OPT round 6, R6.1): grid / shard orientation of the large block-sharded 3x3 convs at their model shapes.

ttnn block-shards e.g. the eight 512 -> 512 3x3 convs of the CustomFPN / DepthNet ([6, 16, 44, 512]) ROW_MAJOR on
8 x 10 cores (16 channel tiles do not split over 12 columns; 132 tile rows -> 14 per grid row). Candidates
``T<x>x<y>``: ``transpose_shards`` (COL_MAJOR: rows over the x columns, channels over the y grid rows) with
``override_sharding_config`` on an x-by-y grid; ``Tauto``: ttnn's own grid with ``transpose_shards``. Each candidate runs
(a) from a DRAM input with a DRAM output (``slicing="l1"``: I2S / S2I inside; the model's form for most convs) and
(b) from an input already in the candidate's own L1 shards with the output left there (the DN_L1 chain case). Per
candidate: output bits against the base (same input), max |diff|, device time per call (a trace of ``reps`` calls,
best of 3 blocking replays).

    bin/devrun -t 900 -- python -u code/scripts/probe_conv512.py --json logs/bevdet/opt_r6/probe_c512.json \
        --conv img_depth_net... --shape 6,16,44,512 --act relu --cands base,T12x8
"""
from __future__ import annotations

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

import numpy as np


def timed(device, fn, reps):
    """Device time per call: a trace of ``reps`` calls, best of 3 blocking replays. The trace capture is always
    ended (an error inside the capture would otherwise leave the device capturing)."""
    import ttnn

    fn().deallocate()
    ttnn.synchronize_device(device)
    outs = []
    tid = ttnn.begin_trace_capture(device, cq_id=0)
    try:
        for _ in range(reps):
            outs.append(fn())
    finally:
        ttnn.end_trace_capture(device, tid, cq_id=0)
    best = 1e9
    try:
        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)
    finally:
        ttnn.release_trace(device, tid)
        for o in outs:
            o.deallocate()
    return best


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--json")
    ap.add_argument("--reps", type=int, default=8)
    ap.add_argument("--conv", default="img_bev_encoder_neck.conv.3",
                    help="module name for params.provider (DN-relative names start with 'dn:')")
    ap.add_argument("--shape", default="6,16,44,512", help="n,h,w,cin")
    ap.add_argument("--act", default="relu")
    ap.add_argument("--cands", default="base,T12x8,T11x8,Tauto")
    ap.add_argument("--abh", type=int, default=0, help="act_block_h_override for every candidate")
    ap.add_argument("--nodb", action="store_true", help="no double buffering")
    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.conv_tune import DB
    from tt_bevdet.tt.params import DN, BEVDetParams
    from tt_bevdet.ttaw.ops.conv import Conv2d, FeatureMap

    cfg = default_config()
    params = BEVDetParams(BEVDetWeights(paths.weights_onnx()), cfg)
    prec = str(KNOBS.read().PRECISION)
    name = DN + a.conv[3:] if a.conv.startswith("dn:") else a.conv
    p = params.provider(name)
    n, h, w, c = (int(v) for v in a.shape.split(","))
    BS = ttnn.TensorMemoryLayout.BLOCK_SHARDED
    base = {} if a.nodb else dict(DB)
    if a.abh:
        base["act_block_h_override"] = a.abh
    act = None if a.act in ("", "none") else a.act

    def grid(x, y):
        return ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(x - 1, y - 1))})

    def cand(label):
        if label in ("base", "auto"):
            return dict(base)
        if label in ("autoT", "h4T", "w2T"):
            return {**base, "shard_layout": BS, "transpose_shards": True}
        if label == "autoBS":
            return {**base, "shard_layout": BS}
        if label == "h4":
            return dict(base)
        if label == "Tauto":
            return {**base, "shard_layout": BS, "transpose_shards": True}
        x, y = (int(v) for v in label[1:].split("x"))
        return {**base, "shard_layout": BS, "override_sharding_config": True, "core_grid": grid(x, y),
                "transpose_shards": label[0] == "T"}

    SLICING = {"auto": "auto", "autoT": "auto", "autoBS": "auto", "h4": ("height", 4), "h4T": ("height", 4),
               "w2T": ("width", 2)}
    dev = open_device(trace_region_size=512 << 20)
    print(f"grid {dev.compute_with_storage_grid_size()}, conv {name} {tuple(p.weight.shape)} in {a.shape} prec {prec}",
          flush=True)
    rng = np.random.default_rng(0)
    rows = n * h * w
    host = torch.from_numpy(rng.standard_normal((1, 1, rows, c)).astype(np.float32)).relu()
    xin = ttnn.from_torch(host, dtype=ttnn.bfloat16, layout=ttnn.TILE_LAYOUT, device=dev,
                          memory_config=ttnn.DRAM_MEMORY_CONFIG)
    fm = FeatureMap(xin, n, h, w, c)
    ref = None
    res = []
    try:
        for label in [s for s in a.cands.split(",") if s]:
            ov = cand(label)
            cc = Conv2d(p.weight, p.bias, stride=p.stride, padding=p.padding, dilation=p.dilation, groups=p.groups,
                        activation=act, precision=prec, slicing=SLICING.get(label, "l1"), conv_config=ov,
                        name="probe." + label)
            row = {"label": label}
            try:
                cc.output_memory_config = None if label not in SLICING else "dram"   # output left in its L1 shards
                o = cc(fm)
                mc_out = o.tensor.memory_config()
                got = ttnn.to_torch(ttnn.from_device(o.tensor)).float()
                o.deallocate()
                if ref is None:
                    ref = got
                d = (got - ref).abs()
                row.update(bad=int((d != 0).sum()), maxdiff=float(d.max()),
                           out_shard=(str(mc_out.shard_spec.shape) + " " + str(mc_out.shard_spec.grid.bounding_box()))
                           if mc_out.is_sharded() else "dram")
                cc.output_memory_config = "dram"
                row["ms_dram_io"] = round(timed(dev, lambda: cc(fm), a.reps), 4)
                cc.output_memory_config = None
                if mc_out.memory_layout == BS and label not in SLICING:
                    xs = ttnn.to_memory_config(xin, mc_out)    # input in this candidate's own sharding (chain)
                    fs = FeatureMap(xs, n, h, w, c)
                    o2 = cc(fs)
                    row["chain_same_bits"] = bool((ttnn.to_torch(ttnn.from_device(o2.tensor)).float() == got).all())
                    o2.deallocate()
                    row["ms_l1_chain"] = round(timed(dev, lambda: cc(fs), a.reps), 4)
                    ttnn.deallocate(xs)
            except Exception as e:  # noqa: BLE001
                row["error"] = str(e).splitlines()[0][:200]
            finally:
                cc.release()
                dev.clear_program_cache()
            print(json.dumps(row), flush=True)
            res.append(row)
            if a.json:
                Path(a.json).write_text(json.dumps(res, indent=1))
    finally:
        close_device(dev)
    return 0


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