# SPDX-License-Identifier: Apache-2.0
"""MB0: prove the ``ttnn.generic_op`` + ``UnifiedKernelDescriptor`` plumbing works on the full 11x10 grid.

Device probe for docs/research/device-stack-smoke.md (target 5). It

1. opens the device and records the compute grid, DRAM grid, the runtime
   ``get_optimal_dram_bank_to_logical_worker_assignment`` list (the harvest-pattern question of
   unified-kernels-11x10-audit.md section 1.3) and any harvesting info exposed by the bindings;
2. builds a HEIGHT_SHARDED uint32 output tensor with one [32, 32] shard per compute core;
3. compiles ``kernels/hello_kernel.cpp`` through ``UnifiedKernelDescriptor`` (the same descriptor path the
   deepseek_v3_b1 micro-ops use) for every core of the grid and runs it via ``ttnn.generic_op``;
4. checks that each core wrote (logical_x, logical_y, row-major linear id, magic) into its own shard;
5. reports cold (JIT compile) and warm (program cache) wall times.

Run through the device lock::

    bin/with-device.sh python models/experimental/gr00t/benchmarks/mb0_generic_op_hello.py
"""

from __future__ import annotations

import os
import sys
import time

import torch

import ttnn
from models.demos.deepseek_v3_b1.unified_kernel_descriptor import (
    UnifiedCompileTimeCoreDescriptor,
    UnifiedKernelDescriptor,
)

KERNEL = "models/experimental/gr00t/benchmarks/kernels/hello_kernel.cpp"
MAGIC = 0x00C0FFEE
SHARD = (32, 32)  # 1024 uint32 words per core; we only use the first 4


def _probe_device_facts(device) -> None:
    grid = device.compute_with_storage_grid_size()
    dram = device.dram_grid_size()
    print(
        f"FACT arch={ttnn.get_arch_name()} compute_grid={grid.x}x{grid.y} ({grid.x * grid.y} cores) dram_grid={dram.x}x{dram.y}"
    )
    workers = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_0)
    print(f"FACT optimal_dram_bank_workers(NOC_0)={[(c.x, c.y) for c in workers]}")
    workers1 = device.get_optimal_dram_bank_to_logical_worker_assignment(ttnn.NOC.NOC_1)
    print(f"FACT optimal_dram_bank_workers(NOC_1)={[(c.x, c.y) for c in workers1]}")
    pinned = [(0, 9), (0, 0), (0, 7), (0, 3), (7, 9), (7, 1), (7, 6), (7, 4)]
    print(f"FACT deepseek_pinned_list_matches_runtime={[(c.x, c.y) for c in workers] == pinned}")
    for name in dir(device):
        if "harvest" in name.lower():
            try:
                print(f"FACT device.{name}={getattr(device, name)()}")
            except Exception as exc:  # noqa: BLE001 - probe only
                print(f"FACT device.{name} -> {exc!r}")
    for x in range(grid.x):
        for y in (0, grid.y - 1):
            phys = device.worker_core_from_logical_core(ttnn.CoreCoord(x, y))
            print(f"FACT logical({x},{y}) -> noc0/virtual ({phys.x},{phys.y})")


