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 time
from typing import Tuple
import copy
import numpy as np
import torch.utils.checkpoint as checkpoint
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.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 timm.models.vision_transformer import Mlp

# 新增 ResNet 相关导入
from torchvision.models import resnet34
import torchvision.models as models

# 更新配置，ResNet34的通道数
BACKBONE_CONFIGS = {
    "tiny": [64, 128, 256, 512],
    "small": [96, 192, 384, 768],
    "base": [128, 256, 512, 1024],
    "resnet34": [64, 128, 256, 512]  # ResNet34的各层通道数
}


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


# ============ ResNet34 Backbone ============

class ResNet34Backbone(nn.Module):
    """ResNet34 作为特征提取器的backbone"""

    def __init__(self, pretrained=True):
        super(ResNet34Backbone, self).__init__()

        # 加载预训练的ResNet34
        if pretrained:
            resnet = models.resnet34(pretrained=True)
        else:
            resnet = models.resnet34(pretrained=False)

        # 提取ResNet的各个阶段
        # Stage 0: stem (conv1 + bn1 + relu + maxpool)
        self.stem = nn.Sequential(
            resnet.conv1,  # 7x7 conv, stride=2, output: /2
            resnet.bn1,
            resnet.relu,
            resnet.maxpool  # 3x3 maxpool, stride=2, output: /4
        )


        # Stage 1: layer1 (3 basic blocks) - output: /4
        self.layer1 = resnet.layer1  # 64 channels

        # Stage 2: layer2 (4 basic blocks) - output: /8
        self.layer2 = resnet.layer2  # 128 channels

        # Stage 3: layer3 (6 basic blocks) - output: /16
        self.layer3 = resnet.layer3  # 256 channels

        # Stage 4: layer4 (3 basic blocks) - output: /32
        self.layer4 = resnet.layer4  # 512 channels

        # 输出通道数
        self.out_channels = [64, 128, 256, 512]

    def forward(self, x):
        """
        输入: x (B, 3, H, W)
        输出: [x4, x8, x16, x32] 对应于不同下采样倍数的特征图
        """
        # Stem: /1 -> /4
        x = self.stem(x)  # (B, 64, H/4, W/4)

        # Layer1: /4 -> /4
        x4 = self.layer1(x)  # (B, 64, H/4, W/4)   对应原代码的x128

        # Layer2: /4 -> /8
        x8 = self.layer2(x4)  # (B, 128, H/8, W/8)   对应原代码的x64

        # Layer3: /8 -> /16
        x16 = self.layer3(x8)  # (B, 256, H/16, W/16) 对应原代码的x32

        # Layer4: /16 -> /32
        x32 = self.layer4(x16)  # (B, 512, H/32, W/32) 对应原代码的x16

        return x4, x8, x16, x32


# ============ 其余组件保持不变 ============

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

        # 使用ResNet34的通道配置
        channel_config =[64, 128, 256, 512]

        # 替换backbone为ResNet34
        self.backbone = ResNet34Backbone(pretrained=True)

        self.backbone_blocks = nn.ModuleList([
            GMambaBlock(
                dim=channel_config[0],  # 64
                num_heads=2,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            GMambaBlock(
                dim=channel_config[1],  # 128
                num_heads=4,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            GMambaBlock(
                dim=channel_config[2],  # 256
                num_heads=8,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            GMambaBlock(
                dim=channel_config[3],  # 512
                num_heads=16,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
        ])

        # 使用 MambaVisionBlock，需要根据ResNet34的通道数调整
        self.blocks = nn.ModuleList([
            GMambaBlock(
                dim=channel_config[2],  # 256
                num_heads=8,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            GMambaBlock(
                dim=channel_config[1],  # 128
                num_heads=8,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
            GMambaBlock(
                dim=channel_config[0],  # 64
                num_heads=8,
                counter=0,
                transformer_blocks=[],
                mlp_ratio=4,
                qkv_bias=True,
                drop_path=0.1,
                layer_scale=1e-6,
                window_size=7
            ),
        ])

        # 上采样层，根据ResNet34的通道数调整
        self.Up5 = up_conv(ch_in=channel_config[3], ch_out=channel_config[3] // 2)  # 512 -> 256
        self.Up_conv5 = conv_block(ch_in=channel_config[3], ch_out=channel_config[3] // 2)  # 512 -> 256

        self.Up4 = up_conv(ch_in=channel_config[2], ch_out=channel_config[2] // 2)  # 256 -> 128
        self.Up_conv4 = conv_block(ch_in=channel_config[2], ch_out=channel_config[2] // 2)  # 256 -> 128

        self.Up3 = up_conv(ch_in=channel_config[1], ch_out=channel_config[1] // 2)  # 128 -> 64
        self.Up_conv3 = conv_block(ch_in=channel_config[1], ch_out=channel_config[1] // 2)  # 128 -> 64

        self.Up4x = LRDU(channel_config[0], 4)  # 64通道，4倍上采样
        self.convout = nn.Conv2d(channel_config[0], n_class, kernel_size=1, stride=1, padding=0)

        # 对新增的backbone_blocks进行随机初始化
        # self._init_backbone_blocks()

    def _init_backbone_blocks(self):
        """初始化新增的backbone_blocks"""
        for block in self.backbone_blocks:
            block.apply(self._init_weights)

    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):
        # ResNet34 backbone输出四个尺度的特征
        # x4: (B, 64, H/4, W/4)    对应原来的x128
        # x8: (B, 128, H/8, W/8)    对应原来的x64
        # x16: (B, 256, H/16, W/16) 对应原来的x32
        # x32: (B, 512, H/32, W/32) 对应原来的x16
        x4, x8, x16, x32 = self.backbone(x)

        x4 = self.backbone_blocks[0](x4) + x4    # 64 channels
        x8 = self.backbone_blocks[1](x8) + x8    # 128 channels
        x16 = self.backbone_blocks[2](x16) + x16  # 256 channels
        x32 = self.backbone_blocks[3](x32) + x32  # 512 channels



        d32 = self.Up5(x32)  # x32 (512, H/32, W/32) -> d32 (256, H/16, W/16)
        d32 = torch.cat([x16, d32], dim=1)  # 拼接 x16 (256) + d32 (256) = (512, H/16, W/16)
        d32 = self.Up_conv5(d32)  # 512 -> 256


        d32 = self.blocks[0](d32) + d32  # 残差连接


        d64 = self.Up4(d32)  # d32 (256, H/16, W/16) -> d64 (128, H/8, W/8)
        d64 = torch.cat([x8, d64], dim=1)  # 拼接 x8 (128) + d64 (128) = (256, H/8, W/8)
        d64 = self.Up_conv4(d64)  # 256 -> 128


        d64 = self.blocks[1](d64) + d64  # 残差连接


        d128 = self.Up3(d64)  # d64 (128, H/8, W/8) -> d128 (64, H/4, W/4)
        d128 = torch.cat([x4, d128], dim=1)  # 拼接 x4 (64) + d128 (64) = (128, H/4, W/4)
        d128 = self.Up_conv3(d128)  # 128 -> 64


        d128 = self.blocks[2](d128) + d128  # 残差连接

        # 最终4倍上采样到原始分辨率
        d512 = self.Up4x(d128)  # (64, H/4, W/4) -> (64, H, W)
        out = self.convout(d512)  # 64 -> n_class

        return out




