src/models/model_configs.py

# Model configurations separated from factory for better organization

MODEL_REGISTRY = {
    # EfficientNet family
    'efficientnet-b0': 'efficientnet_b0',
    'efficientnet-b1': 'efficientnet_b1',
    'efficientnet-b2': 'efficientnet_b2',
    'efficientnet-b3': 'efficientnet_b3',
    'efficientnet-b4': 'efficientnet_b4',
    'efficientnet-b5': 'efficientnet_b5',
    'efficientnet-b6': 'efficientnet_b6',
    'efficientnet-b7': 'efficientnet_b7',

    # ConvNeXt V2 family
    'convnext-tiny': 'convnextv2_tiny.fcmae_ft_in22k_in1k',
    'convnext-small': 'convnextv2_small.fcmae_ft_in22k_in1k',
    'convnext-base': 'convnextv2_base.fcmae_ft_in22k_in1k',
    'convnext-tiny-384': 'convnextv2_tiny.fcmae_ft_in22k_in1k_384',

    # ResNet family
    'resnet18': 'resnet18',
    'resnet34': 'resnet34',
    'resnet50': 'resnet50',
    'resnet101': 'resnet101',
    'resnet152': 'resnet152',

    # ResNeXt family
    'resnext50': 'resnext50_32x4d',
    'resnext101': 'resnext101_32x8d',

    # Vision Transformer family
    'vit-tiny': 'vit_tiny_patch16_224',
    'vit-small': 'vit_small_patch16_224',
    'vit-base': 'vit_base_patch16_224',

    # MobileNet family
    'mobilenet-v3-small': 'mobilenetv3_small_100',
    'mobilenet-v3-large': 'mobilenetv3_large_100',
}

# Training configurations for each model
TRAINING_CONFIGS = {
    # EfficientNet configs
    'efficientnet-b0': {
        'input_size': 224,
        'batch_size': 128,
        'learning_rate': 1e-3,
        'weight_decay': 1e-5,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },
    'efficientnet-b1': {
        'input_size': 240,
        'batch_size': 96,
        'learning_rate': 1e-3,
        'weight_decay': 1e-5,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },
    'efficientnet-b2': {
        'input_size': 260,
        'batch_size': 64,
        'learning_rate': 1e-3,
        'weight_decay': 1e-5,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },
    'efficientnet-b3': {
        'input_size': 300,
        'batch_size': 48,
        'learning_rate': 1e-3,
        'weight_decay': 1e-5,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },
    'efficientnet-b4': {
        'input_size': 380,
        'batch_size': 32,
        'learning_rate': 1e-3,
        'weight_decay': 1e-5,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },

    # ConvNeXt configs - these models benefit from different hyperparameters
    'convnext-tiny': {
        'input_size': 224,
        'batch_size': 64,
        'learning_rate': 5e-4,
        'weight_decay': 0.05,  # Higher weight decay for ConvNeXt
        'epochs': 30,
        'warmup_epochs': 5,
        'min_lr': 1e-6,
        'drop_path_rate': 0.1,
    },
    'convnext-base': {
        'input_size': 224,
        'batch_size': 32,
        'learning_rate': 5e-4,
        'weight_decay': 0.05,
        'epochs': 30,
        'warmup_epochs': 5,
        'min_lr': 1e-6,
        'drop_path_rate': 0.2,
    },
    'convnext-tiny-384': {
        'input_size': 384,
        'batch_size': 16,  # Smaller batch due to larger image size
        'learning_rate': 5e-4,
        'weight_decay': 0.05,
        'epochs': 30,
        'warmup_epochs': 5,
        'min_lr': 1e-6,
        'drop_path_rate': 0.1,
    },

    # ResNet configs
    'resnet50': {
        'input_size': 224,
        'batch_size': 64,
        'learning_rate': 1e-3,
        'weight_decay': 1e-4,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },
    'resnext50': {
        'input_size': 224,
        'batch_size': 48,
        'learning_rate': 1e-3,
        'weight_decay': 1e-4,
        'epochs': 30,
        'warmup_epochs': 3,
        'min_lr': 1e-6,
    },

    # Vision Transformer configs
    'vit-tiny': {
        'input_size': 224,
        'batch_size': 128,
        'learning_rate': 1e-4,  # Lower LR for ViT
        'weight_decay': 0.05,
        'epochs': 30,
        'warmup_epochs': 5,
        'min_lr': 1e-6,
        'drop_rate': 0.1,
    },
}

