# src/mindglow/seg2d/models/encoders/resnext.py
from torchvision.models import resnext50_32x4d, resnext101_32x8d
from torchvision.models import ResNeXt50_32X4D_Weights, ResNeXt101_32X8D_Weights
from .resnet import ResNetEncoder

class ResNeXtEncoder(ResNetEncoder):
    def __init__(self,
                 name='resnext50_32x4d',
                 pretrained=True,
                 in_channels=3,
                 output_stride=32,
                 **kwargs):
        model_fn_lookup = {
            'resnext50_32x4d': (resnext50_32x4d, ResNeXt50_32X4D_Weights.IMAGENET1K_V1),
            'resnext101_32x8d': (resnext101_32x8d, ResNeXt101_32X8D_Weights.IMAGENET1K_V1),
        }
        if name not in model_fn_lookup:
            raise ValueError(f"Unknown ResNeXt variant: {name}")
        model_fn = model_fn_lookup[name]

        super().__init__(
            name=name,               
            pretrained=pretrained,
            in_channels=in_channels,
            output_stride=output_stride,
            model_fn=model_fn,   
            **kwargs
        )

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