# src/mindglow/seg2d/models/unet.py
from typing import List, Dict, Optional, Union, Sequence
import torch
import torch.nn as nn
import torch.nn.functional as F  # needed for interpolation
from .blocks import Up, DoubleConv, Conv2dBlock, AttentionUp

class UNet(nn.Module):
    def __init__(
        self,
        encoder: nn.Module,
        decoder_channels: Sequence[int] = (256, 128, 64, 32),
        out_channels: int = 1,
        bilinear: bool = False,
        dropout: float = 0.0,
        norm: str = 'bn',
        activation: str = 'relu',
        deep_supervision: bool = False,
        attention: bool = False,
    ):
        """
        Generic U-Net decoder that works with any encoder following the Encoder interface.

        The encoder must:
            - Have an attribute `out_channels`: a list of output channels for each stage,
            in order from shallowest to deepest (including the bottleneck as the last element).
            - Return a dictionary of feature maps with keys like 'stage0', 'stage1', …,
            'bottleneck' (or any keys that can be sorted to recover the shallow‑to‑deep order).

        Args:
            encoder (nn.Module): An encoder instance (e.g., ResNetEncoder) that provides
                multi‑scale features via its `forward` method and exposes `out_channels`.
            decoder_channels (Sequence[int]): List of output channels for each decoder stage.
                The length must equal `len(encoder.out_channels) - 1`. Default: (256, 128, 64, 32).
            out_channels (int): Number of output segmentation classes. Default: 1.
            bilinear (bool): If True, use bilinear upsampling in `Up` blocks;
                if False, use transposed convolutions. Default: False.
            dropout (float): Dropout probability applied after each convolutional block.
                Default: 0.0.
            norm (str): Normalization type for internal convolutions. Options: 'bn' (BatchNorm),
                'in' (InstanceNorm), or None. Default: 'bn'.
            activation (str): Activation function for internal convolutions. Options:
                'relu', 'leaky', 'swish', 'gelu', 'mish', or None. Default: 'relu'.
            deep_supervision (bool): If True, the model returns auxiliary outputs from
                intermediate decoder stages (for multi‑scale loss). During training,
                the forward method returns a dict with keys 'out' (final output) and
                'aux' (list of auxiliary outputs). During inference, only the final
                output is returned. Default: False.

        Attributes:
            up_blocks (nn.ModuleList): List of `Up` modules, one per decoder stage.
            seg_head (nn.Conv2d): Final 1×1 convolution to produce the segmentation map.
            aux_heads (nn.ModuleList or None): List of 1×1 convolutions for deep supervision
                (only if `deep_supervision=True`).
            deep_supervision (bool): Stores the deep supervision flag.
            bilinear (bool): Stores the upsampling mode.

        Forward:
            Input: (B, C, H, W) – batch of images.
            Returns:
                - If `deep_supervision=False` or model in eval mode: a single tensor of shape
                (B, out_channels, H, W) with the segmentation logits.
                - If `deep_supervision=True` and model in training mode: a dict with:
                    - 'out': final output tensor (B, out_channels, H, W)
                    - 'aux': list of auxiliary tensors from earlier decoder stages,
                    each of shape (B, out_channels, H_i, W_i) where H_i, W_i are the
                    spatial sizes at that stage (not resized to input size).

        Notes:
            - The final output is always resized to the input spatial dimensions using
            bilinear interpolation, ensuring pixel‑wise alignment regardless of internal
            downsampling factors.
            - The decoder automatically adapts to any input size; no external padding is required.
            - For deep supervision, the auxiliary outputs are kept at their native resolutions
            (typically 1/2, 1/4, … of the input) because losses are usually computed on
            downsampled ground truth. You can modify the code to also resize them if needed.
        """
        super().__init__()
        self.encoder = encoder
        self.deep_supervision = deep_supervision
        self.bilinear = bilinear

        enc_channels = encoder.out_channels # e.g., [64, 256, 512, 1024, 2048]
        self._validate_channels(enc_channels, decoder_channels)

        n_decoder_stages = len(decoder_channels) # 4
        assert n_decoder_stages == len(enc_channels) - 1, \
            f"decoder_channels length must be {len(enc_channels)-1}, got {n_decoder_stages}"

        self.up_blocks = nn.ModuleList()
        up_in_channels = [enc_channels[-1]] + list(decoder_channels[:-1]) # = [2048] + [256, 128, 64] = [2048, 256, 128, 64]

        for i in range(n_decoder_stages):
            skip_ch = enc_channels[-(i+2)]
            up_in = up_in_channels[i]
            out_ch = decoder_channels[i]

            if attention:
                up_block = AttentionUp(in_channels=up_in,
                                    skip_channels=skip_ch,
                                    out_channels=out_ch,
                                    bilinear=bilinear,
                                    dropout=dropout,
                                    norm=norm,
                                    activation=activation,
                                    reduction=2
                                    )
                
            else:
                up_block = Up(
                    in_channels=up_in,
                    skip_channels=skip_ch,
                    out_channels=out_ch,
                    bilinear=bilinear,
                    dropout=dropout,
                    norm=norm,
                    activation=activation
                )
        
            self.up_blocks.append(up_block)

        self.seg_head = nn.Conv2d(decoder_channels[-1], out_channels, kernel_size=1)

        if deep_supervision:
            self.aux_heads = nn.ModuleList([
                nn.Conv2d(ch, out_channels, kernel_size=1)
                for ch in decoder_channels[:-1]
            ])
        else:
            self.aux_heads = None

        self._init_weights()

    def _validate_channels(self, enc_channels, dec_channels):
        if len(enc_channels) < 2:
            raise ValueError(f"Encoder must output at least 2 feature maps, got {len(enc_channels)}")

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, (nn.BatchNorm2d, nn.InstanceNorm2d)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x: torch.Tensor) -> Union[torch.Tensor, Dict[str, Union[torch.Tensor, List[torch.Tensor]]]]:
        enc_features = self.encoder(x)
        keys = sorted(enc_features.keys(), key=lambda k: int(k.replace('stage', '')) if 'stage' in k else 999)
        skips = [enc_features[k] for k in keys[:-1]]
        bottleneck = enc_features[keys[-1]]

        z = bottleneck
        decoder_outputs = []
        for i, up_block in enumerate(self.up_blocks):
            skip = skips[-(i+1)]
            z = up_block(z, skip)
            decoder_outputs.append(z)

        out = self.seg_head(decoder_outputs[-1])
        # Dynamically resize to input size
        out = F.interpolate(out, size=x.shape[-2:], mode='bilinear', align_corners=False)

        if self.deep_supervision and self.training and self.aux_heads is not None:
            aux_outputs = []
            for head, fmap in zip(self.aux_heads, decoder_outputs[:-1]):
                aux = head(fmap)
                aux_outputs.append(aux)
            return {'out': out, 'aux': aux_outputs}
        return out