setup.py

"""
PyTorch Automatic Installation Script
Detects GPU environment and installs the optimal PyTorch version.

IMPORTANT: This script requires official Python (python.org) or conda Python.
MSYS2/MinGW Python is NOT supported by PyTorch.
"""
import subprocess
import platform
import sys


def check_python_compatibility():
    """
    Check if Python build is compatible with PyTorch.
    PyTorch requires official CPython builds (MSVC compiler on Windows).
    MSYS2/MinGW Python (GCC compiler) is NOT supported.
    """
    compiler = platform.python_compiler()

    # Check for MSYS2/MinGW Python (uses GCC)
    if 'GCC' in compiler and platform.system() == 'Windows':
        print("\n" + "="*60)
        print("ERROR: Incompatible Python Build")
        print("="*60)
        print(f"Current Python: {sys.version}")
        print(f"Compiler: {compiler}")
        print("\n[PROBLEM]")
        print("  PyTorch does NOT support MSYS2/MinGW Python (GCC-compiled).")
        print("  PyTorch requires official CPython (MSVC-compiled) on Windows.")
        print("\n[SOLUTION]")
        print("  1. Download official Python from: https://www.python.org/downloads/")
        print("  2. Or use Windows launcher: py -V:3.10 -m venv venv")
        print("     (Check available versions: py -0)")
        print("\n[SUPPORTED VERSIONS]")
        print("  - Python 3.8 to 3.12 (official CPython builds)")
        print("  - NOT MSYS2, NOT MinGW, NOT Cygwin")
        print("="*60)
        sys.exit(1)

    # Check Python version
    version_info = sys.version_info
    if version_info < (3, 8) or version_info >= (3, 13):
        print("\n" + "="*60)
        print("ERROR: Unsupported Python Version")
        print("="*60)
        print(f"Current Python: {sys.version}")
        print("\n[PROBLEM]")
        print(f"  PyTorch requires Python 3.8 to 3.12")
        print(f"  You have Python {version_info.major}.{version_info.minor}")
        print("\n[SOLUTION]")
        print("  Install Python 3.8, 3.9, 3.10, 3.11, or 3.12")
        print("="*60)
        sys.exit(1)

    print(f"[OK] Python {version_info.major}.{version_info.minor} ({compiler})")
    return True


def run_command(command, description):
    """Execute command"""
    print(f"\n{'='*60}")
    print(f"{description}")
    print(f"{'='*60}")
    print(f"Running: {command}\n")

    result = subprocess.run(command, shell=True)

    if result.returncode != 0:
        print(f"\n[ERROR] {description}")
        return False
    return True


def detect_nvidia_gpu():
    """Detect NVIDIA GPU"""
    try:
        result = subprocess.run(
            "nvidia-smi",
            shell=True,
            capture_output=True,
            text=True
        )
        if result.returncode == 0:
            print("[OK] NVIDIA GPU detected")
            # Print GPU info
            lines = result.stdout.split('\n')
            for line in lines:
                if 'NVIDIA' in line and '|' in line:
                    print(f"  {line.strip()}")
                    break
            return True
    except:
        pass

    print("[INFO] No NVIDIA GPU")
    return False


def detect_cuda_version():
    """Detect CUDA version"""
    try:
        result = subprocess.run(
            "nvidia-smi",
            shell=True,
            capture_output=True,
            text=True
        )
        if result.returncode == 0:
            # Extract CUDA Version
            for line in result.stdout.split('\n'):
                if 'CUDA Version' in line:
                    cuda_version = line.split('CUDA Version:')[1].strip().split()[0]
                    print(f"[OK] CUDA version: {cuda_version}")

                    # Return major version
                    major = cuda_version.split('.')[0]
                    if int(major) >= 12:
                        return "12.1"
                    elif int(major) == 11:
                        return "11.8"
                    else:
                        return None
    except:
        pass

    return None


def is_apple_silicon():
    """Detect Apple Silicon"""
    if platform.system() == 'Darwin' and platform.machine() == 'arm64':
        print("[OK] Apple Silicon (M1/M2/M3) detected")
        return True
    return False


def detect_intel_gpu():
    """Detect Intel GPU"""
    try:
        # Windows
        if platform.system() == 'Windows':
            result = subprocess.run(
                "wmic path win32_VideoController get name",
                shell=True,
                capture_output=True,
                text=True
            )
            if result.returncode == 0:
                output = result.stdout.lower()
                if 'intel' in output and ('arc' in output or 'iris' in output or 'uhd' in output or 'hd graphics' in output):
                    print("[OK] Intel GPU detected")
                    # Print GPU info
                    for line in result.stdout.split('\n'):
                        if 'intel' in line.lower() and line.strip() and 'name' not in line.lower():
                            print(f"  {line.strip()}")
                            break
                    return True

        # Linux
        elif platform.system() == 'Linux':
            result = subprocess.run(
                "lspci | grep -i vga",
                shell=True,
                capture_output=True,
                text=True
            )
            if result.returncode == 0:
                output = result.stdout.lower()
                if 'intel' in output:
                    print("[OK] Intel GPU detected")
                    print(f"  {result.stdout.strip()}")
                    return True
    except:
        pass

    return False


