digital-drosophila/models/male-cns-lif/runtime-v1/phase4c.py

#!/usr/bin/env python3
"""Phase 4C: freeze reference, audit visual semantics, and optimize exact runtime."""
from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import platform
import shutil
import subprocess
import sys
import time
from collections import Counter
from pathlib import Path
from typing import Any

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import torch
import yaml
from torch.utils.cpp_extension import load as load_extension

import phase4b as p4b

LOG = logging.getLogger("phase4c")


def sha256(path: Path, chunk: int = 8 << 20) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while data := f.read(chunk):
            h.update(data)
    return h.hexdigest()


def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def copy_immutable(source: Path, destination: Path) -> None:
    destination.parent.mkdir(parents=True, exist_ok=True)
    if destination.exists():
        if sha256(source) != sha256(destination):
            raise RuntimeError(f"immutable destination differs: {destination}")
        return
    shutil.copy2(source, destination)
    destination.chmod(0o444)


def exact_paths(cfg: dict[str, Any]) -> dict[str, Path]:
    root = Path(cfg["paths"]["connectome_root"]).resolve()
    allowed = Path("/data/digital-drosophila/datasets/male-cns-v1.0/processed/neuron-only").resolve()
    if root != allowed:
        raise RuntimeError(f"only allowed input is {allowed}")
    paths = {
        "indptr": root / "gpu/csr-indptr.int64.npy",
        "indices": root / "gpu/csr-indices.int32.npy",
        "raw_weights": root / "gpu/csr-weights.int64.npy",
        "metadata": root / "neuron-index/unified-metadata.parquet",
        "mapping": root / "neuron-index/body-id-neuron-idx.parquet",
        "csr_metadata": root / "gpu/csr-metadata.json",
    }
    for path in paths.values():
        if not path.is_file():
            raise FileNotFoundError(path)
    return paths


def freeze_reference(cfg: dict[str, Any], config_path: Path, paths: dict[str, Path]) -> dict[str, Any]:
    model_root = Path(cfg["paths"]["model_root"])
    ref = model_root / "reference-v1"
    ref.mkdir(parents=True, exist_ok=True)
    source_dir = Path(__file__).resolve().parent
    phase4b_root = Path(cfg["paths"]["phase4b_root"])
    source_files = {
        "model.py": source_dir / "phase4b.py",
        "config.yaml": source_dir / "phase4b.yaml",
    }
    for name, source in source_files.items():
        copy_immutable(source, ref / name)
    phase4b_seal = phase4b_root / "manifests/checksums.sha256"
    phase4b_manifest = phase4b_root / "manifests/manifest.json"
    if not phase4b_seal.is_file() or not phase4b_manifest.is_file():
        raise RuntimeError("Phase 4B manifest/seal missing")
    seal_check = subprocess.run(
        ["sha256sum", "-c", "--quiet", str(phase4b_seal)], cwd=phase4b_root,
        text=True, capture_output=True,
    )
    if seal_check.returncode:
        raise RuntimeError(f"Phase 4B seal failed: {seal_check.stderr}")
    frozen = {
        "name": cfg["names"]["reference"],
        "description": cfg["names"]["description"],
        "not_claimed": "strict Shiu 2024 original-model reproduction",
        "reason": "MaleCNS v1.0 transmitter inventory includes histamine and differs from the original dataset/inventory",
        "parameters": cfg["model"],
        "sign_policy": cfg["sign_policy"],
        "graph": cfg["expected_graph"],
        "source_phase": "Phase 4B PASS",
        "phase4b_manifest_sha256": sha256(phase4b_manifest),
        "phase4b_checksums_sha256": sha256(phase4b_seal),
        "phase4b_checksum_verification": "PASS",
        "local_repository_commit": cfg["reproducibility"]["local_repository_commit"],
        "remote_model_directory_is_git_repository": False,
        "code": {name: {"path": str(ref / name), "sha256": sha256(ref / name)} for name in source_files},
        "inputs": {name: {"path": str(path), "sha256": sha256(path)} for name, path in paths.items()},
    }
    version_path = ref / "VERSION.json"
    if version_path.exists():
        previous = json.loads(version_path.read_text(encoding="utf-8"))
        if previous != frozen:
            raise RuntimeError("existing immutable reference VERSION differs")
    else:
        write_json(version_path, frozen)
        version_path.chmod(0o444)
    return frozen


def full_metadata(path: Path) -> pa.Table:
    table = pq.read_table(path)
    required = [
        "neuron_idx", "body_id", "annotation_superclass", "annotation_class",
        "annotation_subclass", "annotation_type", "annotation_instance", "side",
        "neurotransmitter",
    ]
    missing = [x for x in required if x not in table.column_names]
    if missing:
        raise RuntimeError(f"metadata columns missing: {missing}")
    return table


def list_mask(values: list[Any], accepted: set[str]) -> np.ndarray:
    return np.asarray([x in accepted for x in values], dtype=bool)


def populations(table: pa.Table) -> dict[str, np.ndarray]:
    cls = table["annotation_class"].to_pylist()
    sub = table["annotation_subclass"].to_pylist()
    mech = {"mechanosensory", "mechanosensory_tactile", "mechanosensory_proprioceptive"}
    return {
        "visual": np.asarray([x == "visual" for x in cls]),
        "mechanosensory": np.asarray([(x in mech) and y != "auditory" for x, y in zip(cls, sub)]),
        "olfactory": np.asarray([x == "olfactory" for x in cls]),
        "gustatory": np.asarray([x == "gustatory" for x in cls]),
        "auditory": np.asarray([x == "auditory" for x in sub]),
    }


