import torch
import torch.nn as nn
import torchvision.models as models
from transformers import SegformerForSemanticSegmentation


# ==========================================
# 1. CBAM Components (unchanged)
# ==========================================
class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=4):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        return self.sigmoid(out)


class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        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)

        x_cat = torch.cat([avg_out, max_out], dim=1)
        out = self.conv1(x_cat)
        return self.sigmoid(out)


class CBAM(nn.Module):
    def __init__(self, channels=19, reduction=4, spatial_kernel_size=7):
        super(CBAM, self).__init__()
        self.ca = ChannelAttention(channels, ratio=reduction)
        self.sa = SpatialAttention(kernel_size=spatial_kernel_size)

    def forward(self, x):
        x = x * self.ca(x)
        x = x * self.sa(x)
        return x


# ==========================================
# 2. Main Model → MODIFIED: ONLY CO (1 output)
# ==========================================
class PollutionDifferenceModel(nn.Module):
    def __init__(self, num_classes=19, pollution_dims=1):  # ✅ 5 → 1
        super(PollutionDifferenceModel, self).__init__()

        self.backbone = SegformerForSemanticSegmentation.from_pretrained(
            "nvidia/segformer-b5-finetuned-cityscapes-1024-1024",
            use_safetensors=True
        )

        for param in self.backbone.parameters():
            param.requires_grad = False

        self.cbam = CBAM(channels=num_classes)

        self.convnext = models.convnext_tiny(weights=models.ConvNeXt_Tiny_Weights.DEFAULT)
        for param in self.convnext.parameters():
            param.requires_grad = True

        original_stem = self.convnext.features[0][0]
        self.convnext.features[0][0] = nn.Conv2d(
            in_channels=num_classes,
            out_channels=original_stem.out_channels,
            kernel_size=original_stem.kernel_size,
            stride=original_stem.stride,
            padding=original_stem.padding,
            bias=(original_stem.bias is not None)
        )
        nn.init.kaiming_normal_(self.convnext.features[0][0].weight, mode='fan_out', nonlinearity='relu')

        self.convnext.classifier[2] = nn.Identity()
        convnext_out_dim = 768

        # ✅ Final output = 1 (CO only)
        self.mlp_decoder = nn.Sequential(
            nn.Linear(convnext_out_dim, 256),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(256, 64),
            nn.GELU(),
            nn.Dropout(0.3),
            nn.Linear(64, pollution_dims)
        )

    def get_semantic_map(self, x):
        outputs = self.backbone(pixel_values=x)
        seg_logits = outputs.logits
        probs = torch.nn.functional.softmax(seg_logits, dim=1)
        return probs

    def forward(self, img1, img2):
        map1 = self.get_semantic_map(img1)
        map2 = self.get_semantic_map(img2)
        diff_map = map1 - map2
        attended_diff_map = self.cbam(diff_map)
        z_diff = self.convnext(attended_diff_map)
        pred_pollution_delta = self.mlp_decoder(z_diff)
        return pred_pollution_delta