#!/usr/bin/env python3
"""
SB V2 Benchmark — Compare CPU vs GPU colony simulation performance.

Runs the same 10D synth decision (500 cells, 150 generations) through:
  1. SB V2 Python CPU (pure Python, matching JS V2 physics)
  2. SB V2 Python GPU (using SBColonyGPUKernel)

Logs results to docs/AtomataFieldTheory/SoundBlock-V2-GPU-Results.md

Usage:
    cd afs-engine
    python3 -m afs.gpu.sb_v2_benchmark
"""

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

# Ensure afs-engine is on the Python path regardless of cwd
_AFS_ENGINE_DIR = str(Path(__file__).resolve().parent.parent.parent)
if _AFS_ENGINE_DIR not in sys.path:
    sys.path.insert(0, _AFS_ENGINE_DIR)
from datetime import datetime

# ─── Reference Decision (matches JS test_sb_second_opinion.js SYNTH_DECISION) ──

SYNTH_DECISION = {
    "question": "Which synthesizer should I buy?",
    "options": [
        {"name": "Prophet 6", "price": 3500, "analog_warmth": 85, "versatility": 45,
         "build_quality": 80, "resale_value": 75, "learning_curve": 35,
         "patch_memory": 50, "effects_quality": 60, "mod_routing": 55, "portability": 40},
        {"name": "OB-8x", "price": 4999, "analog_warmth": 92, "versatility": 55,
         "build_quality": 90, "resale_value": 85, "learning_curve": 50,
         "patch_memory": 80, "effects_quality": 70, "mod_routing": 65, "portability": 30},
        {"name": "Nord Stage 4", "price": 4000, "analog_warmth": 30, "versatility": 95,
         "build_quality": 85, "resale_value": 60, "learning_curve": 70,
         "patch_memory": 95, "effects_quality": 80, "mod_routing": 40, "portability": 85},
    ],
    "factors": [
        {"name": "price", "range": [2000, 6000], "unit": "$", "goal": "minimize"},
        {"name": "analog_warmth", "range": [0, 100], "unit": "score"},
        {"name": "versatility", "range": [0, 100], "unit": "score"},
        {"name": "build_quality", "range": [0, 100], "unit": "score"},
        {"name": "resale_value", "range": [0, 100], "unit": "%", "goal": "maximize"},
        {"name": "learning_curve", "range": [0, 100], "unit": "hours", "goal": "minimize"},
        {"name": "patch_memory", "range": [0, 100], "unit": "score"},
        {"name": "effects_quality", "range": [0, 100], "unit": "score"},
        {"name": "mod_routing", "range": [0, 100], "unit": "score"},
        {"name": "portability", "range": [0, 100], "unit": "score"},
    ],
    "relationships": [
        {"from": "price", "to": "build_quality", "strength": 0.8},
        {"from": "analog_warmth", "to": "versatility", "strength": -0.6},
        {"from": "learning_curve", "to": "versatility", "strength": 0.5},
        {"from": "build_quality", "to": "resale_value", "strength": 0.7},
        {"from": "versatility", "to": "effects_quality", "strength": 0.6},
        {"from": "analog_warmth", "to": "mod_routing", "strength": 0.5},
        {"from": "portability", "to": "build_quality", "strength": -0.4},
        {"from": "patch_memory", "to": "learning_curve", "strength": -0.3},
        {"from": "price", "to": "analog_warmth", "strength": 0.5},
        {"from": "mod_routing", "to": "versatility", "strength": 0.4},
    ],
}

# ─── Simulation Config (matches JS DEFAULT_CONFIG_V2) ───────────────────────

SIM_CONFIG = {
    "cells": 500,
    "generations": 150,
    "ticks_per_generation": 1,
    "force_scale": 0.10,
    "gravity_strength": 5.0,
    "gravity_range": None,  # auto-computed
    "gravity_sharpness": 3.0,
    "softening": 0.05,
    "thermal_noise": 0.02,
    "momentum": 0.95,  # capped to 0.7 in integration
    "max_speed": 0.3,
    "boundary_margin": 0.08,
    "boundary_strength": 0.05,
    "tradeoff_amplifier": 1.5,
    "spawn_spread": 0.3,
    "dt": 16.67,
    # V2 metabolism
    "initial_energy": 50.0,
    "energy_min": 0.1,
    "energy_max": 100.0,
    "energy_gain_range": 0.3,
    "energy_gain_scale": 5.0,
    "movement_cost_velocity": 0.002,
    "movement_cost_base": 0.001,
    # V2 reproduction
    "reproduction_threshold": 75.0,
    "child_energy_fraction": 0.4,
    "child_spawn_distance": 0.1,
    "child_mutation": 0.2,
    "population_cap": 500,
}