def layer_edge_stats(indptr: np.ndarray, indices: np.ndarray, raw: np.ndarray,
                     layers: dict[str, np.ndarray]) -> dict[str, Any]:
    n = len(indptr) - 1
    label = np.full(n, -1, dtype=np.int16)
    names = list(layers)
    for j, name in enumerate(names):
        label[layers[name]] = j
    counts = np.zeros((len(names), len(names)), dtype=np.int64)
    weights = np.zeros_like(counts)
    for src in range(n):
        a, b = int(indptr[src]), int(indptr[src + 1])
        x = label[src]
        if x < 0 or a == b:
            continue
        y = label[indices[a:b]]
        valid = y >= 0
        if np.any(valid):
            counts[x] += np.bincount(y[valid], minlength=len(names))
            weights[x] += np.bincount(y[valid], weights=raw[a:b][valid], minlength=len(names)).astype(np.int64)
    return {
        "layer_order": names,
        "edge_counts": counts.tolist(),
        "synapse_weight_sums": weights.tolist(),
    }


def visual_audit(exp: Path, table: pa.Table, indptr: np.ndarray,
                 indices: np.ndarray, raw: np.ndarray) -> tuple[dict[str, Any], dict[str, np.ndarray]]:
    out = exp / "visual-audit"
    out.mkdir(parents=True, exist_ok=True)
    pops = populations(table)
    visual = pops["visual"]
    nodes = np.flatnonzero(visual).astype(np.int32)
    fields = [
        "neuron_idx", "body_id", "annotation_superclass", "annotation_class",
        "annotation_subclass", "annotation_type", "annotation_instance", "side", "neurotransmitter",
    ]
    visual_table = table.select(fields).filter(pa.array(visual))
    pq.write_table(visual_table, out / "phase4b-visual-6091.parquet", compression="zstd")
    arrays = {name: table[name].to_pylist() for name in fields}
    breakdown: dict[str, Any] = {}
    for field in fields[2:]:
        c = Counter("<null>" if arrays[field][i] is None else str(arrays[field][i]) for i in nodes)
        breakdown[field] = dict(sorted(c.items(), key=lambda x: (-x[1], x[0])))
    visual_types = sorted({str(arrays["annotation_type"][i]) for i in nodes if arrays["annotation_type"][i] is not None})
    exact_photoreceptor = visual & list_mask(arrays["annotation_superclass"], {"ol_sensory"})

    # These anatomical memberships are derived only from explicit metadata fields.
    super_values = arrays["annotation_superclass"]
    type_values = ["" if x is None else str(x) for x in arrays["annotation_type"]]
    ol_intrinsic = list_mask(super_values, {"ol_intrinsic"})
    visual_projection = list_mask(super_values, {"visual_projection", "visual_projection_tbc"})
    lamina_types = {"L1", "L2", "L3", "L4", "L5", "C2", "C3", "T1", "Lawf1", "Lawf2"}
    lamina = ol_intrinsic & np.asarray([x in lamina_types for x in type_values])
    lobula_plate = ol_intrinsic & np.asarray([
        x.startswith(("T4", "T5", "LPi", "LPLC", "LLPC")) for x in type_values
    ])
    lobula = ol_intrinsic & ~lobula_plate & np.asarray([
        x.startswith(("LC", "LT", "Li", "Lobula")) for x in type_values
    ])
    medulla = ol_intrinsic & ~lamina & ~lobula & ~lobula_plate & np.asarray([
        x.startswith(("Mi", "Dm", "Tm", "TmY", "Pm", "Cm", "Me")) for x in type_values
    ])
    optic_other = ol_intrinsic & ~(lamina | medulla | lobula | lobula_plate)
    anatomy = {
        "photoreceptor": exact_photoreceptor,
        "retinal_sensory": exact_photoreceptor,
        "lamina": lamina,
        "medulla": medulla,
        "lobula": lobula,
        "lobula_plate": lobula_plate,
        "optic_lobe_interneurons": ol_intrinsic,
        "visual_projection_neurons": visual_projection,
        "optic_lobe_interneurons_unassigned_to_named_region": optic_other,
    }
    membership_rows: list[dict[str, Any]] = []
    for role, mask in anatomy.items():
        for i in np.flatnonzero(mask):
            membership_rows.append({"anatomical_role": role, "neuron_idx": int(i), "body_id": int(arrays["body_id"][i])})
    pq.write_table(pa.Table.from_pylist(membership_rows), out / "anatomical-memberships.parquet", compression="zstd")

    # Exact graph hierarchy uses shortest directed distance, not inferred body IDs.
    distance = p4b.directed_distances(indptr, indices, nodes.astype(np.int64), len(visual))
    optic_super = {"ol_intrinsic", "visual_projection", "visual_projection_tbc", "visual_centrifugal"}
    optic = list_mask(super_values, optic_super)
    layers = {
        "visual_sensory": distance == 0,
        "first_order": distance == 1,
        "second_order": distance == 2,
        "higher_visual": (distance >= 3) & optic,
        "central_brain": (distance >= 3) & ~optic,
    }
    for name, mask in layers.items():
        ids = np.flatnonzero(mask).astype(np.int32)
        pq.write_table(pa.table({
            "neuron_idx": pa.array(ids, type=pa.int32()),
            "body_id": pa.array(np.asarray(arrays["body_id"], dtype=np.uint64)[ids], type=pa.uint64()),
            "shortest_visual_path_depth": pa.array(distance[ids], type=pa.int16()),
        }), out / f"hierarchy-{name}.parquet", compression="zstd")
    edge_stats = layer_edge_stats(indptr, indices, raw, layers)
    write_json(out / "hierarchy-edge-statistics.json", edge_stats)
    summary = {
        "phase4b_visual_count": int(visual.sum()),
        "all_histaminergic": all(arrays["neurotransmitter"][i] == "histamine" for i in nodes),
        "all_ol_sensory": all(arrays["annotation_superclass"][i] == "ol_sensory" for i in nodes),
        "raw_visual_types": visual_types,
        "conclusion": "The 6091 annotation_class=visual bodies are photoreceptor/retinal sensory input neurons, not the whole visual processing network.",
        "field_breakdown": breakdown,
        "anatomical_membership_counts": {k: int(v.sum()) for k, v in anatomy.items()},
        "anatomical_rule_provenance": {
            "photoreceptor_retinal_sensory": "annotation_class=visual AND annotation_superclass=ol_sensory; observed types recorded verbatim",
            "optic_lobe_interneurons": "annotation_superclass=ol_intrinsic",
            "visual_projection_neurons": "annotation_superclass in {visual_projection, visual_projection_tbc}",
            "named_optic_regions": "explicit annotation_type prefix/set rules, preserved in this audit; no body_id was guessed",
        },
        "hierarchy_definition": "exclusive shortest directed connectome distance from the 6091 sensory neurons: depth 0 sensory, depth 1 first-order, depth 2 second-order, depth >=3 split into optic/visual-superclass higher_visual versus non-optic downstream central_brain; unreachable bodies are excluded",
        "hierarchy_counts": {k: int(v.sum()) for k, v in layers.items()},
        "edge_statistics_file": str(out / "hierarchy-edge-statistics.json"),
    }
    write_json(out / "summary.json", summary)
    return summary, layers


