#!/usr/bin/env python3
"""
SB Tiled Matmul Benchmark — Compare naive GPU vs tiled GPU vs NumPy BLAS.

Tests at multiple matrix sizes to find the crossover point where
GPU tiled matmul beats NumPy/BLAS.

Usage:
    cd afs-engine
    python3 -m afs.gpu.sb_tiled_matmul_benchmark
"""

import numpy as np
import struct
import time
import sys
from pathlib import Path

try:
    import wgpu
    _WGPU_AVAILABLE = True
except ImportError:
    _WGPU_AVAILABLE = False

_SHADER_DIR = Path(__file__).parent

# ─── Naive matmul WGSL (no shared memory, for comparison) ───────────────────

NAIVE_MATMUL_WGSL = """
struct MatmulParams {
    M: u32,
    N: u32,
    K: u32,
    _pad: u32,
}

@group(0) @binding(0) var<uniform> params: MatmulParams;
@group(0) @binding(1) var<storage, read> A: array<f32>;
@group(0) @binding(2) var<storage, read> B: array<f32>;
@group(0) @binding(3) var<storage, read_write> C: array<f32>;

@compute @workgroup_size(16, 16)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
    let row = global_id.y;
    let col = global_id.x;

    if (row >= params.M || col >= params.N) { return; }

    var sum: f32 = 0.0;
    for (var k = 0u; k < params.K; k++) {
        sum += A[row * params.K + k] * B[k * params.N + col];
    }
    C[row * params.N + col] = sum;
}
"""


class GPUMatmul:
    """GPU matmul wrapper for benchmarking."""

    def __init__(self, tiled: bool = True):
        if not _WGPU_AVAILABLE:
            raise ImportError("wgpu package required: pip install wgpu")

        self._device = wgpu.utils.get_default_device()
        self._tiled = tiled

        if tiled:
            shader_src = (_SHADER_DIR / "sb_tiled_matmul.wgsl").read_text()
        else:
            shader_src = NAIVE_MATMUL_WGSL

        module = self._device.create_shader_module(code=shader_src)

        self._bgl = self._device.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])

        pl = self._device.create_pipeline_layout(bind_group_layouts=[self._bgl])
        self._pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "main"})

    def matmul(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
        """Compute C = A @ B on GPU. Returns C as numpy array."""
        M, K = A.shape
        K2, N = B.shape
        assert K == K2, f"Inner dimensions mismatch: {K} vs {K2}"

        d = self._device

        STORAGE_RW = (wgpu.BufferUsage.STORAGE |
                      wgpu.BufferUsage.COPY_SRC |
                      wgpu.BufferUsage.COPY_DST)
        STORAGE_RO = (wgpu.BufferUsage.STORAGE |
                      wgpu.BufferUsage.COPY_DST)

        buf_a = d.create_buffer(size=M * K * 4, usage=STORAGE_RO)
        buf_b = d.create_buffer(size=K * N * 4, usage=STORAGE_RO)
        buf_c = d.create_buffer(size=M * N * 4, usage=STORAGE_RW)
        buf_params = d.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

        d.queue.write_buffer(buf_a, 0, A.astype(np.float32).tobytes())
        d.queue.write_buffer(buf_b, 0, B.astype(np.float32).tobytes())
        d.queue.write_buffer(buf_params, 0, struct.pack("IIII", M, N, K, 0))

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": buf_params}},
            {"binding": 1, "resource": {"buffer": buf_a}},
            {"binding": 2, "resource": {"buffer": buf_b}},
            {"binding": 3, "resource": {"buffer": buf_c}},
        ])

        wg_x = (N + 15) // 16
        wg_y = (M + 15) // 16

        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._pipeline)
        compute_pass.set_bind_group(0, bg)
        compute_pass.dispatch_workgroups(wg_x, wg_y)
        compute_pass.end()
        d.queue.submit([encoder.finish()])
        d._poll_wait()

        c_data = d.queue.read_buffer(buf_c, size=M * N * 4)
        return np.frombuffer(c_data, dtype=np.float32).reshape(M, N).copy()


