# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
#   https://github.com/fla-org/flash-linear-attention/graphs/contributors

# Code is adapted from flash-attn.bert_padding.py

# NOTE: This is a copy of `fla.layers.utils` but we do not sync the FLA layers in hf-kernels

import torch
from einops import rearrange, repeat

from ..ops.utils.index import (
    prepare_cu_seqlens_from_lens,
    prepare_cu_seqlens_from_mask,
    prepare_lens_from_mask,
)
from . import tensor_cache

_LAYER_IDX_REQUIRED_MSG = "{cls} requires `layer_idx` when `past_key_values` is provided."


class IndexFirstAxis(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, indices):
        ctx.save_for_backward(indices)
        assert x.ndim >= 2
        ctx.first_axis_dim, other_shape = x.shape[0], x.shape[1:]
        second_dim = other_shape.numel()
        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
        # return x[indices]
        return torch.gather(
            rearrange(x, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim),
        ).reshape(-1, *other_shape)

    @staticmethod
    def backward(ctx, do):
        (indices,) = ctx.saved_tensors
        assert do.ndim >= 2
        other_shape = do.shape[1:]
        do = rearrange(do, "b ... -> b (...)")
        dx = torch.zeros(
            [ctx.first_axis_dim, do.shape[1]],
            device=do.device,
            dtype=do.dtype,
        )
        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
        # dx[indices] = do
        dx.scatter_(0, repeat(indices, "z -> z d", d=do.shape[1]), do)
        return dx.reshape(ctx.first_axis_dim, *other_shape), None


index_first_axis = IndexFirstAxis.apply


class IndexPutFirstAxis(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, indices, first_axis_dim):
        ctx.save_for_backward(indices)
        assert indices.ndim == 1
        assert x.ndim >= 2
        y = torch.zeros(first_axis_dim, *x.shape[1:], device=x.device, dtype=x.dtype)
        # TODO [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
        y[indices] = x
        # y.scatter_(0, repeat(indices, 'z -> z d', d=x.shape[1]), x)
        return y

    @staticmethod
    def backward(ctx, do):
        (indices,) = ctx.saved_tensors
        # TODO [2022-03-04] For some reason torch.gather is a bit faster than indexing.
        dx = do[indices]
        # dx = torch.gather(do, 0, repeat(indices, 'z -> z d', d=do.shape[1]))
        return dx, None, None


index_put_first_axis = IndexPutFirstAxis.apply


