"""The low-res reduction that replaces natten, diffed against natten itself.

natten cannot run this model on Apple Silicon (no libnatten -> cutlass-fna raises;
flex-fna is CPU-only, rejects NAF's asymmetric head dims, and was OOM-killed at 512).
So `lr_neighborhood_attention` replaces it. That is only defensible if it is EXACTLY
equivalent, which is what this asserts — against natten's own kernel, at the sizes
where natten does run.

Skips cleanly if natten is absent; the reduction does not need it at runtime.
"""

import sys
from pathlib import Path

import torch

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))

from pixal3d_mlx.naf import lr_neighborhood_attention  # noqa: E402

PASS, FAIL = [], []


def _natten_reference(q, k_lr, v_lr, n, ks):
    """What NAF actually does: upsample K/V, then dilated neighborhood attention."""
    import natten
    import torch.nn.functional as F

    hq, wq = q.shape[-2:]
    hk, wk = k_lr.shape[-2:]

    def to_bhwnd(x):
        return x.reshape(x.shape[0], n, -1, x.shape[-2], x.shape[-1]).permute(0, 3, 4, 1, 2)

    kk = to_bhwnd(F.interpolate(k_lr, size=(hq, wq), mode="nearest-exact"))
    vv = to_bhwnd(F.interpolate(v_lr, size=(hq, wq), mode="nearest-exact"))
    out = natten.na2d(to_bhwnd(q), kk, vv, kernel_size=ks,
                      dilation=(hq // hk, wq // wk), stride=1, backend="flex-fna")
    b = out.shape[0]
    return out.permute(0, 3, 4, 1, 2).reshape(b, -1, hq, wq)


def test_matches_natten():
    """natten requires kernel*dilation <= input, so low-res must be >= kernel."""
    try:
        import natten  # noqa: F401
    except ImportError:
        print("  SKIP  natten not installed — reduction is self-sufficient at runtime")
        return

    torch.manual_seed(0)
    n, ks, dq = 4, (9, 9), 16
    for hk, hq in ((16, 64), (16, 128), (32, 256)):
        q = torch.randn(1, n * dq, hq, hq)
        k = torch.randn(1, n * dq, hk, hk)
        v = torch.randn(1, n * dq, hk, hk)
        mine = lr_neighborhood_attention(q, k, v, n, ks)
        ref = _natten_reference(q, k, v, n, ks)
        diff = float((mine - ref).abs().max())
        ok = diff < 5e-6
        (PASS if ok else FAIL).append(f"natten dil {hq // hk}")
        print(f"  {'PASS' if ok else 'FAIL'}  vs natten  LR {hk:3} -> HR {hq:4} "
              f"(dil {hq // hk:2})   max diff {diff:.3e}")


def test_asymmetric_head_dims():
    """NAF uses qk=64 with v=256 — the exact case flex-fna refuses."""
    torch.manual_seed(1)
    n, hk, hq = 4, 16, 64
    q = torch.randn(1, n * 64, hq, hq)
    k = torch.randn(1, n * 64, hk, hk)
    v = torch.randn(1, n * 256, hk, hk)
    out = lr_neighborhood_attention(q, k, v, n, (9, 9))
    ok = out.shape == (1, n * 256, hq, hq) and bool(torch.isfinite(out).all())
    (PASS if ok else FAIL).append("asymmetric head dims")
    print(f"  {'PASS' if ok else 'FAIL'}  asymmetric qk=64 v=256 -> {tuple(out.shape)}")


def test_rejects_non_integer_ratio():
    """The reduction is only valid when the dilation is exactly the upsample factor."""
    q = torch.randn(1, 4 * 16, 50, 50)
    k = torch.randn(1, 4 * 16, 16, 16)
    try:
        lr_neighborhood_attention(q, k, k, 4, (9, 9))
        ok = False
    except ValueError:
        ok = True
    (PASS if ok else FAIL).append("guards non-integer ratio")
    print(f"  {'PASS' if ok else 'FAIL'}  rejects non-integer target ratio")


if __name__ == "__main__":
    torch.set_grad_enabled(False)
    for fn in (test_matches_natten, test_asymmetric_head_dims, test_rejects_non_integer_ratio):
        fn()
    print(f"\n{len(PASS)}/{len(PASS) + len(FAIL)} passed")
    sys.exit(1 if FAIL else 0)