# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple, Union

import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule
from mmengine.model import BaseModule
from torch import Tensor

from mmdet.registry import MODELS
from mmdet.utils import ConfigType, MultiConfig, OptConfigType


import torch
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import trunc_normal_, DropPath
from timm.models.registry import register_model
import math
import math
import time
from typing import Tuple
import copy
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as checkpoint
# from depthwise_conv2d_implicit_gemm import DepthWiseConv2dImplicitGEMM as DWConv
from einops import einsum, rearrange, repeat
from einops.layers.torch import Rearrange
from fvcore.nn import FlopCountAnalysis, flop_count_table
from mmcv.cnn import ConvModule
from mmengine.runner import load_state_dict
#from natten import NeighborhoodAttention2D, use_fused_na, use_gemm_na
#from natten.functional import na2d, na2d_av, na2d_qk, natten2dav, natten2dqkrpb
# from timm.models.vision_transformer import _cfg
from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.layers import DropPath, to_2tuple, trunc_normal_
from timm.models.registry import register_model
from fvcore.nn import FlopCountAnalysis, flop_count_str, flop_count, parameter_count
#from natten.flops import qk_2d_rpb_flop, av_2d_flop, add_natten_handle

try:
    from csm_triton import CrossScanTriton, CrossMergeTriton
except:
    from mmdet.models.necks.csm_triton import CrossScanTriton, CrossMergeTriton

# import selective_scan_cuda_oflex

# import selective_scan_fn


import numpy as np
import cv2
from einops import einsum, rearrange, repeat
from einops.layers.torch import Rearrange
from timm.models.layers import DropPath, to_2tuple, trunc_normal_

try:
    from csm_triton import CrossScanTriton, CrossMergeTriton
except:
    from mmdet.models.necks.csm_triton import CrossScanTriton, CrossMergeTriton

# import selective_scan_cuda_oflex

# Import MambaVision components from first code
# from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn


from timm.models.vision_transformer import Mlp

config = {
    "tiny": [64, 128, 256, 512],
    "small": [96, 192, 384, 768],
    "base": [128, 256, 512, 1024]
}

import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from einops import rearrange, repeat
from mamba_ssm import selective_scan_fn