# ─── World Model Builder (matches JS _buildWorldModel + _buildInfluenceMatrix) ──

def build_world_model(decision):
    """Convert decision input into concepts, relationships, attractors, influence matrix."""
    factors = decision["factors"]
    relationships = decision["relationships"]
    options = decision["options"]

    factor_names = [f["name"] for f in factors]
    name_to_idx = {name: i for i, name in enumerate(factor_names)}
    n_dims = len(factors)

    # Build relationships with auto-feedback for leaf concepts
    rels = []
    for r in relationships:
        strength = r.get("strength", 0.5)
        direction = -1 if strength < 0 else 1
        rels.append({
            "from": r["from"], "to": r["to"],
            "direction": direction, "strength": abs(strength),
        })

    # Auto-generate weak feedback for leaf concepts
    outgoing = set(r["from"] for r in rels)
    existing_pairs = set(f"{r['from']}|{r['to']}" for r in rels)
    leaves = [fn for fn in factor_names if fn not in outgoing]
    for leaf in leaves:
        incoming = [r for r in rels if r["to"] == leaf]
        for inc in incoming:
            key = f"{leaf}|{inc['from']}"
            if key not in existing_pairs:
                existing_pairs.add(key)
                rels.append({
                    "from": leaf, "to": inc["from"],
                    "direction": inc["direction"],
                    "strength": inc["strength"] * 0.4,
                })

    # Build influence matrix
    influence_matrix = np.zeros((n_dims, n_dims), dtype=np.float64)
    for rel in rels:
        i = name_to_idx.get(rel["from"])
        j = name_to_idx.get(rel["to"])
        if i is not None and j is not None:
            influence_matrix[i, j] = rel["direction"] * rel["strength"]

    # Build concept (attractor) positions normalized to [0, 1]
    concept_positions = np.zeros((len(options), n_dims), dtype=np.float64)
    for a, opt in enumerate(options):
        for d, factor in enumerate(factors):
            val = opt.get(factor["name"])
            lo, hi = factor["range"]
            if val is not None and hi != lo:
                concept_positions[a, d] = max(0, min(1, (val - lo) / (hi - lo)))
            else:
                concept_positions[a, d] = 0.5

    return {
        "factor_names": factor_names,
        "n_dims": n_dims,
        "influence_matrix": influence_matrix,
        "concept_positions": concept_positions,
        "n_concepts": len(options),
    }


def calibrate(config, n_dims, influence_matrix):
    """Lightweight calibration matching JS _calibrate()."""
    dim_factor = np.sqrt(n_dims)

    # Auto gravity range
    if config["gravity_range"] is None:
        config["gravity_range"] = max(2.0, min(6.0, 1.5 * dim_factor))

    # Analyze matrix
    n = n_dims
    pos_count = neg_count = 0
    total_abs = 0.0
    for i in range(n):
        for j in range(i + 1, n):
            v = influence_matrix[i, j]
            if abs(v) > 0.001:
                if v > 0:
                    pos_count += 1
                else:
                    neg_count += 1
                total_abs += abs(v)
    non_zero = pos_count + neg_count
    mean_abs = total_abs / non_zero if non_zero > 0 else 0
    polarity_ratio = neg_count / non_zero if non_zero > 0 else 0

    if polarity_ratio < 0.2 and config["tradeoff_amplifier"] <= 1.5:
        config["tradeoff_amplifier"] = 2.5

    if mean_abs < 0.3 and config["force_scale"] < 0.15:
        config["force_scale"] = 0.15

    if n_dims >= 8:
        config["spawn_spread"] = min(config["spawn_spread"], 0.2)


def initialize_cells(num_cells, n_dims, concept_positions, spawn_spread, rng):
    """Spawn cells near centroid of concept positions."""
    center = concept_positions.mean(axis=0)
    positions = np.zeros((num_cells, n_dims), dtype=np.float64)
    for c in range(num_cells):
        for d in range(n_dims):
            pos = center[d] + (rng.random() - 0.5) * 2 * spawn_spread
            positions[c, d] = max(0.05, min(0.95, pos))
    velocities = np.zeros((num_cells, n_dims), dtype=np.float64)
    return positions, velocities