def run_hello(device, grid_end: ttnn.CoreCoord, tag: str) -> float:
    """Compile+run the hello kernel on the rectangle (0,0)-(grid_end); return the wall time in seconds."""
    grid = ttnn.CoreRangeSet([ttnn.CoreRange(ttnn.CoreCoord(0, 0), grid_end)])
    num_cores = (grid_end.x + 1) * (grid_end.y + 1)
    torch_out = torch.zeros((num_cores * SHARD[0], SHARD[1]), dtype=torch.int32)
    shard_spec = ttnn.ShardSpec(grid, list(SHARD), ttnn.ShardOrientation.ROW_MAJOR)
    mem = ttnn.MemoryConfig(ttnn.TensorMemoryLayout.HEIGHT_SHARDED, ttnn.BufferType.L1, shard_spec)
    out_t = ttnn.from_torch(torch_out, dtype=ttnn.uint32, layout=ttnn.TILE_LAYOUT, device=device, memory_config=mem)
    # generic_op insists on >= 1 input and >= 1 output tensor (generic_op_device_operation.cpp:135); the kernel
    # never reads this dummy input, it only exists to satisfy that precondition.
    in_t = ttnn.from_torch(torch.zeros(32, 32, dtype=torch.int32), dtype=ttnn.uint32, layout=ttnn.TILE_LAYOUT, device=device)

    out_cb = 0
    cb = ttnn.cb_descriptor_from_sharded_tensor(out_cb, out_t)
    num_pages = 1
    cb.format_descriptors[0].page_size = SHARD[0] * SHARD[1] * 4

    ncrisc_args = [
        ("hello_out_cb", out_cb),
        ("hello_num_pages", num_pages),
        ("hello_grid_start_x", 0),
        ("hello_grid_start_y", 0),
        ("hello_grid_end_x", grid_end.x),
        ("hello_grid_end_y", grid_end.y),
    ]
    ukd = UnifiedKernelDescriptor(
        kernel_source=KERNEL,
        core_ranges=grid,
        ncrisc_named_compile_time_args=ncrisc_args,
        brisc_named_compile_time_args=[],
        trisc_named_compile_time_args=[],
        ncrisc_common_runtime_args=[MAGIC],
        trisc_compute_config=ttnn.ComputeConfigDescriptor(
            math_fidelity=ttnn.MathFidelity.LoFi, math_approx_mode=False, fp32_dest_acc_en=False, dst_full_sync_en=False
        ),
        unified_compile_time_core_descriptors=[
            UnifiedCompileTimeCoreDescriptor(
                named_compile_time_arg="is_active_core", core_range=grid, value=1, other_value=0
            )
        ],
    )
    program = ttnn.ProgramDescriptor(kernels=ukd.get_kernel_descriptors().kernels, cbs=[cb])

    t0 = time.perf_counter()
    ttnn.generic_op([in_t, out_t], program)
    ttnn.synchronize_device(device)
    dt = time.perf_counter() - t0

    got = ttnn.to_torch(out_t).view(num_cores, SHARD[0] * SHARD[1])
    bad = []
    for x in range(grid_end.x + 1):
        for y in range(grid_end.y + 1):
            lin = y * (grid_end.x + 1) + x  # row-major id, matches linear_id_in_grid<true>
            # HEIGHT_SHARDED shards are laid out in ROW_MAJOR core order -> shard index == lin
            words = got[lin, :4].tolist()
            if words != [x, y, lin, MAGIC]:
                bad.append(((x, y), words))
    status = "OK" if not bad else f"MISMATCH on {len(bad)} cores, first {bad[:4]}"
    print(
        f"RESULT hello[{tag}] cores={num_cores} grid=(0,0)-({grid_end.x},{grid_end.y}) wall={dt * 1e3:.2f} ms {status}"
    )
    if bad:
        raise SystemExit(f"hello kernel verification failed: {status}")
    ttnn.deallocate(out_t)
    ttnn.deallocate(in_t)
    return dt


def main() -> int:
    print(
        f"ENV TT_METAL_CACHE={os.environ.get('TT_METAL_CACHE', '<default>')} slow_dispatch={os.environ.get('TT_METAL_SLOW_DISPATCH_MODE')}"
    )
    t0 = time.perf_counter()
    device = ttnn.open_device(device_id=0)
    print(f"RESULT open_device {time.perf_counter() - t0:.2f} s")
    try:
        _probe_device_facts(device)
        grid = device.compute_with_storage_grid_size()
        full = ttnn.CoreCoord(grid.x - 1, grid.y - 1)
        cold = run_hello(device, full, "cold")
        warm = run_hello(device, full, "warm")
        warm2 = run_hello(device, full, "warm2")
        # A second grid shape forces a second JIT (different CT args) - measures compile alone (firmware cached)
        cold_small = run_hello(device, ttnn.CoreCoord(grid.x - 1, 0), "cold_1row")
        print(
            f"RESULT compile_estimate full_grid={cold - warm:.3f} s (cold {cold:.3f} s, warm {warm:.4f}/{warm2:.4f} s); "
            f"second-variant cold={cold_small:.3f} s"
        )
    finally:
        ttnn.close_device(device)
    return 0


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