#!/usr/bin/env python3
"""
Training script using best hyperparameters from Optuna optimization.
This script trains the model with the optimized hyperparameters and additional
regularization techniques to reduce overfitting.
"""

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import pandas as pd
import numpy as np
import torch
from torch.utils.data import DataLoader, random_split
from transformers import CLIPModel as CLIPModel_transformers
import warnings
import config
from main_model import CustomDataset, load_models, train_model

warnings.filterwarnings("ignore")

def train_with_best_params(
    learning_rate=1.42e-05,      # Best from Optuna
    temperature=0.0503,           # Best from Optuna
    alignment_weight=0.5639,      # Best from Optuna
    weight_decay=2.76e-05,        # Best from Optuna
    num_epochs=20,
    batch_size=32,
    subset_size=20000,            # Increased for better generalization
    use_early_stopping=True,
    patience=7
):
    """
    Train model with best hyperparameters and anti-overfitting techniques.
    
    Args:
        learning_rate: Learning rate for optimizer (from Optuna)
        temperature: Temperature for contrastive loss (from Optuna)
        alignment_weight: Weight for alignment loss (from Optuna)
        weight_decay: L2 regularization weight (from Optuna)
        num_epochs: Number of training epochs
        batch_size: Batch size for training
        subset_size: Size of dataset subset
        use_early_stopping: Whether to use early stopping
        patience: Patience for early stopping
    """
    print("="*80)
    print("🚀 Training with Optimized Hyperparameters")
    print("="*80)
    
    print(f"\n📋 Configuration:")
    print(f"  Learning rate: {learning_rate:.2e}")
    print(f"  Temperature: {temperature:.4f}")
    print(f"  Alignment weight: {alignment_weight:.4f}")
    print(f"  Weight decay: {weight_decay:.2e}")
    print(f"  Num epochs: {num_epochs}")
    print(f"  Batch size: {batch_size}")
    print(f"  Subset size: {subset_size}")
    print(f"  Early stopping: {use_early_stopping} (patience={patience})")
    
    # Load data
    print(f"\n📂 Loading data...")
    df = pd.read_csv(config.local_dataset_path)
    df_clean = df.dropna(subset=[config.column_local_image_path])
    print(f"  Total samples: {len(df_clean)}")
    
    # Create dataset
    dataset = CustomDataset(df_clean)
    
    # Create subset
    subset_size = min(subset_size, len(dataset))
    train_size = int(0.8 * subset_size)
    val_size = subset_size - train_size
    
    np.random.seed(42)
    subset_indices = np.random.choice(len(dataset), subset_size, replace=False)
    subset_dataset = torch.utils.data.Subset(dataset, subset_indices)
    
    train_dataset, val_dataset = random_split(
        subset_dataset, 
        [train_size, val_size],
        generator=torch.Generator().manual_seed(42)
    )
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset, 
        batch_size=batch_size, 
        shuffle=True, 
        num_workers=2, 
        pin_memory=True if torch.cuda.is_available() else False
    )
    val_loader = DataLoader(
        val_dataset, 
        batch_size=batch_size, 
        shuffle=False, 
        num_workers=2, 
        pin_memory=True if torch.cuda.is_available() else False
    )
    
    print(f"  Train: {len(train_dataset)} samples")
    print(f"  Val: {len(val_dataset)} samples")
    
    # Load feature models
    print(f"\n🔧 Loading feature models...")
    feature_models = load_models()
    
    # Load main model
    print(f"\n📦 Loading main model...")
    clip_model = CLIPModel_transformers.from_pretrained(
        'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'
    )
    # Frozen reference CLIP for text-space regularization (helps cross-domain generalization)
    reference_clip = CLIPModel_transformers.from_pretrained(
        'laion/CLIP-ViT-B-32-laion2B-s34B-b79K'
    )
    
    # Optionally load previous checkpoint
    if os.path.exists(config.main_model_path):
        user_input = input(f"\n⚠️  Found existing checkpoint at {config.main_model_path}. Load it? (y/n): ")
        if user_input.lower() == 'y':
            print(f"  Loading checkpoint...")
            checkpoint = torch.load(config.main_model_path, map_location=config.device)
            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
                clip_model.load_state_dict(checkpoint['model_state_dict'])
                print(f"  ✅ Checkpoint loaded from epoch {checkpoint.get('epoch', '?')}")
            else:
                clip_model.load_state_dict(checkpoint)
                print(f"  ✅ Checkpoint loaded")
        else:
            print(f"  Starting from pretrained model")
    else:
        print(f"  Starting from pretrained model")
    
    clip_model = clip_model.to(config.device)
    reference_clip = reference_clip.to(config.device)
    reference_clip.eval()
    for param in reference_clip.parameters():
        param.requires_grad = False
    
    # Train model with custom training function that uses weight_decay
    print(f"\n🎯 Starting training...")
    print(f"\n" + "="*80)
    
    # We need to modify the train_model function to accept weight_decay
    # For now, we'll use a modified version
    model = clip_model.to(config.device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', patience=3, factor=0.5
    )
    
    from transformers import CLIPProcessor
    from tqdm import tqdm
    from main_model import train_one_epoch, valid_one_epoch
    import matplotlib.pyplot as plt
    
    train_losses = []
    val_losses = []
    best_val_loss = float('inf')
    patience_counter = 0
    
    processor = CLIPProcessor.from_pretrained('laion/CLIP-ViT-B-32-laion2B-s34B-b79K')
    epoch_pbar = tqdm(range(num_epochs), desc="Training Progress", position=0)
    
    for epoch in epoch_pbar:
        epoch_pbar.set_description(f"Epoch {epoch+1}/{num_epochs}")
        
        # Training
        color_model = feature_models[config.color_column]
        hierarchy_model = feature_models[config.hierarchy_column]
        train_loss, align_metrics = train_one_epoch(
            model, train_loader, optimizer, feature_models, color_model, hierarchy_model,
            config.device, processor, temperature, alignment_weight,
            reference_model=reference_clip, reference_weight=0.1
        )
        train_losses.append(train_loss)
        
        # Validation
        val_loss = valid_one_epoch(
            model, val_loader, feature_models, config.device, processor, 
            temperature=temperature, alignment_weight=alignment_weight,
            reference_model=reference_clip, reference_weight=0.1
        )
        val_losses.append(val_loss)
        
        # Learning rate scheduling
        scheduler.step(val_loss)
        
        # Update progress bar
        epoch_pbar.set_postfix({
            'Train Loss': f'{train_loss:.4f}',
            'Val Loss': f'{val_loss:.4f}',
            'LR': f'{optimizer.param_groups[0]["lr"]:.2e}',
            'Best Val': f'{best_val_loss:.4f}'
        })
        
        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            
            # Save checkpoint
            save_path = config.main_model_path.replace('.pt', '_best_optuna.pt')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'train_loss': train_loss,
                'val_loss': val_loss,
                'best_val_loss': best_val_loss,
                'hyperparameters': {
                    'learning_rate': learning_rate,
                    'temperature': temperature,
                    'alignment_weight': alignment_weight,
                    'weight_decay': weight_decay,
                }
            }, save_path)
            print(f"\n💾 Best model saved at epoch {epoch+1}")
        else:
            patience_counter += 1
        
        # Early stopping
        if use_early_stopping and patience_counter >= patience:
            print(f"\n🛑 Early stopping triggered after {patience_counter} epochs without improvement")
            break
    
    # Plot training curves
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss', color='blue', linewidth=2)
    plt.plot(val_losses, label='Val Loss', color='red', linewidth=2)
    plt.title('Training and Validation Loss (Optimized)', fontsize=14, fontweight='bold')
    plt.xlabel('Epoch', fontsize=12)
    plt.ylabel('Loss', fontsize=12)
    plt.legend(fontsize=11)
    plt.grid(True, alpha=0.3)
    
    plt.subplot(1, 2, 2)
    gap = [train_losses[i] - val_losses[i] for i in range(len(train_losses))]
    plt.plot(gap, label='Train-Val Gap', color='purple', linewidth=2)
    plt.axhline(y=0, color='black', linestyle='--', alpha=0.3)
    plt.title('Overfitting Gap (Optimized)', fontsize=14, fontweight='bold')
    plt.xlabel('Epoch', fontsize=12)
    plt.ylabel('Train Loss - Val Loss', fontsize=12)
    plt.legend(fontsize=11)
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('training_curves_optimized.png', dpi=300, bbox_inches='tight')
    plt.close()
    
    print("\n" + "="*80)
    print("✅ Training completed!")
    print(f"  Best model: {save_path}")
    print(f"  Training curves: training_curves_optimized.png")
    print("\n📊 Final results:")
    print(f"  Last train loss: {train_losses[-1]:.4f}")
    print(f"  Last validation loss: {val_losses[-1]:.4f}")
    print(f"  Best validation loss: {best_val_loss:.4f}")
    print(f"  Overfitting gap: {train_losses[-1] - val_losses[-1]:.4f}")
    print("="*80)
    
    return train_losses, val_losses

