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 mmcv.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 .csm_triton import CrossScanTriton, CrossMergeTriton

# import selective_scan_cuda_oflex

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 .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


class ConvNeXtBlock(nn.Module):
    def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
        super().__init__()
        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
        self.norm = LayerNorm(dim, eps=1e-6)
        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * dim, dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()

    def forward(self, x):
        input = x
        x = self.dwconv(x)
        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma * x
        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)

        x = input + self.drop_path(x)
        return x


class ConvNeXt(nn.Module):
    r""" ConvNeXt
        A PyTorch impl of : `A ConvNet for the 2020s`  -
          https://arxiv.org/pdf/2201.03545.pdf

    Args:
        in_chans (int): Number of input image channels. Default: 3
        num_classes (int): Number of classes for classification head. Default: 1000
        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
        drop_path_rate (float): Stochastic depth rate. Default: 0.
        layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
    """

    def __init__(self, in_chans=3, depths=[3, 3, 9, 3],
                 dims=[96, 192, 384, 768], drop_path_rate=0.1,
                 layer_scale_init_value=1e-6):
        super().__init__()

        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
        stem = nn.Sequential(
            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
        )
        self.downsample_layers.append(stem)
        for i in range(3):
            downsample_layer = nn.Sequential(
                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
                nn.Conv2d(dims[i], dims[i + 1], kernel_size=2, stride=2),
            )
            self.downsample_layers.append(downsample_layer)

        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple residual blocks
        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
        cur = 0
        for i in range(4):
            stage = nn.Sequential(
                *[ConvNeXtBlock(dim=dims[i], drop_path=dp_rates[cur + j],
                                layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
            )
            self.stages.append(stage)
            cur += depths[i]

        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer

        # Use MambaVisionBlock instead of original Block
        self.blocks = nn.ModuleList([
            GMambaBlock(
                dim=dims[0],
                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=dims[1],
                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=dims[2],
                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=dims[3],
                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
            ),
        ])

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, (nn.Conv2d, nn.Linear)):
            trunc_normal_(m.weight, std=.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward_features(self, x):
        stages_out = []
        for i in range(4):
            x = self.downsample_layers[i](x)
            x = self.stages[i](x)

            # Use MambaVisionBlock processing
            # 消融实验，基本UNet
            x = self.blocks[i](x)

            stages_out.append(x)
        return stages_out

    def forward(self, x):
        x = self.forward_features(x)
        return x


class LayerNorm(nn.Module):


    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(normalized_shape))
        self.bias = nn.Parameter(torch.zeros(normalized_shape))
        self.eps = eps
        self.data_format = data_format
        if self.data_format not in ["channels_last", "channels_first"]:
            raise NotImplementedError
        self.normalized_shape = (normalized_shape,)

    def forward(self, x):
        if self.data_format == "channels_last":
            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
        elif self.data_format == "channels_first":
            u = x.mean(1, keepdim=True)
            s = (x - u).pow(2).mean(1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.eps)
            x = self.weight[:, None, None] * x + self.bias[:, None, None]
            return x


model_urls = {
    "convnext_tiny_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
    "convnext_small_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth",
    "convnext_base_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth",
    "convnext_large_1k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth",
    "convnext_tiny_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_22k_224.pth",
    "convnext_small_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_small_22k_224.pth",
    "convnext_base_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth",
    "convnext_large_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth",
    "convnext_xlarge_22k": "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth",
}


@register_model
def convnext_tiny(pretrained=False, in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 9, 3], dims=[64, 128, 256, 512], **kwargs)
    if pretrained:
        url = model_urls['convnext_tiny_22k'] if in_22k else model_urls['convnext_tiny_1k']
        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu", check_hash=True)
        # 自动判断是否有 "model" 层
        state_dict = checkpoint["model"] if "model" in checkpoint else checkpoint

        # 打印一下确认 shape 是不是对的
        print("Loaded state_dict keys:", list(state_dict.keys())[:5])
        print("Shape of first weight:", state_dict[list(state_dict.keys())[0]].shape)
        model.load_state_dict(checkpoint["model"], strict=False)
    return model


