import torch.nn as nn
import torch.nn.functional as F

from .dpt_backbone import DPTRefineNetStack


class ProgressiveOutputHead(nn.Module):
    """Progressive channel-reduction head (GroupNorm for AMP/small-batch stability)."""
    def __init__(self, in_features=256, num_classes=21, intermediate_features=32, groups=8):
        super().__init__()
        self.num_classes = num_classes
        self.output_conv1 = nn.Sequential(
            nn.Conv2d(in_features, in_features // 2, 3, padding=1, bias=False),
            nn.GroupNorm(num_groups=groups, num_channels=in_features // 2),
            nn.ReLU(inplace=True),
        )
        self.output_conv2 = nn.Sequential(
            nn.Conv2d(in_features // 2, intermediate_features, 3, padding=1, bias=False),
            nn.GroupNorm(num_groups=max(1, min(groups, intermediate_features)), num_channels=intermediate_features),
            nn.ReLU(inplace=True),
            nn.Conv2d(intermediate_features, num_classes, 1),
        )
        self.final_activation = nn.Identity()  # sigmoid/softmax applied in the loss
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        # Background-prior bias on the final logits to avoid early sigmoid overflow.
        final_conv = self.output_conv2[-1]
        if isinstance(final_conv, nn.Conv2d) and final_conv.bias is not None:
            nn.init.constant_(final_conv.bias, -2.2 if self.num_classes == 1 else -0.5)

    def forward(self, x):
        return self.final_activation(self.output_conv2(self.output_conv1(x)))


class OriginalDPTSegmentationDecoder(DPTRefineNetStack):
    def __init__(self, in_channels, num_classes=21, features=256, target_size=(512, 512)):
        super().__init__(features=features, use_bn=False)
        self.target_size = target_size

        out_channels = [256, 512, 1024, 1024]
        self.projects = nn.ModuleList([nn.Conv2d(in_ch, oc, 1) for in_ch, oc in zip(in_channels, out_channels)])
        self.resize_layers = nn.ModuleList([
            nn.ConvTranspose2d(out_channels[0], out_channels[0], kernel_size=4, stride=4),
            nn.ConvTranspose2d(out_channels[1], out_channels[1], kernel_size=2, stride=2),
            nn.Identity(),
            nn.Conv2d(out_channels[3], out_channels[3], kernel_size=3, stride=2, padding=1),
        ])
        self.output_head = ProgressiveOutputHead(
            in_features=features, num_classes=num_classes, intermediate_features=32, groups=8)

    def forward(self, features):
        proj = [p(f) for p, f in zip(self.projects, features)]
        resized = [r(f) for r, f in zip(self.resize_layers, proj)]
        logits = self.output_head(self.fuse(resized))
        if self.target_size is not None:
            logits = F.interpolate(logits, size=self.target_size, mode='bilinear', align_corners=True)
        return logits