#!/usr/bin/env python3
"""
SBFusedPipeline — Pre-compiled GPU operation sequences. Zero per-op overhead.

The problem: Individual GPU dispatches cost ~0.5ms each in wgpu-py due to
Python→wgpu→Vulkan driver overhead. SBNumPyV2 eliminated buffer allocation
overhead but still does one submit+wait per operation.

The solution: Record a SEQUENCE of operations, compile them into ONE command
encoder with multiple compute passes, and execute the entire sequence with
a single queue.submit() + _poll_wait().

Usage:
    pipe = SBFusedPipeline(device)
    pipe.begin_recording()
    pipe.matmul('A', 'B', 'C', M, K, N)       # recorded, not executed
    pipe.argmax('C', N*N)                       # recorded (multi-pass reduce)
    pipe.search('A', query, N)                  # recorded (elementwise + copy + reduce)
    pipe.compile()                              # builds ONE command buffer

    # Execute 100 times — ONE submit per execution
    for i in range(100):
        pipe.upload('A', new_data)
        pipe.execute()                          # single queue.submit + one _poll_wait
        result = pipe.readback('indices', 4)    # read 4 bytes (one u32 index)

Key insight: wgpu compute passes within a single command encoder execute
sequentially on the GPU. We exploit this to chain matmul→reduce→elementwise
with zero inter-op CPU overhead.

Challenge: Uniform buffers (params) are written via queue.write_buffer() which
executes BEFORE the command buffer. Multiple writes to the same uniform buffer
would only keep the last value. Solution: pre-allocate separate param buffers
for each recorded operation, write all params before submit, execute all passes
in one command buffer.
"""

import numpy as np
import struct
import time
import math

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

