"""The fused projection must agree with the dense one it replaces.

Skipping a zero row is exact, so the only difference allowed is the fp32 atomicAdd
combine's arrival order — ~1 ULP, the split-K GEMM class. The tests therefore compare
against a dense fp32 reference rather than against fp16 `linear`, and check that the
fused path is no further from the truth than the dense path is.
"""

import pytest
import torch

cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")


def _acc(hidden, device="cuda"):
    """The caller-owned fp32 accumulator the op now writes into, one per module."""
    return torch.zeros(hidden, device=device, dtype=torch.float32)


def _sparse_activation(inter, density, device, generator):
    """A squared-ReLU activation: mostly exact zeros, the rest strictly positive."""
    raw = torch.randn(inter, device=device, generator=generator, dtype=torch.float16)
    keep = torch.rand(inter, device=device, generator=generator) < density
    return torch.where(keep, raw.abs(), torch.zeros_like(raw))


@cuda
@pytest.mark.parametrize("density", [0.0, 0.0667, 0.5, 1.0])
def test_matches_dense_reference(density):
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    torch.manual_seed(0)
    g = torch.Generator(device="cuda").manual_seed(0)
    hidden, inter = 512, 1024
    weight = (torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16) * 0.02)
    act = _sparse_activation(inter, density, "cuda", g)

    got = torch.ops.rwkv7_kernel.sparse_cmix(act, _tile_value_weight(weight), _acc(hidden))
    dense = torch.nn.functional.linear(act, weight)
    truth = torch.nn.functional.linear(act.float(), weight.float())

    fused_err = (got.float().squeeze(0) - truth).abs().max()
    dense_err = (dense.float() - truth).abs().max()
    # Not "fused == dense" -- neither is exact in fp16. The bar is that swapping in
    # the kernel does not move the answer further from the fp32 truth than the
    # implementation it replaces already is.
    assert fused_err <= dense_err * 1.5 + 1e-4, (fused_err.item(), dense_err.item())


@cuda
def test_all_zero_activation_gives_zero():
    """The density-0 case has no rows to accumulate; the accumulator must still be
    clean, which is only true if the previous call re-zeroed it."""
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    g = torch.Generator(device="cuda").manual_seed(1)
    hidden, inter = 512, 1024
    weight = torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16)
    tiled = _tile_value_weight(weight)
    acc = _acc(hidden)

    torch.ops.rwkv7_kernel.sparse_cmix(_sparse_activation(inter, 0.5, "cuda", g), tiled, acc)
    zeros = torch.zeros(inter, device="cuda", dtype=torch.float16)
    got = torch.ops.rwkv7_kernel.sparse_cmix(zeros, tiled, acc)
    assert torch.equal(got, torch.zeros_like(got)), got.abs().max().item()


@cuda
def test_repeated_calls_do_not_accumulate():
    """The persistent accumulator is reused across layers and steps. If the finalize
    pass failed to re-zero it, the second identical call would return roughly twice
    the first -- which no correctness test on a single call would notice."""
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    g = torch.Generator(device="cuda").manual_seed(2)
    hidden, inter = 512, 1024
    weight = (torch.randn(hidden, inter, device="cuda", generator=g, dtype=torch.float16) * 0.02)
    tiled = _tile_value_weight(weight)
    acc = _acc(hidden)
    act = _sparse_activation(inter, 0.0667, "cuda", g)

    first = torch.ops.rwkv7_kernel.sparse_cmix(act, tiled, acc).clone()
    for _ in range(5):
        again = torch.ops.rwkv7_kernel.sparse_cmix(act, tiled, acc)
    torch.testing.assert_close(first, again, rtol=0, atol=0)


@cuda
def test_tiled_weight_follows_weight_updates():
    """A cached tile that survives a weight edit projects with the old weights and
    says nothing about it."""
    from rwkv7_kernel.sparse_ffn import _tiled_weight

    hidden, inter = 512, 1024
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()

    before = _tiled_weight(module).clone()
    with torch.no_grad():
        module.value.weight.mul_(2.0)
    after = _tiled_weight(module)

    torch.testing.assert_close(after, before * 2, rtol=1e-3, atol=1e-3)