class FGMM(nn.Module):

    def __init__(self, d_model, reduction=4):
        super().__init__()
        self.d_model = d_model

        self.freq_conv = nn.Sequential(
            nn.Conv2d(d_model, d_model // reduction, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(d_model // reduction, d_model, 1),
            nn.Sigmoid()
        )

        self.adaptive_weight = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(d_model, d_model // reduction, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(d_model // reduction, 2, 1),
            nn.Softmax(dim=1)
        )

    def forward(self, spatial_feat, freq_feat):

        freq_enhanced = freq_feat * self.freq_conv(freq_feat)

        combined_feat = spatial_feat + freq_enhanced
        weights = self.adaptive_weight(combined_feat)
        spatial_weight = weights[:, 0:1, :, :]
        freq_weight = weights[:, 1:2, :, :]
        output = spatial_weight * spatial_feat + freq_weight * freq_enhanced

        return output, freq_weight.squeeze()


class FEM(nn.Module):

    def __init__(self, d_model):
        super().__init__()
        self.d_model = d_model

        self.low_freq_weight = nn.Parameter(torch.ones(1))
        self.high_freq_weight = nn.Parameter(torch.ones(1))

        self.low_freq_proj = nn.Linear(d_model, d_model)
        self.high_freq_proj = nn.Linear(d_model, d_model)

        self.freq_fusion = nn.Sequential(
            nn.Linear(d_model * 3, d_model * 2),
            nn.ReLU(),
            nn.Linear(d_model * 2, d_model)
        )


        self.weight_norm = nn.Sigmoid()

    def forward(self, x, H, W):
        B, L, D = x.shape


        x_img = rearrange(x, 'b (h w) d -> b d h w', h=H, w=W)


        x_freq = torch.fft.rfft2(x_img, dim=(-2, -1))
        x_freq_mag = torch.abs(x_freq)


        spatial_feat = rearrange(x_img, 'b d h w -> b (h w) d')


        low_weight = self.weight_norm(self.low_freq_weight)
        high_weight = self.weight_norm(self.high_freq_weight)

        freq_h, freq_w = x_freq.shape[-2], x_freq.shape[-1]
        low_freq_mask = torch.zeros_like(x_freq_mag)
        center_h, center_w = freq_h // 4, freq_w // 4
        low_freq_mask[:, :, :center_h, :center_w] = 1

        low_freq = x_freq * low_freq_mask
        low_freq_spatial = torch.fft.irfft2(low_freq, s=(H, W), dim=(-2, -1))
        low_freq_feat = self.low_freq_proj(rearrange(low_freq_spatial, 'b d h w -> b (h w) d'))
        low_freq_feat = low_freq_feat * low_weight

        high_freq_mask = torch.ones_like(x_freq_mag)
        high_freq_mask[:, :, :center_h, :center_w] = 0

        high_freq = x_freq * high_freq_mask
        high_freq_spatial = torch.fft.irfft2(high_freq, s=(H, W), dim=(-2, -1))
        high_freq_feat = self.high_freq_proj(rearrange(high_freq_spatial, 'b d h w -> b (h w) d'))
        high_freq_feat = high_freq_feat * high_weight

        combined_freq = torch.cat([spatial_feat, low_freq_feat, high_freq_feat], dim=-1)
        freq_enhanced = self.freq_fusion(combined_freq)


        freq_weights = {
            'low_freq_weight': low_weight.item(),
            'high_freq_weight': high_weight.item()
        }

        return freq_enhanced, x_freq_mag, freq_weights


class MambaVisionMixer(nn.Module):

    def __init__(
            self,
            d_model,
            d_state=16,
            d_conv=4,
            expand=2,
            dt_rank="auto",
            dt_min=0.001,
            dt_max=0.1,
            dt_init="random",
            dt_scale=1.0,
            dt_init_floor=1e-4,
            conv_bias=True,
            bias=False,
            use_fast_path=True,
            layer_idx=None,
            device=None,
            dtype=None,
            use_frequency=True,
            safe_mode=False,
    ):
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.d_conv = d_conv
        self.expand = expand
        self.d_inner = int(self.expand * self.d_model)
        self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
        self.use_fast_path = use_fast_path
        self.layer_idx = layer_idx
        self.use_frequency = use_frequency
        self.safe_mode = safe_mode

        # SSM组件
        self.in_proj = nn.Linear(self.d_model, self.d_inner, bias=bias, **factory_kwargs)
        self.x_proj = nn.Linear(
            self.d_inner // 2, self.dt_rank + self.d_state * 2, bias=False, **factory_kwargs
        )
        self.dt_proj = nn.Linear(self.dt_rank, self.d_inner // 2, bias=True, **factory_kwargs)

        # 初始化dt参数
        dt_init_std = self.dt_rank ** -0.5 * dt_scale
        if dt_init == "constant":
            nn.init.constant_(self.dt_proj.weight, dt_init_std)
        elif dt_init == "random":
            nn.init.uniform_(self.dt_proj.weight, -dt_init_std, dt_init_std)
        else:
            raise NotImplementedError

        dt = torch.exp(
            torch.rand(self.d_inner // 2, **factory_kwargs) * (math.log(dt_max) - math.log(dt_min))
            + math.log(dt_min)
        ).clamp(min=dt_init_floor)
        inv_dt = dt + torch.log(-torch.expm1(-dt))
        with torch.no_grad():
            self.dt_proj.bias.copy_(inv_dt)
        self.dt_proj.bias._no_reinit = True

        # SSM参数A和D
        A = repeat(
            torch.arange(1, self.d_state + 1, dtype=torch.float32, device=device),
            "n -> d n",
            d=self.d_inner // 2,
        ).contiguous()
        A_log = torch.log(A)
        self.A_log = nn.Parameter(A_log)
        self.A_log._no_weight_decay = True
        self.D = nn.Parameter(torch.ones(self.d_inner // 2, device=device))
        self.D._no_weight_decay = True

        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)

        self.conv1d_x = nn.Conv1d(
            in_channels=self.d_inner // 2,
            out_channels=self.d_inner // 2,
            bias=conv_bias,
            kernel_size=d_conv,
            groups=self.d_inner // 2,
            **factory_kwargs,
        )
        self.conv1d_z = nn.Conv1d(
            in_channels=self.d_inner // 2,
            out_channels=self.d_inner // 2,
            bias=conv_bias,
            kernel_size=d_conv,
            groups=self.d_inner // 2,
            **factory_kwargs,
        )

        if self.use_frequency:
            try:

                self.freq_ratio = nn.Parameter(torch.tensor(0.5))
                self.freq_ratio_norm = nn.Sigmoid()

                self.freq_ssm = FEM(d_model=self.d_inner // 2)
                self.freq_adaptive = FGMM(self.d_inner // 2)
                self.freq_gate = nn.Sequential(
                    nn.Linear(self.d_inner // 2, self.d_inner // 4),
                    nn.ReLU(),
                    nn.Linear(self.d_inner // 4, 1),
                    nn.Sigmoid()
                )
                print(f"✅ Frequency enhancement modules initialized successfully")
            except Exception as e:
                print(f"❌ Error initializing frequency modules: {e}")
                print(f"🔄 Falling back to safe mode (frequency disabled)")
                self.use_frequency = False

    def forward(self, hidden_states, H=None, W=None):
        original_shape = hidden_states.shape

        if hidden_states.dim() == 4:
            B, H, W, D = hidden_states.shape
            hidden_states = rearrange(hidden_states, 'b h w d -> b (h w) d')
        elif hidden_states.dim() == 3:
            if len(original_shape) == 3:
                B, L, D = hidden_states.shape

                if H is not None and W is not None:
                    if L == H * W:

                        pass
                    else:
                        if L == H * W:
                            pass
                        else:
                            raise ValueError(f"L={L} doesn't match H*W={H * W}. "
                                             f"Please check if H, W are window sizes for windowed input.")
                else:
                    H = W = int(math.sqrt(L))
                    if H * W != L:
                        raise ValueError(f"Cannot infer square H, W from L={L}. "
                                         f"Please provide H, W explicitly.")
            else:
                raise ValueError(f"Unexpected input shape: {original_shape}")
        else:
            raise ValueError(f"Input should be 3D (B, L, D) or 4D (B, H, W, D), got {hidden_states.dim()}D")

        seqlen = H * W

        current_B, current_L, current_D = hidden_states.shape
        if current_L != seqlen:
            raise ValueError(f"Sequence length {current_L} doesn't match H*W={seqlen}. "
                             f"Input shape: {original_shape}, H={H}, W={W}")
        B = current_B

        A = -torch.exp(self.A_log.float())

        xz = self.in_proj(hidden_states)
        x, z = xz.chunk(2, dim=-1)

        x = rearrange(x, "b l d -> b d l")
        z = rearrange(z, "b l d -> b d l")

        try:
            x = F.silu(F.conv1d(input=x, weight=self.conv1d_x.weight, bias=self.conv1d_x.bias,
                                padding='same', groups=self.d_inner // 2))
            z = F.silu(F.conv1d(input=z, weight=self.conv1d_z.weight, bias=self.conv1d_z.bias,
                                padding='same', groups=self.d_inner // 2))

            if hasattr(self, '_debug_mode') and self._debug_mode:
                print(f"🔍 Debug: After conv1d - x.shape={x.shape}, z.shape={z.shape}")

        except Exception as e:
            print(f"❌ Error in convolution: {e}")
            print(f"   x.shape: {x.shape}, z.shape: {z.shape}")
            print(f"   conv1d_x.weight.shape: {self.conv1d_x.weight.shape}")
            print(f"   d_inner: {self.d_inner}, groups: {self.d_inner // 2}")
            raise e

        freq_weights_info = {}
        if self.use_frequency and hasattr(self, 'freq_ssm'):
            try:
                x_seq = rearrange(x, "b d l -> b l d")
                freq_enhanced_x, freq_mag, freq_weights = self.freq_ssm(x_seq, H, W)
                freq_enhanced_x = rearrange(freq_enhanced_x, "b l d -> b d l")

                adaptive_freq_ratio = self.freq_ratio_norm(self.freq_ratio)

                freq_importance = self.freq_gate(rearrange(freq_enhanced_x, "b d l -> b l d")).mean(dim=1)  # (B, 1)

                if x.shape[-1] != H * W or freq_enhanced_x.shape[-1] != H * W:
                    if hasattr(self, '_debug_mode') and self._debug_mode:
                        print(f"⚠️ Warning: Shape mismatch in frequency processing")
                        print(f"  x.shape={x.shape}, freq_enhanced_x.shape={freq_enhanced_x.shape}")
                        print(
                            f"  Expected L=H*W={H * W}, but got x.L={x.shape[-1]}, freq.L={freq_enhanced_x.shape[-1]}")
                    freq_weights_info = {'status': 'shape_mismatch_skipped'}

                else:
                    x_spatial = rearrange(x, "b d l -> b d h w", h=H, w=W)
                    x_freq = rearrange(freq_enhanced_x, "b d l -> b d h w", h=H, w=W)

                    x_fused, freq_adaptive_weight = self.freq_adaptive(x_spatial, x_freq)
                    x = rearrange(x_fused, "b d h w -> b d l")


                    freq_gate_weight = freq_importance.unsqueeze(-1).unsqueeze(-1)  # (B, 1, 1)
                    final_freq_weight = adaptive_freq_ratio * freq_gate_weight
                    x = x * (1 + final_freq_weight)

                    freq_weights_info = {
                        'status': 'success',
                        'adaptive_freq_ratio': adaptive_freq_ratio.item(),
                        'low_freq_weight': freq_weights['low_freq_weight'],
                        'high_freq_weight': freq_weights['high_freq_weight'],
                        'avg_freq_importance': freq_importance.mean().item(),
                        'avg_adaptive_weight': freq_adaptive_weight.mean().item()
                    }

            except Exception as e:
                if hasattr(self, '_debug_mode') and self._debug_mode:
                    print(f"⚠️ Warning: Error in frequency processing: {e}")
                    print(f"  Skipping frequency enhancement for this layer")
                freq_weights_info = {'status': 'error', 'error_msg': str(e)}

        try:
            x_dbl = self.x_proj(rearrange(x, "b d l -> (b l) d"))
            dt, B_ssm, C_ssm = torch.split(x_dbl, [self.dt_rank, self.d_state, self.d_state], dim=-1)
            dt = rearrange(self.dt_proj(dt), "(b l) d -> b d l", l=seqlen)
            B_ssm = rearrange(B_ssm, "(b l) dstate -> b dstate l", l=seqlen).contiguous()
            C_ssm = rearrange(C_ssm, "(b l) dstate -> b dstate l", l=seqlen).contiguous()

            if hasattr(self, '_debug_mode') and self._debug_mode:
                print(f"🔍 Debug: SSM params - dt.shape={dt.shape}, B.shape={B_ssm.shape}, C.shape={C_ssm.shape}")

            # 选择性扫描
            y = selective_scan_fn(x,
                                  dt,
                                  A,
                                  B_ssm,
                                  C_ssm,
                                  self.D.float(),
                                  z=None,
                                  delta_bias=self.dt_proj.bias.float(),
                                  delta_softplus=True,
                                  return_last_state=None)

        except Exception as e:
            print(f"❌ Error in SSM processing: {e}")
            print(f"   x.shape: {x.shape}")
            print(f"   seqlen: {seqlen}")
            raise e

        try:
            y = torch.cat([y, z], dim=1)
            y = rearrange(y, "b d l -> b l d")
            out = self.out_proj(y)

            if hasattr(self, '_debug_mode') and self._debug_mode:
                print(f"🔍 Debug: Final output shape = {out.shape}")

        except Exception as e:
            print(f"❌ Error in final output processing: {e}")
            print(f"   y.shape: {y.shape}, z.shape: {z.shape}")
            raise e

        if hasattr(self, '_return_weights') and self._return_weights:
            return out, freq_weights_info
        return out



def window_partition(x, window_size):
    B, C, H, W = x.shape
    x = x.view(B, C, H // window_size, window_size, W // window_size, window_size)
    windows = x.permute(0, 2, 4, 3, 5, 1).reshape(-1, window_size * window_size, C)
    return windows


def window_reverse(windows, window_size, H, W):
    B = int(windows.shape[0] / (H * W / window_size / window_size))
    x = windows.reshape(B, H // window_size, W // window_size, window_size, window_size, -1)
    x = x.permute(0, 5, 1, 3, 2, 4).reshape(B, windows.shape[2], H, W)
    return x


class GMambaBlock(nn.Module):
    def __init__(self,
                 dim,
                 num_heads,
                 counter,
                 transformer_blocks,
                 mlp_ratio=4.,
                 qkv_bias=False,
                 qk_scale=False,
                 drop=0.,
                 attn_drop=0.,
                 drop_path=0.,
                 act_layer=nn.GELU,
                 norm_layer=nn.LayerNorm,
                 Mlp_block=Mlp,
                 layer_scale=None,
                 window_size=7,  # Add window_size parameter
                 ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.window_size = window_size

        if counter in transformer_blocks:
            None
        else:
            self.mixer = MambaVisionMixer(d_model=dim,
                                          d_state=8,
                                          d_conv=3,
                                          expand=1,
                                          use_frequency=True,
                                          )





        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp_block(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        use_layer_scale = layer_scale is not None and type(layer_scale) in [int, float]
        self.gamma_1 = nn.Parameter(layer_scale * torch.ones(dim)) if use_layer_scale else 1
        self.gamma_2 = nn.Parameter(layer_scale * torch.ones(dim)) if use_layer_scale else 1

    def forward(self, x):
        # Convert from (B, C, H, W) to windowed format for processing
        B, C, H, W = x.shape

        # Window partition for transformer-like processing
        pad_r = (self.window_size - W % self.window_size) % self.window_size
        pad_b = (self.window_size - H % self.window_size) % self.window_size
        if pad_r > 0 or pad_b > 0:
            x = torch.nn.functional.pad(x, (0, pad_r, 0, pad_b))
            _, _, Hp, Wp = x.shape
        else:
            Hp, Wp = H, W

        x_windowed = window_partition(x, self.window_size)  # (num_windows*B, window_size*window_size, C)

        # Apply mixer and MLP
        x_windowed = x_windowed + self.drop_path(self.gamma_1 * self.mixer(self.norm1(x_windowed)))
        x_windowed = x_windowed + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x_windowed)))

        # Reverse window partition
        x = window_reverse(x_windowed, self.window_size, Hp, Wp)

        # Remove padding if added
        if pad_r > 0 or pad_b > 0:
            x = x[:, :, :H, :W].contiguous()

        return x



@MODELS.register_module()
class FPN(BaseModule):
    r"""Feature Pyramid Network with MambaVisionBlock preprocessing.

    This is an implementation of paper `Feature Pyramid Networks for Object
    Detection <https://arxiv.org/abs/1612.03144>`_ with MambaVisionBlock
    preprocessing on backbone outputs.

    Args:
        in_channels (list[int]): Number of input channels per scale.
        out_channels (int): Number of output channels (used at each scale).
        num_outs (int): Number of output scales.
        dims (list[int]): Dimensions for MambaVisionBlocks, should match in_channels.
        start_level (int): Index of the start input backbone level used to
            build the feature pyramid. Defaults to 0.
        end_level (int): Index of the end input backbone level (exclusive) to
            build the feature pyramid. Defaults to -1, which means the
            last level.
        add_extra_convs (bool | str): If bool, it decides whether to add conv
            layers on top of the original feature maps. Defaults to False.
            If True, it is equivalent to `add_extra_convs='on_input'`.
            If str, it specifies the source feature map of the extra convs.
            Only the following options are allowed

            - 'on_input': Last feat map of neck inputs (i.e. backbone feature).
            - 'on_lateral': Last feature map after lateral convs.
            - 'on_output': The last output feature map after fpn convs.
        relu_before_extra_convs (bool): Whether to apply relu before the extra
            conv. Defaults to False.
        no_norm_on_lateral (bool): Whether to apply norm on lateral.
            Defaults to False.
        conv_cfg (:obj:`ConfigDict` or dict, optional): Config dict for
            convolution layer. Defaults to None.
        norm_cfg (:obj:`ConfigDict` or dict, optional): Config dict for
            normalization layer. Defaults to None.
        act_cfg (:obj:`ConfigDict` or dict, optional): Config dict for
            activation layer in ConvModule. Defaults to None.
        upsample_cfg (:obj:`ConfigDict` or dict, optional): Config dict
            for interpolate layer. Defaults to dict(mode='nearest').
        init_cfg (:obj:`ConfigDict` or dict or list[:obj:`ConfigDict` or \
            dict]): Initialization config dict.

    Example:
        >>> import torch
        >>> in_channels = [2, 3, 5, 7]
        >>> dims = [2, 3, 5, 7]  # Should match in_channels for MambaVisionBlocks
        >>> scales = [340, 170, 84, 43]
        >>> inputs = [torch.rand(1, c, s, s)
        ...           for c, s in zip(in_channels, scales)]
        >>> self = FPN(in_channels, 11, len(in_channels), dims).eval()
        >>> outputs = self.forward(inputs)
        >>> for i in range(len(outputs)):
        ...     print(f'outputs[{i}].shape = {outputs[i].shape}')
        outputs[0].shape = torch.Size([1, 11, 340, 340])
        outputs[1].shape = torch.Size([1, 11, 170, 170])
        outputs[2].shape = torch.Size([1, 11, 84, 84])
        outputs[3].shape = torch.Size([1, 11, 43, 43])
    """

    def __init__(
            self,
            in_channels: List[int],
            out_channels: int,
            num_outs: int,
            dims: List[int],  # New parameter for MambaVisionBlock dimensions
            start_level: int = 0,
            end_level: int = -1,
            add_extra_convs: Union[bool, str] = False,
            relu_before_extra_convs: bool = False,
            no_norm_on_lateral: bool = False,
            conv_cfg: OptConfigType = None,
            norm_cfg: OptConfigType = None,
            act_cfg: OptConfigType = None,
            upsample_cfg: ConfigType = dict(mode='nearest'),
            init_cfg: MultiConfig = dict(
                type='Xavier', layer='Conv2d', distribution='uniform')
    ) -> None:
        super().__init__(init_cfg=init_cfg)
        assert isinstance(in_channels, list)
        assert isinstance(dims, list)
        assert len(dims) == len(in_channels), "dims should match in_channels length"

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = num_outs
        self.dims = dims
        self.relu_before_extra_convs = relu_before_extra_convs
        self.no_norm_on_lateral = no_norm_on_lateral
        self.fp16_enabled = False
        self.upsample_cfg = upsample_cfg.copy()

        if end_level == -1 or end_level == self.num_ins - 1:
            self.backbone_end_level = self.num_ins
            assert num_outs >= self.num_ins - start_level
        else:
            # if end_level is not the last level, no extra level is allowed
            self.backbone_end_level = end_level + 1
            assert end_level < self.num_ins
            assert num_outs == end_level - start_level + 1
        self.start_level = start_level
        self.end_level = end_level
        self.add_extra_convs = add_extra_convs
        assert isinstance(add_extra_convs, (str, bool))
        if isinstance(add_extra_convs, str):
            # Extra_convs_source choices: 'on_input', 'on_lateral', 'on_output'
            assert add_extra_convs in ('on_input', 'on_lateral', 'on_output')
        elif add_extra_convs:  # True
            self.add_extra_convs = 'on_input'



        self.backbone_blocks = nn.ModuleList([
            GMambaBlock(
                dim=96,
                num_heads=8,
                counter=0,  # Use Mamba mixer
                transformer_blocks=[],  # Empty list means use Mamba
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),

            GMambaBlock(
                dim=192,
                num_heads=8,
                counter=0,  # Use Mamba mixer
                transformer_blocks=[],  # Empty list means use Mamba
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            GMambaBlock(
                dim=384,
                num_heads=8,
                counter=0,  # Use Mamba mixer
                transformer_blocks=[],  # Empty list means use Mamba
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            nn.Identity(),
        ])



        self.lateral_convs = nn.ModuleList()
        self.fpn_convs = nn.ModuleList()

        for i in range(self.start_level, self.backbone_end_level):
            l_conv = ConvModule(
                in_channels[i],
                out_channels,
                1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg if not self.no_norm_on_lateral else None,
                act_cfg=act_cfg,
                inplace=False)
            fpn_conv = ConvModule(
                out_channels,
                out_channels,
                3,
                padding=1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg,
                inplace=False)

            self.lateral_convs.append(l_conv)
            self.fpn_convs.append(fpn_conv)

        # add extra conv layers (e.g., RetinaNet)
        extra_levels = num_outs - self.backbone_end_level + self.start_level
        if self.add_extra_convs and extra_levels >= 1:
            for i in range(extra_levels):
                if i == 0 and self.add_extra_convs == 'on_input':
                    in_channels = self.in_channels[self.backbone_end_level - 1]
                else:
                    in_channels = out_channels
                extra_fpn_conv = ConvModule(
                    in_channels,
                    out_channels,
                    3,
                    stride=2,
                    padding=1,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg,
                    inplace=False)
                self.fpn_convs.append(extra_fpn_conv)

    def forward(self, inputs: Tuple[Tensor]) -> tuple:
        """Forward function.

        Args:
            inputs (tuple[Tensor]): Features from the upstream network, each
                is a 4D-tensor.

        Returns:
            tuple: Feature maps, each is a 4D-tensor.
        """
        assert len(inputs) == len(self.in_channels)

        # First, apply MambaVisionBlocks to preprocess backbone outputs with residual connections
        processed_inputs = []
        # for i, (input_tensor, block) in enumerate(zip(inputs, self.blocks)):
        #     # Apply MambaVisionBlock processing with residual connection
        #     processed_input = block(input_tensor) + input_tensor
        #     processed_inputs.append(processed_input)
        for i, (input_tensor, block) in enumerate(zip(inputs, self.backbone_blocks)):
        # Apply CompleteSwinTransformer processing with residual connection
            processed_input = block(input_tensor) + input_tensor
            processed_inputs.append(processed_input)



        # Convert back to tuple for consistency
        processed_inputs = tuple(processed_inputs)

        # build laterals using processed inputs
        laterals = [
            lateral_conv(processed_inputs[i + self.start_level])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        used_backbone_levels = len(laterals)
        for i in range(used_backbone_levels - 1, 0, -1):
            # In some cases, fixing `scale factor` (e.g. 2) is preferred, but
            #  it cannot co-exist with `size` in `F.interpolate`.
            if 'scale_factor' in self.upsample_cfg:
                # fix runtime error of "+=" inplace operation in PyTorch 1.10
                laterals[i - 1] = laterals[i - 1] + F.interpolate(
                    laterals[i], **self.upsample_cfg)
            else:
                prev_shape = laterals[i - 1].shape[2:]
                laterals[i - 1] = laterals[i - 1] + F.interpolate(
                    laterals[i], size=prev_shape, **self.upsample_cfg)

        # build outputs
        # part 1: from original levels
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(used_backbone_levels)
        ]
        # part 2: add extra levels
        if self.num_outs > len(outs):
            # use max pool to get more levels on top of outputs
            # (e.g., Faster R-CNN, Mask R-CNN)
            if not self.add_extra_convs:
                for i in range(self.num_outs - used_backbone_levels):
                    outs.append(F.max_pool2d(outs[-1], 1, stride=2))
            # add conv layers on top of original feature maps (RetinaNet)
            else:
                if self.add_extra_convs == 'on_input':
                    extra_source = processed_inputs[self.backbone_end_level - 1]
                elif self.add_extra_convs == 'on_lateral':
                    extra_source = laterals[-1]
                elif self.add_extra_convs == 'on_output':
                    extra_source = outs[-1]
                else:
                    raise NotImplementedError
                outs.append(self.fpn_convs[used_backbone_levels](extra_source))
                for i in range(used_backbone_levels + 1, self.num_outs):
                    if self.relu_before_extra_convs:
                        outs.append(self.fpn_convs[i](F.relu(outs[-1])))
                    else:
                        outs.append(self.fpn_convs[i](outs[-1]))
        return tuple(outs)
