#!/usr/bin/env python3
"""
Simple Transformer Two-Tower Training Script

Streamlined training without advanced callbacks and complex configurations.
"""

import os
import sys
import time
import pickle
import argparse

# CPU-specific TensorFlow optimizations
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # Reduce TF logging
os.environ['OMP_NUM_THREADS'] = '4'       # Limit OpenMP threads
os.environ['TF_NUM_INTEROP_THREADS'] = '2'  # Inter-op parallelism
os.environ['TF_NUM_INTRAOP_THREADS'] = '4'  # Intra-op parallelism

import tensorflow as tf

# Configure TensorFlow for CPU optimization
tf.config.threading.set_inter_op_parallelism_threads(2)
tf.config.threading.set_intra_op_parallelism_threads(4)

# Disable GPU if available (force CPU)
tf.config.set_visible_devices([], 'GPU')

# Enable CPU optimizations
tf.config.optimizer.set_jit(True)  # Enable XLA compilation for CPU

# Memory management for CPU
import gc
import psutil

def print_memory_usage(stage=""):
    """Print current memory usage with system-wide context"""
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    process_mb = memory_info.rss / 1024 / 1024

    # System memory
    system_memory = psutil.virtual_memory()
    available_gb = system_memory.available / 1024 / 1024 / 1024
    used_percent = system_memory.percent

    stage_str = f" ({stage})" if stage else ""
    print(f"🧠 Memory{stage_str}: {process_mb:.1f}MB process, {used_percent:.1f}% system, {available_gb:.1f}GB available")

    # Warning if memory is getting high
    if used_percent > 85:
        print(f"   ⚠️  HIGH MEMORY USAGE WARNING: {used_percent:.1f}%")
        print(f"   🧹 Running aggressive cleanup...")
        cleanup_memory()
    elif used_percent > 75:
        print(f"   ⚡ Memory usage elevated: {used_percent:.1f}%")

    return process_mb, used_percent

def cleanup_memory():
    """Aggressive garbage collection and memory cleanup"""
    gc.collect()
    tf.keras.backend.clear_session()
    # Force additional cleanup
    gc.collect()

def monitor_memory_limit(limit_gb=25):
    """Check if we're approaching memory limit and take action"""
    system_memory = psutil.virtual_memory()
    used_gb = (system_memory.total - system_memory.available) / 1024 / 1024 / 1024

    if used_gb > limit_gb:
        print(f"⚠️  Memory limit exceeded: {used_gb:.1f}GB > {limit_gb}GB")
        print(f"🧹 Running emergency cleanup...")
        cleanup_memory()
        return True
    return False

# Add src to path
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

from src.training.transformer_item_pretraining import TransformerItemTowerPretrainer
from src.training.transformer_joint_training import TransformerJointTrainer
from src.preprocessing.transformer_user_data_preparation import TransformerUserDatasetCreator
from src.preprocessing.data_loader import DataProcessor


def run_phase_1():
    """Phase 1: Simple item tower pretraining."""
    print("=" * 60)
    print("🔸 PHASE 1: ITEM TOWER PRETRAINING")
    print("=" * 60)

    start_time = time.time()
    pretrainer = TransformerItemTowerPretrainer(
        embedding_dim=128,
        learning_rate=0.001,
        artifacts_prefix="transformer_"
    )

    # Prepare data
    dataset, data_processor, price_normalizer = pretrainer.prepare_data()

    # Build model
    model = pretrainer.build_model(
        item_vocab_size=len(data_processor.item_vocab),
        category_vocab_size=len(data_processor.category_vocab),
        category_code_vocab_size=len(data_processor.category_code_vocab),
        brand_vocab_size=len(data_processor.brand_vocab),
        price_normalizer=price_normalizer
    )

    # Simple training (updated method signature)
    history = pretrainer.train(
        dataset=dataset,
        epochs=15,
        validation_split=0.2
    )

    # Generate and save embeddings
    item_embeddings = pretrainer.generate_item_embeddings(dataset, data_processor)
    pretrainer.save_model()
    pretrainer.save_embeddings(item_embeddings)

    # Save vocabularies with transformer prefix (missing step from original)
    import pickle
    os.makedirs("src/artifacts/transformers", exist_ok=True)
    with open("src/artifacts/transformers/transformer_vocabularies.pkl", 'wb') as f:
        pickle.dump({
            'item_vocab': data_processor.item_vocab,
            'category_vocab': data_processor.category_vocab,
            'category_code_vocab': data_processor.category_code_vocab,
            'brand_vocab': data_processor.brand_vocab,
            'user_vocab': data_processor.user_vocab
        }, f)
    print("✅ Saved transformer vocabularies to src/artifacts/transformers/transformer_vocabularies.pkl")

    phase1_time = time.time() - start_time
    print(f"✅ Phase 1 completed in {phase1_time:.2f} seconds")

    # Cleanup memory after Phase 1
    print_memory_usage("after Phase 1")
    cleanup_memory()
    print("🧹 Memory cleaned up after Phase 1")

    return history


