#!/usr/bin/env python3
"""
Run all models sequentially for comparison
"""

import os
import sys
import runpy
import time
from datetime import datetime

# Add parent directory to path
sys.path.insert(0, '..')

# List of all available models
MODELS = [
    "spikernn",
    "itransformer",
    "spikformer",
    "ispikeformer",
    "timemixer",
    "dlinear"
]

def run_model(model_name: str):
    """Run experiment for a specific model"""
    print(f"\n{'='*60}")
    print(f"🚀 Starting experiment for {model_name.upper()}")
    print(f"⏰ Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{'='*60}")

    start_time = time.time()

    # Configure arguments for this model
    sys.argv = [
        "main.py",
        "--dataset", "electricity",
        "--num_clients", "20",
        "--max_features", "20",
        "--model_name", model_name,
        "--window", "96",
        "--horizon", "48",
        "--batch_size", "128",
        "--global_epochs", "50",
        "--local_epochs", "1",
        "--lr", "1e-3",
        "--optimizer", "adam",
        "--seed", "42",
        "--device", "cuda",
        "--experiment", f"electricity_{model_name}_50global_96-48_dim128",
        "--frac", "1.0",
        "--dim", "128",
        "--early_stop",
        "--patience", "10"
    ]

    try:
        # Change to parent directory and run
        os.chdir("..")
        runpy.run_path("main.py", run_name="__main__")

        elapsed = time.time() - start_time
        print(f"\n✅ {model_name.upper()} completed successfully!")
        print(f"⏱️  Total time: {elapsed:.1f} seconds ({elapsed/60:.1f} minutes)")

    except Exception as e:
        elapsed = time.time() - start_time
        print(f"\n❌ {model_name.upper()} failed after {elapsed:.1f} seconds")
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # Return to experiments directory
        os.chdir("experiments")

def main():
    """Run all models"""
    total_start = time.time()

    print("🔬 FedQAP Multi-Model Comparison Experiment")
    print(f"📅 Started at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"🎯 Models to run: {', '.join(MODELS)}")

    results = {}

    for i, model in enumerate(MODELS, 1):
        print(f"\n📊 Progress: {i}/{len(MODELS)} models")
        model_start = time.time()

        try:
            run_model(model)
            results[model] = "✅ Success"
        except Exception as e:
            results[model] = f"❌ Failed: {str(e)[:50]}..."

        model_time = time.time() - model_start
        print(f"🔄 Model {model} finished in {model_time/60:.1f} minutes")

    # Summary
    total_time = time.time() - total_start
    print(f"\n{'='*80}")
    print("📊 EXPERIMENT SUMMARY")
    print(f"{'='*80}")
    print(f"🕐 Total time: {total_time/60:.1f} minutes ({total_time/3600:.2f} hours)")
    print(f"📅 Finished at: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("\n📋 Results:")
    for model, status in results.items():
        print(f"  {model:15s}: {status}")

    print(f"\n💾 Check the 'output/' directory for detailed results")
    print(f"📊 Use TensorBoard to compare results: tensorboard --logdir=output/")

if __name__ == "__main__":
    main()