#!/usr/bin/env python3
"""
Project setup script for Dynamic Traffic Signal Optimization
Checks dependencies, creates directories, and validates SUMO installation
"""

import os
import sys
import subprocess
import platform
import pkg_resources

def check_python_version():
    """Check if Python version is compatible"""
    version = sys.version_info
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        print("❌ Error: Python 3.8 or higher is required")
        print(f"   Current version: {version.major}.{version.minor}.{version.micro}")
        sys.exit(1)
    print(f"✅ Python {version.major}.{version.minor}.{version.micro}")

def check_required_packages():
    """Check if required packages are installed"""
    required_packages = [
        'torch', 'numpy', 'pandas', 'matplotlib', 'seaborn',
        'opencv-python', 'gym', 'scikit-learn', 'tensorboard',
        'tqdm', 'PyYAML', 'pillow'
    ]
    
    missing_packages = []
    
    for package in required_packages:
        try:
            pkg_resources.get_distribution(package)
            print(f"✅ {package}")
        except pkg_resources.DistributionNotFound:
            missing_packages.append(package)
            print(f"❌ {package}")
    
    if missing_packages:
        print(f"\n⚠️  Missing packages: {', '.join(missing_packages)}")
        print("Install them with: pip install -r requirements.txt")
        return False
    
    return True

def check_sumo_installation():
    """Check if SUMO is properly installed"""
    try:
        result = subprocess.run(["sumo", "--version"], 
                              capture_output=True, text=True, timeout=10)
        if result.returncode == 0:
            version_line = result.stdout.strip().split('\n')[0]
            print(f"✅ SUMO installed: {version_line}")
            return True
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    
    print("❌ SUMO not found in PATH")
    print_sumo_installation_instructions()
    return False

def print_sumo_installation_instructions():
    """Print SUMO installation instructions"""
    system = platform.system()
    
    print("\n📖 SUMO Installation Instructions:")
    print("="*50)
    
    if system == "Windows":
        print("1. Download SUMO from: https://eclipse.org/sumo/")
        print("2. Install the executable")
        print("3. Add SUMO/bin to your PATH environment variable")
        print("4. Set SUMO_HOME environment variable to SUMO installation directory")
    
    elif system == "Darwin":  # macOS
        print("Option 1 - Homebrew:")
        print("  brew install sumo")
        print("\nOption 2 - Manual:")
        print("  1. Download from: https://eclipse.org/sumo/")
        print("  2. Follow macOS installation guide")
        print("  3. You may need to install XQuartz")
    
    elif system == "Linux":
        print("Ubuntu/Debian:")
        print("  sudo apt-get install sumo sumo-tools sumo-doc")
        print("\nCentOS/RHEL/Fedora:")
        print("  sudo yum install sumo sumo-tools")
        print("\nFrom source:")
        print("  Follow instructions at: https://sumo.dlr.de/docs/Installing/Linux_Build.html")
    
    print("\nAfter installation, ensure SUMO_HOME is set:")
    print("export SUMO_HOME=/path/to/sumo")

def create_project_structure():
    """Create the project directory structure"""
    directories = [
        "src/environment",
        "src/agents", 
        "src/training",
        "src/evaluation",
        "src/utils",
        "config",
        "sumo_configs",
        "models/checkpoints",
        "data/raw",
        "data/processed", 
        "logs/tensorboard",
        "results/plots",
        "results/analysis",
        "scripts",
        "notebooks",
        "tests"
    ]
    
    print("\n📁 Creating project structure...")
    for directory in directories:
        os.makedirs(directory, exist_ok=True)
        
        # Create __init__.py files for Python packages
        if directory.startswith("src/"):
            init_file = os.path.join(directory, "__init__.py")
            if not os.path.exists(init_file):
                with open(init_file, 'w') as f:
                    f.write("# Package initialization\n")
    
    print("✅ Project structure created")

