setup.py

#!/usr/bin/env python3
"""
PongAI Setup Script

Initializes the project environment, creates necessary directories,
and validates the installation.
"""

import sys
import os
from pathlib import Path
import subprocess

# Colors for terminal output
GREEN = '\033[92m'
YELLOW = '\033[93m'
RED = '\033[91m'
RESET = '\033[0m'


def print_header(msg: str):
    """Print a header message."""
    print(f"\n{GREEN}{'='*60}{RESET}")
    print(f"{GREEN}{msg:^60}{RESET}")
    print(f"{GREEN}{'='*60}{RESET}\n")


def print_success(msg: str):
    """Print a success message."""
    print(f"{GREEN}✓{RESET} {msg}")


def print_warning(msg: str):
    """Print a warning message."""
    print(f"{YELLOW}!{RESET} {msg}")


def print_error(msg: str):
    """Print an error message."""
    print(f"{RED}✗{RESET} {msg}")


def setup():
    """Run setup procedure."""
    print_header("PongAI Setup")
    
    project_root = Path(__file__).parent.resolve()
    
    # Step 1: Validate Python version
    print(f"Checking Python version...")
    version = sys.version_info
    if version.major < 3 or (version.major == 3 and version.minor < 10):
        print_error(f"Python 3.10+ required, found {version.major}.{version.minor}")
        return False
    print_success(f"Python {version.major}.{version.minor}.{version.micro}")
    
    # Step 2: Create necessary directories
    print(f"\nCreating project directories...")
    directories = [
        project_root / "models",
        project_root / "pong_tensorboard",
    ]
    
    for directory in directories:
        directory.mkdir(parents=True, exist_ok=True)
        print_success(f"Directory: {directory.name}/")
    
    # Step 3: Check requirements.txt
    print(f"\nChecking dependencies...")
    req_file = project_root / "requirements.txt"
    if req_file.exists():
        print_success(f"Found requirements.txt")
        print(f"\nTo install dependencies, run:")
        print(f"  {YELLOW}pip install -r requirements.txt{RESET}\n")
    else:
        print_error("requirements.txt not found")
        return False
    
    # Step 4: Verify module structure
    print(f"Verifying module structure...")
    modules = ['engine', 'rl', 'train', 'demo', 'api']
    all_ok = True
    
    for module in modules:
        module_path = project_root / module
        init_file = module_path / "__init__.py"
        if module_path.exists() and init_file.exists():
            print_success(f"Module: {module}/")
        else:
            print_error(f"Module: {module}/ (MISSING)")
            all_ok = False
    
    if not all_ok:
        return False
    
    # Step 5: Check key files
    print(f"\nVerifying key files...")
    key_files = [
        ('engine/pong.py', 'Physics engine'),
        ('rl/env.py', 'Gymnasium environment'),
        ('train/ppo.py', 'Training pipeline'),
        ('demo/play.py', 'Interactive demo'),
        ('api/app.py', 'FastAPI server'),
        ('config.py', 'Configuration'),
        ('utils.py', 'Utilities'),
        ('main.py', 'CLI entry point'),
        ('README.md', 'Documentation'),
    ]
    
    for file_path, description in key_files:
        full_path = project_root / file_path
        if full_path.exists():
            size_kb = full_path.stat().st_size / 1024
            print_success(f"{file_path:<25} ({size_kb:>6.1f} KB) - {description}")
        else:
            print_error(f"{file_path:<25} - MISSING")
            all_ok = False
    
    if not all_ok:
        return False
    
    # Step 6: Verify imports (optional)
    print(f"\nVerifying Python imports...")
    try:
        # Try importing key modules
        sys.path.insert(0, str(project_root))
        
        # Check if standard library imports work
        import numpy
        print_success("numpy available")
        
        # Note: Don't check pygame, gymnasium, etc. until installed
        print_warning("Other dependencies (gymnasium, stable-baselines3, etc.) not yet installed")
        
    except ImportError as e:
        print_warning(f"Optional import check skipped: {e}")
    
    # Step 7: Print next steps
    print_header("Setup Complete!")
    print(f"""
Next steps:

1. Install dependencies:
   {YELLOW}pip install -r requirements.txt{RESET}

2. Train a model (optional):
   {YELLOW}python main.py train{RESET}
   
   Or use custom settings:
   {YELLOW}python main.py train --timesteps 500000 --envs 4{RESET}

3. Run the interactive demo:
   {YELLOW}python main.py demo{RESET}

4. Start the API server (optional):
   {YELLOW}python main.py api{RESET}

5. Monitor training (if training):
   {YELLOW}tensorboard --logdir=./pong_tensorboard/{RESET}

For more details, see:
- README.md          (Full documentation)
- QUICKREF.md        (Quick reference)
- IMPLEMENTATION.md  (Implementation details)

Project directory: {project_root}
""")
    
    return True


if __name__ == "__main__":
    try:
        success = setup()
        sys.exit(0 if success else 1)
    except Exception as e:
        print_error(f"Setup failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)