import torch
import torch.nn as nn
from functools import partial
from mamba_ssm import Mamba
from models.transformer_modules import Mlp, GatedMLP, Block
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from models.utils import flip_padded_sequences

class MambaBlock(nn.Module):
    """
    Wrapper block around the Mamba SSM layer providing the usual
    transformer-style residual + MLP structure.

    If `mamba-ssm` is not installed, the block falls back to an
    identity (no-op) for the SSM portion while still running the
    feed-forward MLP path so networks remain usable.
    """
    def __init__(self, d_model, d_state=16, d_conv=4, expand=2,
                 mlp_ratio=4., drop=0.1, drop_path=0.1, act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(d_model)
        # Mamba SSM layer (state space model)
        if Mamba is not None:
            self.mamba = Mamba(d_model=d_model, d_state=d_state, d_conv=d_conv, expand=expand)
        else:
            self.mamba = nn.Identity()

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()        
        # self.norm2 = norm_layer(d_model)
        # mlp_hidden_dim = int(d_model * mlp_ratio)
        # self.mlp = GatedMLP(in_features=d_model, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        

    def forward(self, x, mask=None):
        if mask is not None: 
            x = x.masked_fill(~mask.unsqueeze(-1), 0.0)
        # x: (B, L, D)
        # SSM (Mamba) path with residual
        x_in = x
        x = self.mamba(self.norm1(x))
        if mask is not None: 
            x = x.masked_fill(~mask.unsqueeze(-1), 0.0)
        x = x_in + self.drop_path(x)        

        return x

class SeqEncoder(nn.Module):
    def __init__(self, 
                 seq_dim, 
                 depth=2,                    
                 d_state=16, 
                 d_conv=4, 
                 expand=2, 
                 bidi=True,
                 use_transformer=False,
                 norm_layer=nn.LayerNorm):
        super().__init__()
        
        self.bidi = bidi
        self.use_transformer = use_transformer
        if use_transformer:
            assert not bidi, "Bidirectional mode is not supported with transformer encoder"
            # encoder_layer = nn.TransformerEncoderLayer(d_model=seq_dim, nhead=8, dropout=0.1, batch_first=True, norm_first=True)
            # self.seq_layers = nn.TransformerEncoder(encoder_layer, num_layers=depth*2)
            self.seq_layers = nn.ModuleList([            
                Block(dim=seq_dim, num_heads=4, drop=0.1, drop_path=0.0, norm_layer=norm_layer)
                for i in range(depth*2)
            ])
            
        else:
            self.seq_layers = nn.ModuleList([            
                MambaBlock(d_model=seq_dim, d_state=d_state, d_conv=d_conv, expand=expand, norm_layer=norm_layer)
                for i in range(depth)
            ])
        self.norm_fwd = norm_layer(seq_dim)
        # self.post_norm_fwd = norm_layer(seq_dim)
        if self.bidi: 
            self.seq_layers_bwd = nn.ModuleList([
                MambaBlock(d_model=seq_dim, d_state=d_state, d_conv=d_conv, expand=expand, norm_layer=norm_layer)
                for i in range(depth)
            ])
            self.norm_bwd = norm_layer(seq_dim)
        
        # encoder_layer = nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dropout=dropout, batch_first=True, norm_first=True)
        # self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        
        

    def forward(self, x, mask=None):
        
        if self.bidi:
            # Process forward and backward paths independently
            xr = flip_padded_sequences(x, mask)
            mr = mask
            
            for blk, blk_bwd in zip(self.seq_layers, self.seq_layers_bwd):
                xf = blk(x, mask)                
                xr = blk_bwd(xr, mr)
            xf = self.norm_fwd(xf)
            xr = self.norm_bwd(xr)
            return xf, xr
        else:
            if self.use_transformer:
                mask = ~mask
            #     # For transformer, we need to provide a key padding mask
            #     x = self.seq_layers(x, src_key_padding_mask=~mask)
            # else:
            for blk in self.seq_layers:                
                x = blk(x, mask)                
                
            x = self.norm_fwd(x)
            return x