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


class MultiFusionNeck(nn.Module):
    """
    AgriFM decoder neck - faithful to GitHub implementation.

    U-Net style: starts from bottleneck, upsamples 3 times,
    fusing skip features from features_list at each level.

    For PASTIS single-source (S2 only):
        embed_dim = 1024
        in_feature_key = ('S2',)
        feature_size = (8, 8)      # 128 // 16
        out_size = (128, 128)
        in_fusion_key_list = (
            {'S2': 1024},          # fuse with features_list[2] -> (16,16)
            {'S2': 1024},          # fuse with features_list[1] -> (32,32)
            {'S2': 1024},          # fuse with features_list[0] -> (64,64)
        )
    """

    def __init__(self, embed_dim, in_feature_key=('S2',),
                 feature_size=(8, 8), out_size=(128, 128),
                 in_fusion_key_list=(
                     {'S2': 1024},
                     {'S2': 1024},
                     {'S2': 1024},
                 )):
        super().__init__()
        self.embed_dim = embed_dim
        self.in_feature_key = in_feature_key
        self.feature_size = feature_size
        self.out_size = out_size
        self.in_fusion_key_list = in_fusion_key_list

        # Initial conv: fuse multiple sources into embed_dim
        # If single source: Identity, else conv to merge
        if len(in_feature_key) == 1:
            self.in_conv = nn.Identity()
        else:
            self.in_conv = nn.Sequential(
                nn.Conv2d(len(in_feature_key) * embed_dim, embed_dim, 3, 1, 1),
                nn.BatchNorm2d(embed_dim),
                nn.ReLU(inplace=True),
                nn.Conv2d(embed_dim, embed_dim, 3, 1, 1),
            )

        # Build 3 fusion blocks (one per upsample level)
        self.fusion_list = nn.ModuleList()
        pre_embed = embed_dim
        for fusion_keys in in_fusion_key_list:
            in_embed = sum(fusion_keys.values())
            fusion = nn.Sequential(
                nn.Conv2d(in_embed + pre_embed, pre_embed, 3, 1, 1),
                nn.BatchNorm2d(pre_embed),
                nn.ReLU(inplace=True),
                nn.Conv2d(pre_embed, embed_dim, 3, 1, 1),
            )
            self.fusion_list.append(fusion)
            pre_embed = embed_dim

        # Final output conv
        self.out_conv = nn.Sequential(
            nn.Conv2d(pre_embed, pre_embed, 3, 1, 1),
            nn.BatchNorm2d(pre_embed),
            nn.ReLU(inplace=True),
            nn.Conv2d(pre_embed, pre_embed, 3, 1, 1),
        )

    def forward(self, inputs):
        """
        inputs: dict of {source_name: {'encoder_features': tensor,
                                        'features_list': [t0, t1, t2, t3]}}
        returns: (B, embed_dim, H_out, W_out)
        """
        # Step 1: collect and interpolate bottleneck features
        in_features = []
        for key in self.in_feature_key:
            feat = inputs[key]['encoder_features']
            feat = F.interpolate(feat, self.feature_size,
                                 mode='bilinear', align_corners=False)
            in_features.append(feat)
        in_features = torch.cat(in_features, dim=1)
        in_features = self.in_conv(in_features)  # (B, embed_dim, H0, W0)

        # Step 2: 3 upsample + skip fusion levels
        # features_list has 4 entries (stages 0,1,2,3)
        # level 0 → skip from features_list[2]  (index = 3-1-0 = 2)
        # level 1 → skip from features_list[1]  (index = 3-1-1 = 1)
        # level 2 → skip from features_list[0]  (index = 3-1-2 = 0)
        for i, fusion_keys in enumerate(self.in_fusion_key_list):
            # Upsample current features x2
            in_features = F.interpolate(in_features, scale_factor=2,
                                        mode='bilinear', align_corners=False)
            H, W = in_features.shape[-2:]

            # Skip connection index (from deep to shallow)
            skip_idx = len(self.in_fusion_key_list) - 1 - i  # 2, 1, 0

            # Collect skip features from each source
            skip_feats = []
            for key in fusion_keys:
                feat = inputs[key]['features_list'][skip_idx]
                feat = F.interpolate(feat, (H, W),
                                     mode='bilinear', align_corners=False)
                skip_feats.append(feat)

            skip_cat = torch.cat(skip_feats, dim=1)
            in_features = torch.cat([in_features, skip_cat], dim=1)
            in_features = self.fusion_list[i](in_features)

        # Step 3: final conv + upsample to output size
        out = self.out_conv(in_features)
        out = F.interpolate(out, self.out_size,
                            mode='bilinear', align_corners=False)
        return out