def generate_forced(steps: int, target_count: int, rate_hz: float, dt: float,
                    seed: int, device: torch.device) -> torch.Tensor:
    probability = rate_hz * dt / 1000.0
    generator = torch.Generator(device=device)
    generator.manual_seed(seed)
    forced = torch.empty((steps, target_count), dtype=torch.uint8, device=device)
    # Deliberately preserve Phase 4B's one rand() call per timestep.
    for step in range(steps):
        forced[step].copy_((torch.rand(target_count, generator=generator, device=device) < probability).to(torch.uint8))
    return forced


def expand_rows(indptr: torch.Tensor, sources: torch.Tensor) -> torch.Tensor:
    if sources.numel() == 0:
        return torch.empty(0, dtype=torch.int64, device=indptr.device)
    starts = indptr[sources]
    counts = indptr[sources + 1] - starts
    cumulative = torch.cumsum(counts, 0)
    total = int(cumulative[-1].item()) if cumulative.numel() else 0
    if total == 0:
        return torch.empty(0, dtype=torch.int64, device=indptr.device)
    seq = torch.arange(total, dtype=torch.int64, device=indptr.device)
    row = torch.searchsorted(cumulative, seq, right=True)
    return starts[row] + seq - (cumulative - counts)[row]


class Runtime:
    def __init__(self, cfg: dict[str, Any], table: pa.Table, indptr_np: np.ndarray,
                 indices_np: np.ndarray, raw_np: np.ndarray, device: torch.device,
                 extension: Any | None):
        self.cfg, self.table, self.device, self.extension = cfg, table, device, extension
        self.n = int(cfg["expected_graph"]["neurons"])
        self.indptr = torch.as_tensor(np.asarray(indptr_np), dtype=torch.int64, device=device)
        self.indices = torch.as_tensor(np.asarray(indices_np), dtype=torch.int32, device=device)
        nt = table["neurotransmitter"].to_pylist()
        policy = cfg["sign_policy"]
        signs = np.empty(self.n, dtype=np.int8)
        for i, value in enumerate(nt):
            key = "" if value is None else str(value).strip().lower()
            if not key:
                signs[i] = int(policy["unknown_sign"])
            elif key == "unclear":
                signs[i] = int(policy["unclear_sign"])
            elif key in policy["neurotransmitter_sign"]:
                signs[i] = int(policy["neurotransmitter_sign"][key])
            else:
                raise RuntimeError(f"unmapped neurotransmitter: {value!r}")
        effective = np.empty(len(raw_np), dtype=np.float32)
        for src in range(self.n):
            a, b = int(indptr_np[src]), int(indptr_np[src + 1])
            effective[a:b] = raw_np[a:b].astype(np.float32) * signs[src] * float(cfg["model"]["w_syn_mv_per_anatomical_synapse"])
        self.weights = torch.as_tensor(effective, dtype=torch.float32, device=device)
        self.signs = signs
        self.superclass = np.asarray(table["annotation_superclass"].to_pylist(), dtype=object)

    def compiled_state(self, mode: str):
        m = self.cfg["model"]
        gd = 1.0 - float(m["dt_ms"]) / float(m["synaptic_tau_ms"])
        alpha = float(m["dt_ms"]) / float(m["membrane_tau_ms"])
        rest, reset, threshold = (float(m[x]) for x in ("resting_potential_mv", "reset_potential_mv", "threshold_mv"))

        def fn(v, g, refractory, arrival):
            g2 = g * gd
            g2 = g2 + arrival
            r2 = torch.clamp_min(refractory, 1) - 1
            available = r2 == 0
            leak = v - rest
            drive = g2 - leak
            integrated = v + alpha * drive
            v2 = torch.where(available, integrated, torch.full_like(v, reset))
            natural = available & (v2 >= threshold)
            return v2, g2, r2, natural, available

        kwargs: dict[str, Any] = {"backend": "inductor", "fullgraph": True}
        if mode == "compile-reduce-overhead":
            kwargs["mode"] = "reduce-overhead"
        return torch.compile(fn, **kwargs), fn

    def run(self, backend: str, targets_np: np.ndarray, rate_hz: float, seed: int,
            duration_ms: float, capture_events: bool, monitor: bool = False,
            forced_override: torch.Tensor | None = None) -> dict[str, Any]:
        m = self.cfg["model"]
        dt = float(m["dt_ms"]); steps = int(round(duration_ms / dt))
        delay_steps = int(round(float(m["synaptic_delay_ms"]) / dt)); qlen = delay_steps + 1
        ref_steps = int(round(float(m["refractory_ms"]) / dt))
        alpha = dt / float(m["membrane_tau_ms"]); gd = 1.0 - dt / float(m["synaptic_tau_ms"])
        rest = float(m["resting_potential_mv"]); reset = float(m["reset_potential_mv"]); threshold = float(m["threshold_mv"])
        targets = torch.as_tensor(targets_np, dtype=torch.int64, device=self.device)
        forced = forced_override if forced_override is not None else generate_forced(steps, len(targets_np), rate_hz, dt, seed, self.device)
        if forced.shape != (steps, len(targets_np)):
            raise RuntimeError("forced input shape mismatch")
        v = torch.full((self.n,), rest, dtype=torch.float32, device=self.device)
        g = torch.zeros(self.n, dtype=torch.float32, device=self.device)
        refractory = torch.zeros(self.n, dtype=torch.int32, device=self.device)
        queue = torch.zeros((qlen, self.n), dtype=torch.float32, device=self.device)
        counts = torch.zeros(self.n, dtype=torch.int32, device=self.device)
        event_steps: list[np.ndarray] = []; event_nodes: list[np.ndarray] = []
        compiled = eager_fn = None
        if backend.startswith("compile-"):
            compiled, eager_fn = self.compiled_state(backend)
            # Exclude compile latency and record a graph-break diagnostic separately.
            tv, tg, tr, ta = v.clone(), g.clone(), refractory.clone(), queue[0].clone()
            compiled(tv, tg, tr, ta); torch.cuda.synchronize()
            del tv, tg, tr, ta
        if backend.startswith("custom") or backend in ("cudagraph-custom-state", "hybrid-exact"):
            if self.extension is None:
                raise RuntimeError("custom extension unavailable")
            spike_u8 = torch.zeros(self.n, dtype=torch.uint8, device=self.device)
        graph = staging = None
        if backend == "cudagraph-custom-state":
            staging = torch.zeros(self.n, dtype=torch.float32, device=self.device)
            graph = torch.cuda.CUDAGraph()
            # Warm-up on disposable state, then capture only the graph-safe fixed-shape state kernel.
            sv, sg = v.clone(), g.clone(); sr = refractory.clone(); ss = spike_u8.clone(); sc = counts.clone(); sa = staging.clone()
            self.extension.state_step(sv, sg, sr, sa, ss, sc, gd, alpha, rest, reset, threshold, ref_steps)
            torch.cuda.synchronize()
            del sv, sg, sr, ss, sc, sa
            with torch.cuda.graph(graph):
                self.extension.state_step(v, g, refractory, staging, spike_u8, counts, gd, alpha, rest, reset, threshold, ref_steps)
        sampler = p4b.GPUSampler(self.cfg["device"]["launcher_visible_gpu_uuid"], float(self.cfg["benchmark"]["gpu_sample_interval_seconds"])) if monitor else None
        if sampler: sampler.start()
        torch.cuda.synchronize(); torch.cuda.reset_peak_memory_stats(); start = time.perf_counter()
        cpu_start = time.process_time()
        for step in range(steps):
            slot = step % qlen
            if backend in ("eager", "hybrid-exact"):
                g.mul_(gd).add_(queue[slot]); queue[slot].zero_()
                refractory.clamp_min_(1).sub_(1)
                available = refractory == 0
                v[available] += alpha * (g[available] - (v[available] - rest))
                v[~available] = reset
                spikes = available & (v >= threshold)
                accepted = forced[step].bool() & available[targets]
                forced_sources = targets[accepted]
                if forced_sources.numel(): spikes[forced_sources] = True
                sources = torch.nonzero(spikes, as_tuple=False).flatten()
                counts.add_(spikes.to(torch.int32))
                if sources.numel():
                    v[sources] = reset; refractory[sources] = ref_steps
                    if backend == "hybrid-exact":
                        starts = self.indptr[sources]
                        degrees = self.indptr[sources + 1] - starts
                        cumulative = torch.cumsum(degrees, 0)
                        total = int(cumulative[-1].item()) if cumulative.numel() else 0
                        edges = torch.empty(total, dtype=torch.int64, device=self.device)
                        if total:
                            self.extension.fill_edges(self.indptr, sources, cumulative - degrees, edges)
                    else:
                        edges = expand_rows(self.indptr, sources)
                    if edges.numel():
                        queue[(step + delay_steps) % qlen].scatter_add_(0, self.indices[edges].long(), self.weights[edges])
            elif backend.startswith("compile-"):
                assert compiled is not None
                arrival = queue[slot].clone(); queue[slot].zero_()
                v, g, refractory, spikes, available = compiled(v, g, refractory, arrival)
                accepted = forced[step].bool() & available[targets]
                forced_sources = targets[accepted]
                if forced_sources.numel(): spikes[forced_sources] = True
                sources = torch.nonzero(spikes, as_tuple=False).flatten()
                counts.add_(spikes.to(torch.int32))
                if sources.numel():
                    v[sources] = reset; refractory[sources] = ref_steps
                    edges = expand_rows(self.indptr, sources)
                    if edges.numel(): queue[(step + delay_steps) % qlen].scatter_add_(0, self.indices[edges].long(), self.weights[edges])
            else:
                assert self.extension is not None
                if backend == "cudagraph-custom-state":
                    assert staging is not None and graph is not None
                    staging.copy_(queue[slot]); queue[slot].zero_(); graph.replay()
                else:
                    self.extension.state_step(v, g, refractory, queue[slot], spike_u8, counts, gd, alpha, rest, reset, threshold, ref_steps)
                self.extension.force_step(v, refractory, spike_u8, counts, targets, forced[step], reset, ref_steps)
                self.extension.fanout(spike_u8, self.indptr, self.indices, self.weights, queue[(step + delay_steps) % qlen])
                sources = torch.nonzero(spike_u8, as_tuple=False).flatten() if capture_events else None
            if capture_events:
                if not (backend.startswith("custom") or backend == "cudagraph-custom-state"):
                    pass
                assert sources is not None
                src = sources.detach().cpu().numpy().astype(np.int32, copy=True)
                event_nodes.append(src); event_steps.append(np.full(len(src), step, dtype=np.int32))
        torch.cuda.synchronize(); wall = time.perf_counter() - start; cpu = time.process_time() - cpu_start
        resources = sampler.stop() if sampler else {}
        v_np = v.cpu().numpy().copy(); g_np = g.cpu().numpy().copy(); counts_np = counts.cpu().numpy().copy()
        es = np.concatenate(event_steps) if event_steps else np.empty(0, dtype=np.int32)
        en = np.concatenate(event_nodes) if event_nodes else np.empty(0, dtype=np.int32)
        active = counts_np > 0
        super_values = self.superclass
        desc = np.isin(super_values, ["descending_neuron", "descending_neuron_tbc"])
        motor = np.isin(super_values, ["vnc_motor", "cb_motor"])
        bio_s = duration_ms / 1000.0
        result = {
            "backend": backend, "duration_ms": duration_ms, "steps": steps,
            "wall_clock_seconds": wall, "steps_per_second": steps / wall,
            "biological_realtime_speed": bio_s / wall, "realtime_factor_wall_over_biological": wall / bio_s,
            "total_spikes": int(counts_np.sum(dtype=np.int64)), "active_neurons": int(active.sum()),
            "mean_firing_rate_hz": float(counts_np.sum(dtype=np.int64) / (self.n * bio_s)),
            "descending_active_neurons": int(np.count_nonzero(active & desc)),
            "descending_spikes": int(counts_np[desc].sum(dtype=np.int64)),
            "motor_active_neurons": int(np.count_nonzero(active & motor)),
            "motor_spikes": int(counts_np[motor].sum(dtype=np.int64)),
            "peak_vram_bytes": int(torch.cuda.max_memory_allocated()),
            "cpu_process_utilization_percent_of_one_core": 100.0 * cpu / wall,
            "resource_sampling": resources,
            "event_count": int(len(en)),
            "event_sha256": hashlib.sha256(es.tobytes() + en.tobytes()).hexdigest() if capture_events else None,
            "spike_counts_sha256": hashlib.sha256(counts_np.tobytes()).hexdigest(),
            "final_voltage_sha256": hashlib.sha256(v_np.tobytes()).hexdigest(),
            "final_g_sha256": hashlib.sha256(g_np.tobytes()).hexdigest(),
        }
        return {"metrics": result, "event_steps": es, "event_nodes": en, "counts": counts_np, "v": v_np, "g": g_np}