def benchmark_size(size: int, naive_gpu: GPUMatmul, tiled_gpu: GPUMatmul,
                   warmup: int = 2, trials: int = 5):
    """Benchmark one matrix size across all three methods."""
    np.random.seed(42)
    A = np.random.randn(size, size).astype(np.float32)
    B = np.random.randn(size, size).astype(np.float32)

    # --- NumPy/BLAS ---
    # Warmup
    for _ in range(warmup):
        _ = A @ B
    times = []
    for _ in range(trials):
        t0 = time.perf_counter()
        C_np = A @ B
        times.append(time.perf_counter() - t0)
    numpy_time = np.median(times) * 1000

    # --- Naive GPU ---
    for _ in range(warmup):
        _ = naive_gpu.matmul(A, B)
    times = []
    for _ in range(trials):
        t0 = time.perf_counter()
        C_naive = naive_gpu.matmul(A, B)
        times.append(time.perf_counter() - t0)
    naive_time = np.median(times) * 1000

    # --- Tiled GPU ---
    for _ in range(warmup):
        _ = tiled_gpu.matmul(A, B)
    times = []
    for _ in range(trials):
        t0 = time.perf_counter()
        C_tiled = tiled_gpu.matmul(A, B)
        times.append(time.perf_counter() - t0)
    tiled_time = np.median(times) * 1000

    # Verify correctness (allow f32 tolerance)
    naive_err = np.max(np.abs(C_np - C_naive))
    tiled_err = np.max(np.abs(C_np - C_tiled))

    return {
        "size": size,
        "numpy_ms": numpy_time,
        "naive_ms": naive_time,
        "tiled_ms": tiled_time,
        "naive_err": naive_err,
        "tiled_err": tiled_err,
    }


if __name__ == "__main__":
    if not _WGPU_AVAILABLE:
        print("ERROR: wgpu not available. Install with: pip install wgpu")
        sys.exit(1)

    print("=" * 75)
    print("SB Tiled Matmul Benchmark — Naive GPU vs Tiled GPU vs NumPy/BLAS")
    print("=" * 75)

    try:
        naive_gpu = GPUMatmul(tiled=False)
        tiled_gpu = GPUMatmul(tiled=True)
        info = dict(naive_gpu._device.adapter.info)
        print(f"GPU: {info.get('description', 'unknown')}")
    except Exception as e:
        print(f"GPU init failed: {e}")
        sys.exit(1)

    sizes = [64, 128, 256, 512, 1024]

    print(f"\n{'Size':>6} | {'NumPy/BLAS':>10} | {'Naive GPU':>10} | {'Tiled GPU':>10} | {'Tiled/NumPy':>11} | {'Tiled err':>9}")
    print("-" * 75)

    crossover = None

    for size in sizes:
        try:
            result = benchmark_size(size, naive_gpu, tiled_gpu)
            ratio = result['tiled_ms'] / result['numpy_ms'] if result['numpy_ms'] > 0 else float('inf')
            tiled_vs_naive = result['tiled_ms'] / result['naive_ms'] if result['naive_ms'] > 0 else float('inf')

            print(f"{result['size']:>6} | {result['numpy_ms']:>8.2f}ms | {result['naive_ms']:>8.2f}ms | "
                  f"{result['tiled_ms']:>8.2f}ms | {ratio:>10.2f}x | {result['tiled_err']:>8.2e}")

            if crossover is None and ratio < 1.0:
                crossover = size

        except Exception as e:
            print(f"{size:>6} | ERROR: {e}")

    print()
    if crossover:
        print(f"Crossover point: GPU tiled beats NumPy/BLAS at size {crossover}")
    else:
        print("No crossover found — NumPy/BLAS is faster at all tested sizes.")
        print("This is expected: NumPy uses optimized BLAS (OpenBLAS/MKL) which is")
        print("extremely fast for matmul. GPU wins at larger sizes or when data is")
        print("already GPU-resident (no transfer overhead).")

    # Quick correctness check at small size
    print("\n--- Correctness check (64x64) ---")
    np.random.seed(123)
    A = np.random.randn(64, 64).astype(np.float32)
    B = np.random.randn(64, 64).astype(np.float32)
    C_ref = A @ B
    C_tiled = tiled_gpu.matmul(A, B)
    max_err = np.max(np.abs(C_ref - C_tiled))
    rel_err = max_err / np.max(np.abs(C_ref))
    print(f"  Max absolute error: {max_err:.2e}")
    print(f"  Max relative error: {rel_err:.2e}")
    print(f"  {'PASS' if rel_err < 1e-4 else 'FAIL'}: Tiled matmul matches NumPy within f32 tolerance")
