# Copyright (c) OpenMMLab. All rights reserved.
import torch
import torch.nn as nn
from mmcv.cnn import (ConvModule, DepthwiseSeparableConvModule,
                      build_activation_layer, build_norm_layer)
from mmengine.model import BaseModule
from mmseg.registry import MODELS
from ..utils import resize
import torch.nn.functional as F
from einops import rearrange
from mmcv.cnn import ConvModule, DepthwiseSeparableConvModule
# ===================================================================
# Part 1: CAFM 模块 (稍作修改以简化输出)
# ===================================================================
class CAFM(nn.Module):
    def __init__(self, channels):
        super(CAFM, self).__init__()
        self.conv1_spatial = nn.Conv2d(2, 1, 3, stride=1, padding=1, groups=1)
        self.conv2_spatial = nn.Conv2d(1, 1, 3, stride=1, padding=1, groups=1)
        self.avg1 = nn.Conv2d(channels, 64, 1, stride=1, padding=0)
        self.avg2 = nn.Conv2d(channels, 64, 1, stride=1, padding=0)
        self.max1 = nn.Conv2d(channels, 64, 1, stride=1, padding=0)
        self.max2 = nn.Conv2d(channels, 64, 1, stride=1, padding=0)
        self.avg11 = nn.Conv2d(64, channels, 1, stride=1, padding=0)
        self.avg22 = nn.Conv2d(64, channels, 1, stride=1, padding=0)
        self.max11 = nn.Conv2d(64, channels, 1, stride=1, padding=0)
        self.max22 = nn.Conv2d(64, channels, 1, stride=1, padding=0)

    def forward(self, f1, f2):
        b, c, h, w = f1.size()
        f1_reshaped = f1.reshape([b, c, -1])
        f2_reshaped = f2.reshape([b, c, -1])

        avg_1 = torch.mean(f1_reshaped, dim=-1, keepdim=True).unsqueeze(-1)
        max_1, _ = torch.max(f1_reshaped, dim=-1, keepdim=True)
        max_1 = max_1.unsqueeze(-1)
        avg_1 = F.relu(self.avg1(avg_1))
        max_1 = F.relu(self.max1(max_1))
        avg_1 = self.avg11(avg_1).squeeze(-1)
        max_1 = self.max11(max_1).squeeze(-1)
        a1 = avg_1 + max_1

        avg_2 = torch.mean(f2_reshaped, dim=-1, keepdim=True).unsqueeze(-1)
        max_2, _ = torch.max(f2_reshaped, dim=-1, keepdim=True)
        max_2 = max_2.unsqueeze(-1)
        avg_2 = F.relu(self.avg2(avg_2))
        max_2 = F.relu(self.max2(max_2))
        avg_2 = self.avg22(avg_2).squeeze(-1)
        max_2 = self.max22(max_2).squeeze(-1)
        a2 = avg_2 + max_2

        cross = torch.matmul(a1, a2.transpose(1, 2))
        a1 = torch.matmul(F.softmax(cross, dim=-1), f1_reshaped)
        a2 = torch.matmul(F.softmax(cross.transpose(1, 2), dim=-1), f2_reshaped)

        # Spatial attention for f1
        a1_spatial = a1.reshape([b, c, h, w])
        avg_out = torch.mean(a1_spatial, dim=1, keepdim=True)
        max_out, _ = torch.max(a1_spatial, dim=1, keepdim=True)
        a1_spatial = torch.cat([avg_out, max_out], dim=1)
        a1_spatial = F.relu(self.conv1_spatial(a1_spatial))
        a1_spatial = self.conv2_spatial(a1_spatial)
        a1_spatial = a1_spatial.reshape([b, 1, -1])
        a1_spatial = F.softmax(a1_spatial, dim=-1)

        # Spatial attention for f2
        a2_spatial = a2.reshape([b, c, h, w])
        avg_out = torch.mean(a2_spatial, dim=1, keepdim=True)
        max_out, _ = torch.max(a2_spatial, dim=1, keepdim=True)
        a2_spatial = torch.cat([avg_out, max_out], dim=1)
        a2_spatial = F.relu(self.conv1_spatial(a2_spatial))
        a2_spatial = self.conv2_spatial(a2_spatial)
        a2_spatial = a2_spatial.reshape([b, 1, -1])
        a2_spatial = F.softmax(a2_spatial, dim=-1)

        f1_out = f1_reshaped * a1_spatial + f1_reshaped
        f2_out = f2_reshaped * a2_spatial + f2_reshaped
        
        # 修改点: 直接返回 [B, C, H, W] 格式
        return f1_out.reshape(b, c, h, w), f2_out.reshape(b, c, h, w)


