#!/usr/bin/env python3
"""
SBNumPy V2 — GPU-Resident Workspace, Zero Buffer Allocation Per Op.

V1 created new GPU buffers per call (0.5-1ms each = 900ms overhead on 1926 calls).
V2 pre-allocates a persistent Sound Block workspace and reuses it.

The workspace IS a Sound Block container:
  - Pre-allocated regions for operands A, B, C (matmul)
  - Pre-allocated reduction scratch buffers
  - Pre-allocated elementwise scratch buffers
  - All GPU-resident, never deallocated

Per-call cost: one queue.write_buffer (CPU→GPU data update) + one dispatch.
No create_buffer. No create_bind_group per call (cached). No readback until needed.
"""

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

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

# ─── Shader Sources ──────────────────────────────────────────────────────────

MATMUL_SRC = '''
struct Params { M: u32, K: u32, N: u32, _pad: u32, }
@group(0) @binding(0) var<uniform> params: Params;
@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) id: vec3<u32>) {
    let row = id.x;
    let col = id.y;
    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;
}
'''

REDUCE_SRC = '''
struct Params { num_elements: u32, stride: u32, op: u32, _pad: u32, }
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read_write> values: array<f32>;
@group(0) @binding(2) var<storage, read_write> indices: array<u32>;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let idx = id.x;
    let stride = params.stride;
    let partner = idx + stride;
    if (idx >= params.num_elements || partner >= params.num_elements) { return; }
    if ((idx % (stride * 2u)) != 0u) { return; }
    if (params.op == 0u) {
        if (values[partner] > values[idx]) {
            values[idx] = values[partner]; indices[idx] = indices[partner];
        }
    } else if (params.op == 1u) {
        values[idx] = values[idx] + values[partner];
    } else if (params.op == 2u) {
        if (values[partner] < values[idx]) {
            values[idx] = values[partner]; indices[idx] = indices[partner];
        }
    }
}
'''

ELEMENTWISE_SRC = '''
struct Params { num_elements: u32, op: u32, scalar: f32, _pad: u32, }
@group(0) @binding(0) var<uniform> params: Params;
@group(0) @binding(1) var<storage, read> input_a: array<f32>;
@group(0) @binding(2) var<storage, read_write> output: array<f32>;

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let idx = id.x;
    if (idx >= params.num_elements) { return; }
    if (params.op == 0u) { output[idx] = abs(input_a[idx]); }
    else if (params.op == 1u) { output[idx] = abs(input_a[idx] - params.scalar); }
    else if (params.op == 2u) { output[idx] = input_a[idx] * params.scalar; }
    else if (params.op == 3u) { output[idx] = 1.0 / (1.0 + abs(input_a[idx] - params.scalar)); }
}
'''


