# 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

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange


def fft_conv(
    u: torch.Tensor,
    k: torch.Tensor,
    dropout_mask: torch.Tensor | None,
    gelu: bool = True,
    k_rev: torch.Tensor | None = None,
) -> torch.Tensor:
    seqlen = u.shape[-1]
    fft_size = 2 * seqlen
    k_f = torch.fft.rfft(k, n=fft_size) / fft_size
    if k_rev is not None:
        k_rev_f = torch.fft.rfft(k_rev, n=fft_size) / fft_size
        k_f = k_f + k_rev_f.conj()
    u_f = torch.fft.rfft(u.to(dtype=k.dtype), n=fft_size)

    if len(u.shape) > 3:
        k_f = k_f.unsqueeze(1)
    y = torch.fft.irfft(u_f * k_f, n=fft_size, norm="forward")[..., :seqlen]

    out = y + u
    if gelu:
        out = F.gelu(out)
    if dropout_mask is not None:
        return (out * rearrange(dropout_mask, "b H -> b H 1")).to(dtype=u.dtype)
    else:
        return out.to(dtype=u.dtype)


class LongConvolution(nn.Module):
    """
    Apply FFT convolution with a learned filter of length max_len.

    Args:
        hidden_size (int):
            The number of expected features in the input and output.
        max_len (int):
            The maximum sequence length.

    Returns:
        y (torch.Tensor):
            Output of shape [batch_size, seq_len, hidden_size].
    """

    def __init__(self, hidden_size: int, max_len: int, **kwargs) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.filter = nn.Parameter(torch.randn(self.hidden_size, max_len), requires_grad=True)

    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor):
                Input of shape [batch_size, seq_len, hidden_size].

        Returns:
            y (torch.Tensor):
                Output of shape [batch_size, seq_len, hidden_size].
        """
        x = x.transpose(1, 2)
        y = fft_conv(u=x, k=self.filter, dropout_mask=None, gelu=False)
        y = y.transpose(1, 2)
        return y.to(dtype=x.dtype)


class PositionalEmbedding(nn.Module):
    def __init__(self, emb_dim: int, seq_len: int, **kwargs) -> None:
        """Complex exponential positional embeddings for implicit long convolution filters."""
        super().__init__()

        self.seq_len = seq_len
        # the time embedding fed to the filters is normalized so that t_f = 1
        t = torch.linspace(0, 1, self.seq_len)[None, :, None]  # 1, L, 1

        if emb_dim > 1:
            bands = (emb_dim - 1) // 2
        # use integer time steps for the complex exponential embeddings
        t_rescaled = torch.linspace(0, seq_len - 1, seq_len)[None, :, None]
        w = 2 * math.pi * t_rescaled / seq_len  # 1, L, 1

        f = torch.linspace(1e-4, bands - 1, bands)[None, None]
        z = torch.exp(-1j * f * w)
        z = torch.cat([t, z.real, z.imag], dim=-1)
        self.z = nn.Parameter(z, requires_grad=False)

    def forward(self, L: int) -> torch.Tensor:
        return self.z[:, :L]


class ImplicitLongConvolution(nn.Module):
    """
    Long convolution with implicit filter parameterized by an MLP.

    Args:
        hidden_size (int):
            The number of expected features in the input and output.
        max_len (int):
            The maximum sequence length.
        d_emb (int, Optional):
            The dimension of the positional embeddings. Must be odd and greater or equal to 3 (time, sine and cosine).
            Default: 3.
        d_hidden (int, Optional):
            The number of features in the hidden layer of the MLP. Default: 16.

    Attributes:
        pos_emb (PositionalEmbedding):
            The positional embedding layer.
        mlp (nn.Sequential):
            The MLP that parameterizes the implicit filter.
    """

    def __init__(self, hidden_size: int, max_len: int, d_emb: int = 3, d_hidden: int = 16, **kwargs) -> None:
        super().__init__()
        self.hidden_size = hidden_size
        self.d_emb = d_emb

        assert d_emb % 2 != 0 and d_emb >= 3, "d_emb must be odd and greater or equal to 3 (time, sine and cosine)"
        self.pos_emb = PositionalEmbedding(emb_dim=d_emb, seq_len=max_len)

        self.mlp = nn.Sequential(nn.Linear(d_emb, d_hidden), torch.nn.ReLU(), nn.Linear(d_hidden, hidden_size))

    def filter(self, seq_len: int, *args, **kwargs) -> torch.Tensor:
        return self.mlp(self.pos_emb(seq_len)).transpose(1, 2)

    def forward(self, x: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        """
        Args:
            x (torch.Tensor):
                Input of shape [batch_size, seq_len, hidden_size].

        Returns:
            y (torch.Tensor):
                Output of shape [batch_size, seq_len, hidden_size].
        """
        x = x.transpose(1, 2)
        k = self.filter(x.shape[-1])
        y = fft_conv(u=x, k=k, dropout_mask=None, gelu=False)

        y = y.transpose(1, 2)
        return y.to(dtype=x.dtype)