def install_pytorch():
    """Install PyTorch"""
    print("\n" + "="*60)
    print("PyTorch Installation Started")
    print("="*60)

    # 1. Apple Silicon
    if is_apple_silicon():
        print("\n[Installation Mode] Apple Silicon (MPS support)")
        command = "pip install torch torchvision torchaudio"
        return run_command(command, "PyTorch (MPS) Installation")

    # 2. NVIDIA GPU
    if detect_nvidia_gpu():
        cuda_version = detect_cuda_version()

        if cuda_version == "12.1":
            print("\n[Installation Mode] CUDA 12.1")
            command = "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121"
        elif cuda_version == "11.8":
            print("\n[Installation Mode] CUDA 11.8")
            command = "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118"
        else:
            print("\n[Installation Mode] CUDA version unclear, installing latest CUDA 12.1")
            command = "pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121"

        return run_command(command, "PyTorch (CUDA) Installation")

    # 3. Intel GPU
    if detect_intel_gpu():
        print("\n[Installation Mode] Intel GPU detected")
        print("[INFO] Installing PyTorch (CPU base)")
        print("[INFO] Note: Intel Extension for PyTorch may not be available for all Python versions")

        # Install base PyTorch first
        if not run_command("pip install torch torchvision torchaudio", "PyTorch (CPU) Installation"):
            return False

        # Try to install Intel Extension (may fail on unsupported Python versions)
        print("\n[INFO] Attempting to install Intel Extension for PyTorch...")
        result = subprocess.run(
            "pip install intel-extension-for-pytorch",
            shell=True,
            capture_output=True
        )

        if result.returncode == 0:
            print("[OK] Intel Extension for PyTorch installed successfully")
            print("[INFO] You can use torch.xpu for Intel GPU acceleration")
        else:
            print("[WARNING] Intel Extension for PyTorch installation failed")
            print("[INFO] This is normal for Python 3.12+ or certain configurations")
            print("[INFO] Fallback to CPU mode - training will still work")

        return True

    # 4. CPU only
    print("\n[Installation Mode] CPU only")
    command = "pip install torch torchvision torchaudio"
    return run_command(command, "PyTorch (CPU) Installation")


def install_requirements():
    """Install requirements.txt"""
    command = "pip install -r requirements.txt"
    return run_command(command, "Other dependencies installation (requirements.txt)")


def verify_installation():
    """Verify installation"""
    print("\n" + "="*60)
    print("Installation Verification")
    print("="*60)

    try:
        import torch
        print(f"\n[OK] PyTorch version: {torch.__version__}")

        print("\n[Available Devices]")

        # Check CUDA
        if torch.cuda.is_available():
            print(f"  [OK] CUDA: {torch.cuda.get_device_name(0)}")
            print(f"    - CUDA version: {torch.version.cuda}")
            print(f"    - Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")

        # Check MPS (Apple Silicon)
        elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
            print(f"  [OK] MPS (Apple Silicon)")

        # Check Intel XPU
        else:
            try:
                import intel_extension_for_pytorch as ipex
                print(f"  [OK] Intel Extension for PyTorch version: {ipex.__version__}")

                # Try to detect XPU devices
                if hasattr(torch, 'xpu') and torch.xpu.is_available():
                    print(f"  [OK] Intel XPU available")
                    print(f"    - Device count: {torch.xpu.device_count()}")
                else:
                    print(f"  [INFO] Intel Extension installed but XPU not available")
                    print(f"  [INFO] Fallback to CPU mode")
            except ImportError:
                print(f"  [OK] CPU only")

        return True

    except ImportError as e:
        print(f"\n[ERROR] PyTorch installation failed: {e}")
        return False


def main():
    """Main function"""
    import sys

    print("\n" + "="*60)
    print("CIFAR-10 Project Automatic Setup")
    print("="*60)
    print(f"Python version: {sys.version.split()[0]}")
    print(f"OS: {platform.system()} {platform.machine()}")

    # Check Python compatibility FIRST
    print("\n[0/3] Python compatibility check")
    check_python_compatibility()

    # pip upgrade
    print("\n[1/3] pip upgrade")
    run_command("pip install --upgrade pip", "pip upgrade")

    # PyTorch installation
    print("\n[2/3] PyTorch installation")
    if not install_pytorch():
        print("\n[ERROR] PyTorch installation failed. Manual installation required.")
        print("\nManual installation guide:")
        print("  https://pytorch.org/get-started/locally/")
        sys.exit(1)

    # Other dependencies installation
    print("\n[3/3] Other dependencies installation")
    if not install_requirements():
        print("\n[WARNING] Some packages failed to install. Manual check required.")

    # Verification
    if verify_installation():
        print("\n" + "="*60)
        print("Installation Complete!")
        print("="*60)
        print("\nStart the project with the following commands:")
        print("  cd src")
        print("  python main.py")
    else:
        print("\n[ERROR] Installation verification failed")
        sys.exit(1)


if __name__ == "__main__":
    main()