@register_model
def convnext_small(pretrained=False, in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[96, 192, 384, 768], **kwargs)
    if pretrained:
        url = model_urls['convnext_small_22k'] if in_22k else model_urls['convnext_small_1k']
        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
        model.load_state_dict(checkpoint["model"], strict=False)
    return model


@register_model
def convnext_base(pretrained=False, in_22k=False, **kwargs):
    model = ConvNeXt(depths=[3, 3, 27, 3], dims=[128, 256, 512, 1024], **kwargs)
    if pretrained:
        url = model_urls['convnext_base_22k'] if in_22k else model_urls['convnext_base_1k']
        checkpoint = torch.hub.load_state_dict_from_url(url=url, map_location="cpu")
        model.load_state_dict(checkpoint["model"], strict=False)
    return model


class LRDU(nn.Module):
    """
    large receptive detailed upsample
    """

    def __init__(self, in_c, factor):
        super(LRDU, self).__init__()

        self.up_factor = factor
        self.factor1 = factor * factor // 2
        self.factor2 = factor * factor
        self.up = nn.Sequential(
            nn.Conv2d(in_c, self.factor1 * in_c, (1, 7), padding=(0, 3), groups=in_c),
            nn.Conv2d(self.factor1 * in_c, self.factor2 * in_c, (7, 1), padding=(3, 0), groups=in_c),
            nn.PixelShuffle(factor),
            nn.Conv2d(in_c, in_c, 3, groups=in_c // 4, padding=1),
            nn.BatchNorm2d(in_c),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x


class up_conv(nn.Module):
    def __init__(self, ch_in, ch_out):
        super(up_conv, self).__init__()

        self.up = nn.Sequential(
            LRDU(ch_in, 2),
            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.up(x)
        return x


class conv_block(nn.Module):
    def __init__(self, ch_in, ch_out):
        super(conv_block, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(ch_in, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True),
            nn.Conv2d(ch_out, ch_out, kernel_size=3, stride=1, padding=1, bias=True),
            nn.BatchNorm2d(ch_out),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv(x)
        return x




class Model(nn.Module):
    def __init__(self, n_class=6, pretrained=True):
        super(Model, self).__init__()
        self.n_class = n_class
        self.in_channel = 3
        config = [96, 192, 384, 768]  # channels of convnext-small
        self.backbone = convnext_small(pretrained, True)

        # # Use MambaVisionBlock instead of original Block
        self.blocks = nn.ModuleList([
            GMambaBlock(
                dim=config[2],
                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=config[1],
                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=config[0],
                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
            ),
        ])

        self.Up5 = up_conv(ch_in=config[3], ch_out=config[3] // 2)
        self.Up_conv5 = conv_block(ch_in=config[3], ch_out=config[3] // 2)

        self.Up4 = up_conv(ch_in=config[2], ch_out=config[2] // 2)
        self.Up_conv4 = conv_block(ch_in=config[2], ch_out=config[2] // 2)

        self.Up3 = up_conv(ch_in=config[1], ch_out=config[1] // 2)
        self.Up_conv3 = conv_block(ch_in=config[1], ch_out=config[1] // 2)

        self.Up4x = LRDU(config[0], 4)
        self.convout = nn.Conv2d(config[0], n_class, kernel_size=1, stride=1, padding=0)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            if m.bias is not None:
                nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)
        elif isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out')
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.ones_(m.weight)
            nn.init.zeros_(m.bias)

    def forward(self, x):
        x128, x64, x32, x16 = self.backbone(x)

        d32 = self.Up5(x16)
        d32 = torch.cat([x32, d32], dim=1)
        d32 = self.Up_conv5(d32)

        d32 = self.blocks[0](d32) + d32

        d64 = self.Up4(d32)
        d64 = torch.cat([x64, d64], dim=1)
        d64 = self.Up_conv4(d64)

        d64 = self.blocks[1](d64) + d64

        d128 = self.Up3(d64)
        d128 = torch.cat([x128, d128], dim=1)
        d128 = self.Up_conv3(d128)


        d128 = self.blocks[2](d128) + d128

        d512 = self.Up4x(d128)
        out = self.convout(d512)

        return out


