# SPDX-License-Identifier: Apache-2.0
"""Shared helpers for the ``mb1_*`` p150a micro-benchmarks (GR00T port, research §6 unknowns 1/4/5/9/20).

Every ``mb1_*.py`` script is a stand-alone argparse program that opens the device itself; this module only holds
the timing / configuration helpers they share so that all numbers are produced by the same code path.

Timing methodology
------------------
* ``time_eager``  : ``warmup`` untimed calls, ``ttnn.synchronize_device``, then ``iters`` back-to-back calls followed by
  a final ``synchronize_device``.  The per-call figure therefore includes Python + host dispatch overhead
  ("all-in eager", the level-A cost of blackhole-perf-model-v2 §6.1).
* ``time_traced`` : the same ``iters`` calls are captured into one Metal trace and the trace is replayed ``reps`` times
  (blocking, each replay wrapped in ``synchronize_device``).  Per-call time = replay wall time / ``iters``; the
  minimum over ``reps`` is reported as the device-side cost (kernel time + trace-replay dispatch, the level-B cost).
  When the op sequence is long enough (>= 16 ops per trace) the replay figure is within a few µs of the raw kernel
  duration, which is what the perf model needs.

Weight-byte accounting follows blackhole-perf-model-v2 §2.1 (tile bytes: bf16 2048, bfp8_b 1088, bfp4_b 576).
Peak compute follows the same table: 4096 FLOP/cycle/core (LoFi) x 110 cores x AICLK.
"""

from __future__ import annotations

import json
import math
import os
import statistics
import sys
import time
import traceback
from dataclasses import asdict, dataclass, field
from typing import Any, Callable, Dict, List, Optional, Sequence

TILE = 32
TILE_BYTES = {"bf16": 2048, "bfp8": 1088, "bfp4": 576, "fp32": 4096}
DRAM_PEAK_GBS = 512.0  # GDDR6 spec, perf-model-v2 §2.1
PEAK_FLOP_PER_CYCLE_PER_CORE_LOFI = 4096  # perf-model-v2 §2.1
FIDELITY_DIVISOR = {"LoFi": 1, "HiFi2": 2, "HiFi3": 3, "HiFi4": 4}
DEFAULT_AICLK_MHZ = 1350.0
NUM_CORES = 110

# Blackhole DRAM page alignment pads a bfp8 tile (1088 B) to 1152 B in DRAM-sharded layouts; interleaved buffers
# keep the raw tile size.  Only used for the reporting of "bytes actually moved" when the caller asks for it.
DRAM_ALIGN = 64


def dtype_name(dtype_str: str) -> str:
    return {"bf16": "bfloat16", "bfp8": "bfloat8_b", "bfp4": "bfloat4_b", "fp32": "float32"}[dtype_str]


def ttnn_dtype(ttnn, dtype_str: str):
    return {
        "bf16": ttnn.bfloat16,
        "bfp8": ttnn.bfloat8_b,
        "bfp4": ttnn.bfloat4_b,
        "fp32": ttnn.float32,
    }[dtype_str]


def weight_bytes(K: int, N: int, dtype_str: str) -> int:
    """Bytes of a tile-padded [K, N] weight in the given format (perf-model-v2 §2.1 tile sizes)."""
    kt = math.ceil(K / TILE)
    nt = math.ceil(N / TILE)
    return kt * nt * TILE_BYTES[dtype_str]


def act_bytes(M: int, K: int, dtype_str: str = "bf16") -> int:
    return math.ceil(M / TILE) * math.ceil(K / TILE) * TILE_BYTES[dtype_str]


def peak_tflops(fidelity: str, aiclk_mhz: float = DEFAULT_AICLK_MHZ, cores: int = NUM_CORES) -> float:
    return cores * PEAK_FLOP_PER_CYCLE_PER_CORE_LOFI * aiclk_mhz * 1e6 / FIDELITY_DIVISOR[fidelity] / 1e12


def divisors(n: int) -> List[int]:
    return [d for d in range(1, n + 1) if n % d == 0]