# Augmentation configurations
AUGMENTATION_CONFIGS = {
    'light': {
        'rotation': 5,
        'brightness': 0.1,
        'contrast': 0.1,
        'saturation': 0.1,
        'hue': 0.02,
        'perspective': 0.1,
        'scale': (0.9, 1.0),
        'h_flip_p': 0.5,
    },
    'medium': {
        'rotation': 10,
        'brightness': 0.2,
        'contrast': 0.2,
        'saturation': 0.2,
        'hue': 0.05,
        'perspective': 0.2,
        'scale': (0.8, 1.0),
        'h_flip_p': 0.5,
    },
    'heavy': {
        'rotation': 15,
        'brightness': 0.3,
        'contrast': 0.3,
        'saturation': 0.3,
        'hue': 0.1,
        'perspective': 0.3,
        'scale': (0.7, 1.0),
        'h_flip_p': 0.5,
        'mixup_alpha': 0.2,
        'cutmix_alpha': 1.0,
    }
}

# Model-specific fine-tuning strategies
FINETUNE_CONFIGS = {
    'efficientnet': {
        'freeze_stages': ['stem', 'blocks.0', 'blocks.1'],  # Freeze early layers
        'unfreeze_epoch': 10,  # Unfreeze all layers after this epoch
        'discriminative_lr': True,  # Use different LRs for different layers
        'lr_mult': [0.1, 0.3, 0.5, 0.7, 1.0],  # LR multipliers from bottom to top
    },
    'convnext': {
        'freeze_stages': ['stem', 'stages.0'],
        'unfreeze_epoch': 10,
        'discriminative_lr': True,
        'lr_mult': [0.1, 0.5, 1.0],
    },
    'resnet': {
        'freeze_stages': ['conv1', 'bn1', 'layer1'],
        'unfreeze_epoch': 10,
        'discriminative_lr': False,
    },
    'vit': {
        'freeze_stages': ['patch_embed', 'blocks.0', 'blocks.1'],
        'unfreeze_epoch': 15,
        'discriminative_lr': True,
        'lr_mult': [0.1, 0.5, 1.0],
    }
}


def get_model_family(model_name: str) -> str:
    """Get the family of a model (e.g., 'efficientnet', 'convnext')"""
    if 'efficientnet' in model_name:
        return 'efficientnet'
    elif 'convnext' in model_name:
        return 'convnext'
    elif 'resnext' in model_name:
        return 'resnet'  # ResNeXt uses similar strategy as ResNet
    elif 'resnet' in model_name:
        return 'resnet'
    elif 'vit' in model_name:
        return 'vit'
    elif 'mobilenet' in model_name:
        return 'efficientnet'  # Similar strategy
    else:
        return 'default'


def get_optimal_batch_size(model_name: str, gpu_memory_gb: int = 8) -> int:
    """
    Get optimal batch size based on model and GPU memory

    Args:
        model_name: Name of the model
        gpu_memory_gb: GPU memory in GB

    Returns:
        Recommended batch size
    """
    base_config = TRAINING_CONFIGS.get(model_name, TRAINING_CONFIGS['resnet50'])
    base_batch_size = base_config['batch_size']

    # Adjust based on GPU memory (assuming base is for 8GB)
    memory_multiplier = gpu_memory_gb / 8.0

    # Some models scale better than others
    if 'efficientnet' in model_name:
        memory_multiplier *= 1.2  # EfficientNets are memory efficient
    elif 'vit' in model_name:
        memory_multiplier *= 0.8  # ViTs use more memory

    adjusted_batch_size = int(base_batch_size * memory_multiplier)

    # Round to nearest power of 2 or multiple of 8
    return max(8, min(256, (adjusted_batch_size // 8) * 8))


# Default config for unknown models
DEFAULT_CONFIG = {
    'input_size': 224,
    'batch_size': 32,
    'learning_rate': 1e-3,
    'weight_decay': 1e-4,
    'epochs': 30,
    'warmup_epochs': 3,
    'min_lr': 1e-6,
}