"""
Create GPU-enabled OpenCLIP environment with complete dependency isolation
"""

import os
import sys
import subprocess
import platform

def create_gpu_environment():
    """Create isolated GPU environment for OpenCLIP"""
    print("Creating GPU-enabled OpenCLIP environment...")
    
    venv_path = "openclip_gpu_env"
    
    # Remove existing environment
    if os.path.exists(venv_path):
        print(f"Removing existing environment: {venv_path}")
        import shutil
        shutil.rmtree(venv_path)
    
    # Create virtual environment
    print(f"Creating virtual environment: {venv_path}")
    result = subprocess.run([
        sys.executable, '-m', 'venv', venv_path
    ], capture_output=True, text=True)
    
    if result.returncode != 0:
        print(f"Virtual environment creation failed: {result.stderr}")
        return False
    
    # Get virtual environment Python
    if platform.system() == "Windows":
        venv_python = os.path.join(venv_path, "Scripts", "python.exe")
    else:
        venv_python = os.path.join(venv_path, "bin", "python")
    
    # Install PyTorch with CUDA 11.8
    print("Installing PyTorch with CUDA 11.8...")
    torch_cmd = [
        venv_python, '-m', 'pip', 'install',
        'torch==2.0.1+cu118', 'torchvision==0.15.2+cu118',
        '--index-url', 'https://download.pytorch.org/whl/cu118'
    ]
    
    result = subprocess.run(torch_cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"PyTorch installation failed: {result.stderr}")
        return False
    
    # Install OpenCLIP
    print("Installing OpenCLIP...")
    openclip_cmd = [venv_python, '-m', 'pip', 'install', 'open-clip-torch']
    result = subprocess.run(openclip_cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"OpenCLIP installation failed: {result.stderr}")
        return False
    
    # Install Pillow
    print("Installing Pillow...")
    pillow_cmd = [venv_python, '-m', 'pip', 'install', 'Pillow']
    result = subprocess.run(pillow_cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"Pillow installation failed: {result.stderr}")
        return False
    
    # Test the environment
    print("Testing GPU environment...")
    test_script = """
import torch
import open_clip

print("=== GPU ENVIRONMENT TEST ===")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
if torch.cuda.is_available():
    print(f"CUDA version: {torch.version.cuda}")
    print(f"GPU count: {torch.cuda.device_count()}")
    for i in range(torch.cuda.device_count()):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")

# Test OpenCLIP
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='laion2b_s34b_b79k')
print("OpenCLIP model loaded successfully")

if torch.cuda.is_available():
    model = model.cuda()
    print("Model moved to GPU - SUCCESS")
else:
    print("Running on CPU")
"""
    
    result = subprocess.run([venv_python, '-c', test_script], capture_output=True, text=True)
    print(result.stdout)
    
    if result.returncode == 0:
        print("GPU environment setup successful!")
        return True
    else:
        print("GPU environment test failed")
        print(result.stderr)
        return False

if __name__ == "__main__":
    success = create_gpu_environment()
    sys.exit(0 if success else 1)