# ─── Gravity ramp / settling (matches JS) ───────────────────────────────────

def gravity_ramp(tick_count, total_ticks):
    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 settling_damp(tick_count, total_ticks):
    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


# ─── CPU Simulation ─────────────────────────────────────────────────────────

def run_cpu(wm, config, rng):
    """Pure Python V2 colony simulation matching JS SoundBlockSecondOpinion._tickV2."""
    n_dims = wm["n_dims"]
    n_concepts = wm["n_concepts"]
    concept_pos = wm["concept_positions"]
    influence_matrix = wm["influence_matrix"]
    num_cells = config["cells"]
    generations = config["generations"]
    ticks_per_gen = config["ticks_per_generation"]
    dt_norm = config["dt"] / 1000.0
    epsilon = 1e-6

    # Calibrate
    calibrate(config, n_dims, influence_matrix)

    # Initialize
    positions, velocities = initialize_cells(
        num_cells, n_dims, concept_pos, config["spawn_spread"], rng)
    energy = np.full(num_cells, config["initial_energy"], dtype=np.float64)

    # Amplified matrix
    amp_matrix = influence_matrix.copy()
    if config["tradeoff_amplifier"] != 1.0:
        mask = amp_matrix < 0
        amp_matrix[mask] *= config["tradeoff_amplifier"]
        np.fill_diagonal(amp_matrix, 0.0)

    total_ticks = generations * ticks_per_gen
    tick_count = 0

    timings = {"init": 0.0, "sim": 0.0, "metabolism": 0.0, "reproduction": 0.0}

    t_sim_start = time.perf_counter()

    for gen in range(generations):
        for tick in range(ticks_per_gen):
            tick_count += 1
            g_ramp = gravity_ramp(tick_count, total_ticks)
            s_damp = settling_damp(tick_count, total_ticks)
            effective_gravity = config["gravity_strength"] * g_ramp
            effective_momentum = min(config["momentum"], 0.7) * s_damp

            for c in range(num_cells):
                c_off = c  # row index

                # 1. Concept gravity + energy gain
                grav_force = np.zeros(n_dims)
                energy_gain = 0.0

                for k in range(n_concepts):
                    delta = concept_pos[k] - positions[c]
                    dist_sq = np.sum(delta ** 2)
                    dist = np.sqrt(dist_sq)

                    if 0.001 < dist < config["gravity_range"]:
                        mag = (effective_gravity * dist
                               / (config["softening"] ** 2 + dist_sq * config["gravity_sharpness"])
                               * dt_norm)
                        inv_dist = 1.0 / (dist + epsilon)
                        grav_force += delta * inv_dist * mag

                    if dist < config["energy_gain_range"]:
                        energy_gain += config["energy_gain_scale"] / (1 + dist_sq)

                # 2. Matrix forces
                mat_force = np.dot(velocities[c], amp_matrix) * config["force_scale"]

                # 3. Thermal noise
                noise = (rng.random(n_dims) * 2 - 1) * config["thermal_noise"]

                # 4. Boundary forces
                bound_force = np.zeros(n_dims)
                margin = config["boundary_margin"]
                b_str = config["boundary_strength"]
                for d in range(n_dims):
                    pos = positions[c, d]
                    if pos < margin:
                        bound_force[d] = b_str * (margin - pos) / (margin + epsilon)
                    elif pos > 1.0 - margin:
                        bound_force[d] = -b_str * (pos - (1.0 - margin)) / (margin + epsilon)

                # 5. Integration
                total_force = grav_force + mat_force + noise + bound_force
                vel_mag_sq = 0.0
                for d in range(n_dims):
                    vel = velocities[c, d] * effective_momentum + total_force[d] * dt_norm
                    vel = max(-config["max_speed"], min(config["max_speed"], vel))
                    velocities[c, d] = vel

                    pos = positions[c, d] + vel * dt_norm
                    if pos < 0.0:
                        velocities[c, d] = abs(vel) * 0.5
                        pos = 0.0
                    elif pos > 1.0:
                        velocities[c, d] = -abs(vel) * 0.5
                        pos = 1.0
                    positions[c, d] = max(0.0, min(1.0, pos))
                    vel_mag_sq += vel * vel

                # 6. Metabolism
                vel_mag = np.sqrt(vel_mag_sq)
                movement_cost = (config["movement_cost_velocity"] * vel_mag * n_dims
                                 + config["movement_cost_base"] * n_dims)
                energy[c] += energy_gain - movement_cost
                energy[c] = max(config["energy_min"], min(config["energy_max"], energy[c]))

        # Reproduction (end of generation)
        positions, velocities, energy, num_cells = reproduce_cpu(
            positions, velocities, energy, num_cells, n_dims, config, rng)

    sim_time = time.perf_counter() - t_sim_start

    return {
        "positions": positions[:num_cells],
        "velocities": velocities[:num_cells],
        "energy": energy[:num_cells],
        "num_cells": num_cells,
        "sim_time_ms": sim_time * 1000,
    }