def open_device(ttnn, l1_small_size: int = 32768, trace_region_size: int = 96 * 1024 * 1024, num_cqs: int = 1):
    """Open device 0 with program cache enabled (all scripts use the same sizes so program-cache behaviour is equal)."""
    kwargs = dict(device_id=0, l1_small_size=l1_small_size, trace_region_size=trace_region_size)
    if num_cqs != 1:
        kwargs["num_command_queues"] = num_cqs
    device = ttnn.open_device(**kwargs)
    device.enable_program_cache()
    return device


def compute_config(
    ttnn, fidelity: str = "HiFi2", fp32_acc: bool = False, approx: bool = False, packer_l1_acc: bool = False
):
    fid = {
        "LoFi": ttnn.MathFidelity.LoFi,
        "HiFi2": ttnn.MathFidelity.HiFi2,
        "HiFi3": ttnn.MathFidelity.HiFi3,
        "HiFi4": ttnn.MathFidelity.HiFi4,
    }[fidelity]
    return ttnn.WormholeComputeKernelConfig(
        math_fidelity=fid,
        math_approx_mode=approx,
        fp32_dest_acc_en=fp32_acc,
        packer_l1_acc=packer_l1_acc,
    )


@dataclass
class Timing:
    """Per-call timings in microseconds."""

    eager_us: Optional[float] = None
    trace_us: Optional[float] = None  # min over replays
    trace_us_mean: Optional[float] = None
    trace_reps_us: List[float] = field(default_factory=list)
    iters: int = 0
    error: Optional[str] = None

    def best_us(self) -> Optional[float]:
        if self.trace_us is not None:
            return self.trace_us
        return self.eager_us


def _sync(ttnn, device):
    ttnn.synchronize_device(device)


def time_eager(ttnn, device, fn: Callable[[], Any], iters: int, warmup: int = 3, dealloc: bool = True) -> float:
    """Return µs per call for ``iters`` back-to-back eager calls (host dispatch included)."""
    for _ in range(warmup):
        out = fn()
        if dealloc:
            _dealloc(ttnn, out)
    _sync(ttnn, device)
    t0 = time.perf_counter()
    for _ in range(iters):
        out = fn()
        if dealloc:
            _dealloc(ttnn, out)
    _sync(ttnn, device)
    t1 = time.perf_counter()
    return (t1 - t0) / iters * 1e6


def _dealloc(ttnn, out) -> None:
    if out is None:
        return
    if isinstance(out, (list, tuple)):
        for o in out:
            _dealloc(ttnn, o)
        return
    try:
        ttnn.deallocate(out)
    except Exception:  # noqa: BLE001 - output may already be freed by the callee
        pass


def time_traced(
    ttnn, device, fn: Callable[[], Any], iters: int, reps: int = 5, dealloc: bool = True, cq_id: int = 0
) -> Dict[str, Any]:
    """Capture ``iters`` calls of ``fn`` in one trace, replay ``reps`` times; return per-call µs stats."""
    # compile / program-cache warm-up (must happen outside capture)
    out = fn()
    if dealloc:
        _dealloc(ttnn, out)
    _sync(ttnn, device)
    tid = ttnn.begin_trace_capture(device, cq_id=cq_id)
    try:
        for _ in range(iters):
            out = fn()
            if dealloc:
                _dealloc(ttnn, out)
    finally:
        ttnn.end_trace_capture(device, tid, cq_id=cq_id)
    _sync(ttnn, device)
    times = []
    try:
        # one untimed replay
        ttnn.execute_trace(device, tid, cq_id=cq_id, blocking=False)
        _sync(ttnn, device)
        for _ in range(reps):
            t0 = time.perf_counter()
            ttnn.execute_trace(device, tid, cq_id=cq_id, blocking=False)
            _sync(ttnn, device)
            t1 = time.perf_counter()
            times.append((t1 - t0) / iters * 1e6)
    finally:
        ttnn.release_trace(device, tid)
    return {"min": min(times), "mean": statistics.fmean(times), "reps": times}


