import torch.nn as nn
from torchvision import models

from src.models import register


# Loads pretrained ResNet and replaces the 1000-class head with a single logit for binary detection
def build(cfg: dict) -> nn.Module:
    backbone = cfg.get("backbone", "resnet18")
    pretrained = cfg.get("pretrained", True)

    if backbone == "resnet18":
        weights = models.ResNet18_Weights.DEFAULT if pretrained else None
        model = models.resnet18(weights=weights)
    elif backbone == "resnet34":
        weights = models.ResNet34_Weights.DEFAULT if pretrained else None
        model = models.resnet34(weights=weights)
    elif backbone == "resnet50":
        weights = models.ResNet50_Weights.DEFAULT if pretrained else None
        model = models.resnet50(weights=weights)
    else:
        raise ValueError(f"Unsupported backbone: {backbone!r}")

    model.fc = nn.Linear(model.fc.in_features, 1)
    return model


register("resnet18", build)
register("resnet34", build)
register("resnet50", build)
