# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
# SPDX-License-Identifier: Apache-2.0

"""WP-K4 timing / zone bench of the DiT megakernel (IMPLEMENTATION_PLAN.md §5.7 gate G2, §5.4 timing chain).

Modes::

    # G2: the steps4 program captured in a trace, N replays -> ms per step vs 1.3 x arena bytes / measured GB/s
    bench_mk_step.py --version n16 --dtype bf16 --replays 20
    # per-phase DeviceZoneScopedN zones of the block variant over a few blocks (TT_METAL_DEVICE_PROFILER=1 is set
    # before ttnn is imported; the profiler cannot be combined with the watcher on Blackhole)
    bench_mk_step.py --version n16 --dtype bf16 --profile --blocks 4

Every run writes ``benchmarks/results/bench_mk_step_<version>_<stamp>.json`` (schema ``gr00t-mk-k4/1``).  The
arena is packed from the real plan tensors (``dit_program.real_arena_plan``), the K/V are hoisted on device from the
canonical golden sample and the masks are the call's (``tests/tt/test_mk_dit_block.K4Session``).  Run through
``bin/with-device.sh`` from the tt-metal root; ``ttnn`` is imported lazily in :func:`main`.
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import statistics
import sys
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import torch

TT_METAL_HOME = Path(os.environ.get("TT_METAL_HOME", str(Path(__file__).resolve().parents[4])))
PROFILE_CSV = TT_METAL_HOME / "generated" / "profiler" / ".logs" / "profile_log_device.csv"
RESULTS_DIR = Path(__file__).resolve().parent / "results"
SCHEMA = "gr00t-mk-k4/1"
MEASURED_GBPS: Dict[str, float] = {"bf16": 464.0, "bfp8_b": 414.0}  # mk_k1_summary.md §3 (direct mode)
G2_FACTOR = 1.3


def parse_args(argv: Optional[Sequence[str]] = None) -> argparse.Namespace:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--version", default="n16", choices=("n15", "n16", "n17"))
    ap.add_argument("--dtype", default="bf16", choices=("bf16", "bfp8_b"), help="arena stream dtype")
    ap.add_argument("--replays", type=int, default=20, help="traced steps4 replays for the G2 timing")
    ap.add_argument("--profile", action="store_true", help="collect the per-phase zones of a short block launch")
    ap.add_argument(
        "--profile-variant",
        default="block",
        choices=("block", "step"),
        help="profiled launch: `block` = blocks [0, --blocks) of the block variant; `step` = one whole denoising step "
        "(all blocks + output head; set GR00T_MK_ZONE_ROUNDS=0 so the hub BRISC's 250 profiler markers hold every block)",
    )
    ap.add_argument("--blocks", type=int, default=4, help="blocks of the profiled block launch")
    ap.add_argument(
        "--first-block",
        type=int,
        default=0,
        help="first block of the profiled block launch (the launch's first block streams its weights cold: profile "
        "from block 1 or 2 to see the steady-state prefetch)",
    )
    ap.add_argument("--step", type=int, default=0, help="denoising step of the profiled block launch")
    ap.add_argument(
        "--cb-w-slots", type=int, default=None, help="weight-ring depth in K-block slots (K4Budget default)"
    )
    ap.add_argument(
        "--k5-device",
        action="store_true",
        help="open the device as the K5 model does (tt.model.worker_l1_size_for(TTPolicy(dit_backend='megakernel')), the "
        "64 KiB worker-L1 cut) instead of the K4 test device (96 KiB): the compute-core binaries must fit K5's "
        "kernel-config ring buffer (mk_k4b_summary.md section 5) -- a program too large fails at the first launch",
    )
    ap.add_argument("--results-dir", default=str(RESULTS_DIR))
    return ap.parse_args(argv)


def read_zones(path: Path = PROFILE_CSV) -> Dict[str, Any]:
    """``{"freq_mhz", "zones": {"<RISC>:<zone>": {count, median_us, mean_us, min_us, max_us, cores}}}`` over the
    ``mk_*`` zones of every launch in the profiler CSV (nested zones handled with a per-(run, core, RISC) stack)."""
    if not path.is_file():
        raise FileNotFoundError(f"{path} not found (set TT_METAL_DEVICE_PROFILER=1 before importing ttnn)")
    with open(path) as fh:
        header = fh.readline()
        freq_mhz = float(header.split("CHIP_FREQ[MHz]:")[1].split(",")[0])
        reader = csv.reader(fh)
        cols = [c.strip() for c in next(reader)]
        idx = {c: i for i, c in enumerate(cols)}
        stacks: Dict[tuple, List[Tuple[str, int]]] = defaultdict(list)
        durations: Dict[str, List[int]] = defaultdict(list)
        cores: Dict[str, set] = defaultdict(set)
        for row in reader:
            if len(row) < len(cols) - 1:
                continue
            zone = row[idx["zone name"]].strip()
            if not zone.startswith("mk_"):
                continue
            core = (
                row[idx["run host ID"]].strip(),
                row[idx["core_x"]],
                row[idx["core_y"]],
                row[idx["RISC processor type"]].strip(),
            )
            t = int(row[idx["time[cycles since reset]"]])
            typ = row[idx["type"]].strip()
            if typ == "ZONE_START":
                stacks[core].append((zone, t))
            elif typ == "ZONE_END":
                st = stacks[core]
                for j in range(len(st) - 1, -1, -1):
                    if st[j][0] == zone:
                        name = f"{core[3]}:{zone}"
                        durations[name].append(t - st[j][1])
                        cores[name].add(core[1:3])
                        del st[j]
                        break
    zones = {}
    for name, ds in durations.items():
        us = [d / freq_mhz for d in ds]
        zones[name] = {
            "count": len(us),
            "cores": len(cores[name]),
            "median_us": statistics.median(us),
            "mean_us": sum(us) / len(us),
            "min_us": min(us),
            "max_us": max(us),
        }
    return {"freq_mhz": freq_mhz, "zones": zones}


def g2_record(per_step_ms: float, bytes_per_step: int, dtype: str) -> Dict[str, Any]:
    floor_ms = bytes_per_step / MEASURED_GBPS[dtype] / 1e6
    return {
        "per_step_ms": per_step_ms,
        "bytes_per_step": bytes_per_step,
        "measured_gbps": MEASURED_GBPS[dtype],
        "floor_ms_per_step": floor_ms,
        "gate_ms_per_step": G2_FACTOR * floor_ms,
        "ratio_to_floor": per_step_ms / floor_ms,
        "effective_gbps": bytes_per_step / per_step_ms / 1e6,
        "pass": per_step_ms <= G2_FACTOR * floor_ms,
    }


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = parse_args(argv)
    if args.profile:
        os.environ["TT_METAL_DEVICE_PROFILER"] = "1"
        # the device CSV of the last run and the profiler's accumulated zone-source-location logs: stale entries of
        # earlier builds of the same zone names at other line numbers can collide in its 16-bit hash (profiler.cpp
        # `populateZoneSrcLocations` throws at device close) -- they are rebuilt from this run's compile logs
        for stale in (PROFILE_CSV, *PROFILE_CSV.parent.glob("*zone_src_locations.log")):
            if stale.exists():
                stale.unlink()
    os.environ.setdefault("TT_METAL_CACHE", str(Path.home() / ".cache" / "tt-metal-cache-gr00t"))
    torch.set_grad_enabled(False)
    import ttnn  # noqa: F401  (after the profiler env var)

    from models.experimental.gr00t.common.preprocessing import encode
    from models.experimental.gr00t.reference.model import load_golden_sample
    from models.experimental.gr00t.tests.tt import test_mk_dit_block as TB
    from models.experimental.gr00t.tests.tt.harness import device_facts
    from models.experimental.gr00t.tt.device import close_gr00t_device
    from models.experimental.gr00t.tt.megakernel import descriptors as D
    from models.experimental.gr00t.tt.megakernel import dit_program as P

    TB.STREAM_DTYPE = args.dtype  # the session packs the arena in this dtype and picks the matching oracle policy
    if args.cb_w_slots is not None:
        os.environ["GR00T_MK_CB_W_SLOTS"] = str(args.cb_w_slots)
    gs, obs = load_golden_sample(args.version, "canonical")
    mi = encode(args.version, obs)
    out: Dict[str, Any] = {
        "schema": SCHEMA,
        "created": datetime.now().astimezone().isoformat(timespec="seconds"),
        "version": args.version,
        "dtype": args.dtype,
        "cb_w_slots": args.cb_w_slots,
        "profile": bool(args.profile),
        "profile_variant": args.profile_variant if args.profile else None,
        "env": {k: v for k, v in os.environ.items() if k.startswith("GR00T_MK_")},
        "argv": list(argv) if argv is not None else sys.argv[1:],
    }
    worker_l1 = None
    if args.k5_device:
        from models.experimental.gr00t.tt.model import worker_l1_size_for
        from models.experimental.gr00t.tt.policy import TTPolicy

        worker_l1 = int(worker_l1_size_for(TTPolicy(dit_backend="megakernel")))
    out["k5_device"] = bool(args.k5_device)
    out["worker_l1_size"] = worker_l1 if worker_l1 is not None else P.K4_WORKER_L1_SIZE
    device = P.open_k4_device(worker_l1_size=out["worker_l1_size"])
    try:
        out["device"] = device_facts(device)
        s = TB.K4Session(args.version, (gs, obs, mi), device, variant="block")
        out["arena"] = {
            "n_blocks": s.arena.n_blocks,
            "bytes_per_step": s.mk.arena_bytes_per_step(),
            "stream_bytes_per_step": s.arena.bytes_per_bank("stream") * 8,
            "small_bytes_per_step": s.arena.bytes_per_bank("small") * 8,
            "weights_load_s": s.weights_s,
            "arena_pack_upload_s": s.arena_s,
        }
        if args.profile:
            # one short launch: the zones of every RISC of every core land in the profiler CSV
            tap = "sa_embs" if args.first_block == 0 else f"dit_block{args.first_block - 1}_out"
            x = s.golden_rows(tap, args.step)
            xr = s.mk.make_x_rows(x)
            if args.profile_variant == "step":
                s.mk.set_variant("step")
                t0 = time.perf_counter()
                s.mk.run_step(xr, args.step)
                wall = time.perf_counter() - t0
                n_blk = s.arena.n_blocks
            else:
                t0 = time.perf_counter()
                s.mk.run_block(xr, args.first_block, matmul_only=False, n_blk=args.blocks, step=args.step)
                wall = time.perf_counter() - t0
                n_blk = args.blocks
            s.mk._ttnn.deallocate(xr)
            out["profile_launch"] = {
                "variant": args.profile_variant,
                "blocks": n_blk,
                "first_block": args.first_block if args.profile_variant == "block" else 0,
                "step": args.step,
                "wall_s": wall,
                "zone_rounds": P.K4_ZONE_ROUNDS,
            }
        else:
            s.mk.set_variant("steps4")
            noise = gs.load("initial_noise")
            state = gs.load("state_features").reshape(-1)
            t0 = time.perf_counter()
            ref = s.mk.run_steps4(noise, state)
            out["steps4_untraced_wall_s"] = time.perf_counter() - t0
            s.mk.set_steps4_inputs(noise, state)
            tid = s.mk.capture_steps4()
            try:
                times = [s.mk.replay(tid) for _ in range(args.replays)]
                traced = s.mk.read_steps4_outputs()
            finally:
                s.mk.release_trace(tid)
            ms = sorted(t * 1e3 for t in times)
            out["steps4_trace"] = {
                "replays": args.replays,
                "ms_min": ms[0],
                "ms_median": statistics.median(ms),
                "ms_max": ms[-1],
                "traced_eq_untraced": bool(torch.equal(traced["action_pred"], ref["action_pred"])),
            }
            out["g2"] = g2_record(statistics.median(ms) / P.STEPS4_N_STEPS, s.mk.arena_bytes_per_step(), args.dtype)
            out["g2_min"] = g2_record(ms[0] / P.STEPS4_N_STEPS, s.mk.arena_bytes_per_step(), args.dtype)
        out["stats"] = s.mk.stats().__dict__
        out["k4b"] = {  # the K4b schedule knobs of this run (dit_program / descriptors)
            "mcast_split": P.K4_MCAST_SPLIT,
            "zone_rounds": P.K4_ZONE_ROUNDS,
            "land_pages": D.K4_LAND_PAGES,
            "hub": list(s.cm.hub),
            "hub2": list(s.cm.hub2),
            "split_blocks": [b for b in range(s.arena.n_blocks) if P.is_split_block(s.arena, b)],
            "cb_w_slots": s.mk.k4.cb_w_slots[args.dtype],
        }
        s.release()
    finally:
        close_gr00t_device(device)
    if args.profile:
        zones = read_zones()
        out["zones"] = zones
        print(f"zones ({zones['freq_mhz']} MHz), {args.blocks} blocks:")
        for name, z in sorted(zones["zones"].items()):
            print(
                f"  {name:28s} n={z['count']:5d} cores={z['cores']:3d} median {z['median_us']:8.2f} us  min {z['min_us']:8.2f}  max {z['max_us']:8.2f}"
            )
    else:
        g = out["g2"]
        print(
            f"{args.version} {args.dtype}: steps4 trace median {out['steps4_trace']['ms_median']:.3f} ms "
            f"(min {out['steps4_trace']['ms_min']:.3f}) -> {g['per_step_ms']:.3f} ms/step; floor {g['floor_ms_per_step']:.3f}, "
            f"gate {g['gate_ms_per_step']:.3f}, ratio {g['ratio_to_floor']:.2f}, {g['effective_gbps']:.0f} GB/s -> G2 {'PASS' if g['pass'] else 'FAIL'}"
        )
    Path(args.results_dir).mkdir(parents=True, exist_ok=True)
    path = Path(args.results_dir) / f"bench_mk_step_{args.version}_{datetime.now().strftime('%Y%m%d-%H%M%S')}.json"
    with open(path, "w") as fh:
        json.dump(out, fh, indent=1, default=str)
    print("wrote", path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())