"""
SBColonyGPUKernel — GPU-accelerated colony simulation for Sound Block Second Opinion V2.

Single WGSL compute shader (sb_colony.wgsl) runs all per-cell physics in one dispatch:
matrix forces, concept gravity, thermal noise, boundary forces, velocity/position integration.

Metabolism and reproduction stay on CPU (branching-heavy, once per generation).

Uses wgpu-py (wgpu-native backend) which auto-selects:
  - Vulkan on Linux (AMD/NVIDIA)
  - Metal on macOS
  - DX12 on Windows

Same WGSL shader works in the browser via native WebGPU API.

Usage:
    kernel = SBColonyGPUKernel(num_dims=10, max_cells=512, max_concepts=8)
    kernel.upload_matrix(influence_matrix)    # D x D, once
    kernel.upload_concepts(concept_positions) # K x D, once
    kernel.upload_cells(positions, velocities) # C x D each

    for gen in range(150):
        noise = np.random.uniform(-1, 1, (num_cells, num_dims)).astype(np.float32)
        positions, velocities = kernel.run_generation(
            ticks_per_gen=1, params_dict={...}, noise=noise)
        # CPU: metabolism, reproduction
        kernel.upload_cells(positions, velocities)  # if population changed

Follows the exact pattern of afs-engine/afs/gpu/wgpu_force_kernel.py.
"""

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

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


def gpu_available() -> bool:
    """Check if WebGPU compute is available on this system."""
    if not _WGPU_AVAILABLE:
        return False
    try:
        device = wgpu.utils.get_default_device()
        return device is not None
    except Exception:
        return False


_SHADER_DIR = Path(__file__).parent

# Params struct layout (20 x f32 = 80 bytes, must match WGSL struct Params)
# Alignment: WGSL uniform buffers require 16-byte alignment.
# Struct fields:
#   u32 num_cells, u32 num_dims, u32 num_concepts, u32 _pad0     (16 bytes)
#   f32 gravity_strength, f32 gravity_sharpness, f32 gravity_range, f32 softening  (16 bytes)
#   f32 force_scale, f32 tradeoff_amplifier, f32 thermal_noise_level, f32 momentum (16 bytes)
#   f32 max_speed, f32 boundary_margin, f32 boundary_strength, f32 dt_norm         (16 bytes)
#   f32 gravity_ramp, f32 settling_damp, f32 _pad1, f32 _pad2                      (16 bytes)
# Total: 80 bytes
_PARAMS_SIZE = 80
_PARAMS_FORMAT = "IIIIffffffffffffffffff"  # 4 u32 + 16 f32 = 20 values... let me count
# Actually: 3 u32 + 1 u32 pad + 16 f32 = 20 values
# struct.pack format: I I I I f f f f  f f f f  f f f f  f f f f = 4I + 16f


