#!/usr/bin/env python3
"""
ÜBERMENSCHETIEN QUICK START
============================
One-command setup and training.

Usage:
  python quickstart.py --full          # Run full pipeline
  python quickstart.py --train-dense   # Just dense training
  python quickstart.py --train-cfhot   # Just CF-HoT heads
  python quickstart.py --improve       # Just self-improvement
  python quickstart.py --test          # Test current model

"From zero to self-improving in one command"
"""

import os
import sys
import argparse
import subprocess
from pathlib import Path

ROOT = os.path.dirname(os.path.abspath(__file__))

def run_command(cmd, description):
    """Run a command with nice output."""
    print(f"\n{'='*70}")
    print(f"🚀 {description}")
    print(f"{'='*70}")
    print(f"$ {cmd}\n")
    
    result = subprocess.run(cmd, shell=True, cwd=ROOT)
    
    if result.returncode != 0:
        print(f"\n❌ Failed: {description}")
        return False
    
    print(f"\n✓ Complete: {description}")
    return True


def check_dependencies():
    """Check required packages are installed."""
    print("\n🔍 Checking dependencies...")
    
    required = [
        "torch",
        "transformers",
        "peft",
        "bitsandbytes",
        "accelerate",
    ]
    
    missing = []
    for pkg in required:
        try:
            __import__(pkg)
            print(f"  ✓ {pkg}")
        except ImportError:
            print(f"  ✗ {pkg}")
            missing.append(pkg)
    
    if missing:
        print(f"\n❌ Missing packages: {', '.join(missing)}")
        print("Install with: pip install " + " ".join(missing))
        return False
    
    print("\n✓ All dependencies installed")
    return True


def train_dense(steps=100):
    """Run THE CONDENSATOR dense training."""
    return run_command(
        f"python the_condensator.py --stages sft,dpo,rl --steps {steps}",
        "THE CONDENSATOR - Dense Response Training"
    )


def train_cfhot(steps=3000):
    """Train CF-HoT behavior heads."""
    success = True
    
    for behavior in ["repetition", "hedging", "verbosity"]:
        if not run_command(
            f"python train_cfhot_head.py --behavior {behavior} --steps {steps}",
            f"CF-HoT {behavior.upper()} Head Training"
        ):
            success = False
    
    return success


def train_self_improve(iterations=5):
    """Run stable self-improvement."""
    return run_command(
        f"python train_self_improve.py --iterations {iterations}",
        "Stable Self-Improvement Loop"
    )


def test_model(checkpoint=None):
    """Test the model."""
    cmd = "python the_condensator.py --eval-only"
    if checkpoint:
        cmd += f" --checkpoint {checkpoint}"
    return run_command(cmd, "Model Evaluation")


def full_pipeline():
    """Run the complete training pipeline."""
    print("\n" + "="*70)
    print("🔥 ÜBERMENSCHETIEN FULL TRAINING PIPELINE")
    print("="*70)
    print("""
    This will run:
    1. THE CONDENSATOR (SFT → DPO → RL)
    2. CF-HoT Head Training (repetition, hedging, verbosity)
    3. Stable Self-Improvement Loop
    
    Estimated time: 2-4 hours on RTX 3090
    """)
    
    if not check_dependencies():
        return False
    
    # Step 1: Dense training
    if not train_dense(100):
        return False
    
    # Step 2: CF-HoT heads
    if not train_cfhot(1000):  # Fewer steps for quick start
        return False
    
    # Step 3: Self-improvement
    if not train_self_improve(3):
        return False
    
    print("\n" + "="*70)
    print("✓ ÜBERMENSCHETIEN TRAINING COMPLETE!")
    print("="*70)
    print("""
    Your model is ready! Run:
    
      python ubermenschetien_v2_full.py
    
    Commands:
      > hello                  # Chat
      > !eval                  # Evaluate quality
      > !improve               # Continue self-improvement
    """)
    
    return True


def main():
    parser = argparse.ArgumentParser(description="Übermenschetien Quick Start")
    parser.add_argument("--full", action="store_true", help="Run full pipeline")
    parser.add_argument("--train-dense", action="store_true", help="Run dense training only")
    parser.add_argument("--train-cfhot", action="store_true", help="Run CF-HoT training only")
    parser.add_argument("--improve", action="store_true", help="Run self-improvement only")
    parser.add_argument("--test", action="store_true", help="Test current model")
    parser.add_argument("--steps", type=int, default=100, help="Training steps")
    parser.add_argument("--checkpoint", type=str, default=None, help="Checkpoint path for testing")
    
    args = parser.parse_args()
    
    if args.full:
        full_pipeline()
    elif args.train_dense:
        train_dense(args.steps)
    elif args.train_cfhot:
        train_cfhot(args.steps)
    elif args.improve:
        train_self_improve()
    elif args.test:
        test_model(args.checkpoint)
    else:
        parser.print_help()
        print("\n💡 Try: python quickstart.py --full")


if __name__ == "__main__":
    main()