src/training_pipeline.py

"""
ViralTrack Predictor - Spotify Popularity Prediction
Predicts track popularity (0-100) using audio features + metadata
"""

import os
import logging
import torch
from pathlib import Path
from typing import Dict, Any, List

from dotenv import load_dotenv
from omegaconf import OmegaConf

from datasets import load_dataset, DatasetDict
from transformers import (
    AutoTokenizer,
    AutoModelForSequenceClassification,
    TrainingArguments,
    Trainer,
    DataCollatorWithPadding,
)
from transformers.trainer_callback import TrainerCallback
import numpy as np
from tqdm import tqdm

# Suppress HTTP logs from transformers/datasets
logging.getLogger("filelock").setLevel(logging.ERROR)
logging.getLogger("urllib3").setLevel(logging.ERROR)
logging.getLogger("huggingface_hub").setLevel(logging.ERROR)
logging.getLogger("datasets").setLevel(logging.ERROR)
logging.getLogger("transformers").setLevel(logging.ERROR)
logging.getLogger("torch").setLevel(logging.ERROR)

load_dotenv()

# Setup logging - cleaner format
logging.basicConfig(
    level=logging.ERROR,
    format='%(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)


class PerformanceCallback(TrainerCallback):
    """Track metrics per epoch with clean output"""
    def __init__(self):
        self.epoch_metrics = []

    def on_epoch_end(self, args, state, control, metrics=None, **kwargs):
        if metrics:
            self.epoch_metrics.append({'epoch': state.epoch, 'metrics': metrics.copy()})
            # Clean epoch summary
            print(f"\n{'='*50}")
            print(f"✅ Epoch {state.epoch:.0f}/{args.num_train_epochs:.0f} Complete")
            print(f"{'='*50}")
            key_metrics = ['loss', 'mae', 'r2']
            for k in key_metrics:
                full_key = f'eval_{k}' if k != 'loss' else k
                if full_key in metrics:
                    val = metrics[full_key]
                    if isinstance(val, (int, float)):
                        print(f"  {k.upper():<15} {val:.4f}")
            print(f"{'='*50}\n")
        return control


def load_config(config_name: str = 'config'):
    """Load YAML config"""
    conf = OmegaConf.load(f'configs/{config_name}.yaml')
    return OmegaConf.to_container(conf, resolve=True)


def compute_metrics(eval_pred, metric_names=['mse', 'mae', 'r2']):
    """Compute regression metrics using scikit-learn"""
    from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
    
    predictions, labels = eval_pred

    if isinstance(predictions, tuple):
        predictions = predictions[0]

    predictions = predictions.squeeze(-1)
    labels = labels.squeeze(-1)

    results = {
        'mse': mean_squared_error(labels, predictions),
        'mae': mean_absolute_error(labels, predictions),
        'r2': r2_score(labels, predictions),
    }

    return results


def get_feature_importance(model, tokenizer, feature_columns, device='cpu'):
    """
    Analyze feature importance by perturbing inputs
    Returns recommendations for improving popularity
    """
    logger.info("\n🔍 Analyzing Feature Importance...")
    
    # Baseline feature importance (correlation-based approximation)
    importance = {}
    for col in feature_columns:
        if col in ['danceability', 'energy', 'valence', 'acousticness', 
                   'instrumentalness', 'liveness', 'speechiness']:
            # These are audio features - we'll use statistical analysis
            importance[col] = {
                'type': 'audio_feature',
                'range': [0.0, 1.0],
                'description': get_feature_description(col)
            }
        elif col in ['tempo', 'duration_ms']:
            importance[col] = {
                'type': 'audio_feature',
                'range': [0, float('inf')],
                'description': get_feature_description(col)
            }
        else:
            importance[col] = {
                'type': 'text_feature',
                'description': get_feature_description(col)
            }
    
    return importance


def get_feature_description(feature: str) -> str:
    """Get human-readable description of audio features"""
    descriptions = {
        'track_name': 'Song title text',
        'artists': 'Artist name(s)',
        'danceability': 'How suitable for dancing (0-1)',
        'energy': 'Intensity and activity level (0-1)',
        'valence': 'Musical positiveness/happiness (0-1)',
        'tempo': 'Speed in BPM',
        'duration_ms': 'Song length in milliseconds',
        'acousticness': 'Acoustic vs electronic (0-1)',
        'instrumentalness': 'No vocals (0-1)',
        'liveness': 'Live performance probability (0-1)',
        'speechiness': 'Spoken word probability (0-1)',
    }
    return descriptions.get(feature, 'Unknown feature')


def generate_recommendations(prediction: float, features: Dict[str, float]) -> List[str]:
    """Generate actionable recommendations based on prediction and features"""
    recommendations = []
    
    if prediction < 50:
        recommendations.append("⚠️ Predicted popularity is LOW - consider these changes:")
    elif prediction < 70:
        recommendations.append("📈 Predicted popularity is MODERATE - optimization opportunities:")
    else:
        recommendations.append("🔥 Predicted popularity is HIGH - track has viral potential!")
    
    # Feature-specific recommendations
    if features.get('duration_ms', 0) > 200000:  # > 3:20
        recommendations.append("  📍 Song is long (>3:20) - consider shorter version for TikTok/Reels")
    
    if features.get('energy', 0) < 0.4:
        recommendations.append("  ⚡ Low energy - consider adding more dynamic elements")
    
    if features.get('danceability', 0) < 0.5:
        recommendations.append("  💃 Low danceability - may not perform well on social platforms")
    
    if features.get('valence', 0) > 0.8:
        recommendations.append("  😊 Very positive mood - great for playlists/morning vibes")
    
    if features.get('acousticness', 0) > 0.7:
        recommendations.append("  🎸 Highly acoustic - consider production polish for mainstream appeal")
    
    if features.get('speechiness', 0) > 0.3:
        recommendations.append("  🎤 High speechiness - may work well for podcast/hip-hop audiences")
    
    return recommendations


def train(config_name: str = 'config', epochs: int = None, batch_size: int = None, num_samples: int = None):
    """Main training function for regression"""
    print(f"\n{'🎵'*30}")
    print("  VIRALTRACK PREDICTOR - Spotify Popularity Prediction")
    print(f"{'🎵'*30}\n")

    # Load config
    cfg = load_config(config_name)
    print(f"📋 Config: {config_name}\n")

    # Override config with CLI args if provided
    if epochs is not None:
        cfg['training']['epochs'] = epochs
    if batch_size is not None:
        cfg['training']['batch_size'] = batch_size
    if num_samples is not None:
        cfg['dataset']['num_samples'] = num_samples

    # Setup HF auth
    hf_token = os.getenv("HF_TOKEN")
    if hf_token:
        print("✓ Hugging Face token loaded\n")

    # Load dataset
    ds_cfg = cfg['dataset']
    print(f"📊 Dataset: {ds_cfg['name']}")

    load_kwargs = {'path': ds_cfg['name']}
    if ds_cfg.get('config'):
        load_kwargs['name'] = ds_cfg['config']

    dataset = load_dataset(**load_kwargs)

    if not isinstance(dataset, DatasetDict):
        dataset = dataset.train_test_split(test_size=0.2)
        tv = dataset['train'].train_test_split(test_size=0.1)
        dataset = DatasetDict({
            'train': tv['train'],
            'validation': tv['test'],
            'test': dataset['test']
        })

    # Subsample if requested
    num_samples = ds_cfg.get('num_samples')
    if num_samples is not None:
        print(f"⚡ Using subset: {num_samples} samples (for faster testing)")
        if len(dataset['train']) > num_samples:
            dataset['train'] = dataset['train'].select(range(num_samples))
        if 'validation' in dataset and len(dataset['validation']) > num_samples // 10:
            dataset['validation'] = dataset['validation'].select(range(min(num_samples // 10, len(dataset['validation']))))
        if 'test' in dataset and len(dataset['test']) > num_samples // 10:
            dataset['test'] = dataset['test'].select(range(min(num_samples // 10, len(dataset['test']))))

    print(f"   ├─ Train:      {len(dataset['train']):,} samples")
    if 'validation' in dataset:
        print(f"   ├─ Validation: {len(dataset['validation']):,} samples")
    if 'test' in dataset:
        print(f"   └─ Test:       {len(dataset['test']):,} samples")
    print()

    # Load tokenizer and model
    model_cfg = cfg['model']
    feature_columns = ds_cfg.get('feature_columns', ['text'])
    target_col = ds_cfg.get('target_column', 'label')
    max_length = ds_cfg.get('max_length', 512)

    print(f"🤖 Model: {model_cfg['name']}")
    print(f"   Target: {target_col} (regression)")
    print(f"   Features: {len(feature_columns)} columns\n")

    print("⏳ Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(model_cfg['name'])

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    print("⏳ Loading model weights...\n")

    with tqdm(total=100, desc="Loading weights", bar_format='{desc}: |{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]') as pbar:
        model = AutoModelForSequenceClassification.from_pretrained(
            model_cfg['name'],
            num_labels=1,
            problem_type="regression",
            trust_remote_code=model_cfg.get('trust_remote_code', False),
            ignore_mismatched_sizes=True,
        )
        pbar.update(100)

    print(f"\n📦 Model: {model.__class__.__name__}")
    print(f"   Source: {model_cfg['name']}")
    print(f"   Params: {sum(p.numel() for p in model.parameters()):,}")
    print(f"   ✓ Ready for training\n")

    # Tokenize - combine text features and normalize audio features
    print("🔧 Preprocessing data...")

    def normalize_features(ex):
        # Combine text features
        text_parts = []
        for col in ['track_name', 'artists']:
            if col in ex and ex[col] is not None:
                text_parts.append(str(ex[col]))
        combined_text = ' '.join(text_parts) if text_parts else ""
        
        # Get numerical features
        numerical = []
        for col in feature_columns:
            if col in ex and col not in ['track_name', 'artists']:
                val = ex[col]
                if val is not None:
                    numerical.append(f"{col}:{float(val):.3f}")
        
        # Combine all into text for the model
        full_text = f"{combined_text} | {' '.join(numerical)}"
        
        tokenized = tokenizer(full_text, padding='max_length', truncation=True, max_length=max_length)
        
        # Set regression target (normalize to 0-1 range for stability)
        tokenized['labels'] = [float(ex[target_col]) / 100.0]
        
        return tokenized

    tokenized = {}
    for split in dataset.keys():
        tokenized[split] = dataset[split].map(
            normalize_features, batched=False, remove_columns=dataset[split].column_names
        )

    dataset = DatasetDict(tokenized)
    print("✓ Preprocessing complete\n")

    # Training args
    train_cfg = cfg['training']
    hw_cfg = cfg.get('hardware', {})
    out_cfg = cfg.get('output', {})

    output_dir = Path(out_cfg.get('dir', './outputs'))
    output_dir.mkdir(parents=True, exist_ok=True)

    print(f"{'='*50}")
    print("🚀 TRAINING CONFIGURATION")
    print(f"{'='*50}")
    print(f"  Epochs:        {train_cfg['epochs']}")
    print(f"  Batch size:    {train_cfg['batch_size']}")
    print(f"  Learning rate: {train_cfg['learning_rate']}")
    print(f"  Output dir:    {output_dir}")
    print(f"{'='*50}\n")

    # Split data for validation
    has_validation = 'validation' in dataset
    if not has_validation:
        print("  Creating validation split...")
        train_val = dataset['train'].train_test_split(test_size=0.1)
        dataset = DatasetDict({
            'train': train_val['train'],
            'validation': train_val['test']
        })
        has_validation = True

    print(f"📈 Training: {len(dataset['train']):,} samples")
    print(f"   Validating: {len(dataset['validation']):,} samples\n")

    training_args = TrainingArguments(
        output_dir=str(output_dir),
        num_train_epochs=train_cfg['epochs'],
        per_device_train_batch_size=train_cfg['batch_size'],
        per_device_eval_batch_size=train_cfg['batch_size'],
        learning_rate=train_cfg['learning_rate'],
        weight_decay=train_cfg.get('weight_decay', 0.01),
        warmup_steps=100,
        fp16=hw_cfg.get('mixed_precision', 'fp16') == 'fp16',
        save_strategy='epoch',
        logging_steps=out_cfg.get('logging_steps', 10),
        eval_strategy='epoch',
        load_best_model_at_end=True,
        metric_for_best_model='loss',
        greater_is_better=False,
        report_to='none',
        disable_tqdm=False,
        dataloader_pin_memory=False,
    )

    # Train
    print("⏳ Starting training...\n")
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset['train'],
        eval_dataset=dataset['validation'],
        processing_class=tokenizer,
        data_collator=DataCollatorWithPadding(tokenizer),
        compute_metrics=lambda x: compute_metrics(x, cfg.get('evaluation', {}).get('metrics', ['mse', 'mae', 'r2'])),
        callbacks=[PerformanceCallback()],
    )

    trainer.train()

    # Evaluate
    print(f"\n{'='*50}")
    print("📈 EVALUATION")
    print(f"{'='*50}")
    if 'test' in dataset:
        eval_dataset = dataset['test']
    else:
        eval_dataset = dataset['validation']
    metrics = trainer.evaluate(eval_dataset)

    print(f"\n=== Final Metrics ===")
    for k, v in metrics.items():
        if isinstance(v, (int, float)):
            if k in ['eval_mse', 'eval_mae']:
                print(f"  {k:<15} {v * 100:.4f} (on 0-100 scale)")
            elif k == 'eval_r2':
                print(f"  {k:<15} {v:.4f}")
            else:
                print(f"  {k:<15} {v:.4f}")
    print(f"{'='*50}\n")

    # Save
    model_path = output_dir  # Save directly to output_dir (e.g., ./model)
    model.save_pretrained(str(model_path))
    tokenizer.save_pretrained(str(model_path))
    print(f"💾 Model saved to: {model_path}\n")

    # Feature importance analysis
    feature_importance = get_feature_importance(model, tokenizer, feature_columns)
    print(f"{'='*50}")
    print("📊 FEATURE ANALYSIS")
    print(f"{'='*50}")
    for feat, info in feature_importance.items():
        print(f"  {feat}: {info['description']}")
    print(f"{'='*50}\n")

    print(f"{'🎵'*30}")
    print("  ✅ TRAINING COMPLETE!")
    print(f"{'🎵'*30}")
    print("   Model can predict track popularity and provide recommendations\n")

    return {'metrics': metrics, 'model_path': str(model_path)}


if __name__ == '__main__':
    import argparse
    
    parser = argparse.ArgumentParser(description='Train ViralTrack Predictor')
    parser.add_argument('config', nargs='?', default='config', help='Config file name (default: config)')
    parser.add_argument('--epochs', type=int, default=None, help='Number of training epochs')
    parser.add_argument('--batch_size', type=int, default=None, help='Training batch size')
    parser.add_argument('--num_samples', type=int, default=None, help='Number of samples to use (for faster testing)')
    
    args = parser.parse_args()
    train(args.config, epochs=args.epochs, batch_size=args.batch_size, num_samples=args.num_samples)