class SBColonyGPUKernel:
    """GPU colony kernel using sb_colony.wgsl compute shader."""

    def __init__(self, num_dims: int, max_cells: int = 512, max_concepts: int = 8):
        if not _WGPU_AVAILABLE:
            raise ImportError("wgpu package required: pip install wgpu")

        self.num_dims = num_dims
        self.max_cells = max_cells
        self.max_concepts = max_concepts

        self._num_cells = 0
        self._num_concepts = 0
        self._bind_group_dirty = True

        self._device = wgpu.utils.get_default_device()
        self._load_shader()
        self._create_pipeline()
        self._create_buffers()

        # Cached bind group
        self._bg = None

    def _load_shader(self):
        self._shader_src = (_SHADER_DIR / "sb_colony.wgsl").read_text()
        self._metabolism_src = (_SHADER_DIR / "sb_metabolism.wgsl").read_text()
        self._pack_src = (_SHADER_DIR / "sb_pack_transfer.wgsl").read_text()
        self._noise_src = (_SHADER_DIR / "sb_noise_gen.wgsl").read_text()

    def _create_pipeline(self):
        d = self._device
        module = d.create_shader_module(code=self._shader_src)

        # 7 bindings: params(uniform), positions(storage), velocities(storage),
        #             forces(storage), influence_matrix(read_only_storage),
        #             concept_positions(read_only_storage), noise(read_only_storage)
        self._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}},
            {"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            {"binding": 4, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 5, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 6, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
        ])

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

        # ── Metabolism pipeline ──
        metab_module = d.create_shader_module(code=self._metabolism_src)
        self._metab_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}},
            {"binding": 4, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            {"binding": 5, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
        ])
        metab_pl = d.create_pipeline_layout(bind_group_layouts=[self._metab_bgl])
        self._metab_pipeline = d.create_compute_pipeline(
            layout=metab_pl, compute={"module": metab_module, "entry_point": "main"})

        self._metab_bg = None

        # ── Pack/Unpack pipeline (for Sound Block I/O) ──
        pack_module = d.create_shader_module(code=self._pack_src)
        self._pack_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}},
        ])
        pack_pl = d.create_pipeline_layout(bind_group_layouts=[self._pack_bgl])
        self._unpack_pipeline = d.create_compute_pipeline(
            layout=pack_pl, compute={"module": pack_module, "entry_point": "unpack_main"})
        self._pack_pipeline = d.create_compute_pipeline(
            layout=pack_pl, compute={"module": pack_module, "entry_point": "pack_main"})

        # ── Noise generator pipeline ──
        noise_module = d.create_shader_module(code=self._noise_src)
        self._noise_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}},
        ])
        noise_pl = d.create_pipeline_layout(bind_group_layouts=[self._noise_bgl])
        self._noise_pipeline = d.create_compute_pipeline(
            layout=noise_pl, compute={"module": noise_module, "entry_point": "main"})

        self._noise_bg = None
        self._buf_noise_params = None  # created in _create_buffers

    def _create_buffers(self):
        d = self._device
        D = self.num_dims
        C_max = self.max_cells
        K_max = self.max_concepts

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

        # Per-cell buffers (max_cells * num_dims * 4 bytes)
        cell_buf_size = C_max * D * 4
        self._buf_positions = d.create_buffer(size=cell_buf_size, usage=STORAGE_RW)
        self._buf_velocities = d.create_buffer(size=cell_buf_size, usage=STORAGE_RW)
        self._buf_forces = d.create_buffer(size=cell_buf_size, usage=STORAGE_RW)
        self._buf_noise = d.create_buffer(size=cell_buf_size, usage=STORAGE_RO)

        # Energy and cell state buffers (per cell, 1 value each)
        self._buf_energy = d.create_buffer(size=C_max * 4, usage=STORAGE_RW)
        self._buf_cell_states = d.create_buffer(size=C_max * 4, usage=STORAGE_RW)

        # Matrix and concept buffers
        self._buf_influence_matrix = d.create_buffer(size=D * D * 4, usage=STORAGE_RO)
        self._buf_concept_positions = d.create_buffer(
            size=K_max * D * 4, usage=STORAGE_RO)

        # Params uniform (physics)
        self._buf_params = d.create_buffer(
            size=_PARAMS_SIZE,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

        # Metabolism params uniform (16 values × 4 bytes = 64 bytes)
        self._buf_metab_params = d.create_buffer(
            size=64,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

        # Noise gen params uniform (3 u32 + 1 f32 = 16 bytes)
        self._buf_noise_params = d.create_buffer(
            size=16,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

        # Pack/unpack transfer buffers — packed is 1/3 the size of unpacked
        packed_cell_size = ((C_max * D + 5) // 6) * 2 * 4  # ceil(N/6) * 2 f32s
        self._buf_packed_transfer = d.create_buffer(
            size=max(packed_cell_size, 16),
            usage=STORAGE_RW)
        self._buf_pack_params = d.create_buffer(
            size=16,  # 4 u32s
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

    def _rebuild_bind_group(self):
        """Rebuild bind group when cell/concept count changes."""
        C = self._num_cells
        K = self._num_concepts
        D = self.num_dims
        d = self._device

        self._bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_params}},
            {"binding": 1, "resource": {"buffer": self._buf_positions,
                                        "size": C * D * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_velocities,
                                        "size": C * D * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_forces,
                                        "size": C * D * 4}},
            {"binding": 4, "resource": {"buffer": self._buf_influence_matrix}},
            {"binding": 5, "resource": {"buffer": self._buf_concept_positions,
                                        "size": K * D * 4}},
            {"binding": 6, "resource": {"buffer": self._buf_noise,
                                        "size": C * D * 4}},
        ])

        self._bind_group_dirty = False

        # Metabolism bind group (shares positions, velocities, concept_positions buffers)
        self._metab_bg = d.create_bind_group(layout=self._metab_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_metab_params}},
            {"binding": 1, "resource": {"buffer": self._buf_positions,
                                        "size": C * D * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_velocities,
                                        "size": C * D * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_energy,
                                        "size": C * 4}},
            {"binding": 4, "resource": {"buffer": self._buf_cell_states,
                                        "size": C * 4}},
            {"binding": 5, "resource": {"buffer": self._buf_concept_positions,
                                        "size": K * D * 4}},
        ])

    # ─── Data Upload ───────────────────────────────────────────

    def upload_energy(self, energy: np.ndarray):
        """Upload cell energy values. Call after initialization and reproduction."""
        C = energy.shape[0]
        self._device.queue.write_buffer(
            self._buf_energy, 0, energy.astype(np.float32).tobytes())

    def readback_energy(self) -> tuple:
        """Read energy and cell states from GPU."""
        C = self._num_cells
        self._device._poll_wait()
        energy_data = self._device.queue.read_buffer(self._buf_energy, size=C * 4)
        states_data = self._device.queue.read_buffer(self._buf_cell_states, size=C * 4)
        energy = np.frombuffer(energy_data, dtype=np.float32)[:C].copy()
        states = np.frombuffer(states_data, dtype=np.uint32)[:C].copy()
        return energy, states

    def tick_metabolism(self, metab_params: dict):
        """Run one metabolism tick on GPU. Energy/states updated in-place on GPU."""
        if self._metab_bg is None:
            self._rebuild_bind_group()

        d = self._device
        C = self._num_cells
        workgroups = (C + 255) // 256

        # Pack metabolism params: 3 u32 + 1 u32 pad + 9 f32 + 3 f32 pad = 16 values, 64 bytes
        packed = struct.pack("IIIIffffffffffff",
            metab_params.get("num_cells", C),
            metab_params.get("num_dims", self.num_dims),
            metab_params.get("num_concepts", self._num_concepts),
            0,  # pad
            metab_params.get("absorption_range", 0.3),
            metab_params.get("energy_gain_factor", 5.0),
            metab_params.get("energy_falloff_exp", 2.0),
            metab_params.get("movement_cost_factor", 0.002),
            metab_params.get("base_cost_per_dim", 0.001),
            metab_params.get("min_energy", 0.1),
            metab_params.get("max_energy", 100.0),
            metab_params.get("hot_threshold", 75.0),
            metab_params.get("warm_threshold", 30.0),
            0.0, 0.0, 0.0,  # pad
        )

        d.queue.write_buffer(self._buf_metab_params, 0, packed)

        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._metab_pipeline)
        compute_pass.set_bind_group(0, self._metab_bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])

    def upload_matrix(self, influence_matrix: np.ndarray):
        """Upload D x D influence matrix to GPU. Call once at init."""
        assert influence_matrix.shape == (self.num_dims, self.num_dims), \
            f"Expected ({self.num_dims}, {self.num_dims}), got {influence_matrix.shape}"
        self._device.queue.write_buffer(
            self._buf_influence_matrix, 0,
            influence_matrix.astype(np.float32).tobytes())

    def upload_concepts(self, concept_positions: np.ndarray):
        """Upload K x D concept positions to GPU. Call once at init."""
        K, D = concept_positions.shape
        assert D == self.num_dims and K <= self.max_concepts, \
            f"Expected (<=  {self.max_concepts}, {self.num_dims}), got {concept_positions.shape}"
        if K != self._num_concepts:
            self._num_concepts = K
            self._bind_group_dirty = True
        self._device.queue.write_buffer(
            self._buf_concept_positions, 0,
            concept_positions.astype(np.float32).tobytes())

    def upload_cells(self, positions: np.ndarray, velocities: np.ndarray):
        """Upload cell positions and velocities. Call before tick loop and after reproduction."""
        C, D = positions.shape
        assert D == self.num_dims and C <= self.max_cells, \
            f"Expected (<= {self.max_cells}, {self.num_dims}), got {positions.shape}"
        assert velocities.shape == (C, D)
        if C != self._num_cells:
            self._num_cells = C
            self._bind_group_dirty = True
        d = self._device
        d.queue.write_buffer(self._buf_positions, 0,
                             positions.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_velocities, 0,
                             velocities.astype(np.float32).tobytes())

    def upload_noise(self, noise: np.ndarray):
        """Upload C x D noise buffer. Call once per generation."""
        C, D = noise.shape
        assert C == self._num_cells and D == self.num_dims
        self._device.queue.write_buffer(
            self._buf_noise, 0,
            noise.astype(np.float32).tobytes())

    # ─── GPU Noise Generation (zero CPU transfer) ──────────────

    def generate_noise_gpu(self, generation: int):
        """Generate noise entirely on GPU. No CPU→GPU transfer at all.

        Uses PCG hash per cell per dimension with generation as seed.
        The noise buffer is written directly on GPU — never touches CPU.
        """
        C = self._num_cells
        D = self.num_dims
        N = C * D
        d = self._device

        # Build noise bind group if needed
        if self._noise_bg is None:
            self._noise_bg = d.create_bind_group(layout=self._noise_bgl, entries=[
                {"binding": 0, "resource": {"buffer": self._buf_noise_params}},
                {"binding": 1, "resource": {"buffer": self._buf_noise,
                                            "size": N * 4}},
            ])

        # Upload noise params (just the seed changes per generation)
        params = struct.pack("IIIf", C, D, generation * 31337 + 42, 1.0)
        d.queue.write_buffer(self._buf_noise_params, 0, params)

        # Dispatch noise generator
        workgroups = (N + 255) // 256
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._noise_pipeline)
        compute_pass.set_bind_group(0, self._noise_bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])

    # ─── Sound Block Packed I/O ─────────────────────────────────

    def upload_noise_packed(self, noise: np.ndarray):
        """Upload noise via mantissa packing — 2.67x less data over the bus.

        CPU: pack C*D floats into ceil(C*D/6)*2 f32s
        Transfer: packed data over PCIe bus (2.67x smaller)
        GPU: unpack shader expands to full C*D array
        """
        C, D = noise.shape
        assert C == self._num_cells and D == self.num_dims
        N = C * D

        # CPU-side mantissa pack (vectorized NumPy — no Python loops)
        scaled = ((noise.flatten() + 1.0) * 127.5).clip(0, 255).astype(np.uint32)
        bins_needed = (N + 5) // 6
        padded = np.zeros(bins_needed * 6, dtype=np.uint32)
        padded[:N] = scaled

        # Reshape to (bins, 6) and pack 3 bytes per f32
        reshaped = padded.reshape(bins_needed, 6)
        re_packed = (reshaped[:, 0] << 16) | (reshaped[:, 1] << 8) | reshaped[:, 2]
        im_packed = (reshaped[:, 3] << 16) | (reshaped[:, 4] << 8) | reshaped[:, 5]

        # Interleave re, im into packed array
        packed = np.zeros(bins_needed * 2, dtype=np.float32)
        packed[0::2] = re_packed.astype(np.float32)
        packed[1::2] = im_packed.astype(np.float32)

        # Transfer packed data (2.67x smaller)
        d = self._device
        d.queue.write_buffer(self._buf_packed_transfer, 0, packed.tobytes())

        # GPU unpack: dispatch shader to expand into noise buffer
        pack_params = struct.pack("IIII", N, D, 0, 0)
        d.queue.write_buffer(self._buf_pack_params, 0, pack_params)

        bg = d.create_bind_group(layout=self._pack_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_pack_params}},
            {"binding": 1, "resource": {"buffer": self._buf_packed_transfer,
                                        "size": bins_needed * 2 * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_noise,
                                        "size": N * 4}},
        ])

        workgroups = (bins_needed + 255) // 256
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._unpack_pipeline)
        compute_pass.set_bind_group(0, bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])

    def readback_packed(self) -> tuple:
        """Read positions and velocities via GPU-side packing — 2.67x less data back.

        GPU: pack shader compresses positions into packed buffer
        Transfer: packed data over PCIe bus (2.67x smaller)
        CPU: unpack to full arrays
        """
        C = self._num_cells
        D = self.num_dims
        N = C * D
        bins_needed = (N + 5) // 6
        d = self._device

        # Pack positions on GPU
        pack_params = struct.pack("IIII", N, D, 0, 0)
        d.queue.write_buffer(self._buf_pack_params, 0, pack_params)

        bg = d.create_bind_group(layout=self._pack_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_pack_params}},
            {"binding": 1, "resource": {"buffer": self._buf_packed_transfer,
                                        "size": bins_needed * 2 * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_positions,
                                        "size": N * 4}},
        ])

        workgroups = (bins_needed + 255) // 256
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._pack_pipeline)
        compute_pass.set_bind_group(0, bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])

        d._poll_wait()

        # Read packed data (2.67x smaller transfer)
        packed_size = bins_needed * 2 * 4
        packed_data = d.queue.read_buffer(self._buf_packed_transfer, size=packed_size)
        packed = np.frombuffer(packed_data, dtype=np.float32).copy()

        # CPU-side unpack (vectorized NumPy — no Python loops)
        re_vals = packed[0::2].astype(np.uint32)
        im_vals = packed[1::2].astype(np.uint32)

        unpacked = np.zeros((bins_needed, 6), dtype=np.float32)
        unpacked[:, 0] = ((re_vals >> 16) & 0xFF).astype(np.float32) / 255.0
        unpacked[:, 1] = ((re_vals >> 8) & 0xFF).astype(np.float32) / 255.0
        unpacked[:, 2] = (re_vals & 0xFF).astype(np.float32) / 255.0
        unpacked[:, 3] = ((im_vals >> 16) & 0xFF).astype(np.float32) / 255.0
        unpacked[:, 4] = ((im_vals >> 8) & 0xFF).astype(np.float32) / 255.0
        unpacked[:, 5] = (im_vals & 0xFF).astype(np.float32) / 255.0

        positions = unpacked.flatten()[:N]

        # Also readback velocities (raw — velocity range needs different packing)
        vel_data = d.queue.read_buffer(self._buf_velocities, size=N * 4)
        velocities = np.frombuffer(vel_data, dtype=np.float32).reshape(C, D).copy()

        return positions.reshape(C, D), velocities

    # ─── GPU Tick ──────────────────────────────────────────────

    def _pack_params(self, p: dict) -> bytes:
        """Pack params dict into uniform buffer bytes matching WGSL struct layout."""
        return struct.pack("IIIIffffffffffffffff",
            p["num_cells"],
            p["num_dims"],
            p["num_concepts"],
            0,  # _pad0
            p["gravity_strength"],
            p["gravity_sharpness"],
            p["gravity_range"],
            p["softening"],
            p["force_scale"],
            p["tradeoff_amplifier"],
            p["thermal_noise_level"],
            p["momentum"],
            p["max_speed"],
            p["boundary_margin"],
            p["boundary_strength"],
            p["dt_norm"],
            p["gravity_ramp"],
            p["settling_damp"],
            0.0,  # _pad1
            0.0,  # _pad2
        )

    def tick(self, params_dict: dict):
        """Run one tick on GPU. No readback. Positions/velocities stay GPU-resident."""
        if self._bind_group_dirty:
            self._rebuild_bind_group()

        d = self._device
        C = self._num_cells
        workgroups = (C + 255) // 256  # workgroup_size(256)

        # Upload params uniform
        d.queue.write_buffer(self._buf_params, 0, self._pack_params(params_dict))

        # Single compute dispatch
        encoder = d.create_command_encoder()
        compute_pass = encoder.begin_compute_pass()
        compute_pass.set_pipeline(self._pipeline)
        compute_pass.set_bind_group(0, self._bg)
        compute_pass.dispatch_workgroups(workgroups)
        compute_pass.end()
        d.queue.submit([encoder.finish()])

    def run_generation(self, ticks_per_gen: int, params_dict: dict,
                       noise: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
        """
        Run N ticks with one noise upload, N dispatches, one readback.

        This is the batched dispatch pattern from wgpu_force_kernel.py:
        upload once, dispatch N times, readback once.

        Args:
            ticks_per_gen: number of physics ticks per generation
            params_dict: physics parameters (gravity_ramp/settling_damp pre-computed)
            noise: (C, D) f32 noise buffer [-1, 1] range

        Returns:
            (positions, velocities) as (C, D) float32 arrays
        """
        # Upload noise for this generation
        self.upload_noise(noise)

        # Dispatch N ticks — all GPU-resident, no per-tick transfer
        for _ in range(ticks_per_gen):
            self.tick(params_dict)

        # Single readback at end of generation
        return self.readback()

    # ─── Readback ──────────────────────────────────────────────

    def readback(self) -> tuple[np.ndarray, np.ndarray]:
        """Read positions and velocities back from GPU."""
        C = self._num_cells
        D = self.num_dims
        size = C * D * 4

        # Ensure GPU work is complete
        self._device._poll_wait()

        pos_data = self._device.queue.read_buffer(self._buf_positions, size=size)
        vel_data = self._device.queue.read_buffer(self._buf_velocities, size=size)

        positions = np.frombuffer(pos_data, dtype=np.float32).reshape(C, D).copy()
        velocities = np.frombuffer(vel_data, dtype=np.float32).reshape(C, D).copy()
        return positions, velocities

    @property
    def adapter_info(self) -> dict:
        """Return GPU adapter information."""
        return dict(self._device.adapter.info)


# ─── Gravity Ramp + Settling (matches JS gravityScale/settlingDamping) ─────

def compute_gravity_ramp(tick_count: int, total_ticks: int) -> float:
    """
    Gravity ramping: smoothstep from 5% to 100% over generation budget.
    Maps progress [0.1, 0.9] to [0, 1], then smoothstep, then 0.05 + 0.95 * smoothstep.
    Matches Python _compute_gravity_scale() and JS gravityScale().
    """
    if total_ticks <= 0:
        return 1.0
    progress = min(1.0, tick_count / total_ticks)
    t = max(0.0, min(1.0, (progress - 0.1) / 0.8))
    return 0.05 + 0.95 * t * t * (3 - 2 * t)


def compute_settling_damp(tick_count: int, total_ticks: int) -> float:
    """
    Settling: progressively slow cells in final 15%.
    Returns momentum multiplier (1.0 = no damping, lower = more damping).
    Matches Python _compute_settling_damping() and JS settlingDamping().
    """
    if total_ticks <= 0:
        return 1.0
    progress = min(1.0, tick_count / total_ticks)
    settling_start = 0.85
    if progress <= settling_start:
        return 1.0
    settling = (progress - settling_start) / (1.0 - settling_start)
    return 1.0 - 0.7 * settling
