#!/usr/bin/env python3
"""
PyTorch MPS Flash Attention Benchmark - Baseline for RustyTorch comparison.

This script benchmarks PyTorch's scaled_dot_product_attention on Apple MPS backend.
Use this as a baseline to compare against RustyTorch Metal Flash Attention.

Usage:
    python3 scripts/bench_pytorch_mps.py

Run Rust benchmark with:
    cargo bench -p rtx-flash-metal-attention --bench metal_vs_pytorch
"""

import torch
import torch.nn.functional as F
import time
import sys
import platform


def check_mps():
    """Check if MPS backend is available."""
    if not torch.backends.mps.is_available():
        print("ERROR: MPS not available on this system")
        print("This benchmark requires macOS with Apple Silicon")
        sys.exit(1)

    print("=" * 70)
    print("PyTorch MPS Benchmark - scaled_dot_product_attention")
    print("=" * 70)
    print(f"PyTorch version:  {torch.__version__}")
    print(f"MPS available:    {torch.backends.mps.is_available()}")
    print(f"Python version:   {platform.python_version()}")
    print(f"macOS version:    {platform.mac_ver()[0]}")
    print(f"Chip:             {platform.processor()}")
    print()


def benchmark(
    batch_size: int,
    seq_len: int = 128,
    head_dim: int = 64,
    num_heads: int = 8,
    name: str = "Test",
    iterations: int = 1000,
    warmup: int = 50,
    dtype=torch.float32,
):
    """
    Run a single benchmark scenario.

    Args:
        batch_size: Number of sequences in batch
        seq_len: Sequence length
        head_dim: Dimension per attention head
        num_heads: Number of attention heads
        name: Scenario name for display
        iterations: Number of timed iterations
        warmup: Number of warmup iterations
        dtype: Data type (float32 to match Rust implementation)
    """
    device = torch.device("mps")

    # Create input tensors [batch, heads, seq, head_dim]
    q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)

    # Warmup: wake up GPU and JIT compile kernels
    for _ in range(warmup):
        _ = F.scaled_dot_product_attention(q, k, v)
    torch.mps.synchronize()

    # Timed run with synchronization for fair comparison
    start = time.perf_counter()

    for _ in range(iterations):
        _ = F.scaled_dot_product_attention(q, k, v)
        torch.mps.synchronize()  # Block until GPU done (matches Rust waitUntilCompleted)

    end = time.perf_counter()

    # Calculate metrics
    avg_time_ms = ((end - start) / iterations) * 1000
    total_elements = batch_size * num_heads * seq_len * head_dim
    throughput = total_elements / (avg_time_ms / 1000)  # elements per second

    print(f"| {name:<20} | {avg_time_ms:>10.4f} ms | {throughput/1e6:>10.2f} M/s |")


def benchmark_causal(
    batch_size: int,
    seq_len: int = 128,
    head_dim: int = 64,
    num_heads: int = 8,
    name: str = "Test",
    iterations: int = 1000,
    warmup: int = 50,
    dtype=torch.float32,
):
    """Run benchmark with causal masking (autoregressive attention)."""
    device = torch.device("mps")

    q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    k = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)
    v = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype)

    # Warmup
    for _ in range(warmup):
        _ = F.scaled_dot_product_attention(q, k, v, is_causal=True)
    torch.mps.synchronize()

    # Timed run
    start = time.perf_counter()

    for _ in range(iterations):
        _ = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.mps.synchronize()

    end = time.perf_counter()

    avg_time_ms = ((end - start) / iterations) * 1000
    total_elements = batch_size * num_heads * seq_len * head_dim
    throughput = total_elements / (avg_time_ms / 1000)

    print(f"| {name:<20} | {avg_time_ms:>10.4f} ms | {throughput/1e6:>10.2f} M/s |")


def main():
    check_mps()

    # Standard attention benchmarks (matching Rust scenarios)
    print("Standard Attention (non-causal)")
    print("-" * 70)
    print(f"| {'Scenario':<20} | {'Avg Time':>10} | {'Throughput':>12} |")
    print("-" * 70)

    benchmark(1, name="Latency_BS1")
    benchmark(64, name="Throughput_BS64")
    benchmark(256, name="Heavy_BS256")

    print("-" * 70)
    print()

    # Causal attention benchmarks
    print("Causal Attention (autoregressive)")
    print("-" * 70)
    print(f"| {'Scenario':<20} | {'Avg Time':>10} | {'Throughput':>12} |")
    print("-" * 70)

    benchmark_causal(1, name="Causal_BS1")
    benchmark_causal(32, seq_len=256, name="Causal_BS32")

    print("-" * 70)
    print()

    # Instructions
    print("To compare with RustyTorch Metal Flash Attention:")
    print("  cargo bench -p rtx-flash-metal-attention --bench metal_vs_pytorch")
    print()
    print("Expected Results:")
    print("  - Latency_BS1:    RustyTorch should win (zero-copy advantage)")
    print("  - Throughput_BS64: Parity expected (GPU-bound)")
    print("  - Heavy_BS256:    PyTorch may win (Apple-optimized tiling)")


if __name__ == "__main__":
    main()