def compare(reference: dict[str, Any], candidate: dict[str, Any], tolerance: float) -> dict[str, Any]:
    events = np.array_equal(reference["event_steps"], candidate["event_steps"]) and np.array_equal(reference["event_nodes"], candidate["event_nodes"])
    counts = np.array_equal(reference["counts"], candidate["counts"])
    vdiff = np.abs(reference["v"] - candidate["v"])
    gdiff = np.abs(reference["g"] - candidate["g"])
    return {
        "spike_events_exact": bool(events), "per_neuron_spike_counts_exact": bool(counts),
        "final_voltage_max_abs_error": float(vdiff.max(initial=0)),
        "final_voltage_mean_abs_error": float(vdiff.mean()),
        "final_g_max_abs_error": float(gdiff.max(initial=0)), "final_g_mean_abs_error": float(gdiff.mean()),
        "state_within_tolerance": bool(vdiff.max(initial=0) <= tolerance and gdiff.max(initial=0) <= tolerance),
        "eligible_for_runtime_v1": bool(events and counts),
    }


def save_reference_trial(exp: Path, spec: dict[str, Any], run: dict[str, Any]) -> None:
    base = exp / "regression/reference" / spec["name"]
    base.mkdir(parents=True, exist_ok=True)
    pq.write_table(pa.table({
        "step": pa.array(run["event_steps"], type=pa.int32()),
        "timestamp_ms": pa.array(run["event_steps"].astype(np.float64) * 0.1, type=pa.float64()),
        "neuron_idx": pa.array(run["event_nodes"], type=pa.int32()),
    }), base / "spike-events.parquet", compression="zstd")
    np.save(base / "per-neuron-spike-counts.int32.npy", run["counts"], allow_pickle=False)
    np.save(base / "final-voltage.float32.npy", run["v"], allow_pickle=False)
    np.save(base / "final-g.float32.npy", run["g"], allow_pickle=False)
    write_json(base / "statistics.json", run["metrics"] | {"trial": spec})


