"""
SBHomeostasisKernel — GPU Homeostasis Calibrator as Sound Block.

Runs N probe colonies in parallel on GPU using sb_homeostasis.wgsl.
Each probe partition has its own physics config. All share the same
influence matrix and attractor positions.

Stage 1: 4 wide-scan probes → pick top 2
Stage 2: 3 bisection probes → pick best overall
Total: 2 GPU dispatches (one per stage), ~1-2ms each.

Usage:
    kernel = SBHomeostasisKernel(num_dims=10, max_attractors=8)
    best_cfg = kernel.calibrate(influence_matrix, attractor_positions, base_cfg)
"""

import math
import struct
import time
import numpy as np
from pathlib import Path
from typing import Optional

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

_SHADER_DIR = Path(__file__).parent

# SharedParams: 8 u32 = 32 bytes
_SHARED_PARAMS_SIZE = 32
_SHARED_PARAMS_FORMAT = "IIIIIIII"

# ConfigEntry: 16 f32 = 64 bytes (matches WGSL struct)
_CONFIG_ENTRY_SIZE = 64
_CONFIG_ENTRY_FORMAT = "f" * 16


class SBHomeostasisKernel:
    """GPU-accelerated homeostasis calibrator.

    Runs probe colonies in parallel: one GPU dispatch per stage,
    all configs in a single pass via partitioned probe buffers.
    """

    # Probe sizing (matches Python HomeostasisCalibrator)
    PROBE_CELLS_LOW_D = 32
    PROBE_CELLS_HIGH_D = 48
    PROBE_GENS_LOW_D = 10
    PROBE_GENS_MID_D = 15
    PROBE_GENS_HIGH_D = 20
    MAX_CONFIGS_PER_STAGE = 4

    def __init__(self, num_dims: int, max_attractors: int = 64):
        if not _WGPU_AVAILABLE:
            raise ImportError("wgpu required: pip install wgpu")

        self.num_dims = num_dims
        self._cells_per_config = (self.PROBE_CELLS_HIGH_D if num_dims >= 10
                                  else self.PROBE_CELLS_LOW_D)
        self._probe_gens = (self.PROBE_GENS_HIGH_D if num_dims >= 10
                            else (self.PROBE_GENS_MID_D if num_dims >= 7
                                  else self.PROBE_GENS_LOW_D))
        self._max_configs = self.MAX_CONFIGS_PER_STAGE
        self._max_probes = self._max_configs * self._cells_per_config
        self._max_attractors = max_attractors

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

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

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

        # 8 bindings: shared params, amplitudes, velocities, forces,
        # influence_matrix, attractor_positions, deltas, configs
        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.storage}},
            {"binding": 7, "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"})

    def _create_buffers(self):
        d = self._device
        D = self.num_dims
        P = self._max_probes
        A = self._max_attractors

        RW = (wgpu.BufferUsage.STORAGE |
              wgpu.BufferUsage.COPY_SRC |
              wgpu.BufferUsage.COPY_DST)
        RO = wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_DST

        probe_size = P * D * 4
        self._buf_amplitudes = d.create_buffer(size=probe_size, usage=RW)
        self._buf_velocities = d.create_buffer(size=probe_size, usage=RW)
        self._buf_forces = d.create_buffer(size=probe_size, usage=RW)
        self._buf_deltas = d.create_buffer(size=P * 4, usage=RW)

        self._buf_influence_matrix = d.create_buffer(size=D * D * 4, usage=RO)
        self._buf_attractor_positions = d.create_buffer(size=A * D * 4, usage=RO)

        self._buf_run_params = d.create_buffer(
            size=_SHARED_PARAMS_SIZE,
            usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)

        # Config buffer: max_configs × 64 bytes
        self._buf_configs = d.create_buffer(
            size=self._max_configs * _CONFIG_ENTRY_SIZE, usage=RO)

    # ─── Field Analysis ──────────────────────────────────────────────────────

    def _analyze_matrix(self, matrix: np.ndarray) -> dict:
        n = self.num_dims
        upper = matrix[np.triu_indices(n, k=1)]
        non_zero = upper[np.abs(upper) > 0.001]

        if len(non_zero) == 0:
            return {"polarity_ratio": 0, "mean_abs": 0, "spread": 0}

        return {
            "polarity_ratio": float(np.sum(non_zero < 0) / len(non_zero)),
            "mean_abs": float(np.mean(np.abs(non_zero))),
            "spread": float(np.max(upper) - np.min(upper)),
        }

    def _param_space(self, analysis: dict) -> dict:
        n = self.num_dims
        df = math.sqrt(n)
        return {
            "force_scale":        (0.03, 0.35),
            "gravity_strength":   (1.5, 8.0),
            "gravity_range":      (max(1.5, df * 0.8), min(8.0, df * 2.5)),
            "gravity_sharpness":  (2.0, 8.0),
            "thermal_noise":      (0.01, 0.12),
            "tradeoff_amplifier": (1.0 if analysis["polarity_ratio"] > 0.4 else 2.0, 5.0),
        }

    def _sample_configs(self, space: dict, centers, width: float,
                        n_configs: int) -> list[dict]:
        keys = ["force_scale", "gravity_strength", "gravity_range",
                "gravity_sharpness", "thermal_noise", "tradeoff_amplifier"]
        configs = []

        if centers is None:
            for i in range(n_configs):
                t = i / max(1, n_configs - 1)
                cfg = {"label": f"wide-{i+1}"}
                for k in keys:
                    lo, hi = space[k]
                    if k in ("tradeoff_amplifier", "gravity_strength"):
                        cfg[k] = lo + (hi - lo) * (0.2 + 0.6 * ((i * 0.618) % 1.0))
                    else:
                        cfg[k] = lo + (hi - lo) * t
                configs.append(cfg)
        else:
            if len(centers) >= 2:
                mid = {"label": "bisect-mid"}
                for k in keys:
                    lo, hi = space[k]
                    a = centers[0].get(k, (lo + hi) / 2)
                    b = centers[1].get(k, (lo + hi) / 2)
                    val = math.sqrt(a * b) if (a > 0 and b > 0) else (a + b) / 2
                    mid[k] = max(lo, min(hi, val))
                configs.append(mid)

            best = centers[0]
            for i in range(n_configs - len(configs)):
                perturb = {"label": f"narrow-{i+1}"}
                rng_offset = (i + 1) * 0.618
                for k in keys:
                    lo, hi = space[k]
                    base = best.get(k, (lo + hi) / 2)
                    h = (len(k) * 7 + ord(k[0])) % 7 + 1
                    delta = (hi - lo) * width * (0.3 + 0.2 * ((rng_offset * h) % 1.0))
                    sign = 1 if i % 2 == 0 else -1
                    perturb[k] = max(lo, min(hi, base + sign * delta))
                configs.append(perturb)

        return configs[:n_configs]

    # ─── GPU Dispatch ────────────────────────────────────────────────────────

    def _pack_config_entry(self, cfg: dict) -> bytes:
        return struct.pack(_CONFIG_ENTRY_FORMAT,
            cfg.get("gravity_strength", 5.0),
            cfg.get("gravity_sharpness", 3.0),
            cfg.get("gravity_range", 3.0),
            cfg.get("softening", 0.05),
            cfg.get("force_scale", 0.1),
            cfg.get("tradeoff_amplifier", 1.5),
            cfg.get("thermal_noise", 0.02),
            cfg.get("momentum", 0.7),
            cfg.get("max_amplitude", 1.0),
            cfg.get("boundary_margin", 0.08),
            cfg.get("boundary_strength", 0.05),
            cfg.get("dt_norm", 0.01667),
            cfg.get("settling_start", 0.85),
            cfg.get("settling_strength", 0.7),
            0.0, 0.0,  # pads
        )

    def _init_probes_partitioned(self, n_configs: int,
                                  attractor_positions: np.ndarray) -> tuple:
        """Initialize probe partitions — 60% near attractors, 40% exploration."""
        C = self._cells_per_config
        D = self.num_dims
        total = n_configs * C
        n_att = attractor_positions.shape[0]

        amps = np.full((total, D), 0.5, dtype=np.float32)
        vels = np.zeros((total, D), dtype=np.float32)

        for p in range(n_configs):
            offset = p * C
            per_att = max(2, int(C * 0.6 / max(1, n_att)))
            att_total = per_att * n_att
            explore_total = C - att_total

            idx = 0
            for a in range(n_att):
                for _ in range(per_att):
                    if idx >= C:
                        break
                    amps[offset + idx] = attractor_positions[a] + \
                        (np.random.rand(D).astype(np.float32) - 0.5) * 0.5
                    idx += 1

            centroid = attractor_positions.mean(axis=0)
            for _ in range(explore_total):
                if idx >= C:
                    break
                amps[offset + idx] = centroid + \
                    (np.random.rand(D).astype(np.float32) - 0.5) * 0.6
                idx += 1

        amps = np.clip(amps, 0.05, 0.95)
        return amps, vels

    def _run_stage_gpu(self, probe_configs: list[dict],
                       influence_matrix: np.ndarray,
                       attractor_positions: np.ndarray) -> list[dict]:
        """Run one stage of probe configs on GPU. Returns scored results."""
        n_configs = len(probe_configs)
        if n_configs == 0:
            return []

        C = self._cells_per_config
        D = self.num_dims
        total_probes = n_configs * C
        total_passes = self._probe_gens
        A = attractor_positions.shape[0]
        d = self._device

        # Initialize partitioned probes
        amps, vels = self._init_probes_partitioned(n_configs, attractor_positions)

        # Upload shared data
        d.queue.write_buffer(self._buf_influence_matrix, 0,
                             influence_matrix.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_attractor_positions, 0,
                             attractor_positions.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_amplitudes, 0,
                             amps.astype(np.float32).tobytes())
        d.queue.write_buffer(self._buf_velocities, 0,
                             vels.astype(np.float32).tobytes())

        # Upload per-config physics
        config_data = b""
        for pc in probe_configs:
            config_data += self._pack_config_entry(pc)
        # Pad to max_configs
        while len(config_data) < self._max_configs * _CONFIG_ENTRY_SIZE:
            config_data += b"\x00" * _CONFIG_ENTRY_SIZE
        d.queue.write_buffer(self._buf_configs, 0, config_data)

        # Upload shared params (pass_index=0, updated per dispatch)
        shared_data = struct.pack(_SHARED_PARAMS_FORMAT,
            total_probes, D, A, 0,  # pass_index
            total_passes, C, n_configs, 0)
        d.queue.write_buffer(self._buf_run_params, 0, shared_data)

        # Build bind group
        bg = d.create_bind_group(layout=self._bgl, entries=[
            {"binding": 0, "resource": {"buffer": self._buf_run_params}},
            {"binding": 1, "resource": {"buffer": self._buf_amplitudes,
                                        "size": total_probes * D * 4}},
            {"binding": 2, "resource": {"buffer": self._buf_velocities,
                                        "size": total_probes * D * 4}},
            {"binding": 3, "resource": {"buffer": self._buf_forces,
                                        "size": total_probes * D * 4}},
            {"binding": 4, "resource": {"buffer": self._buf_influence_matrix}},
            {"binding": 5, "resource": {"buffer": self._buf_attractor_positions,
                                        "size": A * D * 4}},
            {"binding": 6, "resource": {"buffer": self._buf_deltas,
                                        "size": total_probes * 4}},
            {"binding": 7, "resource": {"buffer": self._buf_configs,
                                        "size": n_configs * _CONFIG_ENTRY_SIZE}},
        ])

        workgroups = (total_probes + 255) // 256

        # ── Fused multi-pass dispatch: one encoder, all passes ──
        encoder = d.create_command_encoder()
        for p in range(total_passes):
            d.queue.write_buffer(self._buf_run_params, 12,
                                 struct.pack("I", p))
            cp = encoder.begin_compute_pass()
            cp.set_pipeline(self._pipeline)
            cp.set_bind_group(0, bg)
            cp.dispatch_workgroups(workgroups)
            cp.end()
        d.queue.submit([encoder.finish()])

        # Readback final amplitudes
        d._poll_wait()
        data = d.queue.read_buffer(self._buf_amplitudes,
                                   size=total_probes * D * 4)
        final_amps = np.frombuffer(data, dtype=np.float32).reshape(
            total_probes, D).copy()

        # Score each partition
        results = []
        for i in range(n_configs):
            partition = final_amps[i * C:(i + 1) * C]
            score = self._score_partition(partition, attractor_positions)
            results.append({
                "label": probe_configs[i].get("label", f"probe-{i}"),
                "score": score["total"],
                "detail": score,
            })

        return results

    def _score_partition(self, positions: np.ndarray,
                         attractor_positions: np.ndarray) -> dict:
        """Score a probe partition on differentiation quality.
        Matches HomeostasisCalibrator._score_colony() weights."""
        n_cells = positions.shape[0]
        D = positions.shape[1]

        if n_cells < 4:
            return {"total": 0, "cluster_count": 0}

        # Check for NaN
        if np.any(np.isnan(positions)):
            return {"total": 0, "unstable": True}

        # Simple density clustering
        clusters = self._cluster_partition(positions)
        cluster_count = len(clusters)

        # Separation — mean pairwise centroid distance
        separation = 0.0
        if cluster_count >= 2:
            centroids = np.array([c["centroid"] for c in clusters])
            pair_count = 0
            for i in range(cluster_count):
                for j in range(i + 1, cluster_count):
                    separation += float(np.linalg.norm(centroids[i] - centroids[j]))
                    pair_count += 1
            separation /= max(1, pair_count)

        # Spread — RMS variance
        variance = np.var(positions, axis=0)
        spread = float(np.sqrt(np.mean(variance)))

        # Energy health (proxy from spread)
        energy_health = 1.0 - 2.0 * abs(min(1.0, spread / 0.3) - 0.5)

        # Pair distance variance
        pair_dist_var = 0.0
        n_att = attractor_positions.shape[0]
        if n_att >= 2:
            dists = []
            for i in range(n_att):
                for j in range(i + 1, n_att):
                    dists.append(float(np.linalg.norm(
                        attractor_positions[i] - attractor_positions[j])))
            if dists:
                mean_d = sum(dists) / len(dists)
                pair_dist_var = sum((d - mean_d) ** 2 for d in dists) / len(dists)

        total = (cluster_count * 25
                 + separation * 20
                 + spread * 30
                 + energy_health * 10
                 + 1.0 * 10  # growth = 1.0 (fixed population)
                 + pair_dist_var * 15)

        return {
            "total": round(total, 2),
            "cluster_count": cluster_count,
            "separation": round(separation, 4),
            "spread": round(spread, 4),
            "energy_health": round(energy_health, 3),
            "pair_dist_variance": round(pair_dist_var, 4),
        }

    def _cluster_partition(self, positions: np.ndarray) -> list:
        """Simple density clustering matching production logic."""
        n = positions.shape[0]
        D = positions.shape[1]
        radius = 0.10 + 0.04 * math.log(max(1, D))
        min_size = max(3, int(n * 0.05))
        assigned = np.zeros(n, dtype=bool)
        clusters = []

        for _ in range(8):  # max 8 clusters
            # Find densest unassigned point
            best_idx = -1
            best_count = 0
            for i in range(n):
                if assigned[i]:
                    continue
                count = 0
                for j in range(n):
                    if assigned[j]:
                        continue
                    dist = float(np.linalg.norm(positions[i] - positions[j]))
                    if dist < radius:
                        count += 1
                if count > best_count:
                    best_count = count
                    best_idx = i

            if best_idx < 0 or best_count < min_size:
                break

            # Assign neighbors
            members = []
            for j in range(n):
                if assigned[j]:
                    continue
                dist = float(np.linalg.norm(positions[best_idx] - positions[j]))
                if dist < radius:
                    assigned[j] = True
                    members.append(j)

            if len(members) >= min_size:
                centroid = positions[members].mean(axis=0)
                clusters.append({
                    "centroid": centroid,
                    "cell_count": len(members),
                })

        return clusters

    # ─── Public: Calibrate ───────────────────────────────────────────────────

    def calibrate(self, influence_matrix: np.ndarray,
                  attractor_positions: np.ndarray,
                  base_cfg: dict) -> dict:
        """Run full adaptive bisection calibration on GPU.

        Returns dict with winning config + calibration metadata.
        """
        t0 = time.perf_counter()
        analysis = self._analyze_matrix(influence_matrix)
        space = self._param_space(analysis)

        # ── Stage 1: 4 wide-scan probes (single GPU dispatch) ──
        stage1_cfgs = self._sample_configs(space, None, 1.0, 4)
        # Ensure gravity_range is set
        default_gr = base_cfg.get("gravity_range",
            max(2.0, min(6.0, 1.5 * math.sqrt(self.num_dims))))
        for c in stage1_cfgs:
            c.setdefault("gravity_range", default_gr)

        t1 = time.perf_counter()
        stage1_results = self._run_stage_gpu(
            stage1_cfgs, influence_matrix, attractor_positions)
        t_stage1 = (time.perf_counter() - t1) * 1000

        stage1_results.sort(key=lambda r: r["score"], reverse=True)

        # ── Stage 2: 3 bisection probes (single GPU dispatch) ──
        top2_cfgs = []
        for r in stage1_results[:2]:
            for c in stage1_cfgs:
                if c["label"] == r["label"]:
                    top2_cfgs.append(c)
                    break

        stage2_cfgs = self._sample_configs(space, top2_cfgs, 0.5, 3)
        for c in stage2_cfgs:
            c.setdefault("gravity_range", default_gr)

        t2 = time.perf_counter()
        stage2_results = self._run_stage_gpu(
            stage2_cfgs, influence_matrix, attractor_positions)
        t_stage2 = (time.perf_counter() - t2) * 1000

        # Pick best across both stages
        all_results = stage1_results + stage2_results
        all_cfgs = stage1_cfgs + stage2_cfgs
        all_results.sort(key=lambda r: r["score"], reverse=True)
        best = all_results[0]

        # Find winning config
        win_cfg = None
        for c in all_cfgs:
            if c["label"] == best["label"]:
                win_cfg = c
                break

        total_ms = (time.perf_counter() - t0) * 1000

        return {
            "winning_config": {
                "force_scale": win_cfg["force_scale"],
                "gravity_strength": win_cfg["gravity_strength"],
                "gravity_range": win_cfg["gravity_range"],
                "gravity_sharpness": win_cfg["gravity_sharpness"],
                "thermal_noise": win_cfg["thermal_noise"],
                "tradeoff_amplifier": win_cfg["tradeoff_amplifier"],
            } if win_cfg else {},
            "calibration": {
                "method": "gpu-adaptive-bisection",
                "label": best["label"],
                "score": best["score"],
                "detail": best["detail"],
                "probe_count": len(all_results),
                "stage1_ms": round(t_stage1, 2),
                "stage2_ms": round(t_stage2, 2),
                "total_ms": round(total_ms, 2),
                "field_analysis": analysis,
                "all_scores": [{"label": r["label"], "score": r["score"]}
                               for r in all_results],
            },
        }
