#!/usr/bin/env python3
"""Test script to verify PyTorch MPS (Metal Performance Shaders) GPU acceleration."""

import torch
import time

def test_mps_acceleration():
    """Compare performance between CPU and MPS (GPU) for matrix operations."""
    
    # Check MPS availability
    if not torch.backends.mps.is_available():
        print("❌ MPS is not available on this system")
        return
    
    print("✅ MPS (Metal Performance Shaders) is available!")
    print(f"PyTorch version: {torch.__version__}")
    print("-" * 50)
    
    # Test parameters
    size = 4096
    iterations = 100
    
    # Create random matrices
    print(f"\n📊 Testing matrix multiplication ({size}x{size})...")
    a = torch.randn(size, size)
    b = torch.randn(size, size)
    
    # CPU benchmark
    print("\n🖥️  CPU Performance:")
    start = time.time()
    for _ in range(iterations):
        c_cpu = torch.matmul(a, b)
    cpu_time = time.time() - start
    print(f"   Time: {cpu_time:.3f} seconds")
    
    # MPS (GPU) benchmark
    print("\n🚀 MPS (GPU) Performance:")
    device = torch.device("mps")
    a_mps = a.to(device)
    b_mps = b.to(device)
    
    # Warm up GPU
    for _ in range(10):
        _ = torch.matmul(a_mps, b_mps)
    torch.mps.synchronize()  # Ensure GPU operations complete
    
    start = time.time()
    for _ in range(iterations):
        c_mps = torch.matmul(a_mps, b_mps)
    torch.mps.synchronize()  # Ensure all operations complete
    mps_time = time.time() - start
    print(f"   Time: {mps_time:.3f} seconds")
    
    # Results
    speedup = cpu_time / mps_time
    print("\n📈 Results:")
    print(f"   Speedup: {speedup:.2f}x faster on MPS")
    print(f"   {'🎉 MPS acceleration is working!' if speedup > 1 else '⚠️  No acceleration detected'}")
    
    # Memory info
    print("\n💾 MPS Memory Info:")
    print(f"   Allocated: {torch.mps.current_allocated_memory() / 1024**2:.2f} MB")
    print(f"   Driver: {torch.mps.driver_allocated_memory() / 1024**2:.2f} MB")
    
    # Test moving tensors between devices
    print("\n🔄 Testing tensor movement between devices...")
    test_tensor = torch.randn(1000, 1000)
    
    # To MPS
    mps_tensor = test_tensor.to('mps')
    print(f"   ✅ Moved to MPS: {mps_tensor.device}")
    
    # Back to CPU
    cpu_tensor = mps_tensor.cpu()
    print(f"   ✅ Moved to CPU: {cpu_tensor.device}")
    
    print("\n✨ MPS setup complete and working correctly!")

if __name__ == "__main__":
    test_mps_acceleration()