"""Reference sparse attention ported from vendor/FlashMLA/tests/ref.py (no TestParam wrappers)."""

import torch


def ref_sparse_attn_fwd(q, kv, indices, sm_scale, topk_length=None, attn_sink=None, d_v=512, rms=None):
    """q [s_q,h_q,d], kv [s_kv,1,d], indices [s_q,1,topk] -> (out bf16, max_logits, lse)."""
    s_q, h_q, d_qk = q.shape
    s_kv = kv.shape[0]
    topk = indices.shape[-1]
    idx = indices.clone().squeeze(1)
    if topk_length is not None:
        mask = torch.arange(topk).unsqueeze(0).broadcast_to(s_q, topk) >= topk_length.unsqueeze(1)
        idx.masked_fill_(mask, -1)
    invalid = (idx < 0) | (idx >= s_kv)
    idx = idx.masked_fill(invalid, 0)

    qf = q.float()
    gathered = kv.view(s_kv, d_qk).index_select(0, idx.flatten().long()).reshape(s_q, topk, d_qk).float()
    P = qf @ gathered.transpose(1, 2)
    if rms is not None:
        P = P * (sm_scale * rms.unsqueeze(-1))
    else:
        P = P * sm_scale
    P.masked_fill_(invalid.unsqueeze(1).broadcast_to(P.shape), float("-inf"))

    orig_lse = torch.logsumexp(P, dim=-1)
    max_logits = P.max(dim=-1).values

    lse_for_o = orig_lse
    if attn_sink is not None:
        lse_for_o = torch.logsumexp(
            torch.stack([orig_lse.view(s_q, h_q), attn_sink.broadcast_to(s_q, h_q)], dim=0), dim=0
        )
    lse_for_o = lse_for_o.clone()
    lse_for_o.masked_fill_(lse_for_o == float("-inf"), float("+inf"))
    s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
    out = s_for_o @ gathered[..., :d_v]

    orig_lse = orig_lse.clone()
    orig_lse.masked_fill_(orig_lse == float("-inf"), float("+inf"))
    return out.to(torch.bfloat16), max_logits, orig_lse


def ref_sparse_attn_decode(q, scopes, sm_scale, attn_sink=None, d_v=512, rms=None):
    """q [b,s_q,h_q,d]; scopes = list of (blocked_k_bf16 [n_tokens, d], indices [b,s_q,topk], topk_length|None).

    Returns (out bf16 [b,s_q,h_q,d_v], lse [b,h_q,s_q]).
    """
    b, s_q, h_q, d_qk = q.shape
    gathered_all, invalid_all = [], []
    for blocked_k, indices, topk_length in scopes:
        topk = indices.shape[-1]
        fixed = torch.clamp_min(indices, 0)
        g = blocked_k.view(-1, d_qk).index_select(0, fixed.reshape(-1).long()).view(b, s_q, topk, d_qk)
        invalid = indices == -1
        if topk_length is not None:
            invalid = invalid | (
                torch.arange(0, topk).view(1, 1, topk).broadcast_to(b, s_q, topk) >= topk_length.view(b, 1, 1)
            )
        gathered_all.append(g)
        invalid_all.append(invalid)
    gathered = torch.cat(gathered_all, dim=2)
    invalid = torch.cat(invalid_all, dim=2)

    gathered = gathered.reshape(b * s_q, -1, d_qk).float()
    gathered.masked_fill_(gathered != gathered, 0.0)
    qf = q.float().view(b * s_q, h_q, d_qk)
    w = qf @ gathered.transpose(-1, -2)
    if rms is not None:
        w = w * (sm_scale * rms.view(b * s_q, h_q).unsqueeze(-1))
    else:
        w = w * sm_scale
    w[invalid.view(b * s_q, 1, -1).broadcast_to(b * s_q, h_q, invalid.size(-1))] = float("-inf")
    lse = w.logsumexp(dim=-1)
    w = torch.exp(w - lse.unsqueeze(-1))
    out = w @ gathered[..., :d_v]
    out = out.view(b, s_q, h_q, d_v)
    lse = lse.view(b, s_q, h_q)

    if attn_sink is not None:
        out = out * (1.0 / (1.0 + torch.exp(attn_sink.view(1, 1, h_q) - lse))).unsqueeze(-1)

    lonely = lse == float("-inf")
    out = out.clone()
    out[lonely.unsqueeze(-1).broadcast_to(b, s_q, h_q, d_v)] = 0.0
    lse = lse.clone()
    lse[lonely] = float("+inf")
    return out.to(torch.bfloat16), lse.transpose(1, 2)