def create_sample_files():
    """Create sample configuration and documentation files"""
    
    # Create .gitignore if it doesn't exist
    gitignore_content = """# Python
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg

# Virtual Environment
venv/
env/
ENV/

# IDE
.vscode/
.idea/
*.swp
*.swo

# Jupyter Notebook
.ipynb_checkpoints

# Models and Data
models/*.pth
data/raw/*
data/processed/*
!data/raw/.gitkeep
!data/processed/.gitkeep

# Logs
logs/*.log
logs/tensorboard/*

# Results
results/plots/*
results/analysis/*

# SUMO outputs
*.xml
*.csv
!sumo_configs/*.xml

# OS
.DS_Store
Thumbs.db

# Temporary files
*.tmp
*.temp
"""
    
    if not os.path.exists(".gitignore"):
        with open(".gitignore", 'w') as f:
            f.write(gitignore_content)
        print("✅ .gitignore created")
    
    # Create README if it doesn't exist
    if not os.path.exists("README.md"):
        readme_content = """# Dynamic Traffic Signal Optimization using Reinforcement Learning

## M.Tech Project

### Overview
This project implements an intelligent traffic signal control system using Deep Reinforcement Learning (DRL) to optimize traffic flow at urban intersections.

### Quick Start

1. **Setup environment**
   ```bash
   python scripts/setup_project.py
   ```

2. **Train the model**
   ```bash
   python main.py --mode train
   ```

3. **Test the model**
   ```bash
   python main.py --mode test --model models/final_model.pth
   ```

### Project Structure
```
├── src/                    # Source code
│   ├── environment/        # Traffic environment
│   ├── agents/            # RL agents
│   ├── training/          # Training framework
│   └── evaluation/        # Evaluation tools
├── config/                # Configuration files
├── sumo_configs/          # SUMO network files
├── models/                # Trained models
├── results/               # Results and analysis
└── scripts/               # Utility scripts
```

### Configuration
Edit `config/config.yaml` to customize:
- Training parameters
- Network architecture
- Evaluation settings

### Requirements
- Python 3.8+
- PyTorch
- SUMO Traffic Simulator
- See `requirements.txt` for full list
"""
        
        with open("README.md", 'w') as f:
            f.write(readme_content)
        print("✅ README.md created")

def check_gpu_availability():
    """Check if GPU is available for PyTorch"""
    try:
        import torch
        if torch.cuda.is_available():
            gpu_count = torch.cuda.device_count()
            gpu_name = torch.cuda.get_device_name(0)
            print(f"✅ GPU available: {gpu_name} ({gpu_count} device(s))")
            return True
        else:
            print("ℹ️  GPU not available, will use CPU")
            return False
    except ImportError:
        print("⚠️  PyTorch not installed, cannot check GPU")
        return False

def run_basic_tests():
    """Run basic functionality tests"""
    print("\n🧪 Running basic tests...")
    
    # Test YAML loading
    try:
        import yaml
        test_config = {
            'test': 'value',
            'nested': {'key': 123}
        }
        yaml.dump(test_config)
        print("✅ YAML functionality")
    except Exception as e:
        print(f"❌ YAML test failed: {e}")
    
    # Test NumPy
    try:
        import numpy as np
        arr = np.random.random((3, 3))
        assert arr.shape == (3, 3)
        print("✅ NumPy functionality")
    except Exception as e:
        print(f"❌ NumPy test failed: {e}")
    
    # Test PyTorch
    try:
        import torch
        tensor = torch.randn(2, 3)
        assert tensor.shape == (2, 3)
        print("✅ PyTorch functionality")
    except Exception as e:
        print(f"❌ PyTorch test failed: {e}")

def main():
    print("🚦 Dynamic Traffic Signal Optimization - Project Setup")
    print("="*60)
    
    # Check system requirements
    print("\n1️⃣  Checking Python version...")
    check_python_version()
    
    print("\n2️⃣  Checking required packages...")
    packages_ok = check_required_packages()
    
    print("\n3️⃣  Checking SUMO installation...")
    sumo_ok = check_sumo_installation()
    
    print("\n4️⃣  Creating project structure...")
    create_project_structure()
    
    print("\n5️⃣  Creating sample files...")
    create_sample_files()
    
    print("\n6️⃣  Checking GPU availability...")
    gpu_available = check_gpu_availability()
    
    if packages_ok:
        print("\n7️⃣  Running basic tests...")
        run_basic_tests()
    
    print("\n" + "="*60)
    print("🎉 Setup completed!")
    
    if packages_ok and sumo_ok:
        print("✅ Ready to run the project")
        print("\nNext steps:")
        print("1. Review config/config.yaml")
        print("2. Run: python main.py --mode train")
    else:
        print("⚠️  Please fix the issues above before running the project")
        if not packages_ok:
            print("   - Install missing Python packages")
        if not sumo_ok:
            print("   - Install and configure SUMO")

if __name__ == "__main__":
    main()