import torch.nn as nn
from torchvision import models

from src.models import register


# EfficientNet's classification head is a Sequential; [-1] targets the final Linear
def build(cfg: dict) -> nn.Module:
    backbone = cfg.get("backbone", "efficientnet_b0")
    pretrained = cfg.get("pretrained", True)

    if backbone == "efficientnet_b0":
        weights = models.EfficientNet_B0_Weights.DEFAULT if pretrained else None
        model = models.efficientnet_b0(weights=weights)
    else:
        raise ValueError(f"Unsupported EfficientNet backbone: {backbone!r}. Supported: efficientnet_b0")

    in_features = model.classifier[-1].in_features
    model.classifier[-1] = nn.Linear(in_features, 1)
    return model


register("efficientnet_b0", build)
