# src/mindglow/seg2d/models/encoders/vgg.py
import torch.nn as nn
from torchvision.models import vgg11_bn, vgg13_bn, vgg16_bn, vgg19_bn
from torchvision.models import VGG11_BN_Weights, VGG13_BN_Weights, VGG16_BN_Weights, VGG19_BN_Weights
from .base import Encoder

class VGGEncoder(Encoder):
    def __init__(self,
                 name='vgg16_bn',
                 pretrained=True,
                 in_channels=3,
                 **kwargs):
        super().__init__()
        
        kwargs.pop('output_stride', None)
        
        model_map = {
            'vgg11': (vgg11_bn, VGG11_BN_Weights.IMAGENET1K_V1),
            'vgg13': (vgg13_bn, VGG13_BN_Weights.IMAGENET1K_V1),
            'vgg16': (vgg16_bn, VGG16_BN_Weights.IMAGENET1K_V1),
            'vgg19': (vgg19_bn, VGG19_BN_Weights.IMAGENET1K_V1),
        }
        model_fn, weights = model_map[name]
        vgg = model_fn(weights=weights if pretrained else None, **kwargs)
        features = vgg.features

        # Adapt first conv for different input channels
        if in_channels != 3:
            old_conv = features[0]
            new_conv = nn.Conv2d(
                in_channels, old_conv.out_channels,
                kernel_size=old_conv.kernel_size,
                stride=old_conv.stride,
                padding=old_conv.padding,
            )
            if pretrained:
                avg_weight = old_conv.weight.mean(dim=1, keepdim=True)
                new_conv.weight.data = avg_weight.repeat(1, in_channels, 1, 1)
            else:
                nn.init.kaiming_normal_(new_conv.weight)
            features[0] = new_conv

        # Split features into stages (each stage ends before a maxpool)
        # VGG has 5 maxpool layers, giving 6 stages including the final convs.
        stage_indices = [0, 5, 10, 17, 24, 31]  # valid for vgg*_bn (counted manually)
        # More robust: find indices of 'MaxPool2d' and split there
        pool_positions = [i for i, m in enumerate(features) if isinstance(m, nn.MaxPool2d)]
        stage_indices = [0] + [p+1 for p in pool_positions]  # start indices of each stage

        self.stages = nn.ModuleList()
        for i in range(len(stage_indices)):
            start = stage_indices[i]
            end = stage_indices[i+1] if i+1 < len(stage_indices) else len(features)
            self.stages.append(features[start:end])

        # Output channels: hardcoded for VGG (64, 128, 256, 512, 512, 512)
        self._out_channels = [64, 128, 256, 512, 512, 512][:len(self.stages)]

    def forward(self, x):
        features = {}
        for i, stage in enumerate(self.stages):
            x = stage(x)
            if i == 0:
                features['stage1'] = x
            elif i == 1:
                features['stage2'] = x
            elif i == 2:
                features['stage3'] = x
            elif i == 3:
                features['stage4'] = x
            elif i == 4:
                features['stage5'] = x   # optional
            elif i == 5:
                features['bottleneck'] = x
        # Ensure we have at least the 4 stages + bottleneck
        if 'bottleneck' not in features:
            features['bottleneck'] = x
        return features

    @property
    def out_channels(self):
        # Return only [stage1, stage2, stage3, stage4, bottleneck]
        return self._out_channels