def run_phase_2():
    """Phase 2: Simple data preparation."""
    print("=" * 60)
    print("🔸 PHASE 2: DATA PREPARATION")
    print("=" * 60)

    start_time = time.time()
    print_memory_usage("start of Phase 2")

    # Check memory limit before starting
    if monitor_memory_limit(20):
        print("⚠️  Memory usage high before Phase 2, continuing with caution...")

    # Load data and vocabularies from Phase 1 first
    data_processor = DataProcessor()
    items_df, users_df, interactions_df = data_processor.load_data()

    # Load vocabularies created in Phase 1 instead of rebuilding
    print("Loading vocabularies from Phase 1...")
    with open("src/artifacts/transformers/transformer_vocabularies.pkl", 'rb') as f:
        vocab_data = pickle.load(f)
    data_processor.item_vocab = vocab_data['item_vocab']
    data_processor.category_vocab = vocab_data['category_vocab']
    data_processor.category_code_vocab = vocab_data['category_code_vocab']
    data_processor.brand_vocab = vocab_data['brand_vocab']
    data_processor.user_vocab = vocab_data['user_vocab']
    print("✅ Loaded transformer vocabularies from Phase 1")

    # Create dataset creator with existing data processor to avoid rebuilding vocabularies
    dataset_creator = TransformerUserDatasetCreator(
        max_history_length=100,  # Limit sequence length to 100
        artifacts_prefix="transformer_"
    )
    # Override the internal data processor with our loaded one
    dataset_creator.data_processor = data_processor

    # Use a sample for memory efficiency
    print("Using sample data for memory efficiency...")
    sample_size = min(50000, len(interactions_df))  # Keep original 50k sample size
    sample_interactions = interactions_df.sample(n=sample_size, random_state=42)

    # Filter users and items to match sample
    user_ids = set(sample_interactions['user_id'])
    item_ids = set(sample_interactions['product_id'])
    sample_users = users_df[users_df['user_id'].isin(user_ids)]
    sample_items = items_df[items_df['product_id'].isin(item_ids)]

    print(f"Sample dataset: {len(sample_items)} items, {len(sample_users)} users, {len(sample_interactions)} interactions")

    # Load item embeddings from Phase 1
    item_embeddings = dataset_creator.load_item_embeddings()

    # Create proper 80/20 random split instead of broken temporal split
    from sklearn.model_selection import train_test_split
    train_interactions, val_interactions = train_test_split(
        sample_interactions,
        test_size=0.2,
        random_state=42
    )

    # Create training dataset
    training_features, max_sequence_length = dataset_creator.create_training_dataset(
        train_interactions, sample_items, sample_users, item_embeddings
    )

    # Create validation dataset
    validation_features, _ = dataset_creator.create_training_dataset(
        val_interactions, sample_items, sample_users, item_embeddings
    )

    # Save datasets
    os.makedirs("src/artifacts/transformers", exist_ok=True)
    with open("src/artifacts/transformers/transformer_training_features.pkl", 'wb') as f:
        pickle.dump(training_features, f)
    with open("src/artifacts/transformers/transformer_validation_features.pkl", 'wb') as f:
        pickle.dump(validation_features, f)

    phase2_time = time.time() - start_time
    print(f"✅ Phase 2 completed in {phase2_time:.2f} seconds")

    # Cleanup memory after Phase 2
    print_memory_usage("after Phase 2")
    cleanup_memory()
    print("🧹 Memory cleaned up after Phase 2")

    return training_features, validation_features


