"""Simple, minimal implementation of Mamba in one file of PyTorch.

Suggest reading the following before/while reading the code:
    [1] Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Albert Gu and Tri Dao)
        https://arxiv.org/abs/2312.00752
    [2] The Annotated S4 (Sasha Rush and Sidd Karamcheti)
        https://srush.github.io/annotated-s4

Glossary:
    b: batch size                       (`B` in Mamba paper [1] Algorithm 2)
    l: sequence length                  (`L` in [1] Algorithm 2)
    d or d_model: hidden dim
    n or d_state: latent state dim      (`N` in [1] Algorithm 2)
    expand: expansion factor            (`E` in [1] Section 3.4)
    d_in or d_inner: d * expand         (`D` in [1] Algorithm 2)
    A, B, C, D: state space parameters  (See any state space representation formula)
                                        (B, C are input-dependent (aka selective, a key innovation in Mamba); A, D are not)
    Δ or delta: input-dependent step size
    dt_rank: rank of Δ                  (See [1] Section 3.6 "Parameterization of ∆")

"""
from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from einops import rearrange, repeat, einsum


class Mamba(nn.Module):
    def __init__(
        self,
        d_model: int,
        n_layer: int,
        input_dim: int = 1,
        output_dim: int = 1,
        d_state: int = 16,
        expand: int = 2,
        dt_rank="auto",
        d_conv: int = 4,
        conv_bias: bool = True,
        bias: bool = False,
        return_last: bool = False,
    ):
        """
        Mamba for continuous time-series regression.

        Args:
            input_dim: 每个时间步的输入特征数，例如单变量就是 1，多变量就是 len(features)
            output_dim: 每个时间步的输出维度，一般回归预测为 1
            return_last:
                False -> 输出 [B, L, output_dim]
                True  -> 输出 [B, output_dim]，只取最后一个时间步
        """
        super().__init__()

        self.d_model = d_model
        self.n_layer = n_layer
        self.input_dim = input_dim
        self.output_dim = output_dim
        self.d_state = d_state
        self.expand = expand
        self.dt_rank = dt_rank
        self.d_conv = d_conv
        self.conv_bias = conv_bias
        self.bias = bias
        self.return_last = return_last

        self.d_inner = int(self.expand * self.d_model)

        if self.dt_rank == "auto":
            self.dt_rank = math.ceil(self.d_model / 16)

        # 连续气象/光伏特征输入：Linear 投影到 d_model
        self.input_proj = nn.Linear(self.input_dim, self.d_model)

        self.layers = nn.ModuleList([
            ResidualBlock(
                d_model=self.d_model,
                conv_bias=self.conv_bias,
                bias=self.bias,
                d_conv=self.d_conv,
                d_inner=self.d_inner,
                dt_rank=self.dt_rank,
                d_state=self.d_state
            )
            for _ in range(self.n_layer)
        ])

        self.norm_f = RMSNorm(self.d_model)

        # 回归输出层
        self.output_proj = nn.Linear(self.d_model, self.output_dim)

    def forward(self, x):
        """
        Args:
            x: shape [B, L] 或 [B, L, C]

        Returns:
            if return_last=False: [B, L, output_dim]
            if return_last=True:  [B, output_dim]
        """
        x = x.float()

        # 如果输入是 [B, L]，说明是单变量序列，自动扩展成 [B, L, 1]
        if x.dim() == 2:
            x = x.unsqueeze(-1)

        if x.dim() != 3:
            raise ValueError(f"Expected input shape [B, L, C], but got {x.shape}")

        if x.shape[-1] != self.input_dim:
            raise ValueError(
                f"Input feature dim mismatch: model expects input_dim={self.input_dim}, "
                f"but got x.shape[-1]={x.shape[-1]}. Full input shape: {x.shape}"
            )

        # [B, L, input_dim] -> [B, L, d_model]
        x = self.input_proj(x)

        for layer in self.layers:
            x = layer(x)

        x = self.norm_f(x)

        # [B, L, d_model] -> [B, L, output_dim]
        out = self.output_proj(x)

        if self.return_last:
            out = out[:, -1, :]

        B , L , C = out.shape
        out = out.reshape(B,-1)
        return out

