# === Begin: Final corrected code to be appended for EnBiSeNetV2 ===
import torch
import torch.nn as nn
from mmcv.cnn import ConvModule, build_conv_layer, build_norm_layer
from mmengine.model import BaseModule, Sequential
from mmseg.registry import MODELS
# Import existing BiSeNetV2 components to inherit from them
from.bisenetv2 import BiSeNetV2, GELayer, StemBlock

# --- Helper Modules for EnBiSeNetV2 ---

class DepthwiseSeparableConvModule(BaseModule):
    """Depthwise Separable Convolutional Module."""
    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 stride=1,
                 padding=0,
                 dilation=1,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super().__init__(init_cfg)
        self.depthwise_conv = ConvModule(
            in_channels,
            in_channels,
            kernel_size,
            stride=stride,
            padding=padding,
            dilation=dilation,
            groups=in_channels,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)
        self.pointwise_conv = ConvModule(
            in_channels,
            out_channels,
            kernel_size=1,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

    def forward(self, x):
        x = self.depthwise_conv(x)
        x = self.pointwise_conv(x)
        return x

class SpatialAttentionModule(BaseModule):
    """Spatial Attention Module."""
    def __init__(self, init_cfg=None):
        super().__init__(init_cfg)
        self.conv = ConvModule(
            2, 1, kernel_size=7, padding=3, act_cfg=None, norm_cfg=None)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        attention_map = torch.cat([avg_out, max_out], dim=1)
        attention_map = self.conv(attention_map)
        attention_map = self.sigmoid(attention_map)
        return x * attention_map

class ChannelAttentionModule(BaseModule):
    """Channel Attention Module."""
    def __init__(self, in_channels, reduction=16, init_cfg=None):
        super().__init__(init_cfg)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(in_channels, in_channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(in_channels // reduction, in_channels, bias=False),
            nn.Sigmoid()
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y.expand_as(x)

class SEModule(BaseModule):
    """Squeeze-and-Excitation Module."""
    def __init__(self, channels, reduction=16, init_cfg=None):
        super().__init__(init_cfg)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Sequential(
            nn.Linear(channels, channels // reduction, bias=False),
            nn.ReLU(inplace=True),
            nn.Linear(channels // reduction, channels, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        b, c, _, _ = x.size()
        y = self.avg_pool(x).view(b, c)
        y = self.fc(y).view(b, c, 1, 1)
        return x * y

class CEBlockSimAspp(BaseModule):
    """Simplified ASPP for Context Embedding."""
    def __init__(self, in_channels, out_channels, rates=(1, 6, 12), norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU'), init_cfg=None):
        super().__init__(init_cfg)
        self.in_channels = in_channels
        self.out_channels = out_channels
        
        self.global_avg_pool = nn.Sequential(
            nn.AdaptiveAvgPool2d((1, 1)),
            ConvModule(in_channels, out_channels, 1, norm_cfg=norm_cfg, act_cfg=act_cfg)
        )
        
        self.aspp_branches = nn.ModuleList()
        for rate in rates:
            reduced_channels = in_channels // 2
            if reduced_channels == 0:
                reduced_channels = 1
            self.aspp_branches.append(
                nn.Sequential(
                    ConvModule(in_channels, reduced_channels, 1, norm_cfg=norm_cfg, act_cfg=act_cfg),
                    ConvModule(reduced_channels, out_channels, 3, padding=rate, dilation=rate, norm_cfg=norm_cfg, act_cfg=act_cfg)
                )
            )
            
        self.project = ConvModule(out_channels * (len(rates) + 1), out_channels, 1, norm_cfg=norm_cfg, act_cfg=act_cfg)

    def forward(self, x):
        res = []# <-- FIX
        res.append(nn.functional.interpolate(self.global_avg_pool(x), size=x.shape[2:], mode='bilinear', align_corners=False))
        for branch in self.aspp_branches:
            res.append(branch(x))
        
        res = torch.cat(res, dim=1)
        return self.project(res)

# In en_bisenetv2.py, replace the whole GELayerWithSE class with this:

class GELayerWithSE(BaseModule):
    """Gather-and-Expansion Layer with optional SE Module."""
    def __init__(self, in_channels, out_channels, mid_channels, stride, add_se_module=False,
                 norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU'), init_cfg=None):
        super().__init__(init_cfg)

        # 1. Pointwise Conv for channel expansion
        self.pw_conv1 = ConvModule(
            in_channels, mid_channels, kernel_size=1, norm_cfg=norm_cfg, act_cfg=act_cfg)

        # 2. Depthwise Conv
        self.dw_conv = ConvModule(
            mid_channels,
            mid_channels,
            kernel_size=3,
            stride=stride,
            padding=1,
            groups=mid_channels,
            norm_cfg=norm_cfg,
            act_cfg=act_cfg)

        # 3. Squeeze-and-Excitation module
        self.add_se_module = add_se_module
        if self.add_se_module:
            self.se = SEModule(mid_channels)

        # 4. Pointwise Conv for channel projection
        self.pw_conv2 = ConvModule(
            mid_channels, out_channels, kernel_size=1, norm_cfg=norm_cfg, act_cfg=None)

        # 5. Shortcut connection
        if stride == 1 and in_channels == out_channels:
            # Use forward_v1 for identity shortcut
            self.shortcut = nn.Identity()
            self.forward = self.forward_v1 if not add_se_module else self.forward_v1_se
        else:
            # Use forward_v2 for projection shortcut
            self.shortcut = ConvModule(
                in_channels,
                out_channels,
                kernel_size=1,
                stride=stride,
                norm_cfg=norm_cfg,
                act_cfg=None)
            self.forward = self.forward_v2 if not add_se_module else self.forward_v2_se

    def forward_v1(self, x):
        out = self.pw_conv1(x)
        out = self.dw_conv(out)
        out = self.pw_conv2(out)
        out = out + self.shortcut(x)
        return out

    def forward_v2(self, x):
        out = self.pw_conv1(x)
        out = self.dw_conv(out)
        out = self.pw_conv2(out)
        out = out + self.shortcut(x)
        return out

    def forward_v1_se(self, x):
        out = self.pw_conv1(x)
        out = self.se(out)
        out = self.dw_conv(out)
        out = self.pw_conv2(out)
        out = out + self.shortcut(x)
        return out

    def forward_v2_se(self, x):
        out = self.pw_conv1(x)
        out = self.se(out)
        out = self.dw_conv(out)
        out = self.pw_conv2(out)
        out = out + self.shortcut(x)
        return out

class FuseModule(BaseModule):
    """Fuse features from different levels."""
    def __init__(self, in_channels_list, out_channels, norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU'), init_cfg=None):
        super().__init__(init_cfg)
        self.conv = ConvModule(sum(in_channels_list), out_channels, 3, padding=1, norm_cfg=norm_cfg, act_cfg=act_cfg)

    def forward(self, features):
        target_size = features[-1].shape[2:]
        
        upsampled_features = []# <-- FIX
        for feat in features:
            upsampled_features.append(
                nn.functional.interpolate(feat, size=target_size, mode='bilinear', align_corners=False)
            )
        
        fused_feature = torch.cat(upsampled_features, dim=1)
        return self.conv(fused_feature)

class SimBGABlock(BaseModule):
    """Simplified Bilateral Guided Aggregation Block."""
    def __init__(self, detail_channels, semantic_channels, out_channels, norm_cfg=dict(type='BN'), act_cfg=dict(type='ReLU'), init_cfg=None):
        super().__init__(init_cfg)
        self.conv_detail = ConvModule(detail_channels, out_channels, 3, padding=1, norm_cfg=norm_cfg, act_cfg=act_cfg)
        self.conv_semantic = ConvModule(semantic_channels, out_channels, 3, padding=1, norm_cfg=norm_cfg, act_cfg=act_cfg)
        self.conv_out = ConvModule(out_channels, out_channels, 3, padding=1, norm_cfg=norm_cfg, act_cfg=act_cfg)

    def forward(self, detail_feat, semantic_feat):
        semantic_feat_upsampled = nn.functional.interpolate(
            semantic_feat, size=detail_feat.shape[2:], mode='bilinear', align_corners=False
        )
        
        detail_feat_processed = self.conv_detail(detail_feat)
        semantic_feat_processed = self.conv_semantic(semantic_feat_upsampled)
        
        fused_feat = detail_feat_processed + semantic_feat_processed
        return self.conv_out(fused_feat)


# --- Main EnBiSeNetV2 Backbone ---

@MODELS.register_module()
class EnBiSeNetV2(BiSeNetV2):
    """
    This class is the implementation of En_bisenetv2.
    """
    def __init__(self,
                 in_channels,
                 detail_channels_stages=(64, 64, 128),
                 semantic_channels_stages=(16, 32, 64, 128),
                 semantic_expansion_ratio=6,
                 bga_channels=128,
                 out_indices=(0, 1, 2, 3, 4),
                 align_corners=False,
                 norm_cfg=dict(type='BN'),
                 act_cfg=dict(type='ReLU'),
                 init_cfg=None):
        super(BiSeNetV2, self).__init__(init_cfg)
        self.in_channels = in_channels
        self.out_indices = out_indices
        self.align_corners = align_corners
        self.norm_cfg = norm_cfg
        self.act_cfg = act_cfg

        self.detail_branch = self._make_detail_branch(detail_channels_stages)
        self.spatial_attention = SpatialAttentionModule()

        self.semantic_branch = self._make_semantic_branch(
            semantic_channels_stages, semantic_expansion_ratio)
        self.channel_attention = ChannelAttentionModule(semantic_channels_stages[-1])
        
        self.fuse_module = FuseModule(
            in_channels_list=[semantic_channels_stages[3], semantic_channels_stages[3]],
            out_channels=semantic_channels_stages[-1]
        )

        self.bga_layer = SimBGABlock(
            detail_channels=detail_channels_stages[-1],
            semantic_channels=semantic_channels_stages[-1],
            out_channels=bga_channels
        )

    def _make_detail_branch(self, channels_stages):
        layers = []
        # stage 1
        layers.append(
            DepthwiseSeparableConvModule(
                self.in_channels, channels_stages[0], 3, stride=2, padding=1,
                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        layers.append(
            DepthwiseSeparableConvModule(
                channels_stages[0], channels_stages[0], 3, stride=1, padding=1,
                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        # stage 2
        layers.append(
            DepthwiseSeparableConvModule(
                channels_stages[0], channels_stages[1], 3, stride=2, padding=1,
                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        layers.append(
            DepthwiseSeparableConvModule(
                channels_stages[1], channels_stages[1], 3, stride=1, padding=1,
                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        # stage 3
        layers.append(
            DepthwiseSeparableConvModule(
                channels_stages[1], channels_stages[2], 3, stride=2, padding=1,
                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        layers.append(
            DepthwiseSeparableConvModule(
                channels_stages[2], channels_stages[2], 3, stride=1, padding=1,
                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        return Sequential(*layers)

    # def _make_detail_branch(self, channels_stages):
    #     layers = [] # <-- FIX
    #     layers.append(
    #         DepthwiseSeparableConvModule(
    #             self.in_channels, channels_stages, 3, stride=2, padding=1,
    #             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
    #     layers.append(
    #         DepthwiseSeparableConvModule(
    #             channels_stages, channels_stages, 3, stride=1, padding=1,
    #             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
    #     layers.append(
    #         DepthwiseSeparableConvModule(
    #             channels_stages, channels_stages[3], 3, stride=2, padding=1,
    #             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
    #     layers.append(
    #         DepthwiseSeparableConvModule(
    #             channels_stages[3], channels_stages[3], 3, stride=1, padding=1,
    #             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
    #     layers.append(
    #         DepthwiseSeparableConvModule(
    #             channels_stages[3], channels_stages[1], 3, stride=2, padding=1,
    #             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
    #     layers.append(
    #         DepthwiseSeparableConvModule(
    #             channels_stages[1], channels_stages[1], 3, stride=1, padding=1,
    #             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
    #     return Sequential(*layers)

    def _make_semantic_branch(self, channels_stages, expansion_ratio):
        # stage 1
        stage1 = StemBlock(self.in_channels, channels_stages[0],
                        norm_cfg=self.norm_cfg, act_cfg=self.act_cfg)
        # stage 2
        stage2 = self._make_ge_layer(
            channels_stages[0], channels_stages[1], 2, expansion_ratio, add_se=False)
        # stage 3
        stage3 = self._make_ge_layer(
            channels_stages[1], channels_stages[2], 2, expansion_ratio, add_se=True)
        # stage 4
        stage4 = self._make_ge_layer(
            channels_stages[2], channels_stages[3], 2, expansion_ratio, add_se=True)
        # stage 5: Simplified ASPP instead of CEBlock
        stage5 = CEBlockSimAspp(channels_stages[3], channels_stages[3],
                                norm_cfg=self.norm_cfg, act_cfg=self.act_cfg)
        return nn.ModuleDict(
            {'1': stage1, '2': stage2, '3': stage3, '4': stage4, '5': stage5}
        )

    # def _make_semantic_branch(self, channels_stages, expansion_ratio):
    #     stage1 = StemBlock(self.in_channels, channels_stages,
    #                   norm_cfg=self.norm_cfg, act_cfg=self.act_cfg)
    #     stage2 = self._make_ge_layer(
    #         channels_stages, channels_stages[3], 2, expansion_ratio, add_se=False)
    #     stage3 = self._make_ge_layer(
    #         channels_stages[3], channels_stages[1], 2, expansion_ratio, add_se=True)
    #     stage4 = self._make_ge_layer(
    #         channels_stages[1], channels_stages[2], 2, expansion_ratio, add_se=True)
    #     stage5 = CEBlockSimAspp(channels_stages[2], channels_stages[2],
    #                             norm_cfg=self.norm_cfg, act_cfg=self.act_cfg)
    #     return nn.ModuleDict(
    #         {'1': stage1, '2': stage2, '3': stage3, '4': stage4, '5': stage5}
    #     )

    def _make_ge_layer(self, in_channels, out_channels, stride, expansion_ratio, add_se=False):
        mid_channels = in_channels * expansion_ratio
        layers = []# <-- FIX
        layers.append(
            GELayerWithSE(in_channels, out_channels, mid_channels, stride, add_se_module=add_se,
                          norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        for _ in range(1, 3):
            layers.append(
                GELayerWithSE(out_channels, out_channels, mid_channels, 1, add_se_module=add_se,
                              norm_cfg=self.norm_cfg, act_cfg=self.act_cfg))
        return Sequential(*layers)

    def forward(self, x):
        detail_out = self.detail_branch(x)
        detail_out = self.spatial_attention(detail_out)

        semantic_out_s1 = self.semantic_branch['1'](x)
        semantic_out_s2 = self.semantic_branch['2'](semantic_out_s1)
        semantic_out_s3 = self.semantic_branch['3'](semantic_out_s2)
        semantic_out_s4 = self.semantic_branch['4'](semantic_out_s3)
        semantic_out_s5 = self.semantic_branch['5'](semantic_out_s4)
        
        semantic_fused = self.fuse_module([semantic_out_s4, semantic_out_s5])
        semantic_fused = self.channel_attention(semantic_fused)

        out = self.bga_layer(detail_out, semantic_fused)

        outs = [detail_out, semantic_out_s3, semantic_out_s4, semantic_fused, out]
        outs = [outs[i] for i in self.out_indices]
        return tuple(outs)

# === End: Final corrected code to be appended for EnBiSeNetV2 ===