@cuda
def test_two_modules_do_not_share_an_accumulator():
    """The bug that made the accumulator caller-owned.

    It used to be a process-wide map keyed on (device, hidden), so two modules of the
    same width shared one buffer. On a single stream that is still correct -- kernel
    completion orders the reuse -- and off it the partial sums interleave and the
    logits come back wrong with no NaN and no error. Two modules of equal width must
    now hold two distinct buffers; the same object would put the old hazard back.
    """
    from rwkv7_kernel.sparse_ffn import _tiled_weight

    hidden, inter = 512, 1024
    modules = []
    for _ in range(2):
        m = torch.nn.Module()
        m.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
        _tiled_weight(m)
        modules.append(m)

    a, b = (m._rwkv_hf_accumulator for m in modules)
    assert a.data_ptr() != b.data_ptr()
    assert a.shape == (hidden,) and a.dtype == torch.float32


@cuda
def test_mismatched_devices_are_refused():
    """A device_map split would otherwise launch on one GPU with another's pointers."""
    if torch.cuda.device_count() < 2:
        pytest.skip("needs two CUDA devices")
    from rwkv7_kernel._ext import load_extension
    from rwkv7_kernel.sparse_ffn import _tile_value_weight

    load_extension()
    hidden, inter = 512, 1024
    weight = torch.randn(hidden, inter, device="cuda:0", dtype=torch.float16) * 0.02
    act = torch.zeros(inter, device="cuda:1", dtype=torch.float16)
    with pytest.raises(RuntimeError, match="share a device"):
        torch.ops.rwkv7_kernel.sparse_cmix(act, _tile_value_weight(weight), _acc(hidden, "cuda:0"))


@cuda
def test_grad_enabled_forward_is_refused_rather_than_silently_ungraded():
    """There is no backward, so a grad-enabled forward used to return a value whose
    gradient never arrives -- `.backward()` succeeds and the input grad is None."""
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 512, 1024
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
    act = torch.zeros(1, 1, inter, device="cuda", dtype=torch.float16, requires_grad=True)

    with pytest.raises(RuntimeError, match="no backward"):
        sparse_project(module, act)

    with torch.no_grad():  # the documented way through
        sparse_project(module, act)


def test_register_with_a_model_it_cannot_pre_tile_raises():
    """Passing a model means "build the tiles now". Finding none used to return
    quietly, costing the caller the CUDA-graph cliff the eager build exists to avoid
    with nothing said. Runs on CPU: the argument check now fires before load_extension()."""
    from rwkv7_kernel import sparse_ffn

    class NotAnRwkv(torch.nn.Module):
        pass

    with pytest.raises(ValueError, match="no blocks to pre-tile"):
        sparse_ffn.register(NotAnRwkv())


def test_invalidate_forces_a_rebuild():
    """The escape hatch for the edits the fingerprint cannot see."""
    from rwkv7_kernel.sparse_ffn import invalidate

    module = torch.nn.Module()
    module._rwkv_hf_tiled_fp = ("pretend a tile was built",)
    invalidate(module)
    assert module._rwkv_hf_tiled_fp is None


def test_register_that_raises_leaves_the_class_untouched():
    """The raise above exists to guarantee this, and nothing was checking it.

    `register` promises that a rejected call "leaves nothing half-applied". The
    existing test asserts the exception and stops there, so the promise the exception
    is there to keep was never exercised: an implementation that patched the class
    first and validated second would pass it.
    """
    from rwkv7_kernel import sparse_ffn

    # Guard on the module the import below needs; `transformers` alone is satisfied by a
    # stock install that has no RWKV-7, and the test then errors instead of skipping.
    pytest.importorskip("transformers.models.rwkv7")
    from transformers.models.rwkv7.modeling_rwkv7 import Rwkv7FeedForward

    before = Rwkv7FeedForward._project
    with pytest.raises(ValueError, match="no blocks to pre-tile"):
        sparse_ffn.register(torch.nn.Module())

    assert Rwkv7FeedForward._project is before
    assert sparse_ffn._ORIGINAL_PROJECT is None


