"""
Transfer Learning Models - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Transfer learning examples using pre-trained models
"""

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

# Handle compatibility with different torchvision versions
try:
    from torchvision.models import ResNet18_Weights, ResNet50_Weights, VGG16_Weights, DenseNet121_Weights
    TORCHVISION_NEW_API = True
except ImportError:
    TORCHVISION_NEW_API = False


class TransferLearningModel(nn.Module):
    """
    Transfer learning model using pre-trained architectures
    
    Project: PyTorch Neural Networks
    Author: RSK World
    Website: https://rskworld.in
    """
    
    def __init__(self, model_name='resnet18', num_classes=10, pretrained=True, freeze_backbone=False):
        """
        Initialize transfer learning model
        
        Args:
            model_name: Name of pre-trained model ('resnet18', 'resnet50', 'vgg16', etc.)
            num_classes: Number of output classes
            pretrained: Whether to use pre-trained weights
            freeze_backbone: Whether to freeze backbone layers
        """
        super(TransferLearningModel, self).__init__()
        
        # Load pre-trained model
        # Handle compatibility with different torchvision versions
        if TORCHVISION_NEW_API and pretrained:
            # New API (torchvision >= 0.13)
            if model_name == 'resnet18':
                self.backbone = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
                num_features = self.backbone.fc.in_features
                self.backbone.fc = nn.Linear(num_features, num_classes)
            elif model_name == 'resnet50':
                self.backbone = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
                num_features = self.backbone.fc.in_features
                self.backbone.fc = nn.Linear(num_features, num_classes)
            elif model_name == 'vgg16':
                self.backbone = models.vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
                num_features = self.backbone.classifier[6].in_features
                self.backbone.classifier[6] = nn.Linear(num_features, num_classes)
            elif model_name == 'densenet121':
                self.backbone = models.densenet121(weights=DenseNet121_Weights.IMAGENET1K_V1)
                num_features = self.backbone.classifier.in_features
                self.backbone.classifier = nn.Linear(num_features, num_classes)
            else:
                raise ValueError(f"Unsupported model: {model_name}")
        else:
            # Old API (torchvision < 0.13) or pretrained=False
            if model_name == 'resnet18':
                self.backbone = models.resnet18(pretrained=pretrained)
                num_features = self.backbone.fc.in_features
                self.backbone.fc = nn.Linear(num_features, num_classes)
            elif model_name == 'resnet50':
                self.backbone = models.resnet50(pretrained=pretrained)
                num_features = self.backbone.fc.in_features
                self.backbone.fc = nn.Linear(num_features, num_classes)
            elif model_name == 'vgg16':
                self.backbone = models.vgg16(pretrained=pretrained)
                num_features = self.backbone.classifier[6].in_features
                self.backbone.classifier[6] = nn.Linear(num_features, num_classes)
            elif model_name == 'densenet121':
                self.backbone = models.densenet121(pretrained=pretrained)
                num_features = self.backbone.classifier.in_features
                self.backbone.classifier = nn.Linear(num_features, num_classes)
            else:
                raise ValueError(f"Unsupported model: {model_name}")
        
        # Freeze backbone if requested
        if freeze_backbone:
            for param in self.backbone.parameters():
                param.requires_grad = False
            # Unfreeze last layer
            if hasattr(self.backbone, 'fc'):
                for param in self.backbone.fc.parameters():
                    param.requires_grad = True
            elif hasattr(self.backbone, 'classifier'):
                for param in self.backbone.classifier.parameters():
                    param.requires_grad = True
    
    def forward(self, x):
        """Forward pass"""
        return self.backbone(x)
    
    def unfreeze_all(self):
        """Unfreeze all layers for fine-tuning"""
        for param in self.backbone.parameters():
            param.requires_grad = True
    
    def freeze_backbone(self):
        """Freeze backbone layers"""
        for param in self.backbone.parameters():
            param.requires_grad = False
        # Unfreeze classifier
        if hasattr(self.backbone, 'fc'):
            for param in self.backbone.fc.parameters():
                param.requires_grad = True
        elif hasattr(self.backbone, 'classifier'):
            for param in self.backbone.classifier.parameters():
                param.requires_grad = True