"""Reference functions synchronized from PyTorch Helion.

Generated by scripts/sync_from_helion.py. Do not edit directly.
"""

from __future__ import annotations

import torch
import torch.nn.functional as F
from torch.utils.checkpoint import checkpoint


def rel_error(a: torch.Tensor | None, b: torch.Tensor | None) -> float:
    """Relative L2 error between two tensors."""
    assert a is not None and b is not None
    return (a.float() - b.float()).norm().item() / b.float().norm().clamp(
        min=1e-8
    ).item()


def chunked_linear_attn_reference(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor | None = None,
    a: torch.Tensor | None = None,
    C: int = 64,
    initial_state: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Pure PyTorch reference for chunked linear attention forward.

    Supports both scalar decay (g: [B,H,T]) and diagonal decay (g: [B,H,T,D]).
    For correction case (beta provided): uses sequential recurrence within chunks.
    For no-correction case: uses parallel matmul-based intra-chunk computation.
    """
    B, H, T, D = q.shape
    DV = v.shape[-1]
    N = T // C
    assert T % C == 0

    input_dtype = q.dtype
    q, k, v, g = q.float(), k.float(), v.float(), g.float()
    if beta is not None:
        beta = beta.float()
    if a is not None:
        a = a.float()

    diagonal_decay = g.dim() == 4

    qc = q.reshape(B, H, N, C, D)
    kc = k.reshape(B, H, N, C, D)
    vc = v.reshape(B, H, N, C, DV)
    if diagonal_decay:
        gc = g.reshape(B, H, N, C, D)
    else:
        gc = g.reshape(B, H, N, C)

    bc = beta.reshape(B, H, N, C) if beta is not None else None
    ac = a.reshape(B, H, N, C, D) if a is not None else None

    outputs = []
    if initial_state is not None:
        state = initial_state.float().clone()
    else:
        state = q.new_zeros(B, H, D, DV, dtype=torch.float32)

    def correction_chunk(
        qi: torch.Tensor,
        ki: torch.Tensor,
        vi: torch.Tensor,
        gi: torch.Tensor,
        bi: torch.Tensor,
        ai: torch.Tensor,
        state: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """One chunk of the sequential beta-correction recurrence, returning
        (chunk_output, new_state); tensors are explicit args so it can be
        gradient-checkpointed."""
        s = state
        chunk_out = []
        for j in range(C):
            g_j = gi[:, :, j]
            alpha = (
                torch.exp(g_j).unsqueeze(-1)
                if diagonal_decay
                else torch.exp(g_j).unsqueeze(-1).unsqueeze(-1)
            )
            s = alpha * s
            kj = ai[:, :, j]
            vj = vi[:, :, j]
            bj = bi[:, :, j].unsqueeze(-1).unsqueeze(-1)
            kts = torch.einsum("bhd,bhdv->bhv", kj, s)
            s = s - bj * torch.einsum("bhd,bhv->bhdv", kj, kts)
            s = s + bj * torch.einsum("bhd,bhv->bhdv", kj, vj)
            chunk_out.append(torch.einsum("bhd,bhdv->bhv", qi[:, :, j], s))
        return torch.stack(chunk_out, dim=2), s

    for i in range(N):
        qi = qc[:, :, i]
        ki = kc[:, :, i]
        vi = vc[:, :, i]
        gi = gc[:, :, i]

        bi = bc[:, :, i] if bc is not None else None
        ai = ac[:, :, i] if ac is not None else ki

        if bi is not None:
            # Sequential beta-correction recurrence, checkpointed per chunk to
            # avoid OOMs.
            chunk_out, state = checkpoint(  # pyrefly: ignore[not-iterable]
                correction_chunk,
                qi,
                ki,
                vi,
                gi,
                bi,
                ai,
                state,
                use_reentrant=False,
            )
            outputs.append(chunk_out.to(input_dtype))

        elif diagonal_decay:
            g_cs = gi.cumsum(-2)
            g_last = g_cs[:, :, -1]

            q_scaled = qi * torch.exp(g_cs)
            k_scaled_intra = ki * torch.exp(-g_cs)

            o_cross = torch.einsum("bhcd,bhdv->bhcv", q_scaled, state)

            attn = torch.einsum("bhid,bhjd->bhij", q_scaled, k_scaled_intra)
            causal_mask = torch.tril(torch.ones(C, C, device=q.device))
            attn = attn * causal_mask
            o_intra = torch.einsum("bhij,bhjv->bhiv", attn, vi)

            outputs.append((o_cross + o_intra).to(input_dtype))

            k_for_state = ki * torch.exp(g_last[:, :, None, :] - g_cs)
            state = state * torch.exp(g_last).unsqueeze(-1) + torch.einsum(
                "bhcd,bhcv->bhdv", k_for_state, vi
            )

        else:
            g_cs = gi.cumsum(-1)
            g_last = g_cs[:, :, -1]

            o_cross = torch.einsum("bhcd,bhdv->bhcv", qi, state)
            o_cross = o_cross * torch.exp(g_cs).unsqueeze(-1)

            attn = torch.einsum("bhid,bhjd->bhij", qi, ki)
            decay_mask = torch.exp(g_cs.unsqueeze(-1) - g_cs.unsqueeze(-2))
            causal_mask = torch.tril(torch.ones(C, C, device=q.device))
            attn = attn * decay_mask * causal_mask
            o_intra = torch.einsum("bhij,bhjv->bhiv", attn, vi)

            outputs.append((o_cross + o_intra).to(input_dtype))

            k_scaled = ki * torch.exp(g_last[..., None, None] - g_cs.unsqueeze(-1))
            state = state * torch.exp(g_last).unsqueeze(-1).unsqueeze(
                -1
            ) + torch.einsum("bhcd,bhcv->bhdv", k_scaled, vi)

    return torch.stack(outputs, dim=2).reshape(B, H, T, DV).to(input_dtype)


def naive_recurrent_reference(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    decay: torch.Tensor,
    beta: torch.Tensor | None = None,
    q_scale: float = 1.0,
    A_log: torch.Tensor | None = None,
    dt_bias: torch.Tensor | None = None,
    lower_bound: float | None = None,
    use_qk_l2norm_in_kernel: bool = False,
    use_gate_in_kernel: bool = False,
    use_beta_sigmoid_in_kernel: bool = False,
    cu_seqlens: torch.Tensor | None = None,
) -> torch.Tensor:
    """Step-by-step recurrent computation. Slow but correct.

    The optional arguments apply the input preamble first, so the reference can be
    fed the same pre-activation tensors a model produces:
        use_qk_l2norm_in_kernel:    q, k = q / ||q||, k / ||k||
        use_beta_sigmoid_in_kernel: beta = sigmoid(beta)
        use_gate_in_kernel:         decay = lower_bound * sigmoid(exp(A_log) *
                                    (decay + dt_bias)), or -exp(A_log) *
                                    softplus(...) when lower_bound is None
    The names match the flags on the kernels, so a caller forwards one dict to both.

    cu_seqlens [N+1] marks sequence boundaries in a varlen batch (B == 1), where
    sequence n spans tokens cu_seqlens[n] : cu_seqlens[n+1]. The state resets to
    zero at each boundary, so no key from one sequence reaches another's output.
    """
    if use_qk_l2norm_in_kernel:
        q = F.normalize(q.float(), dim=-1).to(q.dtype)
        k = F.normalize(k.float(), dim=-1).to(k.dtype)
    if use_gate_in_kernel:
        assert A_log is not None
        H_, D_ = decay.shape[1], decay.shape[-1]
        gb = decay.float()
        if dt_bias is not None:
            gb = gb + dt_bias.view(1, H_, 1, D_)
        a = torch.exp(A_log.float()).view(1, H_, 1, 1)
        if lower_bound is not None:
            decay = lower_bound * torch.sigmoid(a * gb)
        else:
            decay = -a * F.softplus(gb)
    if use_beta_sigmoid_in_kernel:
        assert beta is not None
        beta = torch.sigmoid(beta.float()).to(beta.dtype)

    B, H, T, D = q.shape
    DV = v.shape[-1]
    S = q.new_zeros(B, H, D, DV, dtype=torch.float32)
    outputs = []
    diagonal = decay.dim() == 4
    starts = set() if cu_seqlens is None else set(cu_seqlens[:-1].tolist())

    for t in range(T):
        if t in starts:
            S = torch.zeros_like(S)
        qt = q[:, :, t].float() * q_scale
        kt = k[:, :, t].float()
        vt = v[:, :, t].float()
        bt = beta[:, :, t] if beta is not None else None

        if diagonal:
            alpha = torch.exp(decay[:, :, t])
            S = alpha.unsqueeze(-1) * S
        else:
            alpha = torch.exp(decay[:, :, t])
            S = alpha.unsqueeze(-1).unsqueeze(-1) * S

        if bt is not None:
            kts = torch.einsum("bhd,bhdv->bhv", kt, S)
            S = S - bt.unsqueeze(-1).unsqueeze(-1) * torch.einsum(
                "bhd,bhv->bhdv", kt, kts
            )
            S = S + bt.unsqueeze(-1).unsqueeze(-1) * torch.einsum(
                "bhd,bhv->bhdv", kt, vt
            )
        else:
            S = S + torch.einsum("bhd,bhv->bhdv", kt, vt)

        out_t = torch.einsum("bhd,bhdv->bhv", qt, S)
        outputs.append(out_t)

    return torch.stack(outputs, dim=2)


def recurrent_step_reference(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    state: torch.Tensor,
    alpha: float | torch.Tensor = 1.0,
    beta_val: torch.Tensor | None = None,
    a_val: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Single-step recurrent update (pure PyTorch reference).

    Returns (output, new_state).
    """
    k_sq = k.squeeze(2)
    v_sq = v.squeeze(2)
    q_sq = q.squeeze(2)

    if isinstance(alpha, torch.Tensor):
        alpha_sq = alpha.squeeze(2)
        if alpha_sq.dim() == 3:
            state = alpha_sq.unsqueeze(-1) * state
        else:
            state = alpha_sq.unsqueeze(-1).unsqueeze(-1) * state
    else:
        state = alpha * state

    if beta_val is not None:
        b = beta_val.squeeze(2)
        a_sq = a_val.squeeze(2) if a_val is not None else k_sq

        if b.dim() == 3:
            r = b.shape[-1]
            for j in range(r):
                bj = b[:, :, j].unsqueeze(-1).unsqueeze(-1)
                aj = a_sq[:, :, j]
                kts = torch.einsum("bhd,bhdv->bhv", aj, state)
                state = state - bj * torch.einsum("bhd,bhv->bhdv", aj, kts)
                state = state + bj * torch.einsum("bhd,bhv->bhdv", aj, v_sq)
        else:
            b_exp = b.unsqueeze(-1).unsqueeze(-1)
            kts = torch.einsum("bhd,bhdv->bhv", a_sq, state)
            state = state - b_exp * torch.einsum("bhd,bhv->bhdv", a_sq, kts)
            state = state + b_exp * torch.einsum("bhd,bhv->bhdv", a_sq, v_sq)
    else:
        state = state + torch.einsum("bhd,bhv->bhdv", k_sq, v_sq)

    o = torch.einsum("bhd,bhdv->bhv", q_sq, state)
    return o.unsqueeze(2), state