"""Smoke test: PyTorch can see the Apple GPU and run a training step on it.

Checks three things:
  1. the MPS backend is available,
  2. forward + backward passes produce finite gradients on MPS,
  3. MPS matmul is actually faster than CPU (i.e., the GPU is really used).

Usage: python scripts/check_mps.py
"""

import time

import torch


def bench_matmul(device: str, n: int = 2048, iters: int = 20) -> float:
    """Mean seconds per (n x n) matmul; .item() forces device sync."""
    x = torch.randn(n, n, device=device)
    y = torch.randn(n, n, device=device)
    for _ in range(3):  # warmup
        (x @ y).sum().item()
    start = time.perf_counter()
    for _ in range(iters):
        (x @ y).sum().item()
    return (time.perf_counter() - start) / iters


def main() -> None:
    print(f"torch {torch.__version__}")
    assert torch.backends.mps.is_available(), "MPS backend not available"
    assert torch.backends.mps.is_built(), "torch was built without MPS support"
    print("MPS backend: available")

    # A miniature training step: forward, loss, backward, finite grads.
    w = torch.randn(64, 64, device="mps", requires_grad=True)
    x = torch.randn(128, 64, device="mps")
    loss = ((x @ w) ** 2).mean()
    loss.backward()
    assert w.grad is not None and torch.isfinite(w.grad).all()
    print(f"backward pass on MPS: ok (loss={loss.item():.4f})")

    cpu = bench_matmul("cpu")
    mps = bench_matmul("mps")
    print(f"2048x2048 matmul: cpu {cpu * 1e3:.1f} ms | mps {mps * 1e3:.1f} ms"
          f" | speedup {cpu / mps:.1f}x")
    print("smoke test passed")


if __name__ == "__main__":
    main()