def reproduce_cpu(positions, velocities, energy, num_cells, n_dims, config, rng):
    """V2 reproduction: cells with energy > threshold spawn children."""
    children = []
    for c in range(num_cells):
        if energy[c] > config["reproduction_threshold"]:
            child_energy = energy[c] * config["child_energy_fraction"]
            energy[c] -= child_energy

            child_pos = np.clip(
                positions[c] + (rng.random(n_dims) - 0.5) * 2 * config["child_spawn_distance"],
                0.0, 1.0)
            child_vel = np.clip(
                velocities[c] + (rng.random(n_dims) - 0.5) * 2 * config["child_mutation"],
                -config["max_speed"], config["max_speed"])
            children.append((child_pos, child_vel, child_energy))

    if not children:
        return positions, velocities, energy, num_cells

    new_total = num_cells + len(children)
    new_pos = np.zeros((new_total, n_dims), dtype=np.float64)
    new_vel = np.zeros((new_total, n_dims), dtype=np.float64)
    new_energy = np.zeros(new_total, dtype=np.float64)

    new_pos[:num_cells] = positions[:num_cells]
    new_vel[:num_cells] = velocities[:num_cells]
    new_energy[:num_cells] = energy[:num_cells]

    for i, (cp, cv, ce) in enumerate(children):
        new_pos[num_cells + i] = cp
        new_vel[num_cells + i] = cv
        new_energy[num_cells + i] = ce

    final_num = new_total
    if final_num > config["population_cap"]:
        # Keep top by energy
        indices = np.argsort(new_energy[:final_num])[::-1][:config["population_cap"]]
        new_pos = new_pos[indices]
        new_vel = new_vel[indices]
        new_energy = new_energy[indices]
        final_num = config["population_cap"]

    return new_pos, new_vel, new_energy, final_num


# ─── GPU Simulation ──────────────────────────────────────────────────────────

