import torch
import numpy as np
import time
import coremltools as ct

def test_mps():
    print("--- Testing PyTorch MPS ---")
    if not torch.backends.mps.is_available():
        print("MPS not available.")
        return
    
    device = torch.device("mps")
    print(f"Using device: {device}")
    
    # Create tensors
    size = 4000
    x = torch.randn(size, size, device=device)
    y = torch.randn(size, size, device=device)
    
    # Warm up
    _ = torch.matmul(x, y)
    torch.mps.synchronize()
    
    # Benchmark
    start = time.time()
    for _ in range(10):
        z = torch.matmul(x, y)
    torch.mps.synchronize()
    end = time.time()
    
    print(f"Matrix multiplication {size}x{size} on MPS took: {(end - start)/10:.4f}s per iteration")
    print("PyTorch MPS check passed.\n")

def test_coreml():
    print("--- Testing CoreML ---")
    # Simple check for CoreML availability and version
    print(f"CoreML Tools version: {ct.__version__}")
    # On macOS, CoreML is always available as a system framework.
    # The actual execution depends on the model's 'compute_units'.
    print("CoreML framework is available on macOS.\n")

if __name__ == "__main__":
    test_mps()
    test_coreml()