def propagation_stats(run: dict[str, Any], targets: np.ndarray, distance: np.ndarray,
                      superclass: np.ndarray) -> dict[str, Any]:
    active = run["counts"] > 0
    nonnegative = distance[active & (distance >= 0)]
    return {
        "active_neurons": int(active.sum()),
        "downstream_active_neurons": int(np.count_nonzero(active & ~np.isin(np.arange(len(active)), targets))),
        "propagation_depth": int(nonnegative.max(initial=0)),
        "descending_active_neurons": int(np.count_nonzero(active & np.isin(superclass, ["descending_neuron", "descending_neuron_tbc"]))),
        "motor_active_neurons": int(np.count_nonzero(active & np.isin(superclass, ["vnc_motor", "cb_motor"]))),
    }


def visual_dynamics(exp: Path, cfg: dict[str, Any], runtime: Runtime,
                    layers: dict[str, np.ndarray], indptr: np.ndarray, indices: np.ndarray) -> dict[str, Any]:
    spec = cfg["visual_dynamics"]; duration = float(spec["duration_ms"]); rate = float(spec["rate_hz"]); seed = int(spec["seed"])
    results: dict[str, Any] = {}
    tests: list[tuple[str, np.ndarray, float, int, torch.Tensor | None, str]] = []
    sensory = np.flatnonzero(layers["visual_sensory"]).astype(np.int64)
    tests.append(("A-forced-histaminergic-sensory", sensory, rate, seed, None, "engineering forcing of sensory output spikes"))
    steps = int(round(duration / float(cfg["model"]["dt_ms"])))
    tonic = generate_forced(steps, len(sensory), float(spec["tonic_histamine_baseline_hz"]), float(cfg["model"]["dt_ms"]), seed + 1, runtime.device)
    mod_start = int(round(float(spec["modulation_start_ms"]) / float(cfg["model"]["dt_ms"])))
    mod = generate_forced(steps - mod_start, len(sensory), float(spec["modulation_rate_hz"]), float(cfg["model"]["dt_ms"]), seed + 2, runtime.device)
    tonic[mod_start:] = mod
    tests.append(("B-tonic-histamine-rate-modulation", sensory, 0.0, seed + 1, tonic, "exploratory tonic output plus rate modulation; not baseline"))
    first = np.flatnonzero(layers["first_order"]).astype(np.int64)
    second = np.flatnonzero(layers["second_order"]).astype(np.int64)
    tests.append(("C-first-order-external-input", first, rate, seed + 3, None, "engineering input interface moved to real first-order targets"))
    tests.append(("C-second-order-external-input", second, rate, seed + 4, None, "engineering input interface moved to real second-order targets"))
    for j, name in enumerate(("visual_sensory", "first_order", "second_order", "higher_visual", "central_brain")):
        targets = np.flatnonzero(layers[name]).astype(np.int64)
        # Keep the audit bounded without inventing body IDs: deterministic full layer up to 25k, otherwise fixed index stride.
        if len(targets) > 25000:
            targets = targets[np.linspace(0, len(targets) - 1, 25000, dtype=np.int64)]
        tests.append((f"D-layer-{name}", targets, rate, seed + 10 + j, None, f"engineering external input at hierarchy layer {name}"))
    for name, targets, hz, trial_seed, override, interpretation in tests:
        LOG.info("visual dynamics %s targets=%d", name, len(targets))
        run = runtime.run("eager", targets, hz, trial_seed, duration, False, forced_override=override)
        distance = p4b.directed_distances(indptr, indices, targets, runtime.n)
        record = run["metrics"] | propagation_stats(run, targets, distance, runtime.superclass) | {
            "test": name, "stimulated_population_size": int(len(targets)), "interpretation": interpretation,
            "biological_interpretation": "histamine remains inhibitory; forced sensory spikes are an engineering interface, not phototransduction/receptor-potential biology",
            "baseline_status": "exploratory only" if name.startswith("B-") else "semantic test; no automatic selection",
        }
        results[name] = record
        write_json(exp / "visual-audit/dynamics" / f"{name}.json", record)
    write_json(exp / "visual-audit/dynamics-summary.json", results)
    return results


