"""
Model definition module.
Provides a function to create a pre-trained ResNet50 model
modified for CIFAR-100 (100 output classes).
"""

import torch
import torch.nn as nn
from torchvision import models

def create_resnet50_model(num_classes=100):
    print(f"Loading ResNet50 model with IMAGENET1K_V1 weights...")
    model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
    
    num_ftrs = model.fc.in_features

    model.fc = nn.Linear(num_ftrs, num_classes)
    print(f"Model final layer replaced with {num_ftrs} -> {num_classes} outputs.")
    
    return model

if __name__ == "__main__":
    print("Testing model module...")
    
    model_100 = create_resnet50_model(num_classes=100)
    test_tensor_100 = torch.randn(1, 3, 32, 32)
    output_100 = model_100(test_tensor_100)
    print(f"Output shape for 100 classes: {output_100.shape}")
    
    model_10 = create_resnet50_for_cifar(num_classes=10)
    test_tensor_10 = torch.randn(1, 3, 32, 32)
    output_10 = model_10(test_tensor_10)
    print(f"Output shape for 10 classes: {output_10.shape}")
    
    assert output_100.shape == (1, 100)
    assert output_10.shape == (1, 10)
    
    print("\nModel module test passed!")