# 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 builds a ResNet backbone and adapts it for use in encoder-decoder
    architectures.

    Args:
        name (str): ResNet variant name. 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 convolutional
            layer is replaced. When pretrained=True, the RGB weights are averaged
            across the channel dimension to initialize the new layer.
            Default: 3.
        output_stride (int): Desired output stride (downsampling factor) of the encoder.
            Must be 8, 16, or 32. If 8 or 16, the strides in certain layers are set to 1
            and dilations are introduced to maintain the receptive field.
            Default: 32.
        model_fn (tuple, optional): A tuple `(builder, weights_enum)` to allow custom
            ResNet variants. If provided, `name` is ignored. `builder` is a function
            that constructs the model (e.g., `torchvision.models.resnet50`), and
            `weights_enum` is the corresponding weights enum for pretrained loading.
            Default: None.
        **kwargs: Additional keyword arguments passed to the torchvision model builder
            (e.g., `norm_layer`, `width_per_group`).

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

    Returns (forward):
        dict: A dictionary containing feature maps from each stage:
            - 'stage1': after first residual block (downsampled by 4)
            - 'stage2': after second residual block (downsampled by 8)
            - 'stage3': after third residual block (downsampled by 16)
            - 'stage4': after fourth residual block (downsampled by 32, or modified by output_stride)
            - 'bottleneck': final output of the encoder

    Raises:
        ValueError: If an unknown ResNet variant is provided via `name`, or if
            `output_stride` is not 8, 16, or 32.
    """
    
    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)

        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
            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
            )
            if pretrained:
                # Average the RGB weights over channel dimension
                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:
                    assert new_conv.bias is not None
                    new_conv.bias.data = old_conv.bias.data
            else:
                nn.init.kaiming_normal_(new_conv.weight)
            self.encoder[0] = new_conv

        if output_stride == 16:
            # Replace stride in layer3 and layer4
            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:
            # Replace stride in layer2, layer3, layer4
            for layer in [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 layer is self.encoder[5] else (4, 4)  # progressive
                            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]
        elif name.startswith('resnet50') or name.startswith('resnet101') or name.startswith('resnet152'):
             self._out_channels = [64, 256, 512, 1024, 2048]
        else:
             ValueError(f"Unknown ResNet variant: {name}")

    def forward(self, x):
        x = self.encoder[0](x)   
        x = self.encoder[1](x)   
        x = self.encoder[2](x)   
        x = self.encoder[3](x)   
        stage1 = x

        x = self.encoder[4](x)  
        stage2 = x

        x = self.encoder[5](x)  
        stage3 = x

        x = self.encoder[6](x) 
        stage4 = x

        x = self.encoder[7](x)
        bottleneck = x

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

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