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

"""Model start-up cost of ``Gr00tTT.from_pretrained`` per DiT backend (megakernel-default pass; WP-K5 review NB 7 /
§8 pre-condition 5): wall time of the weight upload and the module builds, the tensors and bytes ``TTWeights`` puts
on the device, and what the megakernel backend skips (``tt.model.megakernel_weight_names``).

Per invocation the device is opened once through ``tt.model.open_model_device(policy)`` and the model is built
``--rounds`` times for each configuration -- ``skip`` (the default: the megakernel packs the DiT / precompute / encoder
tensors into its arena and ``TTWeights`` leaves them out) and ``keep`` (``load_ttnn_dit_weights=True``: the full
Stage-1 upload next to the arena, what WP-K5 measured) -- alternating ``keep, skip, keep, skip, ...`` so both see the
same OS page-cache state of the ``.tensorbin`` tier; every model is released before the next build.  With
``--dit-backend ttnn`` only the ``keep`` configuration exists (the ttnn head reads those weights).

Output ``benchmarks/results/bench_load_<version>_<ts>.json`` (schema ``gr00t-tt-perf/1``): stages
``from_pretrained/<cfg>`` (median ms over the rounds), ``weights_load/<cfg>``, ``head_build/<cfg>`` (the megakernel
arena plan + pack + upload sit inside it); ``meta.configs`` holds the per-round numbers, the tensor counts, the bytes
on the device and the skipped set, ``meta.device_mb`` the MB of plan tensors per configuration and ``meta.delta`` the
keep - skip savings.

Run::

    DEVICE_LOCK_TIMEOUT=14400 /home/deepgadget/experiments/gr00t/bin/with-device.sh \\
        python -m models.experimental.gr00t.benchmarks.bench_load --version n16 [--rounds 2] [--dit-backend ttnn]
"""

from __future__ import annotations

import argparse
import statistics
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence

from models.experimental.gr00t.benchmarks import bench_e2e as E2E
from models.experimental.gr00t.benchmarks import bench_stage_ops as SO
from models.experimental.gr00t.tt import device as D
from models.experimental.gr00t.tt.policy import DEFAULT_DIT_BACKEND, DIT_BACKENDS, MK_ARENA_DTYPES

BENCH_NAME = "bench_load"
CONFIGS: Dict[str, Optional[bool]] = {"keep": True, "skip": None}  # -> from_pretrained(load_ttnn_dit_weights=)


def configs_for(dit_backend: str) -> Dict[str, Optional[bool]]:
    """The configurations that exist for a backend: both for the megakernel, ``keep`` only for ttnn."""
    if dit_backend == "megakernel":
        return dict(CONFIGS)
    return {"keep": None}


def one_build(
    setup: Dict[str, Any], device: Any, cache_root: Optional[Path], load_ttnn: Optional[bool]
) -> Dict[str, Any]:
    """Build the model once, record the timings / weight facts, release it."""
    t0 = time.perf_counter()
    model = E2E.load_model(setup, device, cache_root, load_ttnn_dit_weights=load_ttnn)
    total = time.perf_counter() - t0
    try:
        timing = {k: float(v) for k, v in model.timing.items() if k.endswith("_s")}
        head = {k: float(v) for k, v in model.head.timing.items()}
        rec = {
            "from_pretrained_s": total,
            "weights_load_s": timing.get("weights_load_s"),
            "backbone_build_s": timing.get("backbone_build_s"),
            "adapter_build_s": timing.get("adapter_build_s"),
            "embed_gather_build_s": timing.get("embed_gather_build_s"),
            "executor_and_head_build_s": timing.get("executor_and_head_build_s"),
            "head_timing": head,
            "n_tensors": len(model.W),
            "device_bytes": float(model.W.device_bytes),
            "load_path": model.W.load_stats.get("path"),
            "n_warm": model.W.load_stats.get("n_warm"),
            "n_cold": model.W.load_stats.get("n_cold"),
            "skipped": dict(model.weights_skipped),
            "dit_backend": model.policy.dit_backend,
            "mk_arena_dtype": model.recipe.mk_arena_dtype,
            "adapter_mem": model.adapter_mem,
        }
    finally:
        model.release()
    import ttnn

    ttnn.synchronize_device(device)
    return rec