def test_unregister_restores_the_projection_and_is_idempotent():
    """`unregister` had no test at all, on either of its two paths.

    Driven by setting `_ORIGINAL_PROJECT` directly rather than by calling `register`,
    which would pull in the CUDA extension; the restore path is plain Python and this
    is how `test_invalidate_forces_a_rebuild` already exercises its own state.
    """
    from rwkv7_kernel import sparse_ffn

    pytest.importorskip("transformers.models.rwkv7")
    from transformers.models.rwkv7.modeling_rwkv7 import Rwkv7FeedForward

    original = Rwkv7FeedForward._project
    try:
        sparse_ffn._ORIGINAL_PROJECT = original
        Rwkv7FeedForward._project = sparse_ffn.sparse_project

        sparse_ffn.unregister()
        assert Rwkv7FeedForward._project is original
        assert sparse_ffn._ORIGINAL_PROJECT is None

        # Calling it again must not re-restore a stale function or raise. This is the
        # branch that returns early, and it is the one a caller hits by unregistering
        # in a `finally` that also runs on the path where register never happened.
        sparse_ffn.unregister()
        assert Rwkv7FeedForward._project is original
    finally:
        Rwkv7FeedForward._project = original
        sparse_ffn._ORIGINAL_PROJECT = None


def test_a_non_conforming_input_takes_the_dense_path_and_builds_no_tile():
    """The fallback branch was uncovered, and it is the one most callers hit first.

    Bit-exact against `linear` rather than close to it: the branch does not
    approximate the dense projection, it *is* the dense projection, so any tolerance
    would also pass an implementation that quietly ran the kernel. The tile assertion
    is what actually pins the branch -- a conforming call would have built one.
    """
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 256, 512
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False)
    act = torch.randn(1, 1, inter)  # CPU, fp32: non-conforming on two counts

    got = sparse_project(module, act)

    torch.testing.assert_close(got, torch.nn.functional.linear(act, module.value.weight), rtol=0, atol=0)
    assert getattr(module, "_rwkv_hf_tiled_fp", None) is None
    assert not hasattr(module, "_rwkv_hf_tiled")


class _FakeAdapterWrapper(torch.nn.Module):
    """Minimal stand-in for what PEFT's `lora.Linear` actually looks like from the
    outside: a `.weight` attribute that proxies to a frozen base layer (so shape
    checks upstream still work), plus a `forward()` that adds something the base
    layer alone does not compute. Not importing peft here on purpose -- the bug
    this guards is about *any* module with its own forward() sitting where a bare
    nn.Linear was assumed, not specifically about LoRA's internals.
    """

    def __init__(self, base: torch.nn.Linear, delta: torch.nn.Linear):
        super().__init__()
        self.base = base
        self.delta = delta

    @property
    def weight(self):
        return self.base.weight

    def forward(self, x):
        return self.base(x) + self.delta(x)


def test_a_wrapped_value_module_goes_through_its_own_forward_not_raw_weight():
    """Regression: an earlier version of both fallback branches did
    `F.linear(activated, module.value.weight)`, which reaches past whatever
    `module.value` actually is straight to the frozen base weight. That is exactly
    what `Rwkv7FeedForward._project` (`return self.value(inner)`) does not do, so it
    silently stopped being a drop-in the moment `module.value` was anything other
    than a bare `nn.Linear` -- confirmed against a real PEFT LoRA adapter on
    `ffn.value`: `lora_A`/`lora_B` came back with `grad=None` (never entered the
    graph at all) on every training-shaped call, because `_conforms` requires
    single-token input and a normal training forward never qualifies. This uses a
    hand-rolled wrapper instead of peft so the test has no extra dependency, but the
    property under test -- "does the module's own forward() get called" -- is the
    same one that broke.
    """
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 256, 512
    base = torch.nn.Linear(inter, hidden, bias=False)
    base.weight.requires_grad_(False)  # PEFT freezes the base layer; mirror that
    delta = torch.nn.Linear(inter, hidden, bias=False)
    module = torch.nn.Module()
    module.value = _FakeAdapterWrapper(base, delta)

    act = torch.randn(1, 1, inter, requires_grad=True)  # non-conforming: CPU + fp32

    got = sparse_project(module, act)
    expected = module.value(act)  # what the original _project would have produced

    torch.testing.assert_close(got, expected, rtol=0, atol=0)
    assert not torch.allclose(got, torch.nn.functional.linear(act, base.weight)), (
        "the fallback produced the same answer as skipping the wrapper entirely -- "
        "delta's contribution did not survive, so this test is not actually checking "
        "what it claims to"
    )

    got.sum().backward()
    assert delta.weight.grad is not None and delta.weight.grad.abs().sum() > 0, (
        "the wrapper's own trainable parameter got no gradient -- the bug this "
        "guards against is back"
    )