def run_phase_3():
    """Phase 3: Simple joint training."""
    print("=" * 60)
    print("🔸 PHASE 3: JOINT TRAINING")
    print("=" * 60)

    start_time = time.time()

    # Load training data with memory monitoring
    print("Loading training data...")
    print_memory_usage("before loading data")

    # Check memory before loading large datasets
    if monitor_memory_limit(15):
        print("⚠️  Memory usage high before loading datasets!")

    with open("src/artifacts/transformers/transformer_training_features.pkl", 'rb') as f:
        training_features = pickle.load(f)
    print_memory_usage("after training data")

    with open("src/artifacts/transformers/transformer_validation_features.pkl", 'rb') as f:
        validation_features = pickle.load(f)

    print("Data loaded:")
    print_memory_usage("after all data loaded")

    # Create simple trainer with improved configuration
    trainer = TransformerJointTrainer(
        embedding_dim=128,
        transformer_layers=2,  # Reduced for efficiency and stability
        transformer_heads=4,   # Reduced for efficiency and stability
        transformer_ff_dim=256,  # Reduced from 512 for better regularization
        user_learning_rate=0.0005,  # Optimized learning rate
        item_learning_rate=0.00005,  # 1/10 ratio for stable fine-tuning
        rating_weight=1.0,
        retrieval_weight=0.5,
        artifacts_prefix="transformer_"
    )

    # Build the model components
    print("🔧 Building transformer model...")

    # Load pre-trained item tower
    trainer.load_pre_trained_item_tower()

    # Get max sequence length from training features
    max_seq_len = training_features['item_history_embeddings'].shape[1]

    # Build user tower
    trainer.build_transformer_user_tower(max_sequence_length=max_seq_len)

    # Build complete two-tower model
    trainer.build_two_tower_model()

    # Simple training configuration with item tower freezing
    print("🚀 Starting joint training with item tower freezing...")
    print("🧊 Item tower will be frozen for first 6 epochs to prevent destabilization")

    try:
        # CPU-optimized training parameters with item tower freezing
        history = trainer.train(
            training_features=training_features,
            validation_features=validation_features,
            epochs=15,
            batch_size=64,  # Batch size reduced to 64 
            
            early_stopping_patience=5,
            item_freeze_epochs=6  # NEW: Freeze item tower for first 6 epochs
        )

        print(f"🏁 Training completed successfully!")
        print(f"   - Epochs trained: {len(history.get('total_loss', []))}")
        if 'val_total_loss' in history:
            print(f"   - Best validation loss: {min(history['val_total_loss']):.4f}")

            # Check overfitting control
            if len(history['val_total_loss']) > 0 and len(history['total_loss']) > 0:
                final_val_loss = history['val_total_loss'][-1]
                final_train_loss = history['total_loss'][-1]
                overfitting_ratio = final_val_loss / final_train_loss if final_train_loss > 0 else float('inf')
                print(f"   - Final train/val ratio: {overfitting_ratio:.2f}")

                if overfitting_ratio < 1.8:
                    print("   ✅ Overfitting successfully controlled!")
                else:
                    print("   ⚠️  Some overfitting detected - but training completed")

    except Exception as e:
        print(f"❌ Training failed: {e}")
        import traceback
        traceback.print_exc()
        return None

    # Save final model
    print("💾 Saving model...")
    trainer.save_model()

    # Save training history
    os.makedirs("src/artifacts/transformers", exist_ok=True)
    with open("src/artifacts/transformers/transformer_training_history.pkl", 'wb') as f:
        pickle.dump(history, f)

    phase3_time = time.time() - start_time
    print(f"✅ Phase 3 completed in {phase3_time:.2f} seconds")

    return history


def main():
    """Simple main training function."""
    parser = argparse.ArgumentParser(description='Simple Transformer Training')
    parser.add_argument('--phase', type=str, choices=['1', '2', '3', 'all'], default='all',
                       help='Training phase to run')

    args = parser.parse_args()

    print("🚀 SIMPLE TRANSFORMER TRAINING")
    print(f"Phase: {args.phase}")
    print("-" * 60)

    total_start_time = time.time()

    try:
        if args.phase in ['1', 'all']:
            run_phase_1()

        if args.phase in ['2', 'all']:
            run_phase_2()

        if args.phase in ['3', 'all']:
            run_phase_3()

        total_time = time.time() - total_start_time
        print(f"\n🎉 All phases completed in {total_time:.2f} seconds")

        print("\n🎯 ARCHITECTURAL IMPROVEMENTS APPLIED:")
        print("  ✅ Balanced demographic embeddings (all 8D)")
        print("  ✅ Fixed attention temperature (1.0 vs 0.1)")
        print("  ✅ Demographic variance loss for diversity")
        print("  ✅ Item tower freezing (6 epochs)")
        print("  ✅ Differential learning rates (1/10 ratio)")
        print("  ✅ Reduced model complexity for stability")
        print("\n🚀 Ready for inference with improved transformer model!")

    except Exception as e:
        print(f"\n❌ Training pipeline failed: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()