def run_gpu(wm, config, rng):
    """GPU-accelerated V2 colony simulation using SBColonyGPUKernel."""
    from afs.gpu.sb_colony_kernel import SBColonyGPUKernel, compute_gravity_ramp, compute_settling_damp

    n_dims = wm["n_dims"]
    n_concepts = wm["n_concepts"]
    concept_pos = wm["concept_positions"]
    influence_matrix = wm["influence_matrix"]
    num_cells = config["cells"]
    generations = config["generations"]
    ticks_per_gen = config["ticks_per_generation"]
    dt_norm = config["dt"] / 1000.0

    # Calibrate
    calibrate(config, n_dims, influence_matrix)

    # Initialize cells
    positions, velocities = initialize_cells(
        num_cells, n_dims, concept_pos, config["spawn_spread"], rng)
    energy = np.full(num_cells, config["initial_energy"], dtype=np.float64)

    total_ticks = generations * ticks_per_gen

    # ── Timing breakdown ──
    t_init_start = time.perf_counter()

    # Create GPU kernel
    kernel = SBColonyGPUKernel(num_dims=n_dims, max_cells=config["population_cap"],
                                max_concepts=n_concepts)
    t_kernel_init = (time.perf_counter() - t_init_start) * 1000
    print(f"  GPU kernel init:      {t_kernel_init:.1f} ms")
    print(f"  GPU adapter:          {kernel.adapter_info.get('device', 'unknown')}")

    # Upload static data
    t_upload_start = time.perf_counter()
    kernel.upload_matrix(influence_matrix.astype(np.float32))
    kernel.upload_concepts(concept_pos.astype(np.float32))
    kernel.upload_cells(positions.astype(np.float32), velocities.astype(np.float32))
    kernel.upload_energy(energy.astype(np.float32))
    t_static_upload = (time.perf_counter() - t_upload_start) * 1000
    print(f"  Static upload:        {t_static_upload:.1f} ms")

    # ── Generation loop ──
    tick_count = 0
    total_dispatch_ms = 0.0
    total_readback_ms = 0.0
    total_noise_upload_ms = 0.0
    total_metabolism_ms = 0.0
    total_reproduction_ms = 0.0
    total_cell_upload_ms = 0.0

    t_sim_start = time.perf_counter()

    # ── GPU-RESIDENT LOOP: nothing leaves the GPU until fully baked ──
    # Pre-build metabolism params (constant across generations)
    metab_params = {
        "num_cells": num_cells,
        "num_dims": n_dims,
        "num_concepts": n_concepts,
        "absorption_range": config["energy_gain_range"],
        "energy_gain_factor": config["energy_gain_scale"],
        "energy_falloff_exp": 2.0,
        "movement_cost_factor": config["movement_cost_velocity"],
        "base_cost_per_dim": config["movement_cost_base"],
        "min_energy": config["energy_min"],
        "max_energy": config["energy_max"],
        "hot_threshold": 75.0,
        "warm_threshold": 30.0,
    }

    # Pre-build physics params template (only gravity_ramp and settling_damp change)
    base_params = {
        "num_cells": num_cells,
        "num_dims": n_dims,
        "num_concepts": n_concepts,
        "gravity_strength": config["gravity_strength"],
        "gravity_sharpness": config["gravity_sharpness"],
        "gravity_range": config["gravity_range"],
        "softening": config["softening"],
        "force_scale": config["force_scale"],
        "tradeoff_amplifier": config["tradeoff_amplifier"],
        "thermal_noise_level": config["thermal_noise"],
        "momentum": min(config["momentum"], 0.7),
        "max_speed": config["max_speed"],
        "boundary_margin": config["boundary_margin"],
        "boundary_strength": config["boundary_strength"],
        "dt_norm": dt_norm,
    }

    # ── FUSED GPU LOOP: noise + physics + metabolism in ONE submit per gen ──
    d = kernel._device

    for gen in range(generations):
        tick_count += 1
        g_ramp = compute_gravity_ramp(tick_count, total_ticks)
        s_damp = compute_settling_damp(tick_count, total_ticks)

        # Update params that change per generation
        params = dict(base_params)
        params["gravity_ramp"] = g_ramp
        params["settling_damp"] = s_damp

        # Upload all uniforms (tiny — 16 + 80 + 64 = 160 bytes total)
        noise_params = struct.pack("IIIf", num_cells, n_dims, gen * 31337 + 42, 1.0)
        d.queue.write_buffer(kernel._buf_noise_params, 0, noise_params)
        d.queue.write_buffer(kernel._buf_params, 0, kernel._pack_params(params))
        metab_packed = struct.pack("IIIIffffffffffff",
            num_cells, n_dims, n_concepts, 0,
            config["energy_gain_range"], config["energy_gain_scale"], 2.0,
            config["movement_cost_velocity"], config["movement_cost_base"],
            config["energy_min"], config["energy_max"], 75.0, 30.0,
            0.0, 0.0, 0.0)
        d.queue.write_buffer(kernel._buf_metab_params, 0, metab_packed)

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

        # ONE command encoder, THREE dispatches, ONE submit
        t0 = time.perf_counter()
        encoder = d.create_command_encoder()

        # Dispatch 1: noise generation
        cp1 = encoder.begin_compute_pass()
        cp1.set_pipeline(kernel._noise_pipeline)
        cp1.set_bind_group(0, kernel._noise_bg)
        cp1.dispatch_workgroups((N + 255) // 256)
        cp1.end()

        # Dispatch 2: physics tick
        cp2 = encoder.begin_compute_pass()
        cp2.set_pipeline(kernel._pipeline)
        cp2.set_bind_group(0, kernel._bg)
        cp2.dispatch_workgroups((num_cells + 255) // 256)
        cp2.end()

        # Dispatch 3: metabolism
        cp3 = encoder.begin_compute_pass()
        cp3.set_pipeline(kernel._metab_pipeline)
        cp3.set_bind_group(0, kernel._metab_bg)
        cp3.dispatch_workgroups((num_cells + 255) // 256)
        cp3.end()

        d.queue.submit([encoder.finish()])
        total_dispatch_ms += (time.perf_counter() - t0) * 1000

    # ── SINGLE READBACK via Sound Block packed I/O (2.67x less data) ──
    t0 = time.perf_counter()
    positions_f32, velocities_f32 = kernel.readback_packed()
    energy_f32, states = kernel.readback_energy()
    total_readback_ms = (time.perf_counter() - t0) * 1000

    positions = positions_f32.astype(np.float64)
    velocities = velocities_f32.astype(np.float64)
    energy = energy_f32.astype(np.float64)

    total_sim_ms = (time.perf_counter() - t_sim_start) * 1000

    return {
        "positions": positions[:num_cells],
        "velocities": velocities[:num_cells],
        "energy": energy[:num_cells],
        "num_cells": num_cells,
        "sim_time_ms": total_sim_ms,
        "timings": {
            "kernel_init_ms": t_kernel_init,
            "static_upload_ms": t_static_upload,
            "noise_upload_ms": total_noise_upload_ms,
            "dispatch_ms": total_dispatch_ms,
            "readback_ms": total_readback_ms,
            "metabolism_ms": total_metabolism_ms,
            "reproduction_ms": total_reproduction_ms,
            "cell_upload_ms": total_cell_upload_ms,
            "total_sim_ms": total_sim_ms,
        },
        "adapter_info": kernel.adapter_info,
    }


# ─── Cluster Detection (simple density-based, matches JS _detectClusters) ───

def detect_clusters(positions, max_clusters=8, min_cluster_size=None):
    """Simple density-based clustering for strategy extraction."""
    n_cells, n_dims = positions.shape
    if min_cluster_size is None:
        min_cluster_size = max(3, int(np.ceil(n_cells * 0.05)))

    # Simple grid-based: compute pairwise distances, then greedy merge
    # Use a threshold-based approach matching the JS engine
    assigned = np.full(n_cells, -1, dtype=int)
    cluster_id = 0

    # Compute distance threshold based on dimensionality
    # In D dimensions, "close" means within ~0.15 per dimension
    threshold = 0.15 * np.sqrt(n_dims)

    for c in range(n_cells):
        if assigned[c] >= 0:
            continue
        # Start new cluster
        assigned[c] = cluster_id
        members = [c]
        # Find all cells within threshold
        for c2 in range(c + 1, n_cells):
            if assigned[c2] >= 0:
                continue
            dist = np.linalg.norm(positions[c] - positions[c2])
            if dist < threshold:
                assigned[c2] = cluster_id
                members.append(c2)
        cluster_id += 1

    # Build cluster list
    clusters = []
    for cid in range(cluster_id):
        mask = assigned == cid
        count = np.sum(mask)
        if count >= min_cluster_size:
            centroid = positions[mask].mean(axis=0)
            clusters.append({
                "id": cid,
                "centroid": centroid,
                "cell_count": int(count),
            })

    # Sort by cell count descending, limit to max_clusters
    clusters.sort(key=lambda c: c["cell_count"], reverse=True)
    return clusters[:max_clusters]


# ─── Main ────────────────────────────────────────────────────────────────────

def main():
    print("=" * 70)
    print("SB V2 Benchmark — CPU vs GPU Colony Simulation")
    print("=" * 70)
    print()

    wm = build_world_model(SYNTH_DECISION)
    print(f"Decision:       {SYNTH_DECISION['question']}")
    print(f"Dimensions:     {wm['n_dims']}")
    print(f"Concepts:       {wm['n_concepts']} (attractors)")
    print(f"Cells:          {SIM_CONFIG['cells']}")
    print(f"Generations:    {SIM_CONFIG['generations']}")
    print(f"Ticks/gen:      {SIM_CONFIG['ticks_per_generation']}")
    print()

    # Use fixed seed for reproducibility
    seed = 42

    # ── CPU Run ──
    print("-" * 50)
    print("1. CPU (pure Python)")
    print("-" * 50)
    cpu_config = SIM_CONFIG.copy()
    cpu_rng = np.random.default_rng(seed)

    t0 = time.perf_counter()
    cpu_result = run_cpu(wm, cpu_config, cpu_rng)
    cpu_total_ms = (time.perf_counter() - t0) * 1000

    cpu_clusters = detect_clusters(cpu_result["positions"])
    print(f"  Total time:           {cpu_total_ms:.1f} ms")
    print(f"  Sim time:             {cpu_result['sim_time_ms']:.1f} ms")
    print(f"  Final population:     {cpu_result['num_cells']}")
    print(f"  Clusters found:       {len(cpu_clusters)}")
    print(f"  Avg energy:           {cpu_result['energy'].mean():.1f}")
    for i, cl in enumerate(cpu_clusters):
        print(f"    Cluster {i}: {cl['cell_count']} cells, "
              f"centroid [{', '.join(f'{v:.3f}' for v in cl['centroid'][:3])}...]")
    print()

    # ── GPU Run ──
    print("-" * 50)
    print("2. GPU (SBColonyGPUKernel)")
    print("-" * 50)

    try:
        from afs.gpu.sb_colony_kernel import gpu_available as sb_gpu_available
        if not sb_gpu_available():
            print("  GPU not available — skipping")
            gpu_result = None
            gpu_total_ms = None
        else:
            gpu_config = SIM_CONFIG.copy()
            gpu_rng = np.random.default_rng(seed)

            t0 = time.perf_counter()
            gpu_result = run_gpu(wm, gpu_config, gpu_rng)
            gpu_total_ms = (time.perf_counter() - t0) * 1000

            gpu_clusters = detect_clusters(gpu_result["positions"])
            print(f"  Total time:           {gpu_total_ms:.1f} ms")
            print(f"  Sim time:             {gpu_result['sim_time_ms']:.1f} ms")
            print(f"  Final population:     {gpu_result['num_cells']}")
            print(f"  Clusters found:       {len(gpu_clusters)}")
            print(f"  Avg energy:           {gpu_result['energy'].mean():.1f}")
            for i, cl in enumerate(gpu_clusters):
                print(f"    Cluster {i}: {cl['cell_count']} cells, "
                      f"centroid [{', '.join(f'{v:.3f}' for v in cl['centroid'][:3])}...]")
            print()

            # ── Timing Breakdown ──
            print("-" * 50)
            print("GPU Timing Breakdown")
            print("-" * 50)
            t = gpu_result["timings"]
            print(f"  Kernel init:          {t['kernel_init_ms']:.1f} ms")
            print(f"  Static upload:        {t['static_upload_ms']:.1f} ms")
            print(f"  Noise upload (total): {t['noise_upload_ms']:.1f} ms")
            print(f"  Dispatch (total):     {t['dispatch_ms']:.1f} ms")
            print(f"  Readback (total):     {t['readback_ms']:.1f} ms")
            print(f"  Metabolism (CPU):     {t['metabolism_ms']:.1f} ms")
            print(f"  Reproduction (CPU):   {t['reproduction_ms']:.1f} ms")
            print(f"  Cell re-upload:       {t['cell_upload_ms']:.1f} ms")
            print(f"  Total sim:            {t['total_sim_ms']:.1f} ms")
            print()

            # ── Speedup ──
            print("-" * 50)
            print("Comparison")
            print("-" * 50)
            speedup = cpu_result["sim_time_ms"] / gpu_result["sim_time_ms"]
            print(f"  CPU sim time:         {cpu_result['sim_time_ms']:.1f} ms")
            print(f"  GPU sim time:         {gpu_result['sim_time_ms']:.1f} ms")
            print(f"  Speedup:              {speedup:.1f}x")
            print(f"  GPU pure dispatch:    {t['dispatch_ms']:.1f} ms "
                  f"({t['dispatch_ms'] / SIM_CONFIG['generations']:.2f} ms/gen)")
            print()

    except Exception as e:
        print(f"  GPU run failed: {e}")
        import traceback
        traceback.print_exc()
        gpu_result = None
        gpu_total_ms = None

    # ── Write results to markdown ──
    write_results_md(cpu_result, cpu_total_ms, cpu_clusters,
                     gpu_result, gpu_total_ms,
                     gpu_clusters if gpu_result else None,
                     wm)

    print("=" * 70)
    print("Done.")


def write_results_md(cpu_result, cpu_total_ms, cpu_clusters,
                     gpu_result, gpu_total_ms, gpu_clusters, wm):
    """Write benchmark results to markdown file."""
    results_path = Path(__file__).parent.parent.parent.parent / \
        "docs" / "AtomataFieldTheory" / "SoundBlock-V2-GPU-Results.md"

    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    lines = [
        "# Sound Block V2 GPU Benchmark Results",
        "",
        f"**Generated:** {now}",
        "",
        "## Configuration",
        "",
        f"- **Decision:** {SYNTH_DECISION['question']}",
        f"- **Dimensions:** {wm['n_dims']}",
        f"- **Concepts (attractors):** {wm['n_concepts']}",
        f"- **Cells:** {SIM_CONFIG['cells']}",
        f"- **Generations:** {SIM_CONFIG['generations']}",
        f"- **Ticks per generation:** {SIM_CONFIG['ticks_per_generation']}",
        f"- **Physics:** V2 (metabolism + reproduction)",
        "",
        "## CPU Results (Pure Python)",
        "",
        f"- **Total time:** {cpu_total_ms:.1f} ms",
        f"- **Sim time:** {cpu_result['sim_time_ms']:.1f} ms",
        f"- **Final population:** {cpu_result['num_cells']}",
        f"- **Clusters found:** {len(cpu_clusters)}",
        f"- **Avg energy:** {cpu_result['energy'].mean():.1f}",
        "",
        "### Clusters",
        "",
    ]

    for i, cl in enumerate(cpu_clusters):
        centroid_str = ", ".join(f"{v:.3f}" for v in cl["centroid"][:5])
        lines.append(f"- Cluster {i}: {cl['cell_count']} cells — "
                     f"centroid [{centroid_str}...]")

    lines.append("")

    if gpu_result:
        t = gpu_result["timings"]
        speedup = cpu_result["sim_time_ms"] / gpu_result["sim_time_ms"]
        adapter = gpu_result.get("adapter_info", {})

        lines.extend([
            "## GPU Results (SBColonyGPUKernel)",
            "",
            f"- **GPU:** {adapter.get('device', 'unknown')}",
            f"- **Backend:** {adapter.get('backend_type', 'unknown')}",
            f"- **Total time:** {gpu_total_ms:.1f} ms",
            f"- **Sim time:** {gpu_result['sim_time_ms']:.1f} ms",
            f"- **Final population:** {gpu_result['num_cells']}",
            f"- **Clusters found:** {len(gpu_clusters)}",
            f"- **Avg energy:** {gpu_result['energy'].mean():.1f}",
            "",
            "### Timing Breakdown",
            "",
            "| Phase | Time (ms) |",
            "|-------|-----------|",
            f"| Kernel init | {t['kernel_init_ms']:.1f} |",
            f"| Static upload | {t['static_upload_ms']:.1f} |",
            f"| Noise upload (150 gens) | {t['noise_upload_ms']:.1f} |",
            f"| GPU dispatch (150 gens) | {t['dispatch_ms']:.1f} |",
            f"| Readback (150 gens) | {t['readback_ms']:.1f} |",
            f"| Metabolism (CPU) | {t['metabolism_ms']:.1f} |",
            f"| Reproduction (CPU) | {t['reproduction_ms']:.1f} |",
            f"| Cell re-upload | {t['cell_upload_ms']:.1f} |",
            f"| **Total sim** | **{t['total_sim_ms']:.1f}** |",
            "",
            "### Clusters",
            "",
        ])

        for i, cl in enumerate(gpu_clusters):
            centroid_str = ", ".join(f"{v:.3f}" for v in cl["centroid"][:5])
            lines.append(f"- Cluster {i}: {cl['cell_count']} cells — "
                         f"centroid [{centroid_str}...]")

        lines.extend([
            "",
            "## Speedup",
            "",
            f"- **CPU sim:** {cpu_result['sim_time_ms']:.1f} ms",
            f"- **GPU sim:** {gpu_result['sim_time_ms']:.1f} ms",
            f"- **Speedup:** {speedup:.1f}x",
            f"- **GPU dispatch per generation:** "
            f"{t['dispatch_ms'] / SIM_CONFIG['generations']:.2f} ms",
            "",
            "## Notes",
            "",
            "- GPU physics use f32; CPU uses f64. Cluster counts may differ slightly.",
            "- Metabolism and reproduction run on CPU between GPU generations.",
            "- The same `sb_colony.wgsl` shader works in both Python wgpu-py and browser WebGPU.",
            f"- The JS V2 engine takes ~15-17 seconds for the same workload (single-threaded).",
        ])
    else:
        lines.extend([
            "## GPU Results",
            "",
            "GPU not available or run failed.",
        ])

    lines.append("")

    results_path.parent.mkdir(parents=True, exist_ok=True)
    results_path.write_text("\n".join(lines))
    print(f"\nResults written to: {results_path}")


if __name__ == "__main__":
    main()
