"""
WgpuForceKernel — GPU-accelerated force calculations via WebGPU/WGSL.

Replaces the NumPy force loops in AfsEngine.tick() with three GPU compute shaders:
  1. similarity_forces.wgsl — similarity matrix × velocity coupling
  2. concept_gravity.wgsl  — cell-concept gravitational attraction
  3. position_update.wgsl  — force integration, velocity damping, boundary handling

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

Same WGSL shaders work in the browser via native WebGPU API.

Two usage modes:

  GPU-resident (fast — no per-tick transfer):
    kernel = WgpuForceKernel(num_dims=9)
    kernel.upload_matrix(similarity_matrix)
    kernel.upload_concepts(concept_positions)
    kernel.upload_cells(positions, velocities)
    for tick in range(100):
        kernel.tick(config)          # all on GPU, no readback
    positions, velocities = kernel.readback()  # one readback at end

  Per-tick upload (simpler, for when CPU needs data each tick):
    kernel.compute(positions, velocities, config)
    positions, velocities = kernel.readback()
"""

import numpy as np
import struct
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


class WgpuForceKernel:
    """GPU force kernel using WGSL compute shaders."""

    def __init__(self, num_dims: int, max_cells: int = 1024, max_concepts: int = 32):
        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_groups_dirty = True

        self._device = wgpu.utils.get_default_device()
        self._load_shaders()
        self._create_pipelines()
        self._create_buffers()

        # Cached bind groups (rebuilt when cell/concept count changes)
        self._sim_bg = None
        self._grav_bg = None
        self._upd_bg = None

        # (Force accumulators are zeroed GPU-side — no CPU zero buffer needed)

    def _load_shaders(self):
        self._similarity_src = (_SHADER_DIR / "similarity_forces.wgsl").read_text()
        self._gravity_src = (_SHADER_DIR / "concept_gravity.wgsl").read_text()
        self._update_src = (_SHADER_DIR / "position_update.wgsl").read_text()

    def _create_pipelines(self):
        d = self._device

        # Similarity forces
        sim_module = d.create_shader_module(code=self._similarity_src)
        self._sim_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}},
        ])
        sim_pl = d.create_pipeline_layout(bind_group_layouts=[self._sim_bgl])
        self._sim_pipeline = d.create_compute_pipeline(
            layout=sim_pl, compute={"module": sim_module, "entry_point": "main"})

        # Concept gravity
        grav_module = d.create_shader_module(code=self._gravity_src)
        self._grav_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}},
        ])
        grav_pl = d.create_pipeline_layout(bind_group_layouts=[self._grav_bgl])
        self._grav_pipeline = d.create_compute_pipeline(
            layout=grav_pl, compute={"module": grav_module, "entry_point": "main"})

        # Position update
        upd_module = d.create_shader_module(code=self._update_src)
        self._upd_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}},
        ])
        upd_pl = d.create_pipeline_layout(bind_group_layouts=[self._upd_bgl])
        self._upd_pipeline = d.create_compute_pipeline(
            layout=upd_pl, compute={"module": upd_module, "entry_point": "main"})

    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

        self._buf_positions = d.create_buffer(size=C_max * D * 4, usage=STORAGE_RW)
        self._buf_velocities = d.create_buffer(size=C_max * D * 4, usage=STORAGE_RW)
        self._buf_matrix_forces = d.create_buffer(size=C_max * D * 4, usage=STORAGE_RW)
        self._buf_gravity_forces = d.create_buffer(size=C_max * D * 4, usage=STORAGE_RW)
        self._buf_similarity_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)

        self._buf_sim_params = d.create_buffer(size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        self._buf_grav_params = d.create_buffer(size=48, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
        self._buf_upd_params = d.create_buffer(size=48, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

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

        self._sim_bg = d.create_bind_group(layout=self._sim_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_sim_params}},
            {"binding": 1, "resource": {"buffer": self._buf_velocities, "size": C * D * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_similarity_matrix}},
            {"binding": 3, "resource": {"buffer": self._buf_matrix_forces, "size": C * D * 4}},
        ])

        self._grav_bg = d.create_bind_group(layout=self._grav_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_grav_params}},
            {"binding": 1, "resource": {"buffer": self._buf_positions, "size": C * D * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_concept_positions, "size": K * D * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_gravity_forces, "size": C * D * 4}},
        ])

        self._upd_bg = d.create_bind_group(layout=self._upd_bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_upd_params}},
            {"binding": 1, "resource": {"buffer": self._buf_matrix_forces, "size": C * D * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_gravity_forces, "size": C * D * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_velocities, "size": C * D * 4}},
            {"binding": 4, "resource": {"buffer": self._buf_positions, "size": C * D * 4}},
        ])

        self._bind_groups_dirty = False

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

    def upload_matrix(self, matrix: np.ndarray):
        """Upload similarity matrix to GPU. Call at init and generation boundaries."""
        assert matrix.shape == (self.num_dims, self.num_dims)
        self._device.queue.write_buffer(
            self._buf_similarity_matrix, 0, matrix.astype(np.float32).tobytes())

    def upload_concepts(self, concept_positions: np.ndarray):
        """Upload concept positions to GPU. Call when concepts change."""
        K, D = concept_positions.shape
        assert D == self.num_dims and K <= self.max_concepts
        if K != self._num_concepts:
            self._num_concepts = K
            self._bind_groups_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 to GPU. Call once before tick() loop."""
        C, D = positions.shape
        assert D == self.num_dims and C <= self.max_cells
        assert velocities.shape == (C, D)
        if C != self._num_cells:
            self._num_cells = C
            self._bind_groups_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())

    # ─── GPU-Resident Tick ─────────────────────────────────────

    def tick(self, config: dict):
        """
        Run one tick entirely on GPU. Positions and velocities stay GPU-resident.
        Call readback() only when you need data on CPU (e.g., at generation boundaries).

        This is the fast path — no per-tick data transfer.
        """
        C = self._num_cells
        K = self._num_concepts
        D = self.num_dims
        d = self._device

        if self._bind_groups_dirty:
            self._rebuild_bind_groups()

        dt_norm = config.get("dt", 16.67) / 1000.0
        workgroups = (C + 63) // 64

        # Force accumulators are zeroed GPU-side:
        #   matrix_forces: similarity_forces.wgsl writes with = (full overwrite per cell).
        #   gravity_forces: concept_gravity.wgsl zeros its slice before accumulating.
        # No CPU write_buffer zero pass needed — eliminates 2 transfers per tick.

        # Update uniform params (small — 16-48 bytes each)
        d.queue.write_buffer(self._buf_sim_params, 0, struct.pack("IIff",
            C, D,
            config.get("force_scale", 0.15),
            config.get("tradeoff_amplifier", 1.0)))

        d.queue.write_buffer(self._buf_grav_params, 0, struct.pack("IIIIffffffff",
            C, K, D, 0,
            config.get("gravity_strength", 0.3),
            config.get("gravity_sharpness", 2.0),
            config.get("gravity_range", 1.5),
            config.get("softening", 0.05),
            dt_norm,
            0.0, 0.0, 0.0))

        d.queue.write_buffer(self._buf_upd_params, 0, struct.pack("IIIIffffffff",
            C, D, 0, 0,
            dt_norm,
            min(config.get("momentum", 0.7), 0.7),
            config.get("max_speed", 0.3),
            config.get("boundary_margin", 0.05),
            config.get("boundary_strength", 0.5),
            config.get("thermal_noise", 0.001),
            float(config.get("tick_count", 0)),
            0.0))

        # Dispatch all three shaders in one command buffer
        encoder = d.create_command_encoder()

        p1 = encoder.begin_compute_pass()
        p1.set_pipeline(self._sim_pipeline)
        p1.set_bind_group(0, self._sim_bg)
        p1.dispatch_workgroups(workgroups)
        p1.end()

        p2 = encoder.begin_compute_pass()
        p2.set_pipeline(self._grav_pipeline)
        p2.set_bind_group(0, self._grav_bg)
        p2.dispatch_workgroups(workgroups)
        p2.end()

        p3 = encoder.begin_compute_pass()
        p3.set_pipeline(self._upd_pipeline)
        p3.set_bind_group(0, self._upd_bg)
        p3.dispatch_workgroups(workgroups)
        p3.end()

        d.queue.submit([encoder.finish()])

    # ─── Per-Tick Upload Mode ──────────────────────────────────

    def compute(self, cell_positions: np.ndarray, cell_velocities: np.ndarray, config: dict):
        """Upload cell data + run one tick. Use tick() for GPU-resident mode instead."""
        self.upload_cells(cell_positions, cell_velocities)
        self.tick(config)

    # ─── 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

        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)