# ===================================================================
# Part 2: 全新的 CrossFusedBranches 模块
# ===================================================================
class CrossFusedBranches(BaseModule):
    def __init__(self,
                 detail_channels,
                 in_channels,
                 semantic_channels,
                 semantic_expansion_ratio,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
        
        # 1. --- Porting Detail Branch Layers (with Depthwise Separable Conv) ---
        # 将这里的 ConvModule 全部替换为 DepthwiseSeparableConvModule
        self.d_stage1 = nn.Sequential(
            DepthwiseSeparableConvModule(in_channels, detail_channels[0], 3, stride=2, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg),
            DepthwiseSeparableConvModule(detail_channels[0], detail_channels[0], 3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg)
        )
        self.d_stage2 = nn.Sequential(
            DepthwiseSeparableConvModule(detail_channels[0], detail_channels[1], 3, stride=2, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg),
            DepthwiseSeparableConvModule(detail_channels[1], detail_channels[1], 3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg),
            DepthwiseSeparableConvModule(detail_channels[1], detail_channels[1], 3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg)
        )
        self.d_stage3 = nn.Sequential(
            DepthwiseSeparableConvModule(detail_channels[1], detail_channels[2], 3, stride=2, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg),
            DepthwiseSeparableConvModule(detail_channels[2], detail_channels[2], 3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg),
            DepthwiseSeparableConvModule(detail_channels[2], detail_channels[2], 3, stride=1, padding=1, conv_cfg=conv_cfg, norm_cfg=norm_cfg, act_cfg=act_cfg)
        )

        # 2. --- Porting Semantic Branch Layers ---
        # (这部分保持不变)
        self.s_stage1 = StemBlock(in_channels, semantic_channels[0])
        self.s_stage2 = nn.Sequential(
            GELayer(semantic_channels[0], semantic_channels[1], semantic_expansion_ratio, 2),
            GELayer(semantic_channels[1], semantic_channels[1], semantic_expansion_ratio, 1)
        )
        self.s_stage3 = nn.Sequential(
            GELayer(semantic_channels[1], semantic_channels[2], semantic_expansion_ratio, 2),
            GELayer(semantic_channels[2], semantic_channels[2], semantic_expansion_ratio, 1)
        )
        self.s_stage4 = nn.Sequential(
            GELayer(semantic_channels[2], semantic_channels[3], semantic_expansion_ratio, 2),
            GELayer(semantic_channels[3], semantic_channels[3], semantic_expansion_ratio, 1),
            GELayer(semantic_channels[3], semantic_channels[3], semantic_expansion_ratio, 1),
            GELayer(semantic_channels[3], semantic_channels[3], semantic_expansion_ratio, 1)
        )
        self.s_ce_block = CEBlock(semantic_channels[3], semantic_channels[3])

        # 3. --- Cross Fusion Modules ---
        # (这部分保持不变)
        fusion_channels = detail_channels[1]
        self.adapter_d = ConvModule(detail_channels[1], fusion_channels, 1, norm_cfg=norm_cfg, act_cfg=act_cfg)
        self.adapter_s = ConvModule(semantic_channels[2], fusion_channels, 1, norm_cfg=norm_cfg, act_cfg=act_cfg)
        self.cafm = CAFM(channels=fusion_channels)
        self.injection_conv = ConvModule(fusion_channels, detail_channels[1], 3, padding=1, norm_cfg=norm_cfg, act_cfg=act_cfg)

    def forward(self, x):
        # Forward 逻辑完全保持不变
        semantic_outs = []

        d1 = self.d_stage1(x)
        d2 = self.d_stage2(d1)
        s1 = self.s_stage1(x)
        semantic_outs.append(s1)
        s2 = self.s_stage2(s1)
        semantic_outs.append(s2)
        s3 = self.s_stage3(s2)
        semantic_outs.append(s3)

        s_for_fusion = resize(s3, size=d2.shape[2:], mode='bilinear', align_corners=False)
        d_adapted = self.adapter_d(d2)
        s_adapted = self.adapter_s(s_for_fusion)
        fused_d, _ = self.cafm(d_adapted, s_adapted)
        fused_processed = self.injection_conv(fused_d)
        d2_enhanced = d2 + fused_processed

        detail_final = self.d_stage3(d2_enhanced)
        s4 = self.s_stage4(s3)
        semantic_outs.append(s4)
        s_final = self.s_ce_block(s4)
        semantic_outs.append(s_final)

        return detail_final, semantic_outs


class StemBlock(BaseModule):
    """Stem Block at the beginning of Semantic Branch.

    Args:
        in_channels (int): Number of input channels.
            Default: 3.
        out_channels (int): Number of output channels.
            Default: 16.
        conv_cfg (dict | None): Config of conv layers.
            Default: None.
        norm_cfg (dict | None): Config of norm layers.
            Default: dict(type='BN').
        act_cfg (dict): Config of activation layers.
            Default: dict(type='ReLU').
        init_cfg (dict or list[dict], optional): Initialization config dict.
            Default: None.
    Returns:
        x (torch.Tensor): First feature map in Semantic Branch.
    """

    def __init__(self,
                 in_channels=3,
                 out_channels=16,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)

        self.conv_first = ConvModule(
            in_channels=in_channels,
            out_channels=out_channels,
            kernel_size=3,
            stride=2,
            padding=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)
        self.convs = nn.Sequential(
            ConvModule(
                in_channels=out_channels,
                out_channels=out_channels // 2,
                kernel_size=1,
                stride=1,
                padding=0,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg),
            ConvModule(
                in_channels=out_channels // 2,
                out_channels=out_channels,
                kernel_size=3,
                stride=2,
                padding=1,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=act_cfg))
        self.pool = nn.MaxPool2d(
            kernel_size=3, stride=2, padding=1, ceil_mode=False)
        self.fuse_last = ConvModule(
            in_channels=out_channels * 2,
            out_channels=out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

    def forward(self, x):
        x = self.conv_first(x)
        x_left = self.convs(x)
        x_right = self.pool(x)
        x = self.fuse_last(torch.cat([x_left, x_right], dim=1))
        return x


class GELayer(BaseModule):
    """Gather-and-Expansion Layer.

    Args:
        in_channels (int): Number of input channels.
        out_channels (int): Number of output channels.
        exp_ratio (int): Expansion ratio for middle channels.
            Default: 6.
        stride (int): Stride of GELayer. Default: 1
        conv_cfg (dict | None): Config of conv layers.
            Default: None.
        norm_cfg (dict | None): Config of norm layers.
            Default: dict(type='BN').
        act_cfg (dict): Config of activation layers.
            Default: dict(type='ReLU').
        init_cfg (dict or list[dict], optional): Initialization config dict.
            Default: None.
    Returns:
        x (torch.Tensor): Intermediate feature map in
            Semantic Branch.
    """

    def __init__(self,
                 in_channels,
                 out_channels,
                 exp_ratio=6,
                 stride=1,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
        mid_channel = in_channels * exp_ratio
        self.conv1 = ConvModule(
            in_channels=in_channels,
            out_channels=in_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)
        if stride == 1:
            self.dwconv = nn.Sequential(
                # ReLU in ConvModule not shown in paper
                ConvModule(
                    in_channels=in_channels,
                    out_channels=mid_channel,
                    kernel_size=3,
                    stride=stride,
                    padding=1,
                    groups=in_channels,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg))
            self.shortcut = None
        else:
            self.dwconv = nn.Sequential(
                ConvModule(
                    in_channels=in_channels,
                    out_channels=mid_channel,
                    kernel_size=3,
                    stride=stride,
                    padding=1,
                    groups=in_channels,
                    bias=False,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=None),
                # ReLU in ConvModule not shown in paper
                ConvModule(
                    in_channels=mid_channel,
                    out_channels=mid_channel,
                    kernel_size=3,
                    stride=1,
                    padding=1,
                    groups=mid_channel,
                    conv_cfg=conv_cfg,
                    norm_cfg=norm_cfg,
                    act_cfg=act_cfg),
            )
            self.shortcut = nn.Sequential(
                DepthwiseSeparableConvModule(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    kernel_size=3,
                    stride=stride,
                    padding=1,
                    dw_norm_cfg=norm_cfg,
                    dw_act_cfg=None,
                    pw_norm_cfg=norm_cfg,
                    pw_act_cfg=None,
                ))

        self.conv2 = nn.Sequential(
            ConvModule(
                in_channels=mid_channel,
                out_channels=out_channels,
                kernel_size=1,
                stride=1,
                padding=0,
                bias=False,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=None,
            ))

        self.act = build_activation_layer(act_cfg)

    def forward(self, x):
        identity = x
        x = self.conv1(x)
        x = self.dwconv(x)
        x = self.conv2(x)
        if self.shortcut is not None:
            shortcut = self.shortcut(identity)
            x = x + shortcut
        else:
            x = x + identity
        x = self.act(x)
        return x


class CEBlock(BaseModule):
    """Context Embedding Block for large receptive filed in Semantic Branch.

    Args:
        in_channels (int): Number of input channels.
            Default: 3.
        out_channels (int): Number of output channels.
            Default: 16.
        conv_cfg (dict | None): Config of conv layers.
            Default: None.
        norm_cfg (dict | None): Config of norm layers.
            Default: dict(type='BN').
        act_cfg (dict): Config of activation layers.
            Default: dict(type='ReLU').
        init_cfg (dict or list[dict], optional): Initialization config dict.
            Default: None.
    Returns:
        x (torch.Tensor): Last feature map in Semantic Branch.
    """

    def __init__(self,
                 in_channels=3,
                 out_channels=16,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.gap = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            build_norm_layer(norm_cfg, self.in_channels)[1])
        self.conv_gap = ConvModule(
            in_channels=self.in_channels,
            out_channels=self.out_channels,
            kernel_size=1,
            stride=1,
            padding=0,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)
        # Note: in paper here is naive conv2d, no bn-relu
        self.conv_last = ConvModule(
            in_channels=self.out_channels,
            out_channels=self.out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

    def forward(self, x):
        identity = x
        x = self.gap(x)
        x = self.conv_gap(x)
        x = identity + x
        x = self.conv_last(x)
        return x


class BGALayer(BaseModule):
    """Bilateral Guided Aggregation Layer to fuse the complementary information
    from both Detail Branch and Semantic Branch.

    Args:
        out_channels (int): Number of output channels.
            Default: 128.
        align_corners (bool): align_corners argument of F.interpolate.
            Default: False.
        conv_cfg (dict | None): Config of conv layers.
            Default: None.
        norm_cfg (dict | None): Config of norm layers.
            Default: dict(type='BN').
        act_cfg (dict): Config of activation layers.
            Default: dict(type='ReLU').
        init_cfg (dict or list[dict], optional): Initialization config dict.
            Default: None.
    Returns:
        output (torch.Tensor): Output feature map for Segment heads.
    """

    def __init__(self,
                 out_channels=128,
                 align_corners=False,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
        self.out_channels = out_channels
        self.align_corners = align_corners
        self.detail_dwconv = nn.Sequential(
            DepthwiseSeparableConvModule(
                in_channels=self.out_channels,
                out_channels=self.out_channels,
                kernel_size=3,
                stride=1,
                padding=1,
                dw_norm_cfg=norm_cfg,
                dw_act_cfg=None,
                pw_norm_cfg=None,
                pw_act_cfg=None,
            ))
        self.detail_down = nn.Sequential(
            ConvModule(
                in_channels=self.out_channels,
                out_channels=self.out_channels,
                kernel_size=3,
                stride=2,
                padding=1,
                bias=False,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=None),
            nn.AvgPool2d(kernel_size=3, stride=2, padding=1, ceil_mode=False))
        self.semantic_conv = nn.Sequential(
            ConvModule(
                in_channels=self.out_channels,
                out_channels=self.out_channels,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
                conv_cfg=conv_cfg,
                norm_cfg=norm_cfg,
                act_cfg=None))
        self.semantic_dwconv = nn.Sequential(
            DepthwiseSeparableConvModule(
                in_channels=self.out_channels,
                out_channels=self.out_channels,
                kernel_size=3,
                stride=1,
                padding=1,
                dw_norm_cfg=norm_cfg,
                dw_act_cfg=None,
                pw_norm_cfg=None,
                pw_act_cfg=None,
            ))
        self.conv = ConvModule(
            in_channels=self.out_channels,
            out_channels=self.out_channels,
            kernel_size=3,
            stride=1,
            padding=1,
            inplace=True,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg,
        )

    def forward(self, x_d, x_s):
        detail_dwconv = self.detail_dwconv(x_d)
        detail_down = self.detail_down(x_d)
        semantic_conv = self.semantic_conv(x_s)
        semantic_dwconv = self.semantic_dwconv(x_s)
        semantic_conv = resize(
            input=semantic_conv,
            size=detail_dwconv.shape[2:],
            mode='bilinear',
            align_corners=self.align_corners)
        fuse_1 = detail_dwconv * torch.sigmoid(semantic_conv)
        fuse_2 = detail_down * torch.sigmoid(semantic_dwconv)
        fuse_2 = resize(
            input=fuse_2,
            size=fuse_1.shape[2:],
            mode='bilinear',
            align_corners=self.align_corners)
        output = self.conv(fuse_1 + fuse_2)
        return output


@MODELS.register_module()
class My_BiSeNetV2_A1_add_A2(BaseModule):
    def __init__(self,
                 in_channels=3,
                 detail_channels=(64, 64, 128),
                 semantic_channels=(16, 32, 64, 128),
                 semantic_expansion_ratio=6,
                 bga_channels=128,
                 out_indices=(0, 1, 2, 3, 4),
                 align_corners=False,
                 conv_cfg=None,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        if init_cfg is None:
            init_cfg = [
                dict(type='Kaiming', layer='Conv2d'),
                dict(type='Constant', val=1, layer=['_BatchNorm', 'GroupNorm'])
            ]
        super().__init__(init_cfg=init_cfg)
        self.in_channels = in_channels
        self.out_indices = out_indices
        self.align_corners = align_corners

        # --- 核心修改点 ---
        # 用一个 CrossFusedBranches 模块替换掉旧的 self.detail 和 self.semantic
        self.branches = CrossFusedBranches(
            detail_channels=detail_channels,
            in_channels=in_channels,
            semantic_channels=semantic_channels,
            semantic_expansion_ratio=semantic_expansion_ratio,
            conv_cfg=conv_cfg,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg
        )

        self.bga = BGALayer(bga_channels, self.align_corners)

    def forward(self, x):
        # --- 核心修改点 ---
        # 一次调用即可获得两个分支的最终输出
        x_detail, x_semantic_lst = self.branches(x)
        
        # 后续代码保持不变
        x_head = self.bga(x_detail, x_semantic_lst[-1])
        outs = [x_head] + x_semantic_lst[:-1]
        outs = [outs[i] for i in self.out_indices]
        return tuple(outs)