class ResidualBlock(nn.Module):
    def __init__(self, d_model, conv_bias, bias, d_conv, d_inner, dt_rank, d_state):
        """Simple block wrapping Mamba block with normalization and residual connection."""
        super().__init__()

        self.d_model = d_model  # Fix: store d_model as attribute so it can be used in RMSNorm init below
        self.mixer = MambaBlock(d_model, conv_bias, bias, d_conv, d_inner, dt_rank, d_state)
        self.norm = RMSNorm(self.d_model)
        

    def forward(self, x):
        """
        Args:
            x: shape (b, l, d)    (See Glossary at top for definitions of b, l, d_in, n...)
    
        Returns:
            output: shape (b, l, d)

        Official Implementation:
            Block.forward(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L297
            
            Note: the official repo chains residual blocks that look like
                [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> [Add -> Norm -> Mamba] -> ...
            where the first Add is a no-op. This is purely for performance reasons as this
            allows them to fuse the Add->Norm.

            We instead implement our blocks as the more familiar, simpler, and numerically equivalent
                [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> [Norm -> Mamba -> Add] -> ....
            
        """
        output = self.mixer(self.norm(x)) + x

        return output
            

class MambaBlock(nn.Module):
    def __init__(self, 
        d_model,
        conv_bias,
        bias,
        d_conv,
        d_inner,
        dt_rank,
        d_state
    ):
        """A single Mamba block, as described in Figure 3 in Section 3.4 in the Mamba paper [1]."""
        super().__init__()

        self.d_model = d_model
        self.conv_bias = conv_bias
        self.bias = bias
        self.d_conv = d_conv
        self.d_inner = d_inner
        self.dt_rank = dt_rank
        self.d_state = d_state

        self.in_proj = nn.Linear(self.d_model, self.d_inner * 2, bias=self.bias)

        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            bias=self.conv_bias,
            kernel_size=self.d_conv,
            groups=self.d_inner,
            padding=self.d_conv - 1,
        )

        # x_proj takes in `x` and outputs the input-specific Δ, B, C
        self.x_proj = nn.Linear(self.d_inner, self.dt_rank + self.d_state * 2, bias=False)
        
        # dt_proj projects Δ from dt_rank to d_in
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner, bias=True)

        A = repeat(torch.arange(1, self.d_state + 1), 'n -> d n', d=self.d_inner)
        self.A_log = nn.Parameter(torch.log(A))
        self.D = nn.Parameter(torch.ones(self.d_inner))
        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=self.bias)
        

    def forward(self, x):
        """Mamba block forward. This looks the same as Figure 3 in Section 3.4 in the Mamba paper [1].
    
        Args:
            x: shape (b, l, d)    (See Glossary at top for definitions of b, l, d_in, n...)
    
        Returns:
            output: shape (b, l, d)
        
        Official Implementation:
            class Mamba, https://github.com/state-spaces/mamba/blob/main/mamba_ssm/modules/mamba_simple.py#L119
            mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
            
        """
        (b, l, d) = x.shape
        
        x_and_res = self.in_proj(x)  # shape (b, l, 2 * d_in)
        (x, res) = x_and_res.split(split_size=[self.d_inner, self.d_inner], dim=-1)

        x = rearrange(x, 'b l d_in -> b d_in l')
        x = self.conv1d(x)[:, :, :l]
        x = rearrange(x, 'b d_in l -> b l d_in')
        
        x = F.silu(x)

        y = self.ssm(x)
        
        y = y * F.silu(res)
        
        output = self.out_proj(y)

        return output

    
    def ssm(self, x):
        """Runs the SSM. See:
            - Algorithm 2 in Section 3.2 in the Mamba paper [1]
            - run_SSM(A, B, C, u) in The Annotated S4 [2]

        Args:
            x: shape (b, l, d_in)    (See Glossary at top for definitions of b, l, d_in, n...)
    
        Returns:
            output: shape (b, l, d_in)

        Official Implementation:
            mamba_inner_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L311
            
        """
        (d_in, n) = self.A_log.shape

        # Compute ∆ A B C D, the state space parameters.
        #     A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
        #     ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
        #                                  and is why Mamba is called **selective** state spaces)
        
        A = -torch.exp(self.A_log.float())  # shape (d_in, n)
        D = self.D.float()

        x_dbl = self.x_proj(x)  # (b, l, dt_rank + 2*n)
        
        # BUG FIX: typo, should be self.dt_rank (not self.self.dt_rank)
        (delta, B, C) = x_dbl.split(split_size=[self.dt_rank, n, n], dim=-1)  # delta: (b, l, dt_rank). B, C: (b, l, n)
        delta = F.softplus(self.dt_proj(delta))  # (b, l, d_in)
        
        y = self.selective_scan(x, delta, A, B, C, D)  # This is similar to run_SSM(A, B, C, u) in The Annotated S4 [2]
        
        return y

    
    def selective_scan(self, u, delta, A, B, C, D):
        """Does selective scan algorithm. See:
            - Section 2 State Space Models in the Mamba paper [1]
            - Algorithm 2 in Section 3.2 in the Mamba paper [1]
            - run_SSM(A, B, C, u) in The Annotated S4 [2]

        This is the classic discrete state space formula:
            x(t + 1) = Ax(t) + Bu(t)
            y(t)     = Cx(t) + Du(t)
        except B and C (and the step size delta, which is used for discretization) are dependent on the input x(t).
    
        Args:
            u: shape (b, l, d_in)    (See Glossary at top for definitions of b, l, d_in, n...)
            delta: shape (b, l, d_in)
            A: shape (d_in, n)
            B: shape (b, l, n)
            C: shape (b, l, n)
            D: shape (d_in,)
    
        Returns:
            output: shape (b, l, d_in)
    
        Official Implementation:
            selective_scan_ref(), https://github.com/state-spaces/mamba/blob/main/mamba_ssm/ops/selective_scan_interface.py#L86
            Note: I refactored some parts out of `selective_scan_ref` out, so the functionality doesn't match exactly.
            
        """
        (b, l, d_in) = u.shape
        n = A.shape[1]
        
        # Discretize continuous parameters (A, B)
        # - A is discretized using zero-order hold (ZOH) discretization (see Section 2 Equation 4 in the Mamba paper [1])
        # - B is discretized using a simplified Euler discretization instead of ZOH. From a discussion with authors:
        #   "A is the more important term and the performance doesn't change much with the simplification on B"
        deltaA = torch.exp(einsum(delta, A, 'b l d_in, d_in n -> b l d_in n'))
        deltaB_u = einsum(delta, B, u, 'b l d_in, b l n, b l d_in -> b l d_in n')
        
        # Perform selective scan (see scan_SSM() in The Annotated S4 [2])
        # Note that the below is sequential, while the official implementation does a much faster parallel scan that
        # is additionally hardware-aware (like FlashAttention).
        x = torch.zeros((b, d_in, n), device=deltaA.device)
        ys = []    
        for i in range(l):
            x = deltaA[:, i] * x + deltaB_u[:, i]
            y = einsum(x, C[:, i, :], 'b d_in n, b n -> b d_in')
            ys.append(y)
        y = torch.stack(ys, dim=1)  # shape (b, l, d_in)
        
        y = y + u * D
    
        return y


class RMSNorm(nn.Module):
    def __init__(self,
                 d_model: int,
                 eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(d_model))


    def forward(self, x):
        output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

        return output
        
if __name__ == "__main__":

    # 定义模型参数
    d_model = 32
    n_layer = 2
    vocab_size = 10
    d_state = 16
    expand = 2
    d_conv = 3

    # 实例化Mamba模型
    model = Mamba(
        d_model = d_model,
        n_layer = n_layer,
        input_dim = vocab_size,
        output_dim = 1,
        d_state  = 2,
        expand= 2,
        dt_rank="auto",
        d_conv = 4,
        conv_bias= True,
        bias= False,
        return_last= False,
    )

    # 构建一个随机输入张量 (batch_size, seq_length)
    batch_size = 32
    seq_length = 288

    # 构建一个随机输入张量 (batch_size, seq_length, vocab_size)
    input_ids = torch.randn(batch_size, seq_length, vocab_size)

    logits = model(input_ids)

    print("Input shape :", input_ids.shape)
    print("Logits shape:", logits.shape)  # 应为 (batch_size, seq_length, vocab_size)