"""
SB Fabric Primitives — GPU kernels for the 5 missing Sound Block primitives.

Each primitive wraps a WGSL shader from the sb_fabric_*.wgsl files.
Uses the same wgpu-py pattern as sb_colony_block_kernel.py.

Primitives:
    SBScan       — Harmonic Accumulation (prefix sum)
    SBHistogram  — Spectral Energy Distribution
    SBConvolve   — Resonance Filter (1D convolution)
    SBSegmented  — Multi-Band Processing
    SBCompact    — Active Bin Selection (stream compaction)
"""

import numpy as np
import struct
from pathlib import Path

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

_SHADER_DIR = Path(__file__).parent


def gpu_available() -> bool:
    if not _WGPU_AVAILABLE:
        return False
    try:
        return wgpu.utils.get_default_device() is not None
    except Exception:
        return False


# ─── SBScan: Harmonic Accumulation ────────────────────────────────────────────

class SBScan:
    """Parallel prefix sum (exclusive) on GPU. O(N) work, O(log N) depth."""

    OPS = {"sum": 0, "max": 1, "min": 2, "product": 3}

    def __init__(self, max_size: int = 1_000_000):
        self._device = wgpu.utils.get_default_device()
        self._max_size = max_size
        src = (_SHADER_DIR / "sb_fabric_scan.wgsl").read_text()
        module = self._device.create_shader_module(code=src)

        VIS = wgpu.ShaderStage.COMPUTE
        UNIFORM = wgpu.BufferBindingType.uniform
        STORAGE = wgpu.BufferBindingType.storage

        self._bgl = self._device.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}},
        ])
        pl = self._device.create_pipeline_layout(bind_group_layouts=[self._bgl])
        self._scan_pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "scan_local"})
        self._propagate_pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "propagate"})

        USAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST
        max_blocks = (max_size + 511) // 512
        self._buf_data = self._device.create_buffer(size=max_size * 4, usage=USAGE)
        self._buf_block_sums = self._device.create_buffer(size=max_blocks * 4, usage=USAGE)
        self._buf_params = self._device.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

    def scan(self, data: np.ndarray, op: str = "sum") -> np.ndarray:
        """Exclusive prefix scan. Returns scanned array."""
        N = len(data)
        d = self._device
        op_code = self.OPS.get(op, 0)

        d.queue.write_buffer(self._buf_data, 0, data.astype(np.float32).tobytes())

        num_blocks = (N + 511) // 512
        d.queue.write_buffer(self._buf_params, 0, struct.pack("IIII", N, op_code, 0, 0))

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_params}},
            {"binding": 1, "resource": {"buffer": self._buf_data, "size": N * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_block_sums,
                                        "size": num_blocks * 4}},
        ])

        # Local scan per workgroup
        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(self._scan_pipeline)
        cp.set_bind_group(0, bg)
        cp.dispatch_workgroups(num_blocks)
        cp.end()
        d.queue.submit([encoder.finish()])

        # Multi-block: scan block_sums, then propagate
        if num_blocks > 1:
            # Scan the block sums (recursive, but typically 1 level is enough)
            d.queue.write_buffer(self._buf_params, 0,
                                 struct.pack("IIII", num_blocks, op_code, 0, 0))
            bg2 = d.create_bind_group(layout=self._bgl, entries=[
                {"binding": 0, "resource": {"buffer": self._buf_params}},
                {"binding": 1, "resource": {"buffer": self._buf_block_sums,
                                            "size": num_blocks * 4}},
                {"binding": 2, "resource": {"buffer": self._buf_block_sums,
                                            "size": num_blocks * 4}},
            ])
            encoder = d.create_command_encoder()
            cp = encoder.begin_compute_pass()
            cp.set_pipeline(self._scan_pipeline)
            cp.set_bind_group(0, bg2)
            cp.dispatch_workgroups(1)
            cp.end()
            d.queue.submit([encoder.finish()])

            # Propagate block sums back to data
            # Restore original N in params, use original bg (has data + block_sums)
            d.queue.write_buffer(self._buf_params, 0,
                                 struct.pack("IIII", N, op_code, 2, 0))
            # Propagate uses 256 threads per workgroup, each adds block_sum to its element
            encoder = d.create_command_encoder()
            cp = encoder.begin_compute_pass()
            cp.set_pipeline(self._propagate_pipeline)
            cp.set_bind_group(0, bg)
            cp.dispatch_workgroups((N + 255) // 256)
            cp.end()
            d.queue.submit([encoder.finish()])

        d._poll_wait()
        result = d.queue.read_buffer(self._buf_data, size=N * 4)
        return np.frombuffer(result, dtype=np.float32)[:N].copy()


# ─── SBHistogram: Spectral Energy Distribution ───────────────────────────────

class SBHistogram:
    """GPU histogram via atomic binning."""

    def __init__(self, max_size: int = 1_000_000, max_bins: int = 256):
        self._device = wgpu.utils.get_default_device()
        src = (_SHADER_DIR / "sb_fabric_histogram.wgsl").read_text()
        module = self._device.create_shader_module(code=src)

        VIS = wgpu.ShaderStage.COMPUTE
        self._bgl = self._device.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 2, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])
        pl = self._device.create_pipeline_layout(bind_group_layouts=[self._bgl])
        self._pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "histogram"})

        USAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST
        self._buf_values = self._device.create_buffer(size=max_size * 4, usage=USAGE)
        self._buf_counts = self._device.create_buffer(size=max_bins * 4, usage=USAGE)
        self._buf_params = self._device.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

    def histogram(self, data: np.ndarray, num_bins: int = 64,
                  val_range: tuple = None) -> np.ndarray:
        """Compute histogram of data values."""
        N = len(data)
        d = self._device
        if val_range is None:
            val_range = (float(data.min()), float(data.max()))

        d.queue.write_buffer(self._buf_values, 0, data.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_counts, 0, np.zeros(num_bins, dtype=np.uint32).tobytes())
        d.queue.write_buffer(self._buf_params, 0,
                             struct.pack("IIff", N, num_bins, val_range[0], val_range[1]))

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_params}},
            {"binding": 1, "resource": {"buffer": self._buf_values, "size": N * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_counts, "size": num_bins * 4}},
        ])

        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(self._pipeline)
        cp.set_bind_group(0, bg)
        cp.dispatch_workgroups((N + 255) // 256)
        cp.end()
        d.queue.submit([encoder.finish()])

        d._poll_wait()
        result = d.queue.read_buffer(self._buf_counts, size=num_bins * 4)
        return np.frombuffer(result, dtype=np.uint32)[:num_bins].copy()


# ─── SBConvolve: Resonance Filter ────────────────────────────────────────────

class SBConvolve:
    """1D convolution with shared memory kernel (up to 255 taps)."""

    def __init__(self, max_size: int = 1_000_000):
        self._device = wgpu.utils.get_default_device()
        src = (_SHADER_DIR / "sb_fabric_convolve.wgsl").read_text()
        module = self._device.create_shader_module(code=src)

        VIS = wgpu.ShaderStage.COMPUTE
        self._bgl = self._device.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 2, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.read_only_storage}},
            {"binding": 3, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])
        pl = self._device.create_pipeline_layout(bind_group_layouts=[self._bgl])
        self._pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "convolve_1d"})

        USAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST
        self._buf_input = self._device.create_buffer(size=max_size * 4, usage=USAGE)
        self._buf_output = self._device.create_buffer(size=max_size * 4, usage=USAGE)
        self._buf_kernel = self._device.create_buffer(size=256 * 4, usage=USAGE)
        self._buf_params = self._device.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

    def convolve(self, data: np.ndarray, kernel: np.ndarray) -> np.ndarray:
        """1D convolution. Kernel must have odd length, max 255."""
        N = len(data)
        K = len(kernel)
        assert K % 2 == 1 and K <= 255, f"Kernel size must be odd and <= 255, got {K}"
        d = self._device

        d.queue.write_buffer(self._buf_input, 0, data.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_kernel, 0, kernel.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_params, 0,
                             struct.pack("IIII", N, K, (K - 1) // 2, 0))

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_params}},
            {"binding": 1, "resource": {"buffer": self._buf_input, "size": N * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_kernel, "size": K * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_output, "size": N * 4}},
        ])

        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(self._pipeline)
        cp.set_bind_group(0, bg)
        cp.dispatch_workgroups((N + 255) // 256)
        cp.end()
        d.queue.submit([encoder.finish()])

        d._poll_wait()
        result = d.queue.read_buffer(self._buf_output, size=N * 4)
        return np.frombuffer(result, dtype=np.float32)[:N].copy()


# ─── SBFFT: Spectral Transform ───────────────────────────────────────────────

class SBFFT:
    """GPU FFT via Cooley-Tukey radix-2. The fundamental Sound Block operation."""

    def __init__(self, max_size: int = 1_048_576):
        self._device = wgpu.utils.get_default_device()
        src = (_SHADER_DIR / "sb_fabric_fft.wgsl").read_text()
        module = self._device.create_shader_module(code=src)

        VIS = wgpu.ShaderStage.COMPUTE
        self._bgl = self._device.create_bind_group_layout(entries=[
            {"binding": 0, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.uniform}},
            {"binding": 1, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
            {"binding": 2, "visibility": VIS,
             "buffer": {"type": wgpu.BufferBindingType.storage}},
        ])
        pl = self._device.create_pipeline_layout(bind_group_layouts=[self._bgl])

        self._bit_reverse_pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "bit_reverse"})
        self._butterfly_pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "butterfly"})
        self._scale_pipeline = self._device.create_compute_pipeline(
            layout=pl, compute={"module": module, "entry_point": "fft_scale"})

        USAGE = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST
        self._buf_re = self._device.create_buffer(size=max_size * 4, usage=USAGE)
        self._buf_im = self._device.create_buffer(size=max_size * 4, usage=USAGE)
        self._buf_params = self._device.create_buffer(
            size=16, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

    def fft(self, re: np.ndarray, im: np.ndarray = None,
            inverse: bool = False) -> tuple:
        """Compute FFT (or inverse FFT) on GPU.

        Args:
            re: Real part (length must be power of 2)
            im: Imaginary part (same length, or None for real input)
            inverse: If True, compute inverse FFT (includes 1/N scaling)

        Returns:
            (re_out, im_out) as float32 arrays
        """
        N = len(re)
        assert N > 0 and (N & (N - 1)) == 0, f"FFT size must be power of 2, got {N}"
        d = self._device
        direction = -1 if not inverse else 1

        d.queue.write_buffer(self._buf_re, 0, re.astype(np.float32).tobytes())
        if im is not None:
            d.queue.write_buffer(self._buf_im, 0, im.astype(np.float32).tobytes())
        else:
            d.queue.write_buffer(self._buf_im, 0, np.zeros(N, dtype=np.float32).tobytes())

        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_params}},
            {"binding": 1, "resource": {"buffer": self._buf_re, "size": N * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_im, "size": N * 4}},
        ])

        workgroups = (N + 255) // 256
        half_workgroups = (N // 2 + 255) // 256

        # Step 1: Bit-reversal permutation
        d.queue.write_buffer(self._buf_params, 0,
                             struct.pack("IIiI", N, 0, direction, 0))
        encoder = d.create_command_encoder()
        cp = encoder.begin_compute_pass()
        cp.set_pipeline(self._bit_reverse_pipeline)
        cp.set_bind_group(0, bg)
        cp.dispatch_workgroups(workgroups)
        cp.end()
        d.queue.submit([encoder.finish()])

        # Step 2: log2(N) butterfly stages (one submit per stage for correct params)
        log2n = int(np.log2(N))
        for stage in range(log2n):
            d.queue.write_buffer(self._buf_params, 0,
                                 struct.pack("IIiI", N, stage, direction, 0))
            encoder = d.create_command_encoder()
            cp = encoder.begin_compute_pass()
            cp.set_pipeline(self._butterfly_pipeline)
            cp.set_bind_group(0, bg)
            cp.dispatch_workgroups(half_workgroups)
            cp.end()
            d.queue.submit([encoder.finish()])

        # Step 3: Scale for inverse FFT
        if inverse:
            d.queue.write_buffer(self._buf_params, 0,
                                 struct.pack("IIiI", N, 0, direction, 0))
            encoder = d.create_command_encoder()
            cp = encoder.begin_compute_pass()
            cp.set_pipeline(self._scale_pipeline)
            cp.set_bind_group(0, bg)
            cp.dispatch_workgroups(workgroups)
            cp.end()
            d.queue.submit([encoder.finish()])

        d._poll_wait()
        re_data = d.queue.read_buffer(self._buf_re, size=N * 4)
        im_data = d.queue.read_buffer(self._buf_im, size=N * 4)
        re_out = np.frombuffer(re_data, dtype=np.float32)[:N].copy()
        im_out = np.frombuffer(im_data, dtype=np.float32)[:N].copy()
        return re_out, im_out
