"""
Transfer Learning Example - PyTorch Neural Networks
Project: PyTorch Neural Networks
Author: RSK World
Website: https://rskworld.in
Email: help@rskworld.in
Phone: +91 93305 39277
Description: Example demonstrating transfer learning
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
import sys

sys.path.append('..')

from models.transfer_learning import TransferLearningModel
from training.trainer import Trainer
from training.utils import generate_sample_data


def transfer_learning_example():
    """
    Demonstrate transfer learning
    """
    print("=" * 60)
    print("Transfer Learning Example")
    print("Author: RSK World (https://rskworld.in)")
    print("=" * 60)
    
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"\nUsing device: {device}\n")
    
    # Note: For real transfer learning, you would use actual image data
    # This is a simplified example with synthetic data
    
    print("Creating transfer learning model (ResNet18)...")
    model = TransferLearningModel(
        model_name='resnet18',
        num_classes=10,
        pretrained=True,
        freeze_backbone=True  # Freeze backbone, only train classifier
    ).to(device)
    
    print("Model created!")
    print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")
    print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
    
    # For demonstration, we'll use synthetic data
    # In practice, use real image datasets
    print("\nNote: Using synthetic data for demonstration")
    print("In practice, use torchvision.datasets for real image data")
    
    # Example: Fine-tuning after initial training
    print("\nAfter initial training, you can unfreeze all layers:")
    print("model.unfreeze_all()  # Unfreeze for fine-tuning")
    
    print("\n" + "=" * 60)
    print("Transfer learning setup complete!")
    print("=" * 60)


if __name__ == '__main__':
    transfer_learning_example()