# src/mindglow/seg2d/models/encoders/resnet.py
import torch.nn as nn
from torchvision.models import resnet18, resnet34, resnet50, resnet101, resnet152
from torchvision.models import ResNet18_Weights, ResNet34_Weights, ResNet50_Weights,ResNet101_Weights, ResNet152_Weights
from .base import Encoder

class ResNetEncoder(Encoder):
    """
    ResNet encoder for 2D segmentation tasks, wrapping torchvision ResNet models.

    This encoder outputs feature maps at strides 2, 4, 8, 16, and 32 (bottleneck).
    It is designed to work with a U‑Net decoder that has 4 up blocks followed by
    a final upsampling to produce a full‑resolution segmentation mask.

    Args:
        name (str): ResNet variant. One of: 'resnet18', 'resnet34', 'resnet50',
            'resnet101', 'resnet152'. Default: 'resnet50'.
        pretrained (bool): If True, load ImageNet‑pretrained weights. Default: True.
        in_channels (int): Number of input channels. If not 3, the first conv layer
            is replaced. When pretrained=True, the RGB weights are averaged.
            Default: 3.
        output_stride (int): Desired output stride. Must be 8, 16, or 32.
            If 8 or 16, strides are replaced with dilations to maintain receptive field.
            Default: 32.
        model_fn (tuple, optional): A tuple `(builder, weights_enum)` for custom variants.
        **kwargs: Additional arguments passed to the torchvision model builder.

    Attributes:
        encoder (nn.Sequential): Sequential container of ResNet layers.
        _out_channels (list): Output channels for each stage:
            [stage0, stage1, stage2, stage3, bottleneck].
            For ResNet18/34: [64, 64, 128, 256, 512].
            For ResNet50/101/152: [64, 256, 512, 1024, 2048].

    Returns (forward):
        dict: Feature maps:
            - 'stage0': after first ReLU (stride 2)
            - 'stage1': after layer1 (stride 4)
            - 'stage2': after layer2 (stride 8)
            - 'stage3': after layer3 (stride 16)
            - 'bottleneck': after layer4 (stride 32, or modified by output_stride)
    """

    def __init__(self,
                 name: str = 'resnet50',
                 pretrained: bool = True,
                 in_channels: int = 3,
                 output_stride: int = 32,
                 model_fn=None,
                 **kwargs):
        super().__init__()

        if model_fn is not None:
            builder, weights_enum = model_fn
        else:
            model_fn_lookup = {
                'resnet18': (resnet18, ResNet18_Weights.IMAGENET1K_V1),
                'resnet34': (resnet34, ResNet34_Weights.IMAGENET1K_V1),
                'resnet50': (resnet50, ResNet50_Weights.IMAGENET1K_V1),
                'resnet101': (resnet101, ResNet101_Weights.IMAGENET1K_V1),
                'resnet152': (resnet152, ResNet152_Weights.IMAGENET1K_V1),
            }
            if name not in model_fn_lookup:
                raise ValueError(f"Unknown ResNet variant: {name}")
            builder, weights_enum = model_fn_lookup[name]

        weights = weights_enum if pretrained else None
        resnet = builder(weights=weights, **kwargs)

        # Build encoder as a sequential container
        self.encoder = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4
        )

        if in_channels != 3:
            old_conv = resnet.conv1
            assert old_conv is not None          # conv1 exists
            assert old_conv.weight is not None   # weight is never None
            new_conv = nn.Conv2d(
                in_channels, old_conv.out_channels,
                kernel_size=old_conv.kernel_size,
                stride=old_conv.stride,
                padding=old_conv.padding,
                bias=old_conv.bias is not None
            )
            assert new_conv.weight is not None   # assure type checker that weight exists
            if pretrained:
                # Average RGB weights across channels
                avg_weight = old_conv.weight.mean(dim=1, keepdim=True)
                new_conv.weight.data = avg_weight.repeat(1, in_channels, 1, 1)
                if old_conv.bias is not None:
                    new_conv.bias.data = old_conv.bias.data # type: ignore
            else:
                nn.init.kaiming_normal_(new_conv.weight)
            self.encoder[0] = new_conv

        if output_stride == 16:
            # Remove stride in layer3 and layer4, add dilation
            for layer in [self.encoder[6], self.encoder[7]]:  # layer3, layer4
                for m in layer.modules():
                    if isinstance(m, nn.Conv2d) and m.stride == (2, 2):
                        m.stride = (1, 1)
                        if m.kernel_size[0] == 3:
                            m.dilation = (2, 2)
                            m.padding = (2, 2)
        elif output_stride == 8:
            # Remove stride in layer2, layer3, layer4 with progressive dilation
            for idx, layer in enumerate([self.encoder[5], self.encoder[6], self.encoder[7]]):
                for m in layer.modules():
                    if isinstance(m, nn.Conv2d) and m.stride == (2, 2):
                        m.stride = (1, 1)
                        if m.kernel_size[0] == 3:
                            m.dilation = (2, 2) if idx == 0 else (4, 4)
                            m.padding = (2, 2)
        elif output_stride != 32:
            raise ValueError(f"output_stride must be 8, 16, or 32, got {output_stride}")

        if name.startswith('resnet18') or name.startswith('resnet34'):
            self._out_channels = [64, 64, 128, 256, 512]   # stage0, stage1, stage2, stage3, bottleneck
        elif (name.startswith('resnet50') or name.startswith('resnet101') or name.startswith('resnet152')
                or name.startswith('resnext50') or name.startswith('resnext101')):
            self._out_channels = [64, 256, 512, 1024, 2048] # stage0, stage1, stage2, stage3, bottleneck
        else:
            raise ValueError(f"Unknown ResNet variant: {name}")

    def forward(self, x):
        # Stage 0: after conv1 + bn1 + relu (stride 2)
        x = self.encoder[0](x)
        x = self.encoder[1](x)
        x = self.encoder[2](x)
        stage0 = x

        # Skip maxpool output – we use layer1 as the first real skip
        x = self.encoder[3](x)  # maxpool (stride 4) – not used as skip
        x = self.encoder[4](x)  # layer1 (still stride 4)
        stage1 = x

        # Stage 2: after layer2 (stride 8)
        x = self.encoder[5](x)
        stage2 = x

        # Stage 3: after layer3 (stride 16)
        x = self.encoder[6](x)
        stage3 = x

        # Bottleneck: after layer4 (stride 32 or modified)
        x = self.encoder[7](x)
        bottleneck = x

        return {
            'stage0': stage0,
            'stage1': stage1,
            'stage2': stage2,
            'stage3': stage3,
            'bottleneck': bottleneck,
        }

    @property
    def out_channels(self):
        return self._out_channels