def compiler_report(exp: Path, runtime: Runtime) -> dict[str, Any]:
    report: dict[str, Any] = {"torch_version": torch.__version__, "modes": {}}
    n = runtime.n
    m = runtime.cfg["model"]
    for mode in ("compile-default", "compile-reduce-overhead"):
        try:
            compiled, fn = runtime.compiled_state(mode)
            v = torch.full((n,), float(m["resting_potential_mv"]), device=runtime.device)
            g = torch.zeros(n, device=runtime.device); r = torch.zeros(n, dtype=torch.int32, device=runtime.device); a = torch.zeros(n, device=runtime.device)
            compiled(v, g, r, a); torch.cuda.synchronize()
            report["modes"][mode] = {"status": "compiled", "fullgraph": True, "graph_breaks_in_fused_static_state": 0}
        except Exception as exc:
            report["modes"][mode] = {"status": "failed", "error": repr(exc)}
    report["fullgraph_dynamic_csr"] = {
        "status": "not eligible",
        "graph_break_reasons": [
            "torch.nonzero produces data-dependent dynamic shape",
            "Python branching on sources.numel() and edge_ids.numel()",
            "expand_rows requires cumulative[-1].item() host synchronization",
            "dynamic torch.arange allocation and event-dependent CSR fan-out",
        ],
        "action": "kept eager reference; custom CUDA prototype keeps spike detection/fan-out/accumulation on GPU",
    }
    write_json(exp / "compiler-reports/torch-compile.json", report)
    return report


