"""
SBClusteringGPUKernel — GPU-accelerated density-based clustering.

Ports the CPU clustering (29ms) to GPU using sb_clustering.wgsl.
Two-pass approach:
  Pass 1 (count_neighbors): parallel neighbor count per cell
  Pass 2 (assign_cluster): iterative greedy assignment (one cluster per dispatch)

The iteration (which cluster to assign next) is driven from Python,
but the heavy O(N^2) neighbor counting and O(N) assignment are on GPU.

Usage:
    kernel = SBClusteringGPUKernel(num_dims=10, max_cells=512)
    kernel.upload_positions(positions)
    clusters = kernel.detect_clusters()
    # clusters = list of (cluster_id, member_indices, centroid)
"""

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

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

_SHADER_DIR = Path(__file__).parent


@dataclass
class GPUCluster:
    """GPU-detected cluster result."""
    cluster_id: int
    centroid: np.ndarray
    cell_count: int
    member_indices: list


class SBClusteringGPUKernel:
    """GPU clustering kernel using sb_clustering.wgsl compute shader."""

    def __init__(self, num_dims: int, max_cells: int = 512, max_clusters: 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_clusters = max_clusters

        self._num_cells = 0
        self._positions_np = None  # kept for centroid computation

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

    def _load_shader(self):
        self._shader_src = (_SHADER_DIR / "sb_clustering.wgsl").read_text()

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

        # Bind group layout: params(uniform), positions(read), neighbor_counts(rw), assignments(rw)
        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.read_only_storage}},
            {"binding": 2, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            {"binding": 3, "visibility": wgpu.ShaderStage.COMPUTE,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])

        pl = d.create_pipeline_layout(bind_group_layouts=[self._bgl])

        self._count_pipeline = d.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "count_neighbors"})
        self._assign_pipeline = d.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "assign_cluster"})

    def _create_buffers(self):
        d = self._device
        C = self.max_cells
        D = self.num_dims

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

        # Positions: C * D f32
        self._buf_positions = d.create_buffer(size=C * D * 4, usage=STORAGE_RO)

        # Neighbor counts: C u32
        self._buf_neighbor_counts = d.create_buffer(size=C * 4, usage=STORAGE_RW)

        # Cluster assignments: C u32
        self._buf_assignments = d.create_buffer(size=C * 4, usage=STORAGE_RW)

        # Params uniform: 8 u32/f32 = 32 bytes
        self._buf_params = d.create_buffer(
            size=32,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

    def upload_positions(self, positions: np.ndarray):
        """Upload cell positions. Call before detect_clusters()."""
        C, D = positions.shape
        assert D == self.num_dims and C <= self.max_cells
        self._num_cells = C
        self._positions_np = positions.copy()
        self._device.queue.write_buffer(
            self._buf_positions, 0,
            positions.astype(np.float32).tobytes())

    def detect_clusters(self, max_clusters: int = None) -> list:
        """
        Run GPU-accelerated density-based clustering.

        Returns list of GPUCluster objects sorted by cell_count (descending).
        Matches the algorithm in afs-engine/afs/analysis/cluster_detection.py.
        """
        if max_clusters is None:
            max_clusters = self.max_clusters

        C = self._num_cells
        D = self.num_dims
        d = self._device

        # Compute cluster radius (matches Python detect_clusters)
        cluster_radius = 0.10 + 0.04 * np.log(D)
        cluster_radius_sq = cluster_radius * cluster_radius
        min_cluster_size = max(3, int(C * 0.05))

        # Clear assignments buffer to 0
        zeros = np.zeros(C, dtype=np.uint32)
        d.queue.write_buffer(self._buf_assignments, 0, zeros.tobytes())

        clusters = []
        current_cluster_id = 1

        for _ in range(max_clusters):
            # --- Pass 1: Count neighbors on GPU ---
            params = struct.pack("IIfIIIII",
                C,                      # num_cells
                D,                      # num_dims
                cluster_radius_sq,      # cluster_radius_sq
                min_cluster_size,       # min_cluster_size
                max_clusters,           # max_clusters
                current_cluster_id,     # current_cluster_id (unused in pass 1)
                0,                      # seed_cell (unused in pass 1)
                0,                      # _pad0
            )
            d.queue.write_buffer(self._buf_params, 0, params)

            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_neighbor_counts,
                                            "size": C * 4}},
                {"binding": 3, "resource": {"buffer": self._buf_assignments,
                                            "size": C * 4}},
            ])

            workgroups = (C + 255) // 256
            encoder = d.create_command_encoder()
            compute_pass = encoder.begin_compute_pass()
            compute_pass.set_pipeline(self._count_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 back neighbor counts to find densest seed
            counts_data = d.queue.read_buffer(self._buf_neighbor_counts, size=C * 4)
            counts = np.frombuffer(counts_data, dtype=np.uint32)[:C].copy()

            best_idx = int(np.argmax(counts))
            best_count = int(counts[best_idx])

            if best_count < min_cluster_size:
                break

            # --- Pass 2: Assign neighbors of seed to this cluster on GPU ---
            params = struct.pack("IIfIIIII",
                C,                      # num_cells
                D,                      # num_dims
                cluster_radius_sq,      # cluster_radius_sq
                min_cluster_size,       # min_cluster_size
                max_clusters,           # max_clusters
                current_cluster_id,     # current_cluster_id
                best_idx,               # seed_cell
                0,                      # _pad0
            )
            d.queue.write_buffer(self._buf_params, 0, params)

            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_neighbor_counts,
                                            "size": C * 4}},
                {"binding": 3, "resource": {"buffer": self._buf_assignments,
                                            "size": C * 4}},
            ])

            encoder = d.create_command_encoder()
            compute_pass = encoder.begin_compute_pass()
            compute_pass.set_pipeline(self._assign_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 back assignments to count members and compute centroid
            assign_data = d.queue.read_buffer(self._buf_assignments, size=C * 4)
            assignments = np.frombuffer(assign_data, dtype=np.uint32)[:C].copy()

            member_mask = assignments == current_cluster_id
            member_indices = np.where(member_mask)[0].tolist()

            if len(member_indices) < min_cluster_size:
                break

            centroid = np.mean(self._positions_np[member_mask], axis=0)
            clusters.append(GPUCluster(
                cluster_id=current_cluster_id,
                centroid=centroid,
                cell_count=len(member_indices),
                member_indices=member_indices,
            ))

            current_cluster_id += 1

            # Check if enough unassigned cells remain
            unassigned_count = int(np.sum(assignments == 0))
            if unassigned_count < min_cluster_size:
                break

        clusters.sort(key=lambda c: c.cell_count, reverse=True)
        return clusters

    def readback_assignments(self) -> np.ndarray:
        """Read the full cluster assignment array from GPU."""
        C = self._num_cells
        self._device._poll_wait()
        data = self._device.queue.read_buffer(self._buf_assignments, size=C * 4)
        return np.frombuffer(data, dtype=np.uint32)[:C].copy()

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


# ─── Standalone test: compare GPU vs CPU clustering ─────────────────────────

if __name__ == "__main__":
    import sys
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))
    from afs.analysis.cluster_detection import detect_clusters as cpu_detect_clusters

    print("=" * 60)
    print("SB Clustering GPU Kernel — GPU vs CPU comparison")
    print("=" * 60)

    # Generate test colony data (500 cells, 10 dims with clear clusters)
    np.random.seed(42)
    num_cells = 500
    num_dims = 10

    # Create 4 clusters with known centroids
    centroids = [
        np.full(num_dims, 0.2),
        np.full(num_dims, 0.5),
        np.full(num_dims, 0.8),
        np.random.uniform(0.3, 0.7, num_dims),
    ]
    cluster_sizes = [150, 120, 100, 80]
    noise_cells = num_cells - sum(cluster_sizes)

    positions = []
    for centroid, size in zip(centroids, cluster_sizes):
        cells = centroid + np.random.normal(0, 0.03, (size, num_dims))
        cells = np.clip(cells, 0, 1)
        positions.append(cells)
    # Add noise cells
    positions.append(np.random.uniform(0, 1, (noise_cells, num_dims)))
    positions = np.vstack(positions).astype(np.float32)

    # Shuffle to make it realistic
    perm = np.random.permutation(num_cells)
    positions = positions[perm]

    # --- CPU clustering ---
    t0 = time.perf_counter()
    cpu_clusters = cpu_detect_clusters(positions, max_clusters=8)
    cpu_time = (time.perf_counter() - t0) * 1000

    print(f"\nCPU clustering: {cpu_time:.1f}ms")
    print(f"  Found {len(cpu_clusters)} clusters:")
    for i, c in enumerate(cpu_clusters):
        print(f"    Cluster {i+1}: {c.cell_count} cells")

    # --- GPU clustering ---
    try:
        kernel = SBClusteringGPUKernel(num_dims=num_dims, max_cells=512)
        print(f"\nGPU: {kernel.adapter_info.get('description', 'unknown')}")

        kernel.upload_positions(positions)

        t0 = time.perf_counter()
        gpu_clusters = kernel.detect_clusters()
        gpu_time = (time.perf_counter() - t0) * 1000

        print(f"GPU clustering: {gpu_time:.1f}ms")
        print(f"  Found {len(gpu_clusters)} clusters:")
        for c in gpu_clusters:
            print(f"    Cluster {c.cluster_id}: {c.cell_count} cells")

        # Compare results
        print(f"\n--- Comparison ---")
        print(f"  CPU: {len(cpu_clusters)} clusters, {cpu_time:.1f}ms")
        print(f"  GPU: {len(gpu_clusters)} clusters, {gpu_time:.1f}ms")

        if len(cpu_clusters) == len(gpu_clusters):
            cpu_sizes = sorted([c.cell_count for c in cpu_clusters], reverse=True)
            gpu_sizes = sorted([c.cell_count for c in gpu_clusters], reverse=True)
            print(f"  CPU sizes: {cpu_sizes}")
            print(f"  GPU sizes: {gpu_sizes}")

            # Check if cluster sizes are comparable (exact match unlikely due to tie-breaking)
            total_cpu = sum(cpu_sizes)
            total_gpu = sum(gpu_sizes)
            print(f"  Total assigned: CPU={total_cpu}, GPU={total_gpu}")

            if abs(total_cpu - total_gpu) <= num_cells * 0.1:
                print(f"  PASS: Cluster counts and total assigned cells match within 10%")
            else:
                print(f"  WARN: Total assigned cells differ by more than 10%")
        else:
            print(f"  NOTE: Different cluster counts (CPU={len(cpu_clusters)}, GPU={len(gpu_clusters)})")
            print(f"  This is acceptable — greedy algorithms are sensitive to tie-breaking")

        if gpu_time > 0:
            speedup = cpu_time / gpu_time
            print(f"\n  Speedup: {speedup:.1f}x {'(GPU faster)' if speedup > 1 else '(CPU faster)'}")

    except Exception as e:
        print(f"\nGPU clustering failed: {e}")
        import traceback
        traceback.print_exc()
