import torch 
from torch import nn, Tensor, einsum
from einops import rearrange, repeat

import math
from dataclasses import dataclass
from typing import Union

import torch
from torch import nn 
import torch.nn.functional as F
from typing import List
from zeta.nn import FeedForward 


"""

An implementation of the parallel scan operation in PyTorch (Blelloch version).
This code follows the skeleton proposed by Francois Fleuret in his pscan. However, the keys differences are :
-it has been written in an iterative way (rather than recursive)
-the backward pass has been rewritten

Please see docs/pscan.ipynb for a detailed explanation of what happens here.

"""

# taken straight from https://github.com/johnma2006/mamba-minimal/blob/master/model.py
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


class PScan(torch.autograd.Function):
    @staticmethod
    def pscan(A, X):
        # A : (B, D, L, N)
        # X : (B, D, L, N)

        # modifies X in place by doing a parallel scan.
        # more formally, X will be populated by these values :
        # H[t] = A[t] * H[t-1] + X[t] with H[0] = 0
        # which are computed in parallel (2*log2(T) sequential steps (ideally), instead of T sequential steps)

        B, D, L, _ = A.size()
        num_steps = int(math.log2(L))

        # up sweep or reduction step
        Aa = A
        Xa = X
        for k in range(num_steps):
            T = 2 * (Xa.size(2) // 2)

            Aa = Aa[:, :, :T].view(B, D, T // 2, 2, -1)
            Xa = Xa[:, :, :T].view(B, D, T // 2, 2, -1)

            Xa[:, :, :, 1].add_(Aa[:, :, :, 1].mul(Xa[:, :, :, 0]))
            Aa[:, :, :, 1].mul_(Aa[:, :, :, 0])

            Aa = Aa[:, :, :, 1]
            Xa = Xa[:, :, :, 1]

        # down sweep
        for k in range(num_steps - 1, -1, -1):
            Aa = A[:, :, 2**k - 1 : L : 2**k]
            Xa = X[:, :, 2**k - 1 : L : 2**k]

            T = 2 * (Xa.size(2) // 2)

            if T < Xa.size(2):
                Xa[:, :, -1].add_(Aa[:, :, -1].mul(Xa[:, :, -2]))
                Aa[:, :, -1].mul_(Aa[:, :, -2])

            Aa = Aa[:, :, :T].view(B, D, T // 2, 2, -1)
            Xa = Xa[:, :, :T].view(B, D, T // 2, 2, -1)

            Xa[:, :, 1:, 0].add_(
                Aa[:, :, 1:, 0].mul(Xa[:, :, :-1, 1])
            )
            Aa[:, :, 1:, 0].mul_(Aa[:, :, :-1, 1])

    @staticmethod
    def forward(ctx, A_in, X_in):
        """
        Applies the parallel scan operation, as defined above. Returns a new tensor.

        Args:
            A_in : (B, L, D, N)
            X_in : (B, L, D, N)

        Returns:
            H : (B, L, D, N)
        """

        # clone tensor (in-place ops)
        A = A_in.clone()  # (B, L, D, N)
        X = X_in.clone()  # (B, L, D, N)

        # prepare tensors
        A = A.transpose(2, 1)  # (B, D, L, N)
        X = X.transpose(2, 1)  # (B, D, L, N)

        # parallel scan
        PScan.pscan(A, X)

        ctx.save_for_backward(A_in, X)

        return X.transpose(2, 1)

    @staticmethod
    def backward(ctx, grad_output_in):
        """
        Flows the gradient from the output to the input. Returns two new tensors.

        Args:
            ctx : A_in : (B, L, D, N), X : (B, D, L, N)
            grad_output_in : (B, L, D, N)

        Returns:
            gradA : (B, L, D, N), gradX : (B, L, D, N)
        """

        A_in, X = ctx.saved_tensors

        # clone tensors
        A = A_in.clone()
        # grad_output_in will be cloned with flip()

        # prepare tensors
        A = A.transpose(2, 1)  # (B, D, L, N)
        A = torch.cat((A[:, :, :1], A[:, :, 1:].flip(2)), dim=2)
        grad_output_b = grad_output_in.transpose(2, 1)

        # reverse parallel scan
        grad_output_b = grad_output_b.flip(2)
        PScan.pscan(A, grad_output_b)
        grad_output_b = grad_output_b.flip(2)

        Q = torch.zeros_like(X)
        Q[:, :, 1:].add_(X[:, :, :-1] * grad_output_b[:, :, 1:])

        return Q.transpose(2, 1), grad_output_b.transpose(2, 1)


pscan = PScan.apply


@dataclass
class MambaConfig:
    dim: int  # D
    depth: int
    dt_rank: Union[int, str] = "auto"
    d_state: int = 16  # N in paper/comments
    expand_factor: int = 2  # E in paper/comments
    d_conv: int = 4

    dt_min: float = 0.001
    dt_max: float = 0.1
    dt_init: str = "random"  # "random" or "constant"
    dt_scale: float = 1.0
    dt_init_floor = 1e-4

    bias: bool = False
    conv_bias: bool = True

    pscan: bool = True  # use parallel scan mode or sequential mode when training

    def __post_init__(self):
        self.d_inner = (
            self.expand_factor * self.dim
        )  # E*D = ED in comments

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

class ResidualBlock(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()

        self.mixer = MambaBlock(config)
        self.norm = RMSNorm(config.dim)

    def forward(self, x):
        # x : (B, L, D)

        # output : (B, L, D)

        output = self.mixer(self.norm(x)) + x
        return output

    def step(self, x, cache):
        # x : (B, D)
        # cache : (h, inputs)
        # h : (B, ED, N)
        # inputs: (B, ED, d_conv-1)

        # output : (B, D)
        # cache : (h, inputs)

        output, cache = self.mixer.step(self.norm(x), cache)
        output = output + x
        return output, cache



class MambaBlock(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()

        self.config = config

        # projects block input from D to 2*ED (two branches)
        self.in_proj = nn.Linear(
            config.dim, 2 * config.d_inner, bias=config.bias
        )

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

        # projects x to input-dependent Δ, B, C
        self.x_proj = nn.Linear(
            config.d_inner,
            config.dt_rank + 2 * config.d_state,
            bias=False,
        )

        # projects Δ from dt_rank to d_inner
        self.dt_proj = nn.Linear(
            config.dt_rank, config.d_inner, bias=True
        )

        # dt initialization
        # dt weights
        dt_init_std = config.dt_rank**-0.5 * config.dt_scale
        if config.dt_init == "constant":
            nn.init.constant_(self.dt_proj.weight, dt_init_std)
        elif config.dt_init == "random":
            nn.init.uniform_(
                self.dt_proj.weight, -dt_init_std, dt_init_std
            )
        else:
            raise NotImplementedError

        # dt bias
        dt = torch.exp(
            torch.rand(config.d_inner)
            * (math.log(config.dt_max) - math.log(config.dt_min))
            + math.log(config.dt_min)
        ).clamp(min=config.dt_init_floor)
        inv_dt = dt + torch.log(
            -torch.expm1(-dt)
        )  # inverse of softplus: https://github.com/pytorch/pytorch/issues/72759
        with torch.no_grad():
            self.dt_proj.bias.copy_(inv_dt)
        # self.dt_proj.bias._no_reinit = True # initialization would set all Linear.bias to zero, need to mark this one as _no_reinit
        # todo : explain why removed

        # S4D real initialization
        A = torch.arange(
            1, config.d_state + 1, dtype=torch.float32
        ).repeat(config.d_inner, 1)
        self.A_log = nn.Parameter(
            torch.log(A)
        )  # why store A in log ? to keep A < 0 (cf -torch.exp(...)) ? for gradient stability ?
        self.D = nn.Parameter(torch.ones(config.d_inner))

        # projects block output from ED back to D
        self.out_proj = nn.Linear(
            config.d_inner, config.dim, bias=config.bias
        )

    def forward(self, x):
        # x : (B, L, D)

        # y : (B, L, D)

        _, L, _ = x.shape

        xz = self.in_proj(x)  # (B, L, 2*ED)
        x, z = xz.chunk(2, dim=-1)  # (B, L, ED), (B, L, ED)

        # x branch
        x = x.transpose(1, 2)  # (B, ED, L)
        x = self.conv1d(x)[
            :, :, :L
        ]  # depthwise convolution over time, with a short filter
        x = x.transpose(1, 2)  # (B, L, ED)

        x = F.silu(x)
        y = self.ssm(x)

        # z branch
        z = F.silu(z)

        output = y * z
        output = self.out_proj(output)  # (B, L, D)

        return output

    def ssm(self, x):
        # x : (B, L, ED)

        # y : (B, L, ED)

        A = -torch.exp(self.A_log.float())  # (ED, N)
        D = self.D.float()
        # TODO remove .float()

        deltaBC = self.x_proj(x)  # (B, L, dt_rank+2*N)

        delta, B, C = torch.split(
            deltaBC,
            [
                self.config.dt_rank,
                self.config.d_state,
                self.config.d_state,
            ],
            dim=-1,
        )  # (B, L, dt_rank), (B, L, N), (B, L, N)
        delta = F.softplus(self.dt_proj(delta))  # (B, L, ED)

        if self.config.pscan:
            y = self.selective_scan(x, delta, A, B, C, D)
        else:
            y = self.selective_scan_seq(x, delta, A, B, C, D)

        return y

    def selective_scan(self, x, delta, A, B, C, D):
        # x : (B, L, ED)
        # Δ : (B, L, ED)
        # A : (ED, N)
        # B : (B, L, N)
        # C : (B, L, N)
        # D : (ED)

        # y : (B, L, ED)

        deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (B, L, ED, N)
        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2)  # (B, L, ED, N)

        BX = deltaB * (x.unsqueeze(-1))  # (B, L, ED, N)

        hs = pscan(deltaA, BX)

        y = (hs @ C.unsqueeze(-1)).squeeze(
            3
        )  # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)

        y = y + D * x

        return y

    def selective_scan_seq(self, x, delta, A, B, C, D):
        # x : (B, L, ED)
        # Δ : (B, L, ED)
        # A : (ED, N)
        # B : (B, L, N)
        # C : (B, L, N)
        # D : (ED)

        # y : (B, L, ED)

        _, L, _ = x.shape

        deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (B, L, ED, N)
        deltaB = delta.unsqueeze(-1) * B.unsqueeze(2)  # (B, L, ED, N)

        BX = deltaB * (x.unsqueeze(-1))  # (B, L, ED, N)

        h = torch.zeros(
            x.size(0),
            self.config.d_inner,
            self.config.d_state,
            device=deltaA.device,
        )  # (B, ED, N)
        hs = []

        for t in range(0, L):
            h = deltaA[:, t] * h + BX[:, t]
            hs.append(h)

        hs = torch.stack(hs, dim=1)  # (B, L, ED, N)

        y = (hs @ C.unsqueeze(-1)).squeeze(
            3
        )  # (B, L, ED, N) @ (B, L, N, 1) -> (B, L, ED, 1)

        y = y + D * x

        return y

    # -------------------------- inference -------------------------- #
    """
    Concerning auto-regressive inference

    The cool part of using Mamba : inference is constant wrt to sequence length
    We just have to keep in cache, for each layer, two things :
    - the hidden state h (which is (B, ED, N)), as you typically would when doing inference with a RNN
    - the last d_conv-1 inputs of the layer, to be able to compute the 1D conv which is a convolution over the time dimension
      (d_conv is fixed so this doesn't incur a growing cache as we progress on generating the sequence)
      (and d_conv is usually very small, like 4, so we just have to "remember" the last 3 inputs)

    Concretely, these two quantities are put inside a cache tuple, and are named h and inputs respectively.
    h is (B, ED, N), and inputs is (B, ED, d_conv-1)
    The MambaBlock.step() receives this cache, and, along with outputing the output, alos outputs the updated cache for the next call.

    The cache object is initialized as follows : (None, torch.zeros()).
    When h is None, the selective scan function detects it and start with h=0.
    The torch.zeros() isn't a problem (it's same as just feeding the input, because the conv1d is padded)

    As we need one such cache variable per layer, we store a caches object, which is simply a list of cache object. (See mamba_lm.py)
    """

    def step(self, x, cache):
        # x : (B, D)
        # cache : (h, inputs)
        # h : (B, ED, N)
        # inputs : (B, ED, d_conv-1)

        # y : (B, D)
        # cache : (h, inputs)

        h, inputs = cache

        xz = self.in_proj(x)  # (B, 2*ED)
        x, z = xz.chunk(2, dim=1)  # (B, ED), (B, ED)

        # x branch
        x_cache = x.unsqueeze(2)
        x = self.conv1d(torch.cat([inputs, x_cache], dim=2))[
            :, :, self.config.d_conv - 1
        ]  # (B, ED)

        x = F.silu(x)
        y, h = self.ssm_step(x, h)

        # z branch
        z = F.silu(z)

        output = y * z
        output = self.out_proj(output)  # (B, D)

        # prepare cache for next call
        inputs = torch.cat(
            [inputs[:, :, 1:], x_cache], dim=2
        )  # (B, ED, d_conv-1)
        cache = (h, inputs)

        return output, cache

    def ssm_step(self, x, h):
        # x : (B, ED)
        # h : (B, ED, N)

        # y : (B, ED)
        # h : (B, ED, N)

        A = -torch.exp(
            self.A_log.float()
        )  # (ED, N) # todo : ne pas le faire tout le temps, puisque c'est indépendant de la timestep
        D = self.D.float()
        # TODO remove .float()

        deltaBC = self.x_proj(x)  # (B, dt_rank+2*N)

        delta, B, C = torch.split(
            deltaBC,
            [
                self.config.dt_rank,
                self.config.d_state,
                self.config.d_state,
            ],
            dim=-1,
        )  # (B, dt_rank), (B, N), (B, N)
        delta = F.softplus(self.dt_proj(delta))  # (B, ED)

        deltaA = torch.exp(delta.unsqueeze(-1) * A)  # (B, ED, N)
        deltaB = delta.unsqueeze(-1) * B.unsqueeze(1)  # (B, ED, N)

        BX = deltaB * (x.unsqueeze(-1))  # (B, ED, N)

        if h is None:
            h = torch.zeros(
                x.size(0),
                self.config.d_inner,
                self.config.d_state,
                device=deltaA.device,
            )  # (B, ED, N)

        h = deltaA * h + BX  # (B, ED, N)

        y = (h @ C.unsqueeze(-1)).squeeze(
            2
        )  # (B, ED, N) @ (B, N, 1) -> (B, ED, 1)

        y = y + D * x

        # todo : pq h.squeeze(1) ??
        return y, h.squeeze(1)


class Mamba(nn.Module):
    def __init__(self, config: MambaConfig):
        super().__init__()

        self.config = config

        self.layers = nn.ModuleList(
            [ResidualBlock(config) for _ in range(config.depth)]
        )
        # self.norm_f = RMSNorm(config.dim)

    def forward(self, x):
        # x : (B, L, D)

        # y : (B, L, D)

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

        # x = self.norm_f(x)

        return x

    def step(self, x, caches):
        # x : (B, L, D)
        # caches : [cache(layer) for all layers], cache : (h, inputs)

        # y : (B, L, D)
        # caches : [cache(layer) for all layers], cache : (h, inputs)

        for i, layer in enumerate(self.layers):
            x, caches[i] = layer.step(x, caches[i])

        return x, caches
    
    
class LowLevelMamba(nn.Module):
    """
    LowLevelMamba is a PyTorch module that implements a low-level Mamba model.

    Args:
        dim (int): Dimension of the Mamba model. Default is 4.
        depth (int): Depth of the Mamba model. Default is 3.
        dt_rank (int): Rank of the time step tensor in the Mamba model. Default is 2.
        d_state (int): Dimension of the state tensor in the Mamba model. Default is 2.
        expand_factor (int): Expansion factor of the Mamba model. Default is 2.
        d_conv (int): Dimension of the convolutional kernel in the Mamba model. Default is 3.
        dt_min (float): Minimum value of the time step in the Mamba model. Default is 0.001.
        dt_max (float): Maximum value of the time step in the Mamba model. Default is 0.1.
        dt_init (str): Initialization method for the time step tensor in the Mamba model. Default is "random".
        dt_scale (float): Scaling factor for the time step tensor in the Mamba model. Default is 1.0.
        bias (bool): Whether to include bias terms in the Mamba model. Default is False.
        conv_bias (bool): Whether to include bias terms in the convolutional layers of the Mamba model. Default is True.
        pscan (bool): Whether to use parallel scan operation in the Mamba model. Default is True.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.
    """

    def __init__(
        self,
        dim: int = 4,
        depth: int = 3,
        dt_rank: int = 2,
        d_state: int = 2,
        expand_factor: int = 2,
        d_conv: int = 3,
        dt_min: float = 0.001,
        dt_max: float = 0.1,
        dt_init: str = "random",
        dt_scale: float = 1.0,
        bias: bool = False,
        conv_bias: bool = True,
        pscan: bool = True,
        *args,
        **kwargs
    ):
        super().__init__()
        self.dim = dim
        self.depth = depth
        self.dt_rank = dt_rank
        self.d_state = d_state
        self.expand_factor = expand_factor
        self.d_conv = d_conv
        self.dt_min = dt_min
        self.dt_max = dt_max
        self.dt_init = dt_init
        self.dt_scale = dt_scale
        self.bias = bias
        self.conv_bias = conv_bias
        self.pscan = pscan

        config = MambaConfig(
            dim=dim,
            depth=depth,
            dt_rank=dt_rank,
            d_state=d_state,
            expand_factor=expand_factor,
            d_conv=d_conv,
            dt_min=dt_min,
            dt_max=dt_max,
            dt_init=dt_init,
            dt_scale=dt_scale,
            bias=bias,
            conv_bias=conv_bias,
            pscan=pscan,
            *args,
            **kwargs
        )

        self.model = Mamba(config)

    def forward(self, x: Tensor):
        """
        Forward pass of the LowLevelMamba model.

        Args:
            x (Tensor): Input tensor.

        Returns:
            Tensor: Output tensor.
        """
        return self.model(x)
        
        
        

class HighLevelMamba(nn.Module):
    """
    LowLevelMamba is a PyTorch module that implements a low-level Mamba model.

    Args:
        dim (int): Dimension of the Mamba model. Default is 4.
        depth (int): Depth of the Mamba model. Default is 3.
        dt_rank (int): Rank of the time step tensor in the Mamba model. Default is 2.
        d_state (int): Dimension of the state tensor in the Mamba model. Default is 2.
        expand_factor (int): Expansion factor of the Mamba model. Default is 2.
        d_conv (int): Dimension of the convolutional kernel in the Mamba model. Default is 3.
        dt_min (float): Minimum value of the time step in the Mamba model. Default is 0.001.
        dt_max (float): Maximum value of the time step in the Mamba model. Default is 0.1.
        dt_init (str): Initialization method for the time step tensor in the Mamba model. Default is "random".
        dt_scale (float): Scaling factor for the time step tensor in the Mamba model. Default is 1.0.
        bias (bool): Whether to include bias terms in the Mamba model. Default is False.
        conv_bias (bool): Whether to include bias terms in the convolutional layers of the Mamba model. Default is True.
        pscan (bool): Whether to use parallel scan operation in the Mamba model. Default is True.
        *args: Additional positional arguments.
        **kwargs: Additional keyword arguments.
    """

    def __init__(
        self,
        dim: int = 4,
        depth: int = 3,
        dt_rank: int = 2,
        d_state: int = 2,
        expand_factor: int = 2,
        d_conv: int = 3,
        dt_min: float = 0.001,
        dt_max: float = 0.1,
        dt_init: str = "random",
        dt_scale: float = 1.0,
        bias: bool = False,
        conv_bias: bool = True,
        pscan: bool = True,
        *args,
        **kwargs
    ):
        super().__init__()
        self.dim = dim
        self.depth = depth
        self.dt_rank = dt_rank
        self.d_state = d_state
        self.expand_factor = expand_factor
        self.d_conv = d_conv
        self.dt_min = dt_min
        self.dt_max = dt_max
        self.dt_init = dt_init
        self.dt_scale = dt_scale
        self.bias = bias
        self.conv_bias = conv_bias
        self.pscan = pscan

        config = MambaConfig(
            dim=dim,
            depth=depth,
            dt_rank=dt_rank,
            d_state=d_state,
            expand_factor=expand_factor,
            d_conv=d_conv,
            dt_min=dt_min,
            dt_max=dt_max,
            dt_init=dt_init,
            dt_scale=dt_scale,
            bias=bias,
            conv_bias=conv_bias,
            pscan=pscan,
            *args,
            **kwargs
        )

        self.model = Mamba(config)

    def forward(self, x: Tensor):
        """
        Forward pass of the LowLevelMamba model.

        Args:
            x (Tensor): Input tensor.

        Returns:
            Tensor: Output tensor.
        """
        return self.model(x)
        
        

class HSSSMamba(nn.Module):
    """
    HSSSMamba is a PyTorch module that represents the HSSS Mamba model.

    Args:
        dim_in (int): Input dimension for the low level Mamba module.
        depth_in (int): Depth of the low level Mamba module.
        dt_rank_in (int): Rank of the time step tensor for the low level Mamba module.
        d_state_in (int): Dimension of the state tensor for the low level Mamba module.
        expand_factor_in (int): Expansion factor for the low level Mamba module.
        d_conv_in (int): Dimension of the convolutional filters for the low level Mamba module.
        dt_min_in (float): Minimum value for the time step tensor for the low level Mamba module.
        dt_max_in (float): Maximum value for the time step tensor for the low level Mamba module.
        dt_init_in (str): Initialization method for the time step tensor for the low level Mamba module.
        dt_scale_in (float): Scaling factor for the time step tensor for the low level Mamba module.
        bias_in (bool): Whether to include bias in the low level Mamba module.
        conv_bias_in (bool): Whether to include bias in the convolutional filters for the low level Mamba module.
        pscan_in (bool): Whether to use parallel scan in the low level Mamba module.
        dim (int): Input dimension for the high level Mamba module.
        depth (int): Depth of the high level Mamba module.
        dt_rank (int): Rank of the time step tensor for the high level Mamba module.
        d_state (int): Dimension of the state tensor for the high level Mamba module.
        expand_factor (int): Expansion factor for the high level Mamba module.
        d_conv (int): Dimension of the convolutional filters for the high level Mamba module.
        dt_min (float): Minimum value for the time step tensor for the high level Mamba module.
        dt_max (float): Maximum value for the time step tensor for the high level Mamba module.
        dt_init (str): Initialization method for the time step tensor for the high level Mamba module.
        dt_scale (float): Scaling factor for the time step tensor for the high level Mamba module.
        bias (bool): Whether to include bias in the high level Mamba module.
        conv_bias (bool): Whether to include bias in the convolutional filters for the high level Mamba module.
        pscan (bool): Whether to use parallel scan in the high level Mamba module.
        *args: Variable length argument list.
        **kwargs: Arbitrary keyword arguments.
    """
    def __init__(
        self,
        dim_in: int = 8,
        depth_in: int = 6,
        dt_rank_in: int = 4,
        d_state_in: int = 4,
        expand_factor_in: int = 4,
        d_conv_in: int = 6,
        dt_min_in: float = 0.001,
        dt_max_in: float = 0.1,
        dt_init_in: str = "random",
        dt_scale_in: float = 1.0,
        bias_in: bool = False,
        conv_bias_in: bool = True,
        pscan_in: bool = True,
        dim: int = 4,
        depth: int = 3,
        dt_rank: int = 2,
        d_state: int = 2,
        expand_factor: int = 2,
        d_conv: int = 3,
        dt_min: float = 0.001,
        dt_max: float = 0.1,
        dt_init: str = "random",
        dt_scale: float = 1.0,
        bias: bool = False,
        conv_bias: bool = True,
        pscan: bool = True,
        *args,
        **kwargs
    ):
        super().__init__()
        self.dim_in = dim_in
        self.depth_in = depth_in
        self.dt_rank_in = dt_rank_in
        self.d_state_in = d_state_in
        self.expand_factor_in = expand_factor_in
        self.d_conv_in = d_conv_in
        self.dt_min_in = dt_min_in
        self.dt_max_in = dt_max_in
        self.dt_init_in = dt_init_in
        self.dt_scale_in = dt_scale_in
        self.bias_in = bias_in
        self.conv_bias_in = conv_bias_in
        self.pscan_in = pscan_in
        self.dim = dim
        self.depth = depth
        self.dt_rank = dt_rank
        self.d_state = d_state
        self.expand_factor = expand_factor
        self.d_conv = d_conv
        self.dt_min = dt_min
        self.dt_max = dt_max
        self.dt_init = dt_init
        self.dt_scale = dt_scale
        self.bias = bias
        self.conv_bias = conv_bias
        self.pscan = pscan
        
        self.low_level_mamba = LowLevelMamba(
            dim=dim_in,
            depth=depth_in,
            dt_rank=dt_rank_in,
            d_state=d_state_in,
            expand_factor=expand_factor_in,
            d_conv=d_conv_in,
            dt_min=dt_min_in,
            dt_max=dt_max_in,
            dt_init=dt_init_in,
            dt_scale=dt_scale_in,
            bias=bias_in,
            conv_bias=conv_bias_in,
            pscan=pscan_in,
        )
        
        # Linear projection
        self.ffn = FeedForward(
            dim_in,
            dim,
            dim * 4,
            *args,
            **kwargs
        )
        
        # High level Mamba
        self.high_level_mamba = HighLevelMamba(
            dim=dim,
            depth=depth,
            dt_rank=dt_rank,
            d_state=d_state,
            expand_factor=expand_factor,
            d_conv=d_conv,
            dt_min=dt_min,
            dt_max=dt_max,
            dt_init=dt_init,
            dt_scale=dt_scale,
            bias=bias,
            conv_bias=conv_bias,
            pscan=pscan,
        )
        
    def forward(self, x: Tensor):
        """
        Forward pass of the HSSSMamba model.

        Args:
            x (Tensor): Input tensor.

        Returns:
            Tensor: Output tensor.
        """
        x = self.low_level_mamba(x)
        x = self.ffn(x)
        x = self.high_level_mamba(x)
        return x
    