def build_extension(source_dir: Path, exp: Path) -> tuple[Any | None, dict[str, Any]]:
    cpp, cu = source_dir / "phase4c_ext.cpp", source_dir / "phase4c_ext.cu"
    started = time.perf_counter()
    try:
        ext = load_extension(
            name="male_cns_lif_runtime_v1_ext", sources=[str(cpp), str(cu)],
            extra_cflags=["-O3"], extra_cuda_cflags=["-O3"], verbose=True,
            build_directory=str(exp / "compiler-reports/extension-build"),
        )
        result = {"status": "PASS", "build_seconds": time.perf_counter() - started, "sources": {str(cpp): sha256(cpp), str(cu): sha256(cu)}}
        return ext, result
    except Exception as exc:
        result = {"status": "FAILED", "build_seconds": time.perf_counter() - started, "error": repr(exc), "sources": {str(cpp): sha256(cpp), str(cu): sha256(cu)}}
        return None, result


def seal_outputs(exp: Path, cfg: dict[str, Any], paths: dict[str, Path], code_paths: list[Path], summary: dict[str, Any]) -> None:
    controls = {exp / "manifests/manifest.json", exp / "manifests/checksums.sha256"}
    outputs = sorted(x for x in exp.rglob("*") if x.is_file() and x not in controls)
    manifest = {
        "phase": "4C", "created_unix_time": time.time(), "status": summary["status"],
        "software": {"python": platform.python_version(), "torch": torch.__version__, "numpy": np.__version__, "pyarrow": pa.__version__},
        "code": [{"path": str(x), "bytes": x.stat().st_size, "sha256": sha256(x)} for x in code_paths],
        "inputs": {name: {"path": str(path), "bytes": path.stat().st_size, "sha256": sha256(path)} for name, path in paths.items()},
        "outputs": [{"path": str(x), "bytes": x.stat().st_size, "sha256": sha256(x)} for x in outputs],
        "prohibitions_respected": cfg["prohibitions"],
    }
    write_json(exp / "manifests/manifest.json", manifest)
    targets = outputs + [exp / "manifests/manifest.json"]
    (exp / "manifests/checksums.sha256").write_text("".join(f"{sha256(x)}  {x.relative_to(exp)}\n" for x in targets), encoding="utf-8")


