"""
Simple test to verify OpenCLIP GPU functionality
"""
import subprocess
import os

def test_openclip_gpu():
    """Test OpenCLIP GPU functionality with a simple script"""
    venv_python = "openclip_gpu_env\\Scripts\\python.exe"
    
    if not os.path.exists(venv_python):
        print("ERROR: Virtual environment not found")
        return False
    
    # Simple test script
    test_script = """
import torch
import open_clip
from PIL import Image
import tempfile
import os

print("Testing OpenCLIP GPU...")
print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")

if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
else:
    print("ERROR: CUDA not available")
    exit(1)

# Load model
print("Loading OpenCLIP model...")
model, _, preprocess = open_clip.create_model_and_transforms(
    model_name="ViT-B-32",
    pretrained="laion2b_s34b_b79k"
)

print("Model loaded successfully")

# Move to GPU
model = model.cuda()
print("Model moved to GPU")

# Create a test image (simple red square)
with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
    img_path = f.name

img = Image.new('RGB', (224, 224), color='red')
img.save(img_path)
print(f"Created test image: {img_path}")

# Process image
image = preprocess(img).unsqueeze(0).cuda()
print("Image processed and moved to GPU")

# Test classification
text_labels = ["a photo of a bee", "a photo of a flower", "a photo of a document"]
text_tokens = open_clip.tokenize(text_labels).cuda()

with torch.no_grad():
    image_features = model.encode_image(image)
    text_features = model.encode_text(text_tokens)
    
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)
    
    similarity = (100.0 * image_features @ text_features.T).softmax(dim=-1)
    values, indices = similarity[0].topk(3)
    
    print("Classification results:")
    for i, (val, idx) in enumerate(zip(values, indices)):
        label = text_labels[idx]
        conf = float(val)
        print(f"  {i+1}. {label}: {conf:.4f}")

# Clean up
os.unlink(img_path)
print("Test completed successfully")
"""

    print("Running OpenCLIP GPU test...")
    result = subprocess.run([venv_python, "-c", test_script], capture_output=True, text=True, timeout=60)
    
    print("STDOUT:")
    print(result.stdout)
    
    if result.stderr:
        print("STDERR:")
        print(result.stderr)
    
    if result.returncode == 0:
        print("SUCCESS: OpenCLIP GPU test passed!")
        return True
    else:
        print("ERROR: OpenCLIP GPU test failed")
        return False

if __name__ == "__main__":
    print("SIMPLE OPENCLIP GPU TEST")
    print("=" * 50)
    
    success = test_openclip_gpu()
    
    if success:
        print("\nOpenCLIP GPU is working correctly!")
    else:
        print("\nOpenCLIP GPU test failed. Please check dependencies.")