# ─── Shader Sources (identical to sb_numpy_v2.py) ────────────────────────────

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)); }
}
'''

# Elementwise with add-noise: output[i] = input_a[i] + scalar * input_b[i]
# Uses a separate shader because standard elementwise only has one input.
ELEMENTWISE_ADD_SRC = '''
struct Params { num_elements: u32, scalar: f32, _pad1: u32, _pad2: 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> input_b: array<f32>;
@group(0) @binding(3) 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; }
    output[idx] = input_a[idx] + params.scalar * input_b[idx];
}
'''

# Clamp shader: output[i] = clamp(input[i], lo, hi)
CLAMP_SRC = '''
struct Params { num_elements: u32, lo: f32, hi: 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; }
    output[idx] = clamp(input_a[idx], params.lo, params.hi);
}
'''


# ─── Recorded Operation Types ────────────────────────────────────────────────

OP_MATMUL = 'matmul'
OP_REDUCE = 'reduce'
OP_ELEMENTWISE = 'elementwise'
OP_ELEMENTWISE_ADD = 'elementwise_add'
OP_CLAMP = 'clamp'
OP_COPY = 'copy'          # GPU buffer-to-buffer copy (encoder.copy_buffer_to_buffer)
OP_INIT_INDICES = 'init_indices'  # Upload index array for argmax before reduce


class SBFusedPipeline:
    """Pre-compiled GPU operation sequences. Zero per-op overhead.

    Lifecycle:
        1. __init__    — allocate workspace buffers, build pipelines + bind groups
        2. begin_recording() — start recording operations
        3. matmul/argmax/search/elementwise/clamp — record ops (no GPU calls)
        4. compile()   — build ONE command encoder from all recorded ops
        5. execute()   — single queue.submit() + _poll_wait()
        6. upload/readback — update inputs / read outputs between executions
    """

    # Named workspace regions → buffer index
    REGION_NAMES = {
        'A': 'mat_A', 'B': 'mat_B', 'C': 'mat_C',
        'values': 'vec_values', 'indices': 'vec_indices',
        'scores': 'vec_scores', 'scratch': 'vec_scratch',
    }

    def __init__(self, device=None, max_matrix_dim=1024, max_vector_len=1_000_000,
                 max_reduce_passes=20):
        if not _WGPU:
            raise ImportError("wgpu required")

        self._device = device or wgpu.utils.get_default_device()
        self._max_dim = max_matrix_dim
        self._max_vec = max_vector_len
        self._max_reduce_passes = max_reduce_passes

        # Build compute pipelines (shader modules + pipeline objects)
        self._build_pipelines()

        # Pre-allocate ALL workspace buffers
        self._create_workspace()

        # Pre-build bind groups for each pipeline type
        self._build_bind_groups()

        # Recording state
        self._recording = False
        self._ops = []           # list of recorded operation dicts
        self._compiled = False

        # Compiled state: list of (param_data, param_buffer) for pre-submit writes
        self._param_writes = []  # [(bytes, buffer), ...]
        self._index_inits = []   # [(np.ndarray, buffer), ...] for argmax index setup

        # Stats
        self.stats = {
            'compile_ms': 0.0,
            'execute_ms': 0.0,
            'upload_ms': 0.0,
            'readback_ms': 0.0,
            'executions': 0,
            'ops_per_execution': 0,
        }

    # ─── Pipeline Setup ──────────────────────────────────────────────────────

    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 pipeline
        mm_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": mm_mod, "entry_point": "main"})

        # Reduce pipeline
        red_mod = 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": red_mod, "entry_point": "main"})

        # Elementwise pipeline (single input)
        ew_mod = 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": ew_mod, "entry_point": "main"})

        # Elementwise-add pipeline (two inputs: a + scalar * b)
        ew_add_mod = d.create_shader_module(code=ELEMENTWISE_ADD_SRC)
        self._ew_add_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._ew_add_pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[self._ew_add_bgl]),
            compute={"module": ew_add_mod, "entry_point": "main"})

        # Clamp pipeline
        clamp_mod = d.create_shader_module(code=CLAMP_SRC)
        self._clamp_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._clamp_pipeline = d.create_compute_pipeline(
            layout=d.create_pipeline_layout(bind_group_layouts=[self._clamp_bgl]),
            compute={"module": clamp_mod, "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 x max_dim x 4 bytes each)
        mat_size = dim * dim * 4
        self._buf = {}
        self._buf['mat_A'] = d.create_buffer(size=mat_size, usage=STORAGE)
        self._buf['mat_B'] = d.create_buffer(size=mat_size, usage=STORAGE)
        self._buf['mat_C'] = d.create_buffer(size=mat_size, usage=STORAGE)

        # Vector workspace: values, indices, scores, scratch
        vec_size = vec * 4
        self._buf['vec_values'] = d.create_buffer(size=vec_size, usage=STORAGE)
        self._buf['vec_indices'] = d.create_buffer(size=vec_size, usage=STORAGE)
        self._buf['vec_scores'] = d.create_buffer(size=vec_size, usage=STORAGE)
        self._buf['vec_scratch'] = d.create_buffer(size=vec_size, usage=STORAGE)

        # Pre-allocate param uniform buffers — one per possible recorded op.
        # Each reduce multi-pass needs its own param buffer per stride level.
        # We allocate a pool of param buffers and assign them during compile.
        self._param_pool_size = 64  # enough for complex pipelines
        UNIFORM_USAGE = wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST
        self._param_bufs = [
            d.create_buffer(size=16, usage=UNIFORM_USAGE)
            for _ in range(self._param_pool_size)
        ]
        self._next_param_idx = 0

        self._ws_bytes = mat_size * 3 + vec_size * 4 + 16 * self._param_pool_size

    def _alloc_param_buf(self):
        """Get next available param buffer from the pool."""
        if self._next_param_idx >= self._param_pool_size:
            raise RuntimeError(
                f"Param buffer pool exhausted ({self._param_pool_size}). "
                f"Increase pool size or simplify the pipeline.")
        buf = self._param_bufs[self._next_param_idx]
        self._next_param_idx += 1
        return buf

    def _build_bind_groups(self):
        """Build bind groups are deferred to compile() since each op needs
        its own param buffer. We only build the workspace buffers here."""
        # Bind groups are created during compile() with specific param buffers.
        pass

    def _resolve_region(self, name):
        """Resolve a region name to a buffer key."""
        if name in self._buf:
            return name
        if name in self.REGION_NAMES:
            return self.REGION_NAMES[name]
        raise KeyError(f"Unknown region '{name}'. Available: {list(self.REGION_NAMES.keys())}")

    # ─── Recording API ───────────────────────────────────────────────────────

    def begin_recording(self):
        """Start recording operations. Clears any previous recording."""
        self._recording = True
        self._compiled = False
        self._ops = []
        self._param_writes = []
        self._index_inits = []
        self._next_param_idx = 0

    def matmul(self, src_a, src_b, dst_c, M, K, N):
        """Record a matmul: dst_c = src_a @ src_b.

        Args:
            src_a: region name for input matrix A (M x K)
            src_b: region name for input matrix B (K x N)
            dst_c: region name for output matrix C (M x N)
            M, K, N: matrix dimensions
        """
        assert self._recording, "Call begin_recording() first"
        self._ops.append({
            'type': OP_MATMUL,
            'src_a': self._resolve_region(src_a),
            'src_b': self._resolve_region(src_b),
            'dst': self._resolve_region(dst_c),
            'M': M, 'K': K, 'N': N,
        })

    def argmax(self, src_values, N):
        """Record an argmax operation on a values buffer.

        Reads from src_values, writes result index to 'indices'[0].
        Note: this destructively modifies the values buffer during reduction.
        The source data is first copied to 'values' workspace if it differs.

        Args:
            src_values: region name containing float32 values
            N: number of elements
        """
        assert self._recording, "Call begin_recording() first"
        src_key = self._resolve_region(src_values)

        # If source is not the values buffer, record a copy first
        if src_key != 'vec_values':
            self._ops.append({
                'type': OP_COPY,
                'src': src_key,
                'dst': 'vec_values',
                'size': N * 4,
            })

        # Record index initialization (arange(N) uploaded before submit)
        self._ops.append({
            'type': OP_INIT_INDICES,
            'N': N,
        })

        # Record log2(N) reduce passes with doubling stride
        stride = 1
        while stride < N:
            self._ops.append({
                'type': OP_REDUCE,
                'N': N,
                'stride': stride,
                'op': 0,  # argmax
            })
            stride *= 2

    def search(self, src_data, query, N):
        """Record a resonance search: compute scores, then argmax.

        Computes score[i] = 1.0 / (1.0 + |data[i] - query|) then finds argmax.
        Reads from src_data, writes score to 'scores', result index to 'indices'[0].

        Note: query is baked into the compiled pipeline. To change query between
        executions, recompile.

        Args:
            src_data: region name containing float32 search data
            query: float search target value
            N: number of elements
        """
        assert self._recording, "Call begin_recording() first"
        src_key = self._resolve_region(src_data)

        # Elementwise: scores[i] = 1/(1+|data[i]-query|)  (op=3, resonance)
        self._ops.append({
            'type': OP_ELEMENTWISE,
            'src': src_key,
            'dst': 'vec_scores',
            'N': N,
            'op': 3,  # resonance
            'scalar': query,
        })

        # Copy scores → values for reduce (argmax reads from values)
        self._ops.append({
            'type': OP_COPY,
            'src': 'vec_scores',
            'dst': 'vec_values',
            'size': N * 4,
        })

        # Index init + argmax reduce
        self._ops.append({
            'type': OP_INIT_INDICES,
            'N': N,
        })

        stride = 1
        while stride < N:
            self._ops.append({
                'type': OP_REDUCE,
                'N': N,
                'stride': stride,
                'op': 0,
            })
            stride *= 2

    def elementwise_add(self, src_a, src_b, dst, N, scalar=1.0):
        """Record: dst[i] = src_a[i] + scalar * src_b[i].

        Args:
            src_a, src_b: input region names
            dst: output region name
            N: number of elements
            scalar: multiplier for src_b (default 1.0)
        """
        assert self._recording, "Call begin_recording() first"
        self._ops.append({
            'type': OP_ELEMENTWISE_ADD,
            'src_a': self._resolve_region(src_a),
            'src_b': self._resolve_region(src_b),
            'dst': self._resolve_region(dst),
            'N': N,
            'scalar': scalar,
        })

    def clamp(self, src, dst, N, lo=0.0, hi=1.0):
        """Record: dst[i] = clamp(src[i], lo, hi).

        Args:
            src: input region name
            dst: output region name
            N: number of elements
            lo, hi: clamp bounds
        """
        assert self._recording, "Call begin_recording() first"
        self._ops.append({
            'type': OP_CLAMP,
            'src': self._resolve_region(src),
            'dst': self._resolve_region(dst),
            'N': N,
            'lo': lo,
            'hi': hi,
        })

    def elementwise(self, src, dst, N, op, scalar=0.0):
        """Record a single-input elementwise operation.

        Args:
            src: input region name
            dst: output region name
            N: number of elements
            op: 0=abs, 1=abs(x-s), 2=x*s, 3=resonance(x,s)
            scalar: scalar parameter
        """
        assert self._recording, "Call begin_recording() first"
        self._ops.append({
            'type': OP_ELEMENTWISE,
            'src': self._resolve_region(src),
            'dst': self._resolve_region(dst),
            'N': N,
            'op': op,
            'scalar': scalar,
        })

    # ─── Compile ─────────────────────────────────────────────────────────────

    def compile(self):
        """Compile all recorded ops into a reusable execution recipe.

        wgpu command buffers are single-use (consumed on submit), so we cannot
        cache the command buffer itself. Instead, compile() pre-builds:
        - Bind groups (GPU objects, reusable across submits)
        - Param data (bytes to write to uniform buffers before each submit)
        - Index arrays (for argmax initialization)
        - A recipe list of (pipeline, bind_group, workgroups) tuples and copy ops

        execute() then builds ONE command encoder from this recipe, submits ONCE,
        and waits ONCE. The encoder construction is pure CPU bookkeeping (~0.01ms)
        — the expensive queue.submit+_poll_wait happens exactly once.
        """
        assert not self._recording or len(self._ops) > 0, "No ops recorded"
        self._recording = False

        t0 = time.perf_counter()

        d = self._device
        self._param_writes = []
        self._index_inits = []
        self._recipe = []        # list of dicts describing each pass
        self._next_param_idx = 0

        for op in self._ops:
            if op['type'] == OP_MATMUL:
                self._compile_matmul(op)
            elif op['type'] == OP_REDUCE:
                self._compile_reduce(op)
            elif op['type'] == OP_ELEMENTWISE:
                self._compile_elementwise(op)
            elif op['type'] == OP_ELEMENTWISE_ADD:
                self._compile_elementwise_add(op)
            elif op['type'] == OP_CLAMP:
                self._compile_clamp(op)
            elif op['type'] == OP_COPY:
                self._compile_copy(op)
            elif op['type'] == OP_INIT_INDICES:
                self._compile_init_indices(op)
            else:
                raise ValueError(f"Unknown op type: {op['type']}")

        self._compiled = True

        self.stats['ops_per_execution'] = len(self._ops)
        self.stats['compile_ms'] = (time.perf_counter() - t0) * 1000

    def _compile_matmul(self, op):
        """Pre-build bind group and record recipe for a matmul pass."""
        d = self._device
        param_buf = self._alloc_param_buf()
        param_data = struct.pack("IIII", op['M'], op['K'], op['N'], 0)
        self._param_writes.append((param_data, param_buf))

        bg = d.create_bind_group(layout=self._mm_bgl, entries=[
            {"binding": 0, "resource": {"buffer": param_buf}},
            {"binding": 1, "resource": {"buffer": self._buf[op['src_a']],
                                        "size": op['M'] * op['K'] * 4}},
            {"binding": 2, "resource": {"buffer": self._buf[op['src_b']],
                                        "size": op['K'] * op['N'] * 4}},
            {"binding": 3, "resource": {"buffer": self._buf[op['dst']],
                                        "size": op['M'] * op['N'] * 4}},
        ])

        self._recipe.append({
            'kind': 'compute',
            'pipeline': self._mm_pipeline,
            'bind_group': bg,
            'workgroups': ((op['M'] + 15) // 16, (op['N'] + 15) // 16),
        })

    def _compile_reduce(self, op):
        """Pre-build bind group and record recipe for a reduce pass."""
        d = self._device
        param_buf = self._alloc_param_buf()
        param_data = struct.pack("IIII", op['N'], op['stride'], op['op'], 0)
        self._param_writes.append((param_data, param_buf))

        bg = d.create_bind_group(layout=self._red_bgl, entries=[
            {"binding": 0, "resource": {"buffer": param_buf}},
            {"binding": 1, "resource": {"buffer": self._buf['vec_values'],
                                        "size": op['N'] * 4}},
            {"binding": 2, "resource": {"buffer": self._buf['vec_indices'],
                                        "size": op['N'] * 4}},
        ])

        self._recipe.append({
            'kind': 'compute',
            'pipeline': self._red_pipeline,
            'bind_group': bg,
            'workgroups': ((op['N'] + 255) // 256,),
        })

    def _compile_elementwise(self, op):
        """Pre-build bind group and record recipe for an elementwise pass."""
        d = self._device
        param_buf = self._alloc_param_buf()
        param_data = struct.pack("IIff", op['N'], op['op'], op['scalar'], 0.0)
        self._param_writes.append((param_data, param_buf))

        bg = d.create_bind_group(layout=self._ew_bgl, entries=[
            {"binding": 0, "resource": {"buffer": param_buf}},
            {"binding": 1, "resource": {"buffer": self._buf[op['src']],
                                        "size": op['N'] * 4}},
            {"binding": 2, "resource": {"buffer": self._buf[op['dst']],
                                        "size": op['N'] * 4}},
        ])

        self._recipe.append({
            'kind': 'compute',
            'pipeline': self._ew_pipeline,
            'bind_group': bg,
            'workgroups': ((op['N'] + 255) // 256,),
        })

    def _compile_elementwise_add(self, op):
        """Pre-build bind group and record recipe for an elementwise-add pass."""
        d = self._device
        param_buf = self._alloc_param_buf()
        param_data = struct.pack("IfII", op['N'], op['scalar'], 0, 0)
        self._param_writes.append((param_data, param_buf))

        bg = d.create_bind_group(layout=self._ew_add_bgl, entries=[
            {"binding": 0, "resource": {"buffer": param_buf}},
            {"binding": 1, "resource": {"buffer": self._buf[op['src_a']],
                                        "size": op['N'] * 4}},
            {"binding": 2, "resource": {"buffer": self._buf[op['src_b']],
                                        "size": op['N'] * 4}},
            {"binding": 3, "resource": {"buffer": self._buf[op['dst']],
                                        "size": op['N'] * 4}},
        ])

        self._recipe.append({
            'kind': 'compute',
            'pipeline': self._ew_add_pipeline,
            'bind_group': bg,
            'workgroups': ((op['N'] + 255) // 256,),
        })

    def _compile_clamp(self, op):
        """Pre-build bind group and record recipe for a clamp pass."""
        d = self._device
        param_buf = self._alloc_param_buf()
        # Pack: num_elements(u32), lo(f32), hi(f32), pad(u32)
        param_data = struct.pack("IffI", op['N'], op['lo'], op['hi'], 0)
        self._param_writes.append((param_data, param_buf))

        bg = d.create_bind_group(layout=self._clamp_bgl, entries=[
            {"binding": 0, "resource": {"buffer": param_buf}},
            {"binding": 1, "resource": {"buffer": self._buf[op['src']],
                                        "size": op['N'] * 4}},
            {"binding": 2, "resource": {"buffer": self._buf[op['dst']],
                                        "size": op['N'] * 4}},
        ])

        self._recipe.append({
            'kind': 'compute',
            'pipeline': self._clamp_pipeline,
            'bind_group': bg,
            'workgroups': ((op['N'] + 255) // 256,),
        })

    def _compile_copy(self, op):
        """Record a buffer-to-buffer copy in the recipe."""
        self._recipe.append({
            'kind': 'copy',
            'src': self._buf[op['src']],
            'dst': self._buf[op['dst']],
            'size': op['size'],
        })

    def _compile_init_indices(self, op):
        """Record that we need to upload arange(N) to indices before submit."""
        indices = np.arange(op['N'], dtype=np.uint32)
        self._index_inits.append((indices, self._buf['vec_indices']))

    # ─── Execute ─────────────────────────────────────────────────────────────

    def execute(self):
        """Execute the pre-compiled pipeline. Single queue.submit + _poll_wait.

        Builds a fresh command encoder from the cached recipe (bind groups,
        pipelines, workgroups, copy ops). This is pure CPU bookkeeping — the
        bind groups and pipelines are reused, only the lightweight encoder
        and command buffer objects are recreated.

        Before submitting, writes all param buffers and index arrays via
        queue.write_buffer(). These are queue operations that complete
        before the command buffer begins executing on the GPU.
        """
        assert self._compiled, "Call compile() first"

        d = self._device
        t0 = time.perf_counter()

        # Write all param uniforms (queue ops — execute before command buffer)
        for data, buf in self._param_writes:
            d.queue.write_buffer(buf, 0, data)

        # Write index arrays for argmax
        for arr, buf in self._index_inits:
            d.queue.write_buffer(buf, 0, arr.tobytes())

        # Build command encoder from recipe
        encoder = d.create_command_encoder()
        for step in self._recipe:
            if step['kind'] == 'compute':
                cp = encoder.begin_compute_pass()
                cp.set_pipeline(step['pipeline'])
                cp.set_bind_group(0, step['bind_group'])
                cp.dispatch_workgroups(*step['workgroups'])
                cp.end()
            elif step['kind'] == 'copy':
                encoder.copy_buffer_to_buffer(
                    step['src'], 0, step['dst'], 0, step['size'])

        # Single submit, single wait
        d.queue.submit([encoder.finish()])
        d._poll_wait()

        elapsed = (time.perf_counter() - t0) * 1000
        self.stats['execute_ms'] += elapsed
        self.stats['executions'] += 1

    # ─── Data Upload / Readback ──────────────────────────────────────────────

    def upload(self, region, data):
        """Upload numpy array to a workspace region.

        Args:
            region: region name ('A', 'B', 'C', 'values', 'scores', etc.)
            data: numpy array (will be cast to float32)
        """
        t0 = time.perf_counter()
        buf_key = self._resolve_region(region)
        self._device.queue.write_buffer(
            self._buf[buf_key], 0,
            data.astype(np.float32).tobytes())
        self.stats['upload_ms'] += (time.perf_counter() - t0) * 1000

    def readback(self, region, size_bytes):
        """Read raw bytes from a workspace region.

        Args:
            region: region name
            size_bytes: number of bytes to read

        Returns:
            memoryview of the read data
        """
        t0 = time.perf_counter()
        buf_key = self._resolve_region(region)
        result = self._device.queue.read_buffer(self._buf[buf_key], size=size_bytes)
        self.stats['readback_ms'] += (time.perf_counter() - t0) * 1000
        return result

    def readback_f32(self, region, count):
        """Read float32 array from a workspace region.

        Args:
            region: region name
            count: number of float32 values to read

        Returns:
            numpy float32 array
        """
        raw = self.readback(region, count * 4)
        return np.frombuffer(raw, dtype=np.float32)[:count].copy()

    def readback_u32(self, region, count):
        """Read uint32 array from a workspace region.

        Args:
            region: region name
            count: number of uint32 values to read

        Returns:
            numpy uint32 array
        """
        raw = self.readback(region, count * 4)
        return np.frombuffer(raw, dtype=np.uint32)[:count].copy()

    def reset_stats(self):
        """Reset timing statistics."""
        self.stats = {
            'compile_ms': self.stats.get('compile_ms', 0.0),
            'execute_ms': 0.0,
            'upload_ms': 0.0,
            'readback_ms': 0.0,
            'executions': 0,
            'ops_per_execution': self.stats.get('ops_per_execution', 0),
        }


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

def main():
    print("=" * 70)
    print("SBFusedPipeline — Pre-Compiled GPU Operation Sequences")
    print("=" * 70)
    print()

    pipe = SBFusedPipeline(max_matrix_dim=1024, max_vector_len=1_000_000)
    print(f"GPU: {pipe._device.adapter.info.get('device', 'unknown')}")
    print(f"Workspace: {pipe._ws_bytes / (1024*1024):.1f} MB GPU-resident, "
          f"{pipe._param_pool_size} param buffers")
    print()

    rng = np.random.default_rng(42)

    # ─── Benchmark 1: Fused pipeline 100x (matmul → argmax → search) ─────
    print("=" * 70)
    print("BENCHMARK 1: 100x (matmul 256x256 -> argmax 65536 -> search 10000)")
    print("=" * 70)
    print()

    N_MAT = 256
    N_FLAT = N_MAT * N_MAT  # 65536
    N_SEARCH = 10000
    ITERS = 100

    A = rng.random((N_MAT, N_MAT)).astype(np.float32)
    B = rng.random((N_MAT, N_MAT)).astype(np.float32)
    search_data = rng.random(N_SEARCH).astype(np.float32) * 1000

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

    print(f"  NumPy baseline: {np_total:.1f} ms total, {np_total/ITERS:.3f} ms/iter")
    print()

    # --- Fused pipeline ---
    # Record the sequence once
    pipe.begin_recording()
    pipe.matmul('A', 'B', 'C', N_MAT, N_MAT, N_MAT)
    # argmax over flattened C (C buffer holds N_MAT*N_MAT values)
    # Copy C → values, then argmax
    pipe.argmax('C', N_FLAT)
    # For search, we use a fixed query. In a real pipeline the query would come
    # from the argmax result, but that requires a CPU readback mid-pipeline.
    # To keep the pipeline fully fused, we use a representative query value.
    # The benchmark measures raw throughput — the point is dispatch overhead.
    query_val = float(np.max(A @ B))
    pipe.search('scratch', query_val, N_SEARCH)
    pipe.compile()

    compile_ms = pipe.stats['compile_ms']
    print(f"  Compile time: {compile_ms:.3f} ms")
    print(f"  Recorded ops: {pipe.stats['ops_per_execution']}")
    print()

    # Upload data once
    pipe.upload('A', A)
    pipe.upload('B', B)
    pipe.upload('scratch', search_data)

    # Warmup
    pipe.execute()
    pipe.reset_stats()

    # Execute 100 times
    t0 = time.perf_counter()
    for _ in range(ITERS):
        pipe.execute()
    fused_total = (time.perf_counter() - t0) * 1000

    # Verify correctness: indices[0] holds the SEARCH result (last argmax in pipeline).
    # The matmul argmax result was overwritten by the search argmax.
    # Compare against NumPy search result.
    search_idx_gpu = int(pipe.readback_u32('indices', 1)[0])
    np_search_idx = int(np.argmin(np.abs(search_data - query_val)))
    search_correct = search_idx_gpu == np_search_idx

    # Also verify matmul output
    C_gpu = pipe.readback_f32('C', N_FLAT)
    C_np = (A @ B).flatten()
    matmul_correct = np.allclose(C_gpu, C_np, atol=0.5)

    print(f"  Fused pipeline: {fused_total:.1f} ms total, {fused_total/ITERS:.3f} ms/iter")
    print(f"  NumPy baseline: {np_total:.1f} ms total, {np_total/ITERS:.3f} ms/iter")
    print(f"  Speedup:        {np_total/fused_total:.2f}x")
    print(f"  Correctness:")
    print(f"    matmul:  {'MATCH' if matmul_correct else 'MISMATCH'}")
    print(f"    search:  GPU={search_idx_gpu} NumPy={np_search_idx} {'MATCH' if search_correct else 'MISMATCH'}")
    print()
    print(f"  Per-execute breakdown:")
    print(f"    Execute (submit+wait): {pipe.stats['execute_ms']/ITERS:.3f} ms/iter")
    print(f"    Upload:                {pipe.stats['upload_ms']/ITERS:.3f} ms/iter")
    print(f"    Readback:              {pipe.stats['readback_ms']/ITERS:.3f} ms/iter")
    print()

    # ─── Benchmark 2: Colony-equivalent workload ──────────────────────────
    print("=" * 70)
    print("BENCHMARK 2: Colony-equivalent (matmul 500x10@10x10 -> add noise -> clamp) x150")
    print("=" * 70)
    print()

    N_CELLS = 500
    N_DIMS = 10
    N_ELEMENTS = N_CELLS * N_DIMS  # 5000
    GENS = 150

    positions = rng.random((N_CELLS, N_DIMS)).astype(np.float32)
    influence = rng.random((N_DIMS, N_DIMS)).astype(np.float32) * 0.1
    noise = rng.random((N_CELLS, N_DIMS)).astype(np.float32) * 0.02

    # --- NumPy baseline ---
    pos_np = positions.copy()
    t0 = time.perf_counter()
    for _ in range(GENS):
        forces = pos_np @ influence          # matmul 500x10 @ 10x10
        pos_np = pos_np + 0.02 * forces      # add scaled noise (using forces as "noise")
        pos_np = np.clip(pos_np, 0.0, 1.0)   # clamp
    np_colony = (time.perf_counter() - t0) * 1000

    print(f"  NumPy baseline: {np_colony:.1f} ms ({np_colony/GENS:.3f} ms/gen)")
    print()

    # --- Fused pipeline ---
    # Record: matmul(A=positions, B=influence, C=forces) → add(A + scalar*C, out=A) → clamp(A)
    pipe2 = SBFusedPipeline(max_matrix_dim=1024, max_vector_len=1_000_000)

    pipe2.begin_recording()
    # matmul: C = A @ B (positions @ influence → forces in C)
    pipe2.matmul('A', 'B', 'C', N_CELLS, N_DIMS, N_DIMS)
    # add noise: scratch = A + 0.02 * C (positions + scaled forces)
    pipe2.elementwise_add('A', 'C', 'scratch', N_ELEMENTS, scalar=0.02)
    # clamp: A = clamp(scratch, 0, 1)
    pipe2.clamp('scratch', 'A', N_ELEMENTS, lo=0.0, hi=1.0)
    pipe2.compile()

    compile_ms2 = pipe2.stats['compile_ms']
    print(f"  Compile time: {compile_ms2:.3f} ms")
    print(f"  Recorded ops: {pipe2.stats['ops_per_execution']}")
    print()

    # Upload initial data
    pipe2.upload('A', positions.flatten())
    pipe2.upload('B', influence.flatten())

    # Warmup
    pipe2.execute()
    pipe2.reset_stats()

    # Re-upload fresh positions
    pipe2.upload('A', positions.flatten())

    # Execute 150 generations
    t0 = time.perf_counter()
    for _ in range(GENS):
        pipe2.execute()
    fused_colony = (time.perf_counter() - t0) * 1000

    # Readback final positions for verification
    gpu_positions = pipe2.readback_f32('A', N_ELEMENTS).reshape(N_CELLS, N_DIMS)

    # Verify: compare against NumPy result
    max_diff = float(np.max(np.abs(gpu_positions - pos_np)))
    close = np.allclose(gpu_positions, pos_np, atol=1e-3)

    print(f"  Fused pipeline: {fused_colony:.1f} ms ({fused_colony/GENS:.3f} ms/gen)")
    print(f"  NumPy baseline: {np_colony:.1f} ms ({np_colony/GENS:.3f} ms/gen)")
    print(f"  Speedup:        {np_colony/fused_colony:.2f}x")
    print(f"  Correctness:    max_diff={max_diff:.6f} {'MATCH' if close else 'MISMATCH'}")
    print()
    print(f"  Per-generation breakdown:")
    print(f"    Execute (submit+wait): {pipe2.stats['execute_ms']/GENS:.3f} ms/gen")
    print(f"    Upload:                {pipe2.stats['upload_ms']/GENS:.3f} ms/gen")
    print(f"    Readback:              {pipe2.stats['readback_ms']/GENS:.3f} ms/gen")
    print()

    # ─── Benchmark 3: Mega-fused — ALL 150 gens in ONE submit ──────────
    print("=" * 70)
    print("BENCHMARK 3: Mega-fused — 150 colony gens in ONE queue.submit()")
    print("=" * 70)
    print()

    # Record 150 generations of matmul→add→clamp as one giant recipe.
    # Since each gen's output feeds the next (A is updated in-place via
    # scratch→A clamp), this is a legitimate loop unroll.
    pipe3 = SBFusedPipeline(max_matrix_dim=1024, max_vector_len=1_000_000,
                            max_reduce_passes=20)
    # Need enough param buffers: 3 ops × 150 gens = 450
    pipe3._param_pool_size = 500
    UNIFORM_USAGE = wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST
    pipe3._param_bufs = [
        pipe3._device.create_buffer(size=16, usage=UNIFORM_USAGE)
        for _ in range(pipe3._param_pool_size)
    ]

    pipe3.begin_recording()
    for _ in range(GENS):
        pipe3.matmul('A', 'B', 'C', N_CELLS, N_DIMS, N_DIMS)
        pipe3.elementwise_add('A', 'C', 'scratch', N_ELEMENTS, scalar=0.02)
        pipe3.clamp('scratch', 'A', N_ELEMENTS, lo=0.0, hi=1.0)
    pipe3.compile()

    compile_ms3 = pipe3.stats['compile_ms']
    print(f"  Compile time: {compile_ms3:.3f} ms (150 gens unrolled)")
    print(f"  Recorded ops: {pipe3.stats['ops_per_execution']}")
    print(f"  Param buffers used: {pipe3._next_param_idx}")
    print()

    # Upload initial data
    pipe3.upload('A', positions.flatten())
    pipe3.upload('B', influence.flatten())

    # Warmup
    pipe3.execute()
    pipe3.reset_stats()

    # Re-upload fresh positions
    pipe3.upload('A', positions.flatten())

    # Execute — ALL 150 gens in ONE submit
    t0 = time.perf_counter()
    pipe3.execute()
    mega_fused = (time.perf_counter() - t0) * 1000

    # Readback
    gpu_positions3 = pipe3.readback_f32('A', N_ELEMENTS).reshape(N_CELLS, N_DIMS)
    max_diff3 = float(np.max(np.abs(gpu_positions3 - pos_np)))
    close3 = np.allclose(gpu_positions3, pos_np, atol=1e-3)

    print(f"  Mega-fused:     {mega_fused:.3f} ms (ONE submit for 150 gens)")
    print(f"  Per-gen fused:  {fused_colony:.1f} ms (150 submits)")
    print(f"  NumPy:          {np_colony:.1f} ms")
    print(f"  Speedup vs NumPy:    {np_colony/mega_fused:.2f}x")
    print(f"  Speedup vs per-gen:  {fused_colony/mega_fused:.2f}x")
    print(f"  Correctness:    max_diff={max_diff3:.6f} {'MATCH' if close3 else 'MISMATCH'}")
    print()

    # ─── Benchmark 4: Large matrices — where GPU wins ─────────────────
    print("=" * 70)
    print("BENCHMARK 4: Large matmul sweep (GPU compute vs NumPy)")
    print("=" * 70)
    print()

    print(f"  {'Size':>10} {'NumPy':>10} {'Fused':>10} {'Speedup':>10}")
    print("  " + "-" * 44)

    for N_SIZE in [128, 256, 512, 1024]:
        A_big = rng.random((N_SIZE, N_SIZE)).astype(np.float32)
        B_big = rng.random((N_SIZE, N_SIZE)).astype(np.float32)

        # NumPy
        t0 = time.perf_counter()
        for _ in range(10):
            _ = A_big @ B_big
        np_mm = (time.perf_counter() - t0) * 1000 / 10

        # Fused (single matmul per execute)
        pipe4 = SBFusedPipeline(max_matrix_dim=1024, max_vector_len=100000)
        pipe4.begin_recording()
        pipe4.matmul('A', 'B', 'C', N_SIZE, N_SIZE, N_SIZE)
        pipe4.compile()
        pipe4.upload('A', A_big)
        pipe4.upload('B', B_big)
        pipe4.execute()  # warmup
        pipe4.reset_stats()

        t0 = time.perf_counter()
        for _ in range(10):
            pipe4.execute()
        fused_mm = (time.perf_counter() - t0) * 1000 / 10

        speedup = np_mm / fused_mm if fused_mm > 0 else 0
        print(f"  {N_SIZE}x{N_SIZE:>4}  {np_mm:>8.3f}ms {fused_mm:>8.3f}ms {speedup:>9.2f}x")

    print()

    # ─── Summary ──────────────────────────────────────────────────────────
    print("=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print()
    print("  Benchmark 1: matmul→argmax→search (100 iters, 1 submit/iter)")
    print(f"    NumPy:  {np_total:.1f} ms  |  Fused: {fused_total:.1f} ms  |  {np_total/fused_total:.2f}x")
    print()
    print("  Benchmark 2: colony-equivalent (150 gens, 1 submit/gen)")
    print(f"    NumPy:  {np_colony:.1f} ms  |  Fused: {fused_colony:.1f} ms  |  {np_colony/fused_colony:.2f}x")
    print()
    print("  Benchmark 3: mega-fused colony (150 gens, 1 submit TOTAL)")
    print(f"    NumPy:  {np_colony:.1f} ms  |  Fused: {mega_fused:.3f} ms  |  {np_colony/mega_fused:.2f}x")
    print()
    print("  ANALYSIS:")
    print(f"    Per-submit overhead: ~{fused_colony/GENS:.1f} ms (from Benchmark 2)")
    print(f"    Mega-fused eliminates {GENS-1} of {GENS} submits")
    submit_overhead = fused_colony - mega_fused
    if submit_overhead > 0:
        print(f"    Total submit overhead eliminated: {submit_overhead:.1f} ms")
        print(f"    Pure GPU compute time: ~{mega_fused:.1f} ms for 150 gens")
    print()
    print("  CONCLUSION:")
    if mega_fused < np_colony:
        print("    GPU compute IS faster than NumPy when dispatch overhead is eliminated.")
        print("    The per-dispatch overhead was the only problem.")
    else:
        print("    Even with zero dispatch overhead, GPU is slower at this scale.")
        print(f"    GPU compute: {mega_fused:.1f} ms vs NumPy: {np_colony:.1f} ms")
        print("    The matmul workgroup shader is the bottleneck, not dispatch overhead.")
    print("=" * 70)


if __name__ == '__main__':
    main()
