# 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):
    def __init__(self,
                 name='resnet50',
                 pretrained=True,
                 in_channels=3,
                 output_stride=32,
                 model_fn=None,          # ← NEW: optional pre‑selected (constructor, weights)
                 **kwargs):
        super().__init__()

        # If model_fn is provided, use it; otherwise look up by name
        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)

        # ... rest of the method unchanged ...
        self.encoder = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4
        )

        # (adapt first conv, output stride, out_channels – identical to your code)
        # ... (keep everything after this point exactly as you wrote it)

        # 3. Adapt first conv for different input channels
        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

        # 4. Handle output stride (replace stride with dilation)
        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):
        # We need to run through sequentially to collect intermediate outputs
        x = self.encoder[0](x)   # conv1
        x = self.encoder[1](x)   # bn1
        x = self.encoder[2](x)   # relu
        x = self.encoder[3](x)   # maxpool
        stage1 = x

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

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

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

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

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

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