#!/usr/bin/env python3
"""Debug GPU issues in HuggingFace Spaces environment."""

import os
import sys

print("=== Debugging GPU in HuggingFace Spaces ===")

# Set environment variables BEFORE any imports
os.environ['CUDA_VISIBLE_DEVICES'] = '0'
os.environ['SPACY_PREFER_GPU'] = '1'

# Now import libraries
import torch
import spacy

print("\n1. Environment Check:")
print(f"   Platform: {sys.platform}")
print(f"   Python: {sys.version}")
print(f"   Working dir: {os.getcwd()}")
print(f"   CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
print(f"   SPACY_PREFER_GPU: {os.environ.get('SPACY_PREFER_GPU', 'Not set')}")
print(f"   SPACES: {os.environ.get('SPACES', 'Not set')}")

print("\n2. PyTorch GPU Status:")
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()}")
    print(f"   Current device: {torch.cuda.current_device()}")
    print(f"   GPU 0: {torch.cuda.get_device_name(0)}")
    
    # Force CUDA initialization
    torch.cuda.init()
    print("   ✓ CUDA initialized")
    
    # Set default device
    torch.cuda.set_device(0)
    print("   ✓ Set default CUDA device to 0")

print("\n3. SpaCy GPU Configuration:")
print(f"   SpaCy version: {spacy.__version__}")

# Try multiple methods to enable GPU
print("\n   Attempting spacy.prefer_gpu()...")
gpu_id = spacy.prefer_gpu(gpu_id=0)
print(f"   Result: {gpu_id}")

if torch.cuda.is_available():
    print("\n   Attempting spacy.require_gpu()...")
    try:
        spacy.require_gpu(gpu_id=0)
        print("   ✓ spacy.require_gpu() succeeded")
    except Exception as e:
        print(f"   ✗ spacy.require_gpu() failed: {e}")

print("\n4. Test Model Loading:")
try:
    # Try loading a small model first
    print("   Loading en_core_web_md...")
    nlp_md = spacy.load("en_core_web_md")
    
    # Check if components are on GPU
    print("   Checking MD model components:")
    for name, component in nlp_md.pipeline:
        device = "Unknown"
        if hasattr(component, 'model'):
            if hasattr(component.model, 'device'):
                device = str(component.model.device)
            elif hasattr(component.model, 'parameters'):
                try:
                    param = next(component.model.parameters())
                    device = str(param.device)
                except:
                    pass
        print(f"      {name}: {device}")
    
    # Test processing
    doc = nlp_md("Test sentence")
    print(f"   ✓ MD model processed {len(doc)} tokens")
    
except Exception as e:
    print(f"   ✗ MD model failed: {e}")

print("\n5. Test Transformer Model with GPU:")
try:
    # Force GPU before loading transformer
    if torch.cuda.is_available():
        torch.cuda.set_device(0)
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    
    print("   Loading en_core_web_trf with GPU config...")
    
    # Load with explicit GPU configuration
    config = {
        "nlp": {
            "pipeline": ["transformer", "tagger", "parser", "ner", "lemmatizer"]
        },
        "components": {
            "transformer": {
                "model": {
                    "mixed_precision": True,
                    "@architectures": "spacy-transformers.TransformerModel.v3",
                    "get_spans": {
                        "@span_getters": "spacy-transformers.strided_spans.v1",
                        "window": 128,
                        "stride": 96
                    }
                }
            }
        }
    }
    
    nlp_trf = spacy.load("en_core_web_trf")
    
    # Force components to GPU
    print("   Forcing transformer components to GPU...")
    for name, component in nlp_trf.pipeline:
        if hasattr(component, 'model'):
            if hasattr(component.model, 'to'):
                try:
                    component.model.to('cuda:0')
                    print(f"      ✓ Moved {name} to GPU")
                except Exception as e:
                    print(f"      ✗ Failed to move {name}: {e}")
    
    # Verify GPU usage
    print("\n   Verifying GPU usage:")
    for name, component in nlp_trf.pipeline:
        on_gpu = False
        device_info = "Unknown"
        
        if hasattr(component, 'model'):
            # Check parameters
            if hasattr(component.model, 'parameters'):
                try:
                    for param in component.model.parameters():
                        if param.is_cuda:
                            on_gpu = True
                            device_info = str(param.device)
                            break
                except:
                    pass
            
            # Check device attribute
            if hasattr(component.model, 'device'):
                device_info = str(component.model.device)
                on_gpu = 'cuda' in device_info
        
        status = "✓ GPU" if on_gpu else "✗ CPU"
        print(f"      {name}: {status} ({device_info})")
    
    # Test processing with timing
    print("\n   Testing transformer processing...")
    import time
    
    text = "The quick brown fox jumps over the lazy dog. " * 5
    start = time.time()
    doc = nlp_trf(text)
    end = time.time()
    
    print(f"   ✓ Processed {len(doc)} tokens in {end-start:.2f}s")
    
    # Check memory usage
    if torch.cuda.is_available():
        mem_allocated = torch.cuda.memory_allocated(0) / 1024**3
        mem_reserved = torch.cuda.memory_reserved(0) / 1024**3
        print(f"\n   GPU Memory:")
        print(f"      Allocated: {mem_allocated:.2f} GB")
        print(f"      Reserved: {mem_reserved:.2f} GB")
    
except Exception as e:
    print(f"   ✗ Transformer model failed: {e}")
    import traceback
    traceback.print_exc()

print("\n=== Summary ===")
if torch.cuda.is_available():
    print("✓ CUDA is available")
    print("✓ PyTorch can see GPU")
    print("→ Check if SpaCy models are using GPU above")
else:
    print("✗ No GPU detected in this environment")
    print("→ This script should be run in HuggingFace Spaces with GPU")