#!/usr/bin/env python3
"""
Example: Complete LiveCodeBench evaluation workflow

This script demonstrates the complete workflow:
1. Train steering vectors on LiveCodeBench v1
2. Save model with steering vectors
3. Prepare evaluation instructions for LiveCodeBench v2

Usage:
    python examples/livecodebench_evaluation_example.py
"""

import os
import json
import subprocess
from pathlib import Path

from wisent_guard.core.pipelines import ExperimentRunner
from wisent_guard.core.pipelines.steering_trainer import TrainingConfig
from wisent_guard.core.pipelines.experiment_runner import ExperimentConfig
from wisent_guard.core.models import SteeringCompatibleModel


def main():
    print("🚀 LiveCodeBench Evaluation Example")
    print("===================================\n")
    
    # Step 1: Train steering vectors
    print("Step 1: Training steering vectors on LiveCodeBench v1...")
    
    # Configure training (quick test with 5 problems)
    training_config = TrainingConfig(
        model_name='distilgpt2',
        target_layers=[5],  # DistilGPT2 last layer
        steering_method='CAA',
        batch_size=2,
        device='auto'
    )
    
    experiment_config = ExperimentConfig(
        train_version='release_v1',
        eval_version='release_v2',
        version_split_type='new_only',  # Contamination-free
        data_limit=5,  # Quick test with 5 problems
        pairs_per_problem=1,
        training_config=training_config,
        experiment_name='livecodebench_example',
        output_directory='./example_experiments'
    )
    
    try:
        runner = ExperimentRunner(experiment_config)
        results = runner.run_experiment()
        
        print(f"✅ Training completed!")
        print(f"   Experiment saved to: {experiment_config.output_directory}")
        print(f"   Steering vectors: {len(results.training_results.steering_vectors)} vectors")
        print(f"   Training time: {results.training_results.training_time:.2f} seconds")
        
        # Get the experiment directory
        experiment_dirs = list(Path(experiment_config.output_directory).glob("livecodebench_example_*"))
        if not experiment_dirs:
            print("❌ No experiment directory found!")
            return
        
        latest_experiment = max(experiment_dirs, key=lambda p: p.stat().st_mtime)
        print(f"   Latest experiment: {latest_experiment}")
        
    except Exception as e:
        print(f"❌ Training failed: {e}")
        print("   This might be due to missing dependencies or network issues.")
        print("   For a quick test, you can skip training and use a dummy model.")
        return
    
    # Step 2: Save model for evaluation
    print("\nStep 2: Saving model for evaluation...")
    
    try:
        # Load the trained model
        model = SteeringCompatibleModel.from_pretrained_with_steering(
            'distilgpt2',
            steering_directory=str(latest_experiment)
        )
        
        # Save in evaluation format
        export_path = './example_models/distilgpt2_with_steering'
        model.save_pretrained(export_path)
        
        print(f"✅ Model saved to: {export_path}")
        print(f"   Model info: {model.get_steering_info()}")
        
    except Exception as e:
        print(f"❌ Model saving failed: {e}")
        return
    
    # Step 3: Create evaluation instructions
    print("\nStep 3: Creating evaluation instructions...")
    
    evaluation_script = f"""#!/bin/bash
# LiveCodeBench Evaluation Script
# Auto-generated by livecodebench_evaluation_example.py

echo "🔍 LiveCodeBench Evaluation"
echo "=========================="

# Check if LiveCodeBench is installed
if ! python -c "import lcb_runner" 2>/dev/null; then
    echo "❌ LiveCodeBench not installed. Installing..."
    echo "Please run the following commands:"
    echo "  git clone https://github.com/LiveCodeBench/LiveCodeBench.git"
    echo "  cd LiveCodeBench"
    echo "  pip install -e ."
    echo ""
    echo "Then run this script again."
    exit 1
fi

echo "✅ LiveCodeBench found"

# Create results directory
mkdir -p ./example_results

echo "🚀 Starting evaluation..."
echo "Model: {export_path}"
echo "Dataset: LiveCodeBench v2"
echo "Problems: Limited to 10 for quick test"

# Run evaluation
python -m lcb_runner.runner.main \\
    --model {export_path} \\
    --scenario code_generation \\
    --release_version v2 \\
    --temperature 0.2 \\
    --max_length_generation 512 \\
    --n_samples 1 \\
    --limit 10 \\
    --output_dir ./example_results \\
    --use_cache \\
    --continue_existing \\
    --evaluate

echo "✅ Evaluation completed!"
echo "Results saved to: ./example_results"

# Show quick summary
if [ -f "./example_results/results.json" ]; then
    echo ""
    echo "📊 Quick Results Summary:"
    python -c "
import json
try:
    with open('./example_results/results.json', 'r') as f:
        results = json.load(f)
    print(f'Pass@1: {{results.get(\"pass@1\", \"N/A\")}}')
    print(f'Total problems: {{results.get(\"total\", \"N/A\")}}')
except Exception as e:
    print(f'Could not load results: {{e}}')
"
fi
"""
    
    # Save evaluation script
    script_path = "./run_livecodebench_evaluation.sh"
    with open(script_path, 'w') as f:
        f.write(evaluation_script)
    
    # Make script executable
    os.chmod(script_path, 0o755)
    
    print(f"✅ Evaluation script created: {script_path}")
    
    # Step 4: Show next steps
    print("\n🎯 Next Steps:")
    print("=============")
    print("1. Install LiveCodeBench if not already installed:")
    print("   git clone https://github.com/LiveCodeBench/LiveCodeBench.git")
    print("   cd LiveCodeBench && pip install -e .")
    print("")
    print("2. Run the evaluation:")
    print(f"   ./{script_path}")
    print("")
    print("3. Compare with baseline:")
    print("   python -m lcb_runner.runner.main --model distilgpt2 --scenario code_generation --release_version v2 --limit 10 --evaluate")
    print("")
    print("4. For production use:")
    print("   - Increase data_limit to 100+ for training")
    print("   - Remove --limit flag for full evaluation")
    print("   - Try different target_layers and steering methods")
    print("")
    print("📁 Files created:")
    print(f"   - Model: {export_path}")
    print(f"   - Evaluation script: {script_path}")
    print(f"   - Training results: {latest_experiment}")
    print("")
    print("🎉 Example completed successfully!")


if __name__ == "__main__":
    main()