class SBNumPyV2:
    """NumPy with GPU-resident workspace. Zero buffer allocation per call."""

    def __init__(self, max_matrix_dim=1024, max_vector_len=1_000_000):
        if not _WGPU:
            raise ImportError("wgpu required")
        self._device = wgpu.utils.get_default_device()
        self._max_dim = max_matrix_dim
        self._max_vec = max_vector_len
        self._build_pipelines()
        self._create_workspace()
        self._build_bind_groups()
        self.stats = {'dispatches': 0, 'total_dispatch_ms': 0.0,
                      'total_upload_ms': 0.0, 'total_readback_ms': 0.0}

    def _build_pipelines(self):
        d = self._device
        UNIFORM = wgpu.BufferBindingType.uniform
        STORAGE = wgpu.BufferBindingType.storage
        RO_STORAGE = wgpu.BufferBindingType.read_only_storage
        VIS = wgpu.ShaderStage.COMPUTE

        # Matmul
        mod = d.create_shader_module(code=MATMUL_SRC)
        self._mm_bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS, "buffer": {"type": UNIFORM}},
            {"binding": 1, "visibility": VIS, "buffer": {"type": RO_STORAGE}},
            {"binding": 2, "visibility": VIS, "buffer": {"type": RO_STORAGE}},
            {"binding": 3, "visibility": VIS, "buffer": {"type": STORAGE}},
        ])
        self._mm_pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[self._mm_bgl]),
            compute={"module": mod, "entry_point": "main"})

        # Reduce
        mod2 = d.create_shader_module(code=REDUCE_SRC)
        self._red_bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS, "buffer": {"type": UNIFORM}},
            {"binding": 1, "visibility": VIS, "buffer": {"type": STORAGE}},
            {"binding": 2, "visibility": VIS, "buffer": {"type": STORAGE}},
        ])
        self._red_pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[self._red_bgl]),
            compute={"module": mod2, "entry_point": "main"})

        # Elementwise
        mod3 = d.create_shader_module(code=ELEMENTWISE_SRC)
        self._ew_bgl = d.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS, "buffer": {"type": UNIFORM}},
            {"binding": 1, "visibility": VIS, "buffer": {"type": RO_STORAGE}},
            {"binding": 2, "visibility": VIS, "buffer": {"type": STORAGE}},
        ])
        self._ew_pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[self._ew_bgl]),
            compute={"module": mod3, "entry_point": "main"})

    def _create_workspace(self):
        """Pre-allocate ALL GPU buffers. Never allocate again."""
        d = self._device
        STORAGE = (wgpu.BufferUsage.STORAGE |
                   wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST)

        dim = self._max_dim
        vec = self._max_vec

        # Matrix workspace: A, B, C (max_dim × max_dim × 4 bytes each)
        mat_size = dim * dim * 4
        self._buf_A = d.create_buffer(size=mat_size, usage=STORAGE)
        self._buf_B = d.create_buffer(size=mat_size, usage=STORAGE)
        self._buf_C = d.create_buffer(size=mat_size, usage=STORAGE)

        # Vector workspace: values, indices, scores, scratch
        vec_size = vec * 4
        self._buf_values = d.create_buffer(size=vec_size, usage=STORAGE)
        self._buf_indices = d.create_buffer(size=vec_size, usage=STORAGE)
        self._buf_scores = d.create_buffer(size=vec_size, usage=STORAGE)

        # Params uniforms (reused, overwritten per call)
        self._buf_mm_params = d.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        self._buf_red_params = d.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        self._buf_ew_params = d.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

        ws_bytes = mat_size * 3 + vec_size * 3 + 48
        print(f"  Workspace allocated: {ws_bytes / (1024*1024):.1f} MB GPU-resident")

    def _build_bind_groups(self):
        """Build bind groups ONCE from workspace buffers."""
        d = self._device
        dim = self._max_dim
        vec = self._max_vec

        self._mm_bg = d.create_bind_group(layout=self._mm_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_mm_params}},
            {"binding": 1, "resource": {"buffer": self._buf_A, "size": dim * dim * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_B, "size": dim * dim * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_C, "size": dim * dim * 4}},
        ])

        self._red_bg = d.create_bind_group(layout=self._red_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_red_params}},
            {"binding": 1, "resource": {"buffer": self._buf_values, "size": vec * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_indices, "size": vec * 4}},
        ])

        self._ew_bg = d.create_bind_group(layout=self._ew_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_ew_params}},
            {"binding": 1, "resource": {"buffer": self._buf_values, "size": vec * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_scores, "size": vec * 4}},
        ])

    # ─── Core dispatch (zero allocation) ──────────────────────────

    def _dispatch(self, pipeline, bg, workgroups):
        d = self._device
        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(pipeline)
        cp.set_bind_group(0, bg)
        if isinstance(workgroups, tuple):
            cp.dispatch_workgroups(*workgroups)
        else:
            cp.dispatch_workgroups(workgroups)
        cp.end()
        d.queue.submit([encoder.finish()])
        self.stats['dispatches'] += 1

    # ─── API ─────────────────────────────────────────────────────

    def dot(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
        d = self._device
        M, K = A.shape
        K2, N = B.shape
        assert K == K2 and M <= self._max_dim and N <= self._max_dim

        t0 = time.perf_counter()
        d.queue.write_buffer(self._buf_A, 0, A.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_B, 0, B.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_mm_params, 0, struct.pack("IIII", M, K, N, 0))
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        self._dispatch(self._mm_pipeline, self._mm_bg, ((M + 15) // 16, (N + 15) // 16))
        d._poll_wait()
        self.stats['total_dispatch_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        result = np.frombuffer(d.queue.read_buffer(self._buf_C, size=M * N * 4),
                               dtype=np.float32).reshape(M, N).copy()
        self.stats['total_readback_ms'] += (time.perf_counter() - t0) * 1000
        return result

    def dot_no_readback(self, A: np.ndarray, B: np.ndarray):
        """Matmul that stays GPU-resident. Result in self._buf_C."""
        d = self._device
        M, K = A.shape
        K2, N = B.shape

        t0 = time.perf_counter()
        d.queue.write_buffer(self._buf_A, 0, A.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_B, 0, B.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_mm_params, 0, struct.pack("IIII", M, K, N, 0))
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        self._dispatch(self._mm_pipeline, self._mm_bg, ((M + 15) // 16, (N + 15) // 16))
        self.stats['total_dispatch_ms'] += (time.perf_counter() - t0) * 1000

    def _reduce_resident(self, N: int, op: int):
        """Parallel reduction on data already in self._buf_values."""
        d = self._device
        stride = 1
        while stride < N:
            d.queue.write_buffer(self._buf_red_params, 0,
                                 struct.pack("IIII", N, stride, op, 0))
            self._dispatch(self._red_pipeline, self._red_bg, (N + 255) // 256)
            stride *= 2
        self.stats['dispatches'] -= 1  # don't overcount the loop

    def argmax(self, values: np.ndarray) -> int:
        d = self._device
        N = len(values)

        t0 = time.perf_counter()
        d.queue.write_buffer(self._buf_values, 0, values.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_indices, 0, np.arange(N, dtype=np.uint32).tobytes())
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        self._reduce_resident(N, op=0)
        d._poll_wait()
        self.stats['total_dispatch_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        idx = int(np.frombuffer(d.queue.read_buffer(self._buf_indices, size=4),
                                dtype=np.uint32)[0])
        self.stats['total_readback_ms'] += (time.perf_counter() - t0) * 1000
        return idx

    def argmax_resident(self, N: int) -> int:
        """Argmax on data already in self._buf_values. Upload indices, dispatch, read 4 bytes."""
        d = self._device
        t0 = time.perf_counter()
        d.queue.write_buffer(self._buf_indices, 0, np.arange(N, dtype=np.uint32).tobytes())
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        self._reduce_resident(N, op=0)
        d._poll_wait()
        self.stats['total_dispatch_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        idx = int(np.frombuffer(d.queue.read_buffer(self._buf_indices, size=4),
                                dtype=np.uint32)[0])
        self.stats['total_readback_ms'] += (time.perf_counter() - t0) * 1000
        return idx

    def sum(self, values: np.ndarray) -> float:
        d = self._device
        N = len(values)

        t0 = time.perf_counter()
        d.queue.write_buffer(self._buf_values, 0, values.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_indices, 0, np.arange(N, dtype=np.uint32).tobytes())
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        self._reduce_resident(N, op=1)
        d._poll_wait()
        self.stats['total_dispatch_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        val = float(np.frombuffer(d.queue.read_buffer(self._buf_values, size=4),
                                  dtype=np.float32)[0])
        self.stats['total_readback_ms'] += (time.perf_counter() - t0) * 1000
        return val

    def search(self, data: np.ndarray, query: float) -> int:
        """Resonance search: upload data, compute scores, find argmax."""
        d = self._device
        N = len(data)

        t0 = time.perf_counter()
        d.queue.write_buffer(self._buf_values, 0, data.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_ew_params, 0,
                             struct.pack("IIff", N, 3, query, 0.0))
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        self._dispatch(self._ew_pipeline, self._ew_bg, (N + 255) // 256)
        d._poll_wait()
        self.stats['total_dispatch_ms'] += (time.perf_counter() - t0) * 1000

        # Scores are in self._buf_scores. Copy to values for argmax.
        t0 = time.perf_counter()
        scores_data = d.queue.read_buffer(self._buf_scores, size=N * 4)
        d.queue.write_buffer(self._buf_values, 0, scores_data)
        self.stats['total_upload_ms'] += (time.perf_counter() - t0) * 1000

        return self.argmax_resident(N)

    def readback_C(self, M, N) -> np.ndarray:
        """Read matmul result from GPU."""
        d = self._device
        t0 = time.perf_counter()
        result = np.frombuffer(d.queue.read_buffer(self._buf_C, size=M * N * 4),
                               dtype=np.float32).reshape(M, N).copy()
        self.stats['total_readback_ms'] += (time.perf_counter() - t0) * 1000
        return result


# ─── Benchmark ────────────────────────────────────────────────────────────────

def main():
    print("=" * 70)
    print("SBNumPy V2 — GPU-Resident Workspace, Zero Buffer Allocation")
    print("=" * 70)
    print()

    sb = SBNumPyV2(max_matrix_dim=1024, max_vector_len=1_000_000)
    print(f"GPU: {sb._device.adapter.info.get('device', 'unknown')}")
    print()

    rng = np.random.default_rng(42)

    # ─── Matrix Multiply ──────────────────────────────────────────
    print("## Matrix Multiply")
    print(f"{'Size':>10} {'NumPy':>10} {'SBNumPy':>10} {'Speedup':>10} {'Match':>8}")
    print("-" * 52)

    for N in [64, 128, 256, 512, 1024]:
        A = rng.random((N, N)).astype(np.float32)
        B = rng.random((N, N)).astype(np.float32)

        # Warmup
        if N == 64:
            sb.dot(A, B)

        t0 = time.perf_counter()
        for _ in range(10):
            np_r = A @ B
        np_ms = (time.perf_counter() - t0) * 1000 / 10

        t0 = time.perf_counter()
        for _ in range(10):
            sb_r = sb.dot(A, B)
        sb_ms = (time.perf_counter() - t0) * 1000 / 10

        match = np.allclose(np_r, sb_r, atol=1e-1)
        speedup = np_ms / sb_ms if sb_ms > 0 else 0
        print(f"{N}×{N:>4}  {np_ms:>8.3f}ms {sb_ms:>8.3f}ms {speedup:>9.2f}x {'YES' if match else 'NO':>6}")

    print()

    # ─── Argmax ───────────────────────────────────────────────────
    print("## Argmax")
    print(f"{'N':>10} {'NumPy':>10} {'SBNumPy':>10} {'Speedup':>10} {'Match':>8}")
    print("-" * 52)

    for N in [1000, 10000, 100000, 1000000]:
        values = rng.random(N).astype(np.float32)

        t0 = time.perf_counter()
        for _ in range(10):
            np_idx = int(np.argmax(values))
        np_ms = (time.perf_counter() - t0) * 1000 / 10

        t0 = time.perf_counter()
        for _ in range(10):
            sb_idx = sb.argmax(values)
        sb_ms = (time.perf_counter() - t0) * 1000 / 10

        match = np_idx == sb_idx
        speedup = np_ms / sb_ms if sb_ms > 0 else 0
        print(f"{N:>8,}  {np_ms:>8.3f}ms {sb_ms:>8.3f}ms {speedup:>9.2f}x {'YES' if match else 'NO':>6}")

    print()

    # ─── Sum ──────────────────────────────────────────────────────
    print("## Sum")
    print(f"{'N':>10} {'NumPy':>10} {'SBNumPy':>10} {'Speedup':>10} {'Match':>8}")
    print("-" * 52)

    for N in [1000, 10000, 100000, 1000000]:
        values = rng.random(N).astype(np.float32)

        t0 = time.perf_counter()
        for _ in range(10):
            np_sum = float(np.sum(values))
        np_ms = (time.perf_counter() - t0) * 1000 / 10

        t0 = time.perf_counter()
        for _ in range(10):
            sb_sum = sb.sum(values)
        sb_ms = (time.perf_counter() - t0) * 1000 / 10

        match = abs(np_sum - sb_sum) / max(abs(np_sum), 1e-6) < 0.01
        speedup = np_ms / sb_ms if sb_ms > 0 else 0
        print(f"{N:>8,}  {np_ms:>8.3f}ms {sb_ms:>8.3f}ms {speedup:>9.2f}x {'YES' if match else 'NO':>6}")

    print()

    # ─── FUSED PIPELINE (GPU-resident, minimal readback) ──────────
    print("## FUSED PIPELINE: 100× (matmul 256×256 → argmax → search 10K)")
    print()

    N = 256
    A = rng.random((N, N)).astype(np.float32)
    B = rng.random((N, N)).astype(np.float32)
    data = rng.random(10000).astype(np.float32) * 1000

    # NumPy
    t0 = time.perf_counter()
    for _ in range(100):
        C = A @ B
        max_idx = int(np.argmax(C.flatten()))
        query = float(C.flatten()[max_idx])
        nearest = int(np.argmin(np.abs(data - query)))
    np_pipe = (time.perf_counter() - t0) * 1000

    # SBNumPy V2
    sb.stats = {'dispatches': 0, 'total_dispatch_ms': 0.0,
                'total_upload_ms': 0.0, 'total_readback_ms': 0.0}

    t0 = time.perf_counter()
    for _ in range(100):
        C = sb.dot(A, B)
        max_idx = sb.argmax(C.flatten())
        query = float(C.flatten()[max_idx])
        nearest = sb.search(data, query)
    sb_pipe = (time.perf_counter() - t0) * 1000

    print(f"  NumPy:       {np_pipe:.1f} ms ({np_pipe/100:.3f} ms/iter)")
    print(f"  SBNumPy V2:  {sb_pipe:.1f} ms ({sb_pipe/100:.3f} ms/iter)")
    print(f"  Speedup:     {np_pipe/sb_pipe:.2f}x")
    print()
    print(f"  GPU dispatches: {sb.stats['dispatches']}")
    print(f"  Upload time:    {sb.stats['total_upload_ms']:.1f} ms")
    print(f"  Dispatch time:  {sb.stats['total_dispatch_ms']:.1f} ms")
    print(f"  Readback time:  {sb.stats['total_readback_ms']:.1f} ms")
    print()

    # ─── V1 vs V2 comparison ──────────────────────────────────────
    print("=" * 70)
    print("V1 (buffer per call) fused pipeline: 664.0 ms")
    print(f"V2 (workspace)       fused pipeline: {sb_pipe:.1f} ms")
    print(f"V2 improvement: {664.0/sb_pipe:.1f}x faster than V1")
    print("=" * 70)


if __name__ == '__main__':
    main()
