"""
Example: Using CosineAnnealingLR without float mismatch errors
==============================================================

This example demonstrates how to properly use CosineAnnealingLR with the
refactored training framework, avoiding the common float/int type mismatch.
"""

import torch
import torch.nn as nn
from training.schedulers import SchedulerFactory


def example_cosine_scheduler():
    """
    Example showing proper CosineAnnealingLR usage.
    """
    # Create a simple model and optimizer
    model = nn.Linear(10, 1)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    
    # Configuration for CosineAnnealingLR
    # The key fix: T_max will be automatically converted to int
    scheduler_config = {
        'type': 'cosine',
        'params': {
            'T_max': 50,      # Can be int or float - will be converted
            'eta_min': 1e-6   # Minimum learning rate
        }
    }
    
    # Create scheduler using factory
    scheduler = SchedulerFactory.create_scheduler(
        optimizer=optimizer,
        scheduler_config=scheduler_config,
        num_epochs=50,
        steps_per_epoch=100  # Number of batches per epoch
    )
    
    # Simulate training
    for epoch in range(50):
        # Training loop
        for batch in range(100):
            # Your training code here
            optimizer.zero_grad()
            # ... forward, backward, etc.
            optimizer.step()
        
        # Step the scheduler after each epoch
        scheduler.step()
        
        # Print learning rate
        current_lr = optimizer.param_groups[0]['lr']
        if epoch % 10 == 0:
            print(f"Epoch {epoch}: LR = {current_lr:.6f}")


def example_with_warmup():
    """
    Example showing CosineAnnealingLR with warmup.
    """
    model = nn.Linear(10, 1)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    
    scheduler_config = {
        'type': 'cosine',
        'params': {
            'T_max': 45,      # Reduced because we have 5 warmup epochs
            'eta_min': 1e-6
        }
    }
    
    # Create scheduler with warmup
    scheduler = SchedulerFactory.create_scheduler(
        optimizer=optimizer,
        scheduler_config=scheduler_config,
        num_epochs=50,
        steps_per_epoch=100,
        warmup_epochs=5  # Linear warmup for 5 epochs
    )
    
    print("CosineAnnealingLR with 5 epoch warmup:")
    for epoch in range(10):
        scheduler.step()
        lr = optimizer.param_groups[0]['lr']
        print(f"Epoch {epoch}: LR = {lr:.6f}")


def example_config_yaml():
    """
    Show how to configure in YAML file.
    """
    yaml_content = """
# In your config.yaml file:
TRAIN:
  epochs: 100
  base_lr: 0.001
  
  # New scheduler configuration
  scheduler:
    type: cosine
    params:
      T_max: 100        # Same as epochs
      eta_min: 0.000001 # Final learning rate
  
  # Optional warmup
  warmup_epochs: 10
  
  # For ReduceLROnPlateau, also add:
  scheduler_metric: validation_loss  # Metric to monitor
"""
    print("Example YAML configuration:")
    print(yaml_content)


def common_errors_and_fixes():
    """
    Document common errors and their fixes.
    """
    print("\nCommon CosineAnnealingLR Errors and Fixes:")
    print("=" * 50)
    
    print("\n1. Float mismatch error:")
    print("   Error: 'float' object cannot be interpreted as an integer")
    print("   Cause: T_max parameter is a float")
    print("   Fix: The SchedulerFactory automatically converts to int")
    
    print("\n2. Learning rate not changing:")
    print("   Error: LR stays constant")
    print("   Cause: Not calling scheduler.step()")
    print("   Fix: Call scheduler.step() after each epoch")
    
    print("\n3. Wrong decay pattern:")
    print("   Error: LR decreases too fast/slow")
    print("   Cause: Wrong T_max value")
    print("   Fix: Set T_max to total number of epochs")


if __name__ == "__main__":
    print("CosineAnnealingLR Examples")
    print("=" * 50)
    
    print("\n1. Basic CosineAnnealingLR:")
    example_cosine_scheduler()
    
    print("\n2. CosineAnnealingLR with Warmup:")
    example_with_warmup()
    
    print("\n3. YAML Configuration:")
    example_config_yaml()
    
    print("\n4. Common Errors:")
    common_errors_and_fixes()