# SPDX-License-Identifier: CC-BY-NC-4.0
# Copyright (c) 2026 Jakub Ćwirlej
"""Correctness tests for eugen_sparse.

    python test_eugen_sparse.py                       # GPU: everything; CPU: PyTorch parts only
    TRITON_INTERPRET=1 python test_eugen_sparse.py    # CPU: the Triton kernels in the interpreter (fp32)

Kernel vs reference_sparse_swiglu_ffn with the SAME group selection: y, dX, dW_gate, dW_up, dW_down.
Geometries cover the fused path (group_size <= 64), the unfused path (group_size 128), several token
chunks (atomic weight gradients), a non-power-of-two number of groups and d_model not divisible by 128.
"""
import os
import sys

import torch

import eugen_sparse as es

INTERP = os.environ.get('TRITON_INTERPRET', '0') == '1'


def rel_err(a, b):
    a = a.float(); b = b.float()
    return ((a - b).abs().max() / b.abs().max().clamp(min=1e-12)).item()


def make(M, D, Fd, dev, dt, seed=0):
    torch.manual_seed(seed)
    x = torch.randn(M, D, device=dev, dtype=dt) * 0.5
    gw = torch.randn(Fd, D, device=dev, dtype=dt) * D ** -0.5
    uw = torch.randn(Fd, D, device=dev, dtype=dt) * D ** -0.5
    dw = torch.randn(D, Fd, device=dev, dtype=dt) * Fd ** -0.5
    return x, gw, uw, dw


def parity(tag, M, D, Fd, G, kg, dev, dt, tol_y, tol_g):
    x, gw, uw, dw = make(M, D, Fd, dev, dt)
    idx = es.grouped_topk_select(x, gw, G, kg)
    dy = torch.randn(M, D, device=dev, dtype=dt)
    res = []
    for fn in (lambda *a: es.sparse_swiglu_fwd(*a, idx, G)[0], lambda *a: es.reference_sparse_swiglu_ffn(*a, idx, G)):
        leaves = [t.clone().requires_grad_() for t in (x, gw, uw, dw)]
        y = fn(*leaves)
        y.backward(dy)
        res.append([y.detach()] + [t.grad for t in leaves])
    errs = [rel_err(a, b) for a, b in zip(*res)]
    ok = errs[0] <= tol_y and all(e <= tol_g for e in errs[1:])
    p = es.plan(M, D, Fd, G, kg)
    print(f'  {"PASS" if ok else "FAIL"}  {tag:30s} chunks={p["token_chunks"]} fused={p["fused"]} '
          f'tiles={tuple(p["tiles"].values())} | y/dX/dWg/dWu/dWd ' + ' '.join(f'{e:.1e}' for e in errs))
    return ok