@tensor_cache
def _get_unpad_indices_and_cu(
    attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    lens = prepare_lens_from_mask(attention_mask)
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    cu_seqlens = prepare_cu_seqlens_from_lens(lens)
    return indices, cu_seqlens


@tensor_cache
def _get_last_token_unpad_indices_and_cu(
    attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    return _get_unpad_indices_and_cu(attention_mask[:, -1:])


def get_unpad_indices_and_cu(
    attention_mask: torch.Tensor,
    q_len: int | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return unpadding indices and cumulative sequence lengths."""
    if q_len is None or attention_mask.shape[-1] == q_len:
        return _get_unpad_indices_and_cu(attention_mask)
    if q_len == 1:
        return _get_last_token_unpad_indices_and_cu(attention_mask)
    return _get_unpad_indices_and_cu(attention_mask[:, -q_len:])


@tensor_cache
def get_unpad_data(
    attention_mask: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, int]:
    """
    Retrieves indexing data required to repad unpadded (ragged) tensors.

    Args:
        attention_mask (`torch.Tensor`):
            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.

    Return:
        indices (`torch.Tensor`):
            The indices of non-masked tokens from the flattened input sequence.
        cu_seqlens (`torch.Tensor`):
            The cumulative sequence lengths, used to index into ragged (unpadded) tensors.
            `cu_seqlens` shape is [batch_size + 1].
        max_seqlen_in_batch (`int`):
            Maximum sequence length in batch.
    """
    lens = prepare_lens_from_mask(attention_mask)
    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
    max_seqlen_in_batch = lens.max().item()
    cu_seqlens = prepare_cu_seqlens_from_mask(attention_mask)
    return indices, cu_seqlens, max_seqlen_in_batch


def unpad_hidden_states(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.LongTensor | None,
    attention_mask: torch.Tensor | None,
    q_len: int,
) -> tuple[torch.Tensor, torch.Tensor | None, torch.LongTensor | None]:
    """Unpad hidden states via attention_mask, unless explicit cu_seqlens are given.

    Explicit `cu_seqlens` take precedence: when they are provided, the input is
    returned unchanged and no unpadding happens. Only when `cu_seqlens` is None
    and `attention_mask` is present are the hidden states unpadded via the mask.

    Arguments:
        hidden_states (`torch.Tensor`):
            Input hidden states with padding. Shape: [batch_size, q_len, ...].
        cu_seqlens (`torch.LongTensor | None`):
            Explicit cumulative sequence lengths provided by the caller.
        attention_mask (`torch.Tensor | None`):
            Boolean or int tensor of shape (batch_size, sequence_length), 1 means valid and 0 means not valid.
        q_len (`int`):
            Target length.

    Return:
        hidden_states (`torch.Tensor`):
            The unpadded hidden states of shape [1, total_tokens, ...],
            or the input unchanged when unpadding is skipped.
        indices (`torch.Tensor | None`):
            The indices of non-masked tokens, or None when unpadding is skipped.
        cu_seqlens (`torch.LongTensor | None`):
            The cumulative sequence lengths derived from the mask,
            or the explicit value passed in when unpadding is skipped.
    """
    if cu_seqlens is None and attention_mask is not None:
        indices, cu_seqlens = get_unpad_indices_and_cu(attention_mask, q_len)
        hidden_states = index_first_axis(rearrange(hidden_states, "b s ... -> (b s) ..."), indices).unsqueeze(0)
        return hidden_states, indices, cu_seqlens
    return hidden_states, None, cu_seqlens


def repad_hidden_states(
    hidden_states: torch.Tensor,
    indices: torch.Tensor | None,
    batch_size: int,
    q_len: int,
) -> torch.Tensor:
    """Inverse of `unpad_hidden_states`; a no-op when `indices` is None.

    Arguments:
        hidden_states (`torch.Tensor`):
            Hidden states of shape [1, total_tokens, ...].
        indices (`torch.Tensor | None`):
            The indices returned by `unpad_hidden_states`, or None.
        batch_size (`int`):
            Batch size for the padded sequence.
        q_len (`int`):
            Target length.

    Return:
        hidden_states of shape [batch_size, q_len, ...],
        or the input unchanged when `indices` is None.
    """
    if indices is None:
        return hidden_states
    return pad_input(hidden_states.squeeze(0), indices, batch_size, q_len)


def unpad_input(
    q: torch.Tensor | tuple[torch.Tensor, ...],
    states: tuple[torch.Tensor],
    attention_mask: torch.Tensor,
    q_len: int,
    keepdim: bool = False,
):
    """
    Unpads query, key, and values tensors, using a single dimension for all tokens
    even though they belong to different batches.


    Arguments:
        q (`torch.Tensor` or `Tuple[torch.Tensor]`):
            Query state with padding. Shape: [batch_size, q_len, ...].
            When it is a tuple, do unpadding for each tensor in the tuple.
        states (`Tuple[torch.Tensor]`):
            Attention state with padding. Shape: [batch_size, seq_len, ...].
        attention_mask (`torch.Tensor`):
            Boolean or int tensor of shape [batch_size, sequence_length], 1 means valid and 0 means not valid.
        q_len (`int`):
            Target length.
        keepdim (`bool`):
            Whether to keep the batch dimension. Default: `False`.

    Return:
        q (`torch.Tensor` or `Tuple[torch.Tensor]`):
            Query state without padding.
            Shape: [1, total_target_length, ...] if `keepdim=True` else [total_target_length, ...].
            When the `q` passed in is a tuple, return a tuple of such unpadded tensors.
        states (`Tuple[torch.Tensor]`):
            Attention state without padding.
            Shape: [1, total_source_length, ...] if `keepdim=True` else [total_source_length, ...].
        indices_q (`torch.Tensor`):
            The indices of non-masked tokens from the flattened input target sequence.
        (cu_seqlens_q, cu_seqlens_k) (`Tuple[torch.LongTensor, torch.LongTensor]`):
            The cumulative sequence lengths for the target (query) and source (key, value),
            used to index into ragged (unpadded) tensors.
            `cu_seqlens` shape is [batch_size + 1].
        (max_seqlen_in_batch_q, max_seqlen_in_batch_k) (`Tuple[int, int]`):
            Maximum sequence length in batch (`max_seqlen_in_batch_q` for the target sequence
            i.e. query, `max_seqlen_in_batch_k` for the source sequence i.e. key/value).
    """
    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = get_unpad_data(attention_mask)
    batch_size, seq_len, *_ = states[0].shape

    state = tuple(
        index_first_axis(rearrange(s, "b s ... -> (b s) ..."), indices_k)
        for s in states
    )
    if isinstance(q, torch.Tensor):
        q = (q,)
        cast_tuple = True
    else:
        cast_tuple = False

    if q_len == seq_len:
        q = tuple(index_first_axis(rearrange(q_, "b s ... -> (b s) ..."), indices_k) for q_ in q)
        cu_seqlens_q = cu_seqlens_k
        max_seqlen_in_batch_q = max_seqlen_in_batch_k
        indices_q = indices_k
    elif q_len == 1:
        max_seqlen_in_batch_q = 1
        cu_seqlens_q = torch.arange(batch_size + 1, dtype=torch.int32, device=q[0].device)
        indices_q = cu_seqlens_q[:-1]
        q = tuple(q_.squeeze(1) for q_ in q)
    else:
        raise NotImplementedError("We only support either q_len == k_len (prefilling) or q_len == 1 (decoding)")

    if keepdim:
        q = tuple(q_.unsqueeze(0) for q_ in q)
        state = tuple(s.unsqueeze(0) for s in state)
    if cast_tuple:
        q = q[0]

    return (
        q,
        state,
        indices_q,
        (cu_seqlens_q, cu_seqlens_k),
        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
    )


def pad_input(
    hidden_states: torch.Tensor,
    indices: torch.LongTensor,
    batch_size: int,
    seq_len: int,
) -> torch.Tensor:
    """
    Args:
        hidden_states ([total_tokens, ...]):
            where total_tokens denotes the number of tokens in selected in attention_mask.
        indices ([total_tokens]):
            the indices that represent the non-masked tokens of the original padded input sequence.
        batch_size (int):
            batch_size size for the padded sequence.
        seq_len (int):
            maximum sequence length for the padded sequence.

    Return:
        hidden_states of shape [batch_size, seq_len, ...]
    """
    output = index_put_first_axis(hidden_states, indices, batch_size * seq_len)
    return rearrange(output, "(b s) ... -> b s ...", b=batch_size)


def require_cache_layer_idx(module, past_key_values):
    layer_idx = getattr(module, "layer_idx", None)
    if past_key_values is not None and layer_idx is None:
        raise ValueError(_LAYER_IDX_REQUIRED_MSG.format(cls=module.__class__.__name__))
    return layer_idx


def get_layer_cache(module, past_key_values):
    layer_idx = require_cache_layer_idx(module, past_key_values)
    if past_key_values is not None and len(past_key_values) > layer_idx:
        return past_key_values[layer_idx]
    return None


def update_layer_cache(module, past_key_values, **kwargs):
    layer_idx = require_cache_layer_idx(module, past_key_values)
    if past_key_values is not None:
        return past_key_values.update(layer_idx=layer_idx, **kwargs)
    return None