# src/mindglow/seg2d/models/encoders/__init__.py
from .resnet import ResNetEncoder
from .vgg import VGGEncoder
from .resnext import ResNeXtEncoder

ENCODER_REGISTRY = {
    # ResNet
    'resnet18': ResNetEncoder,
    'resnet34': ResNetEncoder,
    'resnet50': ResNetEncoder,
    'resnet101': ResNetEncoder,
    'resnet152': ResNetEncoder,
    # ResNeXt
    'resnext50_32x4d': ResNeXtEncoder,
    'resnext101_32x8d': ResNeXtEncoder,
    # VGG
    'vgg11': VGGEncoder,
    'vgg13': VGGEncoder,
    'vgg16': VGGEncoder,
    'vgg19': VGGEncoder,
}

def get_encoder(name, **kwargs):
    """Factory: create encoder by name."""
    if name not in ENCODER_REGISTRY:
        raise ValueError(f"Unknown encoder: {name}. Available: {list(ENCODER_REGISTRY.keys())}")
    return ENCODER_REGISTRY[name](name=name, **kwargs)