def test_cpu_reference_and_selection():
    x, gw, uw, dw = make(64, 32, 256, 'cpu', torch.float32)
    sub, spg = es.make_subsample_index(256, 32, 32)                 # full "subsample" = exact energy
    a = es.grouped_topk_select(x, gw, 32, 2, subsample_idx=sub, samples_per_group=spg).sort(-1).values
    b = es.grouped_topk_select(x, gw, 32, 2).sort(-1).values
    sub8, spg8 = es.make_subsample_index(256, 32, 8, seed=1)
    ok_sub = sub8.numel() == 8 * 8 and bool(((sub8 // 32) == torch.arange(8).repeat_interleave(8)).all())
    y = es.sparse_swiglu_ffn(x, gw, uw, dw, es.grouped_topk_select(x, gw, 32, 8), 32)   # all groups active
    dense = torch.nn.functional.linear(torch.nn.functional.silu(x @ gw.t()) * (x @ uw.t()), dw)
    ok = torch.equal(a, b) and ok_sub and rel_err(y, dense) < 1e-6
    print(f'  {"PASS" if ok else "FAIL"}  CPU: full subsample == exact selection, subsample layout, '
          f'all groups == dense SwiGLU ({rel_err(y, dense):.1e})')
    return ok


def test_load_balance_bias():
    ng, target = 8, 0.25
    bias, ema = torch.zeros(ng), torch.full((ng,), target)
    cur = torch.tensor([1.0, 0.0, 0.25, 0.25, 0.25, 0.25, 0.25, 0.25])
    for _ in range(200):
        es.update_load_balance_bias(bias, ema, cur, target, gamma=1e-3, bias_max=0.5, momentum=0.99)
    ok = bias[0] < 0 and bias[1] > 0 and bool((bias[2:] == 0).all()) and float(bias.abs().max()) <= 0.5
    print(f'  {"PASS" if ok else "FAIL"}  load-balance bias: overused -> down, unused -> up, dead zone untouched')
    return ok


def test_plan_invariants():
    """Number of token chunks must not depend on the counting-sort chunk (it decides whether the weight
    gradients are deterministic); in automatic mode every size gets the maximum number of CTAs."""
    if not es._HAS_TRITON:
        return True
    bad = []
    saved = es._CSR_CHUNK
    for chunk in (0, 8192, 4096, 2048, 1024, 512):
        es._CSR_CHUNK = chunk
        for M in (128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536):
            for kg in (1, 2, 3, 4, 8, 16):
                P = M * kg
                ncta, ch, r, c = es._plan_chunks(P)
                nref = -(-P // es._CHUNK)
                ref = min(es._CH, nref)
                while ref > 1 and nref % ref != 0:
                    ref -= 1
                if ch != ref or ncta % ch or ncta // ch != r:
                    bad.append((chunk, M, kg))
                if chunk == 0 and ncta < es._CSR_CTA and c > es._SUB:
                    bad.append(('boost', M, kg))
    es._CSR_CHUNK = saved
    print(f'  {"PASS" if not bad else "FAIL"}  plan: token chunks independent of CSR chunk, CTA boost {bad[:3]}')
    return not bad


def test_interpreter():
    ok = True
    for tag, M, D, Fd, G, kg in (('fused G=32', 256, 64, 512, 32, 4),
                                 ('unfused G=128', 192, 64, 1024, 128, 2),
                                 ('3 token chunks (atomic dW)', 2100, 64, 512, 32, 8),
                                 ('48 groups (not 2^n), d_model 96', 300, 96, 1536, 32, 3)):
        ok = parity(tag, M, D, Fd, G, kg, 'cpu', torch.float32, 1e-4, 1e-4) and ok
    return ok


def test_cuda():
    ok = True
    for G, ng, kg in ((128, 16, 3), (32, 128, 8)):
        for dt, ty, tg in ((torch.float32, 5e-3, 8e-3), (torch.bfloat16, 3e-2, 6e-2)):
            ok = parity(f'{str(dt)[6:]} G={G} ng={ng} kg={kg}', 512, 256, ng * G, G, kg, 'cuda', dt, ty, tg) and ok
    ok = parity('bf16 4 token chunks', 8192, 512, 4096, 32, 8, 'cuda', torch.bfloat16, 3e-2, 6e-2) and ok
    x, gw, uw, dw = make(64, 128, 256, 'cuda', torch.bfloat16)
    idx = es.grouped_topk_select(x, gw, 16, 4).to(torch.int32)
    try:
        torch.library.opcheck(es.sparse_swiglu_fwd, tuple(t.clone().requires_grad_() for t in (x, gw, uw, dw)) + (idx, 16))
        print('  PASS  opcheck')
    except Exception as exc:
        print(f'  FAIL  opcheck: {type(exc).__name__}: {exc}'); ok = False
    import torch._dynamo as dynamo
    x, gw, uw, dw = make(512, 512, 4096, 'cuda', torch.bfloat16)

    def block(a, b, c, d):
        return es.sparse_swiglu_ffn(a, b, c, d, es.grouped_topk_select(a, b, 32, 8), 32).float().pow(2).mean()
    ex = dynamo.explain(block)(x, gw, uw, dw)
    dynamo.reset()
    leaves = [t.clone().requires_grad_() for t in (x, gw, uw, dw)]
    lc = torch.compile(block, fullgraph=True)(*leaves)
    gc = torch.autograd.grad(lc, leaves)
    leaves_e = [t.clone().requires_grad_() for t in (x, gw, uw, dw)]
    le = block(*leaves_e)
    ge = torch.autograd.grad(le, leaves_e)
    d = max(rel_err(a, b) for a, b in zip(gc, ge))
    c_ok = ex.graph_count == 1 and ex.graph_break_count == 0 and d < 3e-2
    print(f'  {"PASS" if c_ok else "FAIL"}  torch.compile: graphs={ex.graph_count} breaks={ex.graph_break_count}, '
          f'compiled vs eager grads {d:.1e}')
    return ok and c_ok


if __name__ == '__main__':
    print(f'torch {torch.__version__} | triton {es._HAS_TRITON} | cuda {torch.cuda.is_available()} | interpreter {INTERP}')
    results = [test_cpu_reference_and_selection(), test_load_balance_bias(), test_plan_invariants()]
    if INTERP and es._HAS_TRITON:
        results.append(test_interpreter())
    elif torch.cuda.is_available() and es._HAS_TRITON:
        results.append(test_cuda())
    else:
        print('  (no CUDA: kernel tests skipped — run with TRITON_INTERPRET=1 to test the kernels on CPU)')
    print(f'\n{sum(results)}/{len(results)} test groups passed')
    sys.exit(0 if all(results) else 1)