def measure(
    ttnn,
    device,
    fn: Callable[[], Any],
    iters: int,
    warmup: int = 3,
    reps: int = 5,
    do_eager: bool = True,
    do_trace: bool = True,
    dealloc: bool = True,
) -> Timing:
    """Run both timing methods, capturing any exception as ``Timing.error`` (never raises)."""
    t = Timing(iters=iters)
    try:
        if do_eager:
            t.eager_us = time_eager(ttnn, device, fn, iters, warmup, dealloc)
        if do_trace:
            tr = time_traced(ttnn, device, fn, iters, reps, dealloc)
            t.trace_us = tr["min"]
            t.trace_us_mean = tr["mean"]
            t.trace_reps_us = [round(x, 3) for x in tr["reps"]]
    except Exception as e:  # noqa: BLE001 - we want to record every failing config, not abort the sweep
        t.error = short_error(e)
    return t


def short_error(e: BaseException, limit: int = 600) -> str:
    msg = str(e).strip().replace("\n", " | ")
    if len(msg) > limit:
        msg = msg[:limit] + " ..."
    return f"{type(e).__name__}: {msg}"


def gbs(bytes_moved: int, us: Optional[float]) -> Optional[float]:
    if us is None or us <= 0:
        return None
    return bytes_moved / (us * 1e-6) / 1e9


def tflops(flop: float, us: Optional[float]) -> Optional[float]:
    if us is None or us <= 0:
        return None
    return flop / (us * 1e-6) / 1e12


def pct(x: Optional[float], ref: float) -> Optional[float]:
    if x is None:
        return None
    return 100.0 * x / ref


def r(x: Optional[float], nd: int = 2) -> Optional[float]:
    return None if x is None else round(x, nd)


def write_json(path: str, payload: Dict[str, Any]) -> None:
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    with open(path, "w") as f:
        json.dump(payload, f, indent=1, default=_json_default)
    print(f"[mb1] wrote {path}")


def _json_default(o):
    if hasattr(o, "__dataclass_fields__"):
        return asdict(o)
    return str(o)


def env_report() -> Dict[str, Any]:
    import platform

    rep = {
        "time": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "python": sys.version.split()[0],
        "platform": platform.platform(),
        "TT_METAL_HOME": os.environ.get("TT_METAL_HOME"),
        "TT_METAL_DEVICE_PROFILER": os.environ.get("TT_METAL_DEVICE_PROFILER"),
    }
    try:
        import torch

        rep["torch"] = torch.__version__
    except Exception:  # noqa: BLE001
        pass
    try:
        import subprocess

        rep["tt_metal_commit"] = (
            subprocess.check_output(["git", "rev-parse", "--short", "HEAD"], cwd=os.environ.get("TT_METAL_HOME", "."))
            .decode()
            .strip()
        )
    except Exception:  # noqa: BLE001
        pass
    return rep


def device_report(ttnn, device) -> Dict[str, Any]:
    g = device.compute_with_storage_grid_size()
    d = device.dram_grid_size()
    rep = {"compute_grid": [g.x, g.y], "num_cores": g.x * g.y, "dram_grid": [d.x, d.y]}
    try:
        rep["arch"] = str(device.arch())
    except Exception:  # noqa: BLE001
        pass
    try:
        workers = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_0)
        rep["dram_bank_workers_noc0"] = [[c.x, c.y] for c in workers]
        workers1 = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_1)
        rep["dram_bank_workers_noc1"] = [[c.x, c.y] for c in workers1]
    except Exception as e:  # noqa: BLE001
        rep["dram_bank_workers_error"] = short_error(e)
    return rep


def print_table(rows: Sequence[Dict[str, Any]], cols: Sequence[str], title: str = "") -> None:
    if title:
        print(f"\n== {title}")
    if not rows:
        print("(no rows)")
        return
    widths = {c: max(len(c), *(len(_fmt(rw.get(c))) for rw in rows)) for c in cols}
    print(" | ".join(c.ljust(widths[c]) for c in cols))
    print("-+-".join("-" * widths[c] for c in cols))
    for rw in rows:
        print(" | ".join(_fmt(rw.get(c)).ljust(widths[c]) for c in cols))


def _fmt(v: Any) -> str:
    if v is None:
        return "-"
    if isinstance(v, float):
        return f"{v:.2f}"
    s = str(v)
    return s if len(s) <= 110 else s[:107] + "..."


def default_out_dir() -> str:
    return os.environ.get(
        "MB1_OUT_DIR",
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "results"),
    )