def run(
    device: Any,
    *,
    version: str,
    policy: str = "mixed_dit",
    trace_layout: str = "per_stage",
    dit_backend: str = DEFAULT_DIT_BACKEND,
    mk_arena_dtype: str = "auto",
    rounds: int = 2,
    cache_root: Optional[Path] = None,
    verbose: bool = True,
) -> Dict[str, Any]:
    if rounds <= 0:
        raise ValueError("rounds must be positive")
    setup = E2E.resolve_setup(
        version, policy=policy, trace_layout=trace_layout, dit_backend=dit_backend, mk_arena_dtype=mk_arena_dtype
    )
    cfgs = configs_for(dit_backend)
    per_cfg: Dict[str, List[Dict[str, Any]]] = {name: [] for name in cfgs}
    for r in range(rounds):
        for name, load_ttnn in cfgs.items():
            rec = one_build(setup, device, cache_root, load_ttnn)
            rec["round"] = r
            per_cfg[name].append(rec)
            if verbose:
                print(
                    f"[bench_load] {version}/{dit_backend}/{name} round {r}: from_pretrained {rec['from_pretrained_s']:.2f} s "
                    f"(weights {rec['weights_load_s']:.2f} s, {rec['n_tensors']} tensors, {rec['device_bytes'] / 1e6:.0f} MB, "
                    f"{rec['load_path']}; head {rec['executor_and_head_build_s']:.2f} s"
                    f"{'; skipped ' + str(rec['skipped'].get('n_tensors')) + ' tensors / ' + format(rec['skipped'].get('device_bytes', 0) / 1e6, '.0f') + ' MB' if rec['skipped'] else ''})"
                )
    stages: Dict[str, Dict[str, Any]] = {}
    for name, recs in per_cfg.items():

        def med(key: str) -> float:
            vals = [float(x[key]) for x in recs if x.get(key) is not None]
            return statistics.median(vals) if vals else float("nan")

        stages[f"from_pretrained/{name}"] = {
            "metric": "ms",
            "value": med("from_pretrained_s") * 1e3,
            "min": min(float(x["from_pretrained_s"]) for x in recs) * 1e3,
            "per_round": [round(float(x["from_pretrained_s"]) * 1e3, 1) for x in recs],
        }
        stages[f"weights_load/{name}"] = {"metric": "ms", "value": med("weights_load_s") * 1e3}
        stages[f"head_build/{name}"] = {
            "metric": "ms",
            "value": med("executor_and_head_build_s") * 1e3,
            "note": "executor + ActionHeadTT build (megakernel: arena plan + pack + upload + io inside)",
        }
    facts = D.device_facts(device)
    facts["worker_l1_size"] = E2E._worker_l1_size(device)
    meta = dict(
        setup["meta"],
        configs={name: recs for name, recs in per_cfg.items()},
        device_mb={
            name: {"mb": recs[0]["device_bytes"] / 1e6, "n_tensors": recs[0]["n_tensors"]}
            for name, recs in per_cfg.items()
        },
        rounds=rounds,
        worker_l1_size=facts["worker_l1_size"],
    )
    if "skip" in per_cfg and "keep" in per_cfg:
        k, s = per_cfg["keep"][0], per_cfg["skip"][0]
        meta["delta"] = {
            "device_mb_saved": (k["device_bytes"] - s["device_bytes"]) / 1e6,
            "n_tensors_saved": k["n_tensors"] - s["n_tensors"],
            "from_pretrained_ms_saved_median": stages["from_pretrained/keep"]["value"]
            - stages["from_pretrained/skip"]["value"],
            "weights_load_ms_saved_median": stages["weights_load/keep"]["value"] - stages["weights_load/skip"]["value"],
        }
    payload = SO.perf_payload(
        BENCH_NAME,
        version,
        dry_run=False,
        args={
            "version": version,
            "policy": policy,
            "trace_layout": trace_layout,
            "dit_backend": dit_backend,
            "mk_arena_dtype": mk_arena_dtype,
            "rounds": rounds,
        },
        method={
            "builds": f"{rounds} round(s) of alternating keep / skip Gr00tTT.from_pretrained on one open device, "
            "each model released before the next build; medians over the rounds",
            "clock": "host wall clock (time.perf_counter)",
        },
        stages=stages,
        rows=[],
        device=facts,
        meta=meta,
    )
    return payload


def build_parser() -> argparse.ArgumentParser:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--version", required=True, choices=E2E.VERSIONS)
    ap.add_argument("--policy", default="mixed_dit", choices=E2E.POLICIES)
    ap.add_argument("--trace-layout", default="per_stage", choices=E2E.TRACE_LAYOUTS)
    ap.add_argument("--dit-backend", default=DEFAULT_DIT_BACKEND, choices=DIT_BACKENDS)
    ap.add_argument("--mk-arena-dtype", default="auto", choices=MK_ARENA_DTYPES)
    ap.add_argument("--rounds", type=int, default=2)
    ap.add_argument("--cache-root", default=None)
    ap.add_argument("--out", default=str(SO.RESULTS_DIR))
    ap.add_argument("--trace-region-mb", type=int, default=64)
    ap.add_argument("--num-cqs", type=int, default=2, choices=(1, 2))
    return ap


def main(argv: Optional[Sequence[str]] = None) -> int:
    args = build_parser().parse_args(argv)
    if E2E.import_gr00t_tt() is None:
        raise ImportError("tt.model.Gr00tTT is not importable")
    SO.ensure_tt_metal_cache()
    from models.experimental.gr00t.tt import model as M
    from models.experimental.gr00t.tt.policy import TTPolicy

    device = M.open_model_device(
        TTPolicy(
            dtype_policy=args.policy,
            trace_layout=args.trace_layout,
            dit_backend=args.dit_backend,
            mk_arena_dtype=args.mk_arena_dtype,
        ),
        trace_region_size=args.trace_region_mb << 20,
        l1_small_size=D.DEFAULT_L1_SMALL_SIZE,
        num_command_queues=args.num_cqs,
    )
    try:
        payload = run(
            device,
            version=args.version,
            policy=args.policy,
            trace_layout=args.trace_layout,
            dit_backend=args.dit_backend,
            mk_arena_dtype=args.mk_arena_dtype,
            rounds=args.rounds,
            cache_root=Path(args.cache_root) if args.cache_root else None,
        )
    finally:
        D.close_gr00t_device(device)
    path = SO.write_perf_json(payload, Path(args.out))
    print(f"[bench_load] wrote {path}")
    return 0


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