def main() -> None:
    ap = argparse.ArgumentParser(); ap.add_argument("--config", required=True, type=Path); args = ap.parse_args()
    cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
    exp = Path(cfg["paths"]["experiment_root"])
    for sub in ("configs", "visual-audit", "benchmarks", "compiler-reports", "profiler", "regression", "manifests", "audit"):
        (exp / sub).mkdir(parents=True, exist_ok=True)
    p4b.setup_logging(exp / "audit/phase4c.log")
    shutil.copy2(args.config, exp / "configs/phase4c.yaml")
    paths = exact_paths(cfg); hashes_before = {k: sha256(v) for k, v in paths.items()}
    indptr, indices, raw = p4b.validate_graph(cfg, paths)
    table = full_metadata(paths["metadata"]); pops = populations(table)
    expected_pop = {"visual": 6091, "auditory": 114}
    for name, count in expected_pop.items():
        if int(pops[name].sum()) != count: raise RuntimeError(f"{name} count mismatch")
    frozen = freeze_reference(cfg, args.config, paths)
    visual_summary, layers = visual_audit(exp, table, indptr, indices, raw)
    if not torch.cuda.is_available() or torch.cuda.device_count() != 1: raise RuntimeError("exactly one visible CUDA GPU required")
    device = torch.device("cuda:0"); props = torch.cuda.get_device_properties(device)
    if cfg["device"]["required_name_substring"] not in props.name: raise RuntimeError(f"wrong GPU: {props.name}")
    source_dir = Path(__file__).resolve().parent
    extension, extension_report = build_extension(source_dir, exp)
    write_json(exp / "compiler-reports/custom-extension.json", extension_report)
    runtime = Runtime(cfg, table, indptr, indices, raw, device, extension)
    compiler = compiler_report(exp, runtime)

    reference_runs: dict[str, dict[str, Any]] = {}
    regression: dict[str, Any] = {}
    tol = float(cfg["regression"]["float_absolute_tolerance"])
    for spec in cfg["reference_trials"]:
        targets = np.flatnonzero(pops[spec["population"]]).astype(np.int64)
        LOG.info("reference trial %s", spec["name"])
        ref = runtime.run("eager", targets, float(spec["rate_hz"]), int(spec["seed"]), float(spec["duration_ms"]), True)
        dist = p4b.directed_distances(indptr, indices, targets, runtime.n)
        ref["metrics"].update(propagation_stats(ref, targets, dist, runtime.superclass))
        save_reference_trial(exp, spec, ref); reference_runs[spec["name"]] = ref
        candidates: dict[str, Any] = {}
        for backend in (["hybrid-exact", "custom-cuda"] if extension is not None else []):
            cand = runtime.run(backend, targets, float(spec["rate_hz"]), int(spec["seed"]), float(spec["duration_ms"]), True)
            comp = compare(ref, cand, tol)
            comp.update({
                "propagation_statistics_exact": propagation_stats(ref, targets, dist, runtime.superclass) == propagation_stats(cand, targets, dist, runtime.superclass),
                "reference_event_sha256": ref["metrics"]["event_sha256"], "candidate_event_sha256": cand["metrics"]["event_sha256"],
            })
            candidates[backend] = comp
        regression[spec["name"]] = candidates
        write_json(exp / "regression" / f"{spec['name']}.json", candidates)
    visual_results = visual_dynamics(exp, cfg, runtime, layers, indptr, indices)

    bench: dict[str, Any] = {}
    auditory = np.flatnonzero(pops["auditory"]).astype(np.int64)
    baseline_by_duration: dict[float, dict[str, Any]] = {}
    for duration in cfg["benchmark"]["durations_ms"]:
        duration = float(duration); bench[str(duration)] = {}
        for backend in cfg["benchmark"]["backends"]:
            if (backend.startswith("custom") or backend in ("cudagraph-custom-state", "hybrid-exact")) and extension is None:
                bench[str(duration)][backend] = {"status": "unavailable", "reason": extension_report.get("error")}; continue
            LOG.info("benchmark backend=%s duration_ms=%s", backend, duration)
            try:
                run = runtime.run(backend, auditory, float(cfg["benchmark"]["rate_hz"]), int(cfg["benchmark"]["seed"]), duration, False, monitor=True)
                record = run["metrics"] | {"status": "PASS"}
                if backend == "eager": baseline_by_duration[duration] = run
                else:
                    ref = baseline_by_duration.get(duration)
                    if ref is not None:
                        state_cmp = compare({**ref, "event_steps": np.empty(0,np.int32), "event_nodes": np.empty(0,np.int32)}, {**run, "event_steps": np.empty(0,np.int32), "event_nodes": np.empty(0,np.int32)}, tol)
                        # Performance runs deliberately disable event capture.  Do not
                        # mislabel two empty arrays as an event-sequence regression.
                        state_cmp["spike_events_exact"] = None
                        state_cmp["event_sequence_checked"] = False
                        record["benchmark_state_regression"] = state_cmp
                bench[str(duration)][backend] = record
            except Exception as exc:
                bench[str(duration)][backend] = {"status": "FAILED", "error": repr(exc)}
    write_json(exp / "benchmarks/runtime-matrix.json", bench)

    eligible = ["eager"]
    if extension is not None and all(regression[x].get("hybrid-exact", {}).get("eligible_for_runtime_v1") for x in regression):
        eligible.append("hybrid-exact")
    if extension is not None and all(regression[x].get("custom-cuda", {}).get("eligible_for_runtime_v1") for x in regression):
        eligible.append("custom-cuda")
    # Compiled/CUDA-graph paths need an exact 1 s event regression before eligibility.
    audit_spec = cfg["reference_trials"][0]; audit_ref = reference_runs[audit_spec["name"]]
    targets = np.flatnonzero(pops[audit_spec["population"]]).astype(np.int64)
    backend_regression: dict[str, Any] = {}
    for backend in ("compile-default", "compile-reduce-overhead", "cudagraph-custom-state"):
        if backend == "cudagraph-custom-state" and extension is None: continue
        try:
            candidate = runtime.run(backend, targets, float(audit_spec["rate_hz"]), int(audit_spec["seed"]), 1000.0, True)
            backend_regression[backend] = compare(audit_ref, candidate, tol)
            if backend_regression[backend]["eligible_for_runtime_v1"]: eligible.append(backend)
        except Exception as exc:
            backend_regression[backend] = {"eligible_for_runtime_v1": False, "error": repr(exc)}
    write_json(exp / "regression/backend-eligibility.json", backend_regression)

    def speed(backend: str) -> float:
        x = bench.get("10000.0", {}).get(backend, {})
        return float(x.get("biological_realtime_speed", 0.0)) if x.get("status") == "PASS" else 0.0
    fastest = max(eligible, key=speed)
    runtime_root = Path(cfg["paths"]["model_root"]) / "runtime-v1"
    runtime_root.mkdir(parents=True, exist_ok=True)
    runtime_sources = (Path(__file__).resolve(), source_dir / "phase4c_ext.cpp", source_dir / "phase4c_ext.cu", source_dir / "requirements.txt", args.config.resolve())
    for source in runtime_sources:
        copy_immutable(source, runtime_root / source.name)
    runtime_version = {
        "name": cfg["names"]["runtime"], "backend": fastest,
        "reference": cfg["names"]["reference"], "eligibility": backend_regression | regression,
        "ten_second_biological_realtime_speed": speed(fastest),
        "spike_behavior_requirement": "exact on all five 1 s frozen reference trials for custom runtime; exact auditory trial for compiler variants",
        "algorithm": "exact sparse event-driven CSR fan-out; no dense full-edge scan; exact 18-step GPU-resident ring buffer",
        "fallback_backend": "eager",
    }
    write_json(runtime_root / "VERSION.json", runtime_version); (runtime_root / "VERSION.json").chmod(0o444)
    hashes_after = {k: sha256(v) for k, v in paths.items()}
    if hashes_before != hashes_after: raise RuntimeError("Phase 3B inputs changed")
    summary = {
        "status": "PASS" if fastest else "FAILED", "reference_frozen": True,
        "reference": frozen, "visual_audit": visual_summary, "visual_dynamics": visual_results,
        "compiler_report": compiler, "custom_extension": extension_report,
        "regression": regression, "backend_regression": backend_regression,
        "runtime_v1_backend": fastest, "benchmarks": bench,
        "device": {"name": props.name, "total_memory_bytes": int(props.total_memory)},
        "input_hashes_unchanged": True, "entered_phase4d": False, "game_rl_web_3d_training": False,
    }
    write_json(exp / "benchmarks/summary.json", summary)
    seal_outputs(exp, cfg, paths, list(runtime_sources), summary)
    LOG.info("PHASE4C_COMPLETE runtime=%s", fastest)


if __name__ == "__main__":
    main()