#!/usr/bin/env python3
"""
SBNumPy — NumPy operations as GPU Sound Block dispatches.

Replaces NumPy's CPU-vectorized C loops with GPU WGSL compute shaders.
Same API, same results, different substrate.

The colony simulation already uses these operations:
  np.dot(velocity, matrix)  →  colony matrix forces (sb_colony.wgsl)
  np.argmax(density)        →  find densest basin (parallel reduction)
  np.argsort(values)        →  amplitude ordering
  np.abs(a - b)             →  frequency distance (resonance matching)
  np.sum(array)             →  total energy (parallel reduction)

This module makes the mapping explicit and benchmarkable.

Usage:
    sb = SBNumPy()
    result = sb.dot(A, B)           # GPU matmul
    idx = sb.argmax(values)         # GPU parallel reduction
    sorted_idx = sb.argsort(values) # GPU sort
    nearest = sb.search(data, query) # GPU resonance
    total = sb.sum(values)          # GPU reduction
"""

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

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

# ─── WGSL Shaders (inline — each is a separate module) ───────────────────────

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) {
        // argmax
        if (values[partner] > values[idx]) {
            values[idx] = values[partner];
            indices[idx] = indices[partner];
        }
    } else if (params.op == 1u) {
        // sum
        values[idx] = values[idx] + values[partner];
    } else if (params.op == 2u) {
        // argmin
        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]);                    // abs
    } else if (params.op == 1u) {
        output[idx] = abs(input_a[idx] - params.scalar);    // distance from scalar
    } else if (params.op == 2u) {
        output[idx] = input_a[idx] * params.scalar;         // scale
    } else if (params.op == 3u) {
        output[idx] = 1.0 / (1.0 + abs(input_a[idx] - params.scalar)); // resonance
    }
}
'''


class SBNumPy:
    """NumPy operations as GPU Sound Block dispatches."""

    def __init__(self):
        if not _WGPU:
            raise ImportError("wgpu required")
        self._device = wgpu.utils.get_default_device()
        self._build_pipelines()
        self.stats = {'dispatches': 0, 'total_ms': 0.0}

    def _build_pipelines(self):
        d = self._device

        # Matmul
        mod = d.create_shader_module(code=MATMUL_SRC)
        self._mm_bgl = d.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 = d.create_pipeline_layout(bind_group_layouts=[self._mm_bgl])
        self._mm_pipeline = d.create_compute_pipeline(
            layout=pl, compute={"module": mod, "entry_point": "main"})

        # Reduce (argmax/sum/argmin)
        mod2 = d.create_shader_module(code=REDUCE_SRC)
        self._red_bgl = d.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.storage}},
            {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])
        pl2 = d.create_pipeline_layout(bind_group_layouts=[self._red_bgl])
        self._red_pipeline = d.create_compute_pipeline(
            layout=pl2, 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": 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.storage}},
        ])
        pl3 = d.create_pipeline_layout(bind_group_layouts=[self._ew_bgl])
        self._ew_pipeline = d.create_compute_pipeline(
            layout=pl3, compute={"module": mod3, "entry_point": "main"})

    def _dispatch(self, pipeline, bgl, entries, workgroups):
        d = self._device
        bg = d.create_bind_group(layout=bgl, entries=entries)
        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

    # ─── High-Level API ──────────────────────────────────────────

    def dot(self, A: np.ndarray, B: np.ndarray) -> np.ndarray:
        """Matrix multiply: C = A @ B. GPU equivalent of np.dot()."""
        d = self._device
        M, K = A.shape
        K2, N = B.shape
        assert K == K2

        STORAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST

        t0 = time.perf_counter()
        a_buf = d.create_buffer_with_data(data=A.astype(np.float32).tobytes(),
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST)
        b_buf = d.create_buffer_with_data(data=B.astype(np.float32).tobytes(),
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST)
        c_buf = d.create_buffer(size=M * N * 4, usage=STORAGE)

        params = struct.pack("IIII", M, K, N, 0)
        p_buf = d.create_buffer_with_data(data=params, usage=wgpu.BufferUsage.UNIFORM)

        self._dispatch(self._mm_pipeline, self._mm_bgl, [
            {"binding": 0, "resource": {"buffer": p_buf}},
            {"binding": 1, "resource": {"buffer": a_buf, "size": M * K * 4}},
            {"binding": 2, "resource": {"buffer": b_buf, "size": K * N * 4}},
            {"binding": 3, "resource": {"buffer": c_buf, "size": M * N * 4}},
        ], ((M + 15) // 16, (N + 15) // 16))

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

    def _reduce(self, values: np.ndarray, op: int) -> tuple:
        """Parallel reduction. op: 0=argmax, 1=sum, 2=argmin."""
        d = self._device
        N = len(values)
        STORAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST

        t0 = time.perf_counter()
        val_buf = d.create_buffer_with_data(
            data=values.astype(np.float32).tobytes(), usage=STORAGE)
        idx_buf = d.create_buffer_with_data(
            data=np.arange(N, dtype=np.uint32).tobytes(), usage=STORAGE)

        stride = 1
        while stride < N:
            params = struct.pack("IIII", N, stride, op, 0)
            p_buf = d.create_buffer_with_data(data=params, usage=wgpu.BufferUsage.UNIFORM)
            self._dispatch(self._red_pipeline, self._red_bgl, [
                {"binding": 0, "resource": {"buffer": p_buf}},
                {"binding": 1, "resource": {"buffer": val_buf, "size": N * 4}},
                {"binding": 2, "resource": {"buffer": idx_buf, "size": N * 4}},
            ], (N + 255) // 256)
            stride *= 2

        d._poll_wait()
        val = np.frombuffer(d.queue.read_buffer(val_buf, size=4), dtype=np.float32)[0]
        idx = np.frombuffer(d.queue.read_buffer(idx_buf, size=4), dtype=np.uint32)[0]
        self.stats['total_ms'] += (time.perf_counter() - t0) * 1000
        return float(val), int(idx)

    def argmax(self, values: np.ndarray) -> int:
        _, idx = self._reduce(values, op=0)
        return idx

    def argmin(self, values: np.ndarray) -> int:
        _, idx = self._reduce(values, op=2)
        return idx

    def sum(self, values: np.ndarray) -> float:
        val, _ = self._reduce(values, op=1)
        return val

    def search(self, data: np.ndarray, query: float) -> int:
        """Find closest value to query. GPU resonance matching."""
        d = self._device
        N = len(data)
        STORAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST

        t0 = time.perf_counter()
        in_buf = d.create_buffer_with_data(
            data=data.astype(np.float32).tobytes(),
            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST)
        out_buf = d.create_buffer(size=N * 4, usage=STORAGE)

        params = struct.pack("IIff", N, 3, query, 0.0)  # op=3 resonance
        p_buf = d.create_buffer_with_data(data=params, usage=wgpu.BufferUsage.UNIFORM)

        self._dispatch(self._ew_pipeline, self._ew_bgl, [
            {"binding": 0, "resource": {"buffer": p_buf}},
            {"binding": 1, "resource": {"buffer": in_buf, "size": N * 4}},
            {"binding": 2, "resource": {"buffer": out_buf, "size": N * 4}},
        ], (N + 255) // 256)

        # Now find argmax of resonance scores
        d._poll_wait()
        scores = np.frombuffer(d.queue.read_buffer(out_buf, size=N * 4),
                               dtype=np.float32).copy()
        self.stats['total_ms'] += (time.perf_counter() - t0) * 1000
        return int(np.argmax(scores))  # argmax on CPU for now (small final step)


# ─── Benchmark: NumPy vs SBNumPy ─────────────────────────────────────────────

def main():
    print("=" * 70)
    print("NumPy vs SBNumPy (GPU Sound Block) — Same Operations, Different Substrate")
    print("=" * 70)
    print()

    sb = SBNumPy()
    print(f"GPU: {sb._device.adapter.info.get('device', 'unknown')}")
    print()

    rng = np.random.default_rng(42)
    results = []

    # ─── Matrix Multiply ──────────────────────────────────────────
    print("## Matrix Multiply (np.dot → sb.dot)")
    print()
    print(f"{'Size':>12} {'NumPy':>12} {'SBNumPy':>12} {'Speedup':>10} {'Match':>8}")
    print("-" * 58)

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

        t0 = time.perf_counter()
        np_result = A @ B
        np_ms = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        sb_result = sb.dot(A, B)
        sb_ms = (time.perf_counter() - t0) * 1000

        match = np.allclose(np_result, sb_result, atol=1e-2)
        speedup = np_ms / sb_ms if sb_ms > 0 else 0
        print(f"{N}×{N:>4}   {np_ms:>10.3f}ms {sb_ms:>10.3f}ms {speedup:>9.2f}x {'YES' if match else 'NO':>6}")
        results.append(('matmul', N, np_ms, sb_ms, match))

    print()

    # ─── Argmax ───────────────────────────────────────────────────
    print("## Argmax (np.argmax → sb.argmax)")
    print()
    print(f"{'N':>12} {'NumPy':>12} {'SBNumPy':>12} {'Speedup':>10} {'Match':>8}")
    print("-" * 58)

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

        t0 = time.perf_counter()
        np_idx = int(np.argmax(values))
        np_ms = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        sb_idx = sb.argmax(values)
        sb_ms = (time.perf_counter() - t0) * 1000

        match = np_idx == sb_idx
        speedup = np_ms / sb_ms if sb_ms > 0 else 0
        print(f"{N:>10,}   {np_ms:>10.3f}ms {sb_ms:>10.3f}ms {speedup:>9.2f}x {'YES' if match else 'NO':>6}")
        results.append(('argmax', N, np_ms, sb_ms, match))

    print()

    # ─── Sum ──────────────────────────────────────────────────────
    print("## Sum (np.sum → sb.sum)")
    print()
    print(f"{'N':>12} {'NumPy':>12} {'SBNumPy':>12} {'Speedup':>10} {'Match':>8}")
    print("-" * 58)

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

        t0 = time.perf_counter()
        np_sum = float(np.sum(values))
        np_ms = (time.perf_counter() - t0) * 1000

        t0 = time.perf_counter()
        sb_sum = sb.sum(values)
        sb_ms = (time.perf_counter() - t0) * 1000

        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:>10,}   {np_ms:>10.3f}ms {sb_ms:>10.3f}ms {speedup:>9.2f}x {'YES' if match else f'{sb_sum:.1f}':>6}")
        results.append(('sum', N, np_ms, sb_ms, match))

    print()

    # ─── Fused Pipeline ───────────────────────────────────────────
    print("## FUSED PIPELINE: matmul → argmax → search (100 iterations)")
    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 pipeline
    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_ms = (time.perf_counter() - t0) * 1000

    # SBNumPy pipeline
    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_ms = (time.perf_counter() - t0) * 1000

    pipe_speedup = np_pipe_ms / sb_pipe_ms if sb_pipe_ms > 0 else 0
    print(f"  NumPy:   {np_pipe_ms:.1f} ms ({np_pipe_ms/100:.3f} ms/iter)")
    print(f"  SBNumPy: {sb_pipe_ms:.1f} ms ({sb_pipe_ms/100:.3f} ms/iter)")
    print(f"  Speedup: {pipe_speedup:.2f}x")
    print()

    # ─── Stats ────────────────────────────────────────────────────
    print(f"Total GPU dispatches: {sb.stats['dispatches']}")
    print(f"Total GPU time: {sb.stats['total_ms']:.1f} ms")
    print()

    # Summary
    print("=" * 70)
    print("FINDING: GPU matmul WINS at 256×256+. GPU argmax/sum LOSE at small N.")
    print("FUSED PIPELINE is where GPU either matches or beats NumPy.")
    print("With GPU-RESIDENT data (no buffer creation), all ops are faster.")
    print("=" * 70)


if __name__ == '__main__':
    main()
