"""Probe whether the Triton in this image can compile and run the kernel features
DeepInfer's H20 Triton backend actually uses.

Each case is a miniature of a real kernel, so a failure names the feature that would
block the port rather than just "triton is broken".
"""

import os
import sys

import torch
import triton
import triton.language as tl

print("=" * 72)
print("python ", sys.version.split()[0])
print("torch  ", torch.__version__, "cuda", torch.version.cuda)
print("triton ", triton.__version__)
print("triton@", os.path.dirname(triton.__file__))
print("device ", torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0))
print("flaggems imported?", "flag_gems" in sys.modules)
try:
    from triton.runtime import driver

    print("driver ", type(driver.active).__name__, "target", getattr(driver.active, "target", None))
except Exception as exc:  # noqa: BLE001
    print("driver probe failed:", exc)
print("=" * 72)


@triton.jit
def _add(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
    off = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = off < n
    tl.store(out_ptr + off, tl.load(x_ptr + off, mask=mask) + tl.load(y_ptr + off, mask=mask), mask=mask)


@triton.jit
def _matmul(a_ptr, b_ptr, c_ptr, M, N, K, BM: tl.constexpr, BN: tl.constexpr, BK: tl.constexpr):
    rm = tl.program_id(0) * BM + tl.arange(0, BM)
    rn = tl.program_id(1) * BN + tl.arange(0, BN)
    rk = tl.arange(0, BK)
    acc = tl.zeros((BM, BN), dtype=tl.float32)
    for k in range(0, tl.cdiv(K, BK)):
        kk = k * BK + rk
        a = tl.load(a_ptr + rm[:, None] * K + kk[None, :], mask=(rm[:, None] < M) & (kk[None, :] < K), other=0.0)
        b = tl.load(b_ptr + kk[:, None] * N + rn[None, :], mask=(kk[:, None] < K) & (rn[None, :] < N), other=0.0)
        acc = tl.dot(a, b, acc)
    tl.store(c_ptr + rm[:, None] * N + rn[None, :], acc.to(tl.float16), mask=(rm[:, None] < M) & (rn[None, :] < N))


@triton.jit
def _gather(pool_ptr, idx_ptr, out_ptr, n, BLOCK: tl.constexpr):
    """Data-dependent gather inside a tile: this is how the compressor picks rows."""
    lanes = tl.arange(0, BLOCK)
    pool = tl.load(pool_ptr + lanes)
    off = tl.program_id(0) * BLOCK + lanes
    mask = off < n
    idx = tl.load(idx_ptr + off, mask=mask, other=0) % BLOCK
    tl.store(out_ptr + off, tl.gather(pool, idx, 0), mask=mask)


@triton.jit
def _rope_pair(x_ptr, out_ptr, n_rows, D: tl.constexpr, BLOCK: tl.constexpr):
    """RoPE pairs (2k, 2k+1) or (k, k+D/2); both are one xor of the lane index."""
    row = tl.program_id(0)
    lanes = tl.arange(0, BLOCK)
    base = x_ptr + row * D
    left = tl.load(base + lanes)
    right = tl.gather(left, lanes ^ 1, 0)
    tl.store(out_ptr + row * D + lanes, left + right)


@triton.jit
def _fp8_quant(x_ptr, out_ptr, sf_ptr, n, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    off = pid * BLOCK + tl.arange(0, BLOCK)
    mask = off < n
    v = tl.load(x_ptr + off, mask=mask, other=0.0).to(tl.float32)
    amax = tl.max(tl.abs(v))
    sf = tl.where(amax > 0.0, amax / 448.0, 1.0)
    tl.store(sf_ptr + pid, sf)
    tl.store(out_ptr + off, (v / sf).to(tl.float8e4nv), mask=mask)


@triton.jit
def _int64_cdf(p_ptr, draw_ptr, out_ptr, n, BLOCK: tl.constexpr):
    """The sampler's CDF is int64 on purpose: uint32 comparisons go signed in Triton."""
    off = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
    mask = off < n
    cdf = tl.load(p_ptr + off, mask=mask, other=0).to(tl.int64)
    draw = tl.load(draw_ptr).to(tl.int64)
    tl.store(out_ptr + off, (cdf > draw).to(tl.int32), mask=mask)


def run(name, fn):
    try:
        fn()
        print(f"[ OK ] {name}")
    except Exception as exc:  # noqa: BLE001
        print(f"[FAIL] {name}: {type(exc).__name__}: {str(exc)[:400]}")


def case_add():
    n = 1 << 20
    x, y = torch.randn(n, device="cuda"), torch.randn(n, device="cuda")
    out = torch.empty_like(x)
    _add[(triton.cdiv(n, 1024),)](x, y, out, n, 1024)
    torch.cuda.synchronize()
    err = (out - (x + y)).abs().max().item()
    assert err < 1e-5, err
    print(f"       max abs err {err:.2e}")


def case_matmul():
    M = N = K = 512
    a = torch.randn(M, K, device="cuda", dtype=torch.float16)
    b = torch.randn(K, N, device="cuda", dtype=torch.float16)
    c = torch.empty(M, N, device="cuda", dtype=torch.float16)
    _matmul[(triton.cdiv(M, 64), triton.cdiv(N, 64))](a, b, c, M, N, K, 64, 64, 32)
    torch.cuda.synchronize()
    ref = a.float() @ b.float()
    rel = ((c.float() - ref).abs().max() / ref.abs().max()).item()
    assert rel < 5e-2, rel
    print(f"       rel err {rel:.4f} (fp16 accumulate in fp32)")


def case_gather():
    n, BLOCK = 4096, 256
    pool = torch.arange(BLOCK, device="cuda", dtype=torch.int32)
    idx = torch.randint(0, BLOCK, (n,), device="cuda", dtype=torch.int32)
    out = torch.empty(n, device="cuda", dtype=torch.int32)
    _gather[(triton.cdiv(n, BLOCK),)](pool, idx, out, n, BLOCK)
    torch.cuda.synchronize()
    assert torch.equal(out, pool[idx.long()]), "gather returned wrong rows"
    print("       rows match pool[idx]")


def case_rope_pair():
    rows, D = 64, 128
    x = torch.randn(rows, D, device="cuda")
    out = torch.empty_like(x)
    _rope_pair[(rows,)](x, out, rows, D, D)
    torch.cuda.synchronize()
    # `lanes ^ 1` swaps adjacent lanes, so the reference is that permutation, not a roll.
    ref = x + x[:, torch.arange(D, device="cuda") ^ 1]
    err = (out - ref).abs().max().item()
    assert err < 1e-5, err
    print("       adjacent-lane pairing exact")


def case_fp8():
    n, BLOCK = 1 << 18, 1024
    x = torch.randn(n, device="cuda")
    out = torch.empty(n, device="cuda", dtype=torch.float8_e4m3fn)
    nblk = triton.cdiv(n, BLOCK)
    sf = torch.empty(nblk, device="cuda")
    _fp8_quant[(nblk,)](x, out, sf, n, BLOCK)
    torch.cuda.synchronize()
    deq = out.float() * sf.repeat_interleave(BLOCK)[:n]
    rel = ((deq - x).abs().max() / x.abs().max()).item()
    assert rel < 1e-1, rel
    print(f"       fp8 e4m3 per-1024 roundtrip rel {rel:.4f}")


def case_int64_cdf():
    n, BLOCK = 4096, 256
    cdf = torch.arange(n, device="cuda", dtype=torch.int64) * 1000
    draw = torch.tensor([123456789], device="cuda", dtype=torch.int64)
    out = torch.empty(n, device="cuda", dtype=torch.int32)
    _int64_cdf[(triton.cdiv(n, BLOCK),)](cdf, draw, out, n, BLOCK)
    torch.cuda.synchronize()
    expect = (cdf > draw[0]).to(torch.int32)
    assert torch.equal(out, expect), "int64 comparison went wrong"
    print("       int64 compare exact")


def report_ptx():
    ck = _add.warmup(torch.empty(1, device="cuda"), torch.empty(1, device="cuda"), torch.empty(1, device="cuda"), 1, BLOCK=1024, grid=(1,))
    head = [ln.strip() for ln in ck.asm["ptx"].splitlines() if ln.strip().startswith((".version", ".target"))]
    print("      ", head)
    print("       num_warps", ck.metadata.num_warps, "shared", ck.metadata.shared)


for label, fn in [
    ("elementwise add", case_add),
    ("tl.dot matmul (tensor core)", case_matmul),
    ("tl.gather on tile (data-dependent)", case_gather),
    ("tl.gather for RoPE pairing", case_rope_pair),
    ("fp8e4nv cast + block scale", case_fp8),
    ("int64 CDF comparison", case_int64_cdf),
    ("PTX target", report_ptx),
]:
    run(label, fn)

print("\nDONE")