import torch.nn as nn
from torchvision.models import (
    resnet50, ResNet50_Weights,
    efficientnet_b0, EfficientNet_B0_Weights,
    vit_b_16, ViT_B_16_Weights,
)


def create_model(model_name: str, num_classes: int = 5, pretrained: bool = True) -> nn.Module:
    """Create a pretrained model with the classification head replaced for num_classes.

    Supported model_name values: 'resnet50', 'efficientnet_b0', 'vit_b_16'
    All backbone layers are trainable (full fine-tuning).
    """
    if model_name == 'resnet50':
        weights = ResNet50_Weights.DEFAULT if pretrained else None
        model = resnet50(weights=weights)
        model.fc = nn.Linear(model.fc.in_features, num_classes)  # 2048 -> num_classes

    elif model_name == 'efficientnet_b0':
        weights = EfficientNet_B0_Weights.DEFAULT if pretrained else None
        model = efficientnet_b0(weights=weights)
        # classifier is Sequential(Dropout(0.2), Linear(1280, 1000))
        model.classifier[1] = nn.Linear(model.classifier[1].in_features, num_classes)

    elif model_name == 'vit_b_16':
        weights = ViT_B_16_Weights.DEFAULT if pretrained else None
        model = vit_b_16(weights=weights)
        # heads is Sequential(head=Linear(768, 1000))
        model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)

    else:
        raise ValueError(f"Unknown model: {model_name!r}. Choose from: resnet50, efficientnet_b0, vit_b_16")

    return model