def main():
    """
    Main function - Uses best parameters from Optuna optimization.
    """
    print("\n" + "="*80)
    print("🚀 Training with Best Optuna Hyperparameters")
    print("="*80)
    
    # Best hyperparameters from Optuna optimization (Trial 29 - Best validation loss: 0.1129)
    # Source: optuna_results.txt
    BEST_PARAMS = {
        'learning_rate': 1.42e-05,      # From Optuna (best trial)
        'temperature': 0.0503,           # From Optuna (best trial)
        'alignment_weight': 0.5639,      # From Optuna (best trial)
        'weight_decay': 2.76e-05,        # From Optuna (best trial)
        'num_epochs': 20,
        'batch_size': 32,
        'subset_size': 20000,           # Increased for better generalization
        'patience': 7
    }
    
    print(f"\n✅ Using optimized hyperparameters from Optuna:")
    print(f"   Learning rate: {BEST_PARAMS['learning_rate']:.2e}")
    print(f"   Temperature: {BEST_PARAMS['temperature']:.4f}")
    print(f"   Alignment weight: {BEST_PARAMS['alignment_weight']:.4f}")
    print(f"   Weight decay: {BEST_PARAMS['weight_decay']:.2e}")
    print(f"   Expected validation loss: ~0.1129 (from Optuna)\n")
    
    train_with_best_params(**BEST_PARAMS)

if __name__ == "__main__":
    main()