"""
Flash attention with learnable softmax support for packed sequences (THD format).

Learnable softmax adds a per-head learnable offset to the attention logits before softmax,
acting as an attention "sink" that can absorb probability mass:
    P = softmax([QK^T/sqrt(d), offset])[:, :, :, :-1]

This is mathematically equivalent to:
    P = softmax(QK^T/sqrt(d)) * sigmoid(logsumexp(QK^T/sqrt(d)) - offset)

We use flash attention for the efficient QK^T computation and apply the sigmoid scaling
as a post-processing step, with a custom backward that correctly propagates gradients
through both the attention output and the logsumexp → sigmoid path.
"""

import math

import torch
from flash_attn.flash_attn_interface import _wrapped_flash_attn_varlen_backward, flash_attn_varlen_func
from torch import Tensor


class _LearnableSoftmaxFlashAttnVarlen(torch.autograd.Function):
    """Custom autograd for flash attention with learnable softmax.

    Forward:
        out_vanilla, lse = flash_attn(q, k, v)
        scale = sigmoid(lse - offset)        # per head per token
        out = scale * out_vanilla

    Backward:
        Pass d_out' = scale * d_out and out' = scale * out_vanilla to flash_attn backward.
        This correctly accounts for the gradient through lse because:
            di' = d_out' . out' = scale^2 * (d_out . out_vanilla)
        which gives the correct total derivative d_S = P * scale * (D - scale * di).
    """

    @staticmethod
    def forward(
        ctx,
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        max_seqlen_q,
        max_seqlen_k,
        softmax_offset,
        softmax_scale,
        causal,
        window_size_left,
        window_size_right,
        dropout_p,
        deterministic,
    ):
        out_vanilla, softmax_lse, _ = flash_attn_varlen_func(
            q,
            k,
            v,
            cu_seqlens_q,
            cu_seqlens_k,
            max_seqlen_q,
            max_seqlen_k,
            dropout_p=dropout_p,
            softmax_scale=softmax_scale,
            causal=causal,
            window_size=(window_size_left, window_size_right),
            deterministic=deterministic,
            return_attn_probs=True,
        )
        # softmax_lse: (nheads, total_q) in fp32
        # softmax_offset: (nheads,)
        scale = torch.sigmoid(softmax_lse - softmax_offset.unsqueeze(1))  # (nheads, total_q)
        scale_t = scale.t().contiguous()  # (total_q, nheads)
        # Keep output in input dtype for flash_attn backward compatibility
        out_learn = (out_vanilla.float() * scale_t.unsqueeze(-1)).to(q.dtype)

        ctx.save_for_backward(
            q,
            k,
            v,
            out_vanilla,
            out_learn,
            softmax_lse,
            softmax_offset,
            scale_t,
            cu_seqlens_q,
            cu_seqlens_k,
        )
        ctx.max_seqlen_q = max_seqlen_q
        ctx.max_seqlen_k = max_seqlen_k
        ctx.softmax_scale = softmax_scale
        ctx.causal = causal
        ctx.window_size_left = window_size_left
        ctx.window_size_right = window_size_right
        ctx.dropout_p = dropout_p
        ctx.deterministic = deterministic
        return out_learn

    @staticmethod
    def backward(ctx, d_out):
        (
            q,
            k,
            v,
            out_vanilla,
            out_learn,
            softmax_lse,
            softmax_offset,
            scale_t,
            cu_seqlens_q,
            cu_seqlens_k,
        ) = ctx.saved_tensors

        # Scale d_out for flash_attn backward: d_out' = scale * d_out
        d_out_modified = (d_out.float() * scale_t.unsqueeze(-1)).to(q.dtype)

        dq = torch.empty_like(q)
        dk = torch.empty_like(k)
        dv = torch.empty_like(v)

        head_size_og = d_out_modified.size(2)
        d_out_padded = d_out_modified
        if head_size_og % 8 != 0:
            d_out_padded = torch.nn.functional.pad(d_out_modified, [0, 8 - head_size_og % 8])

        _wrapped_flash_attn_varlen_backward(
            d_out_padded.contiguous(),
            q,
            k,
            v,
            out_learn.contiguous(),
            softmax_lse,
            dq,
            dk,
            dv,
            cu_seqlens_q,
            cu_seqlens_k,
            ctx.max_seqlen_q,
            ctx.max_seqlen_k,
            ctx.dropout_p,
            ctx.softmax_scale,
            ctx.causal,
            ctx.window_size_left,
            ctx.window_size_right,
            0.0,  # softcap
            None,  # alibi_slopes
            ctx.deterministic,
            rng_state=None,
        )

        dq = dq[..., :head_size_og]
        dk = dk[..., :head_size_og]
        dv = dv[..., :head_size_og]

        # Gradient for softmax_offset
        scale = scale_t.t()  # (nheads, total_q)
        dot_product = (d_out.float() * out_vanilla.float()).sum(-1).t()  # (nheads, total_q)
        sigmoid_grad = scale * (1 - scale)
        d_offset = -(dot_product * sigmoid_grad).sum(1)  # (nheads,)

        return (
            dq,
            dk,
            dv,
            None,
            None,
            None,
            None,
            d_offset,
            None,
            None,
            None,
            None,
            None,
            None,
        )


def learnable_softmax_flash_attn_varlen(
    q: Tensor,
    k: Tensor,
    v: Tensor,
    cu_seqlens_q: Tensor,
    cu_seqlens_k: Tensor,
    max_seqlen_q: int,
    max_seqlen_k: int,
    softmax_offset: Tensor,
    softmax_scale: float | None = None,
    causal: bool = False,
    window_size: tuple[int, int] = (-1, -1),
    dropout_p: float = 0.0,
    deterministic: bool = False,
) -> Tensor:
    """Flash attention with learnable softmax for variable-length packed sequences.

    Args:
        q: (total_q, nheads, headdim) query tensor
        k: (total_k, nheads_k, headdim) key tensor
        v: (total_k, nheads_k, headdim) value tensor
        cu_seqlens_q: (batch_size + 1,) cumulative sequence lengths for queries
        cu_seqlens_k: (batch_size + 1,) cumulative sequence lengths for keys
        max_seqlen_q: max sequence length in queries
        max_seqlen_k: max sequence length in keys
        softmax_offset: (nheads,) learnable per-head offset
        softmax_scale: scaling factor for QK^T (default: 1/sqrt(headdim))
        causal: whether to use causal attention mask
        window_size: (left, right) sliding window sizes, -1 for no window
        dropout_p: dropout probability
        deterministic: use deterministic backward

    Returns:
        out: (total_q, nheads, headdim) attention output
    """
    if softmax_scale is None:
        softmax_scale = 1.0 / math.sqrt(q.size(-1))
    return _LearnableSoftmaxFlashAttnVarlen.apply(
        q,
        k,
        v,
        cu_seqlens_q,
        cu_seqlens_k,
        max_seqlen_q,
        max_seqlen_k,
        softmax_offset,
        softmax_scale,
        causal,
        window_size[0],
        window_size[1],
        dropout_p,
        deterministic,
    )