@cuda
def test_conforming_shape_also_declines_a_wrapped_value_module():
    """The fused-kernel path itself reads `module.value.weight` directly too (for
    tiling), which for a wrapped module is the unadapted base weight -- silently
    wrong under `torch.no_grad()` single-token inference with an unmerged adapter,
    not just a training-time gradient bug. A wrapped module must decline the fused
    path even when the shape/dtype would otherwise conform, and take the module's
    own forward() instead, same as the non-conforming branch."""
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 256, 512
    base = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
    base.weight.requires_grad_(False)
    delta = torch.nn.Linear(inter, hidden, bias=False).half().cuda()
    module = torch.nn.Module()
    module.value = _FakeAdapterWrapper(base, delta)

    act = torch.zeros(1, 1, inter, device="cuda", dtype=torch.float16)  # conforming shape/dtype

    with torch.no_grad():
        got = sparse_project(module, act)
        expected = module.value(act)

    torch.testing.assert_close(got, expected, rtol=0, atol=0)
    assert getattr(module, "_rwkv_hf_tiled_fp", None) is None, (
        "a tile got built for a wrapped module -- the fused path should have "
        "declined before ever reaching _tiled_weight"
    )


@cuda
@pytest.mark.parametrize(
    "reason,shape,dtype",
    [
        ("more than one token", (1, 2), torch.float16),
        ("not fp16", (1, 1), torch.float32),
    ],
)
def test_each_non_conforming_reason_falls_back_on_its_own(reason, shape, dtype):
    """On CPU every input is non-conforming for the same reason, so the individual
    conditions can only be separated on a device where the others are satisfied."""
    from rwkv7_kernel.sparse_ffn import sparse_project

    hidden, inter = 512, 1024
    module = torch.nn.Module()
    module.value = torch.nn.Linear(inter, hidden, bias=False).to("cuda", dtype)
    act = torch.randn(*shape, inter, device="cuda", dtype=dtype)

    with torch.no_grad():
        got = sparse_project(module, act)
        want = torch.nn.functional.linear(act, module.value.weight)

    torch.testing.assert_close(got, want, rtol=0, atol=0), reason
    assert getattr(module, "_rwkv_hf_tiled_fp", None) is None, reason


def test_a_non_default_tile_size_is_read_and_the_layout_still_inverts(monkeypatch):
    """No tile size other than the default was covered anywhere.

    This pins the Python half only: `_tiles` reading the environment at call time,
    and `_tile_value_weight` staying a pure permutation at that size. It deliberately
    does *not* claim the kernel agrees -- the tile shape is a compile-time constant
    on that side, so an env var changed after the extension is built gives the two
    halves different constants, which is exactly why `_tiles` reads late. Checking
    kernel agreement at a second size means rebuilding the extension, which belongs
    in a build matrix rather than here.
    """
    from rwkv7_kernel.sparse_ffn import _tile_value_weight, _tiles

    monkeypatch.setenv("RWKV_HF_KERNELS_C_TILE", "128")
    tile_h, tile_i = _tiles()
    assert (tile_h, tile_i) == (128, 64)

    hidden, inter = 256, 128
    weight = torch.randn(hidden, inter)
    tiled = _tile_value_weight(weight)
    assert tiled.shape == (inter // tile_i, hidden // tile_h, tile_i, tile_h)

    # A permutation loses nothing: undo it and the original must come back exactly.
    back = tiled.permute(0, 2, 1, 3).reshape(inter, hidden).t()
    torch.testing.assert_close(back, weight, rtol=0, atol=0)