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

class UNet(nn.Module):
    """
    Generic U-Net decoder that works with any encoder following the Encoder interface.

    Args:
        encoder: An instance of an Encoder (must provide .out_channels and forward features).
        decoder_channels: List of output channels for each decoder stage.
                           Length should equal len(encoder.out_channels) - 1.
        out_channels: Number of output segmentation classes.
        bilinear: If True, use bilinear upsampling; else use transposed conv.
        dropout: Dropout probability applied after each conv block.
        norm: Normalization type ('bn', 'in', or None).
        activation: Activation type ('relu', 'leaky', 'swish', etc.).
        deep_supervision: If True, return outputs from all decoder stages (for multi-scale loss).
    """
    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,
    ):
        super().__init__()
        self.encoder = encoder
        self.deep_supervision = deep_supervision
        self.bilinear = bilinear

        # Encoder output channels: [skip1, skip2, ..., skipN, bottleneck]
        enc_channels = encoder.out_channels
        self._validate_channels(enc_channels, decoder_channels)

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

        # Build decoder blocks
        self.up_blocks = nn.ModuleList()
        up_in_channels = [enc_channels[-1]] + list(decoder_channels[:-1])  # first = bottleneck, then previous outputs

        for i in range(n_decoder_stages):
            skip_ch = enc_channels[-(i+2)]          # skip connection from encoder
            up_in = up_in_channels[i]                # channels to upsample
            out_ch = decoder_channels[i]             # target output channels
            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)
        self.final_upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=False)
        # Deep supervision heads (optional)
        if deep_supervision:
            self.aux_heads = nn.ModuleList([
                nn.Conv2d(ch, out_channels, kernel_size=1)
                for ch in decoder_channels[:-1]   # exclude last (already has main head)
            ])
        else:
            self.aux_heads = None

        self._init_weights()

    def _validate_channels(self, enc_channels, dec_channels):
        """Ensure encoder and decoder configurations are compatible."""
        if len(enc_channels) < 2:
            raise ValueError(f"Encoder must output at least 2 feature maps, got {len(enc_channels)}")

    def _init_weights(self):
        """Custom weight initialisation."""
        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]]]]:
        # Extract encoder features (ordered from shallow to deep)
        enc_features = self.encoder(x)
        # Convert to list in reverse order: from bottleneck to stage1
        # We assume the encoder returns a dict with keys: stage1, stage2, ..., stageN, bottleneck
        # Sort keys to ensure correct order (stage1 is shallowest, bottleneck deepest)
        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]]   # all skip connections (shallow to deep)
        bottleneck = enc_features[keys[-1]]            # deepest feature

        # Decoder forward pass
        x = bottleneck
        decoder_outputs = []
        for i, up_block in enumerate(self.up_blocks):
            # skip connection is the corresponding feature from the encoder
            skip = skips[-(i+1)]   # take from deep to shallow
            x = up_block(x, skip)
            decoder_outputs.append(x)

        # Final segmentation
        out = self.seg_head(decoder_outputs[-1])
        out = self.final_upsample(out)

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