Build your own fruit fly brain simulation: a tested beginner guide from an empty folder

"""FlyGym 2.1.0 CPU smoke test: build NeuroMechFly on flat ground, step physics, no rendering.
Construction follows FlyGym's own tests/conftest.py at tag v2.1.0."""
import time, resource
t0 = time.time()
import flygym, mujoco
from flygym.anatomy import AxisOrder, JointPreset, ActuatedDOFPreset, Skeleton
from flygym.compose.fly import NeuroMechFly, ActuatorType
from flygym.compose.world import FlatGroundWorld
from flygym.compose.pose import KinematicPosePreset
from flygym.utils.math import Rotation3D
from flygym.simulation import Simulation
print(f"mujoco {mujoco.__version__}; import {time.time()-t0:.1f} s")
pose = KinematicPosePreset.NEUTRAL.get_pose_by_axis_order(AxisOrder.YAW_PITCH_ROLL)
skel = Skeleton(axis_order=AxisOrder.YAW_PITCH_ROLL, joint_preset=JointPreset.LEGS_ONLY)
fly = NeuroMechFly(name="fly")
fly.add_joints(skel, neutral_pose=pose)
fly.add_actuators(skel.get_actuated_dofs_from_preset(ActuatedDOFPreset.LEGS_ACTIVE_ONLY),
                  ActuatorType.POSITION, neutral_input=pose, kp=50)
fly.add_leg_adhesion(gain=1.0)
world = FlatGroundWorld()
world.add_fly(fly, spawn_position=[0, 0, 1.5], spawn_rotation=Rotation3D("quat", [1, 0, 0, 0]))
t1 = time.time()
sim = Simulation(world)
sim.reset()
m = sim.mj_model
print(f"compiled: nq={m.nq} nu={m.nu} nbody={m.nbody} timestep={m.opt.timestep} s in {time.time()-t1:.1f} s")
z0 = float(sim.mj_data.qpos[2]); n = 2000
t2 = time.time()
for _ in range(n):
    sim.step()   # no control written: position actuators hold the neutral pose
dt = time.time() - t2
print(f"{n} steps = {sim.time:.3f} s simulated in {dt:.2f} s wall ({sim.time/dt:.2f}x real time)")
print(f"thorax z {z0:.3f} -> {float(sim.mj_data.qpos[2]):.3f} mm; finite state: {bool(abs(sim.mj_data.qpos).max() < 1e3)}; "
      f"peak RSS {resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024:.0f} MB; total {time.time()-t0:.1f} s")

Run it with MUJOCO_GL=disable in a Python 3.12 environment with flygym==2.1.0. Our output: 2,000 steps = 0.200 s simulated in 0.41 s wall (0.49× real time); the thorax settled from 2.8 to 1.07 mm; peak memory 226 MB.

"""Taste -> action with the Shiu et al. whole-brain LIF model (FlyWire v783).

What this file does
-------------------
1. Loads the FlyWire v783 connectivity published in philshiu/Drosophila_brain_model
   (measured wiring: who connects to whom, how many synapses, excitatory/inhibitory sign).
2. Builds the same leaky integrate-and-fire (LIF) network as upstream model.py
   (modelled dynamics: equations and parameters are imported from upstream, not copied).
3. Drives chosen sugar-sensing gustatory receptor neurons (GRNs) with Poisson spikes
   (hand-made stimulus: which neurons and what rate are our choice, following the paper).
4. Reads out the firing rate of MN9, the motor neuron for proboscis extension
   (hand-made readout choice, following the paper).
5. Maps the rate to an action with a threshold (hand-made rule, not part of the model).

Differences from upstream run_exp (see README "What's real here"):
- Only the three columns the model uses are loaded, as int32/float arrays, to save memory.
- The connectivity file is streamed in chunks (same synapse order) to keep peak memory under 1 GB.
- Each trial runs in a fresh worker process with an explicit random seed (repeatable runs).
- Brian2 runs with the pure-numpy code generation target when no C++ compiler exists.
"""
from __future__ import annotations

import multiprocessing
import os
import shutil
import sys
import time
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

pa.set_memory_pool(pa.system_memory_pool())  # return freed memory to the OS promptly

HERE = Path(__file__).resolve().parent
UPSTREAM = Path(os.environ.get("FLY_UPSTREAM", HERE / "upstream"))
sys.path.insert(0, str(UPSTREAM))

import brian2 as b2  # noqa: E402
from brian2 import Hz, ms, mV  # noqa: E402
from model import default_params  # noqa: E402  (upstream model.py, pinned commit)

COMPLETENESS = UPSTREAM / "Completeness_783.csv"
CONNECTIVITY = UPSTREAM / "Connectivity_783.parquet"

# Right-hemisphere labellar sugar GRNs from upstream example.ipynb (FlyWire v630 IDs).
# One ID changed between v630 and v783; the upstream author gave the new ID in
# https://github.com/philshiu/Drosophila_brain_model/issues/7
SUGAR_GRNS_V783 = [
    720575940624963786, 720575940630233916, 720575940637568838, 720575940638202345,
    720575940617000768, 720575940630797113, 720575940632889389, 720575940621754367,
    720575940621502051, 720575940640649691, 720575940639332736, 720575940616885538,
    720575940639198653, 720575940639259967,  # was 720575940620900446 in v630
    720575940617937543, 720575940632425919, 720575940633143833, 720575940612670570,
    720575940628853239, 720575940629176663, 720575940611875570,
]
MN9 = 720575940660219265  # proboscis motor neuron 9, ID from upstream example.ipynb (same in v783)

# Hand-made action rule. Not part of the model. 20 Hz is our choice (see README).
PROBOSCIS_THRESHOLD_HZ = 20.0


def pick_codegen_target() -> str:
    """Cython (fast) if a C++ compiler is on PATH, else pure numpy (slower, needs no compiler).

    Override with the environment variable FLY_CODEGEN=numpy or FLY_CODEGEN=cython.
    """
    forced = os.environ.get("FLY_CODEGEN")
    if forced:
        return forced
    if any(shutil.which(c) for c in (os.environ.get("CXX", ""), "g++", "clang++", "c++") if c):
        # keep Brian2's compiled-code cache inside the project (works on read-only homes too)
        b2.prefs.codegen.runtime.cython.cache_dir = str(HERE / ".cython_cache")
        return "cython"
    return "numpy"


def load_ids() -> np.ndarray:
    """FlyWire root IDs in model order (row i of the CSV = Brian2 neuron index i)."""
    ids = np.loadtxt(COMPLETENESS, delimiter=",", skiprows=1, usecols=0, dtype=np.int64)
    return ids


def _read_column(name, dtype):
    """Read one parquet column row group by row group into a compact numpy array."""
    f = pq.ParquetFile(CONNECTIVITY, memory_map=False)
    out = np.empty(f.metadata.num_rows, dtype=dtype)
    k = 0
    for g in range(f.num_row_groups):  # streaming keeps peak memory low
        col = f.read_row_group(g, columns=[name]).column(0).to_numpy()
        out[k:k + col.size] = col
        k += col.size
        del col
    assert k == out.size
    return out


def load_edges(shuffle_seed: int | None = None):
    """Return (pre, post) neuron indices of every connection, as int32.

    shuffle_seed: if given, build the degree-preserving control by randomly permuting
    the postsynaptic column in place. Every neuron keeps its number of outgoing and
    incoming connections, and every presynaptic neuron keeps its own weights and sign
    (so excitatory neurons stay excitatory), but who-connects-to-whom is randomised.
    """
    pre = _read_column("Presynaptic_Index", np.int32)
    post = _read_column("Postsynaptic_Index", np.int32)
    if shuffle_seed is not None:
        np.random.default_rng(shuffle_seed).shuffle(post)
    return pre, post


def load_signed_counts():
    """Upstream's 'Excitatory x Connectivity' column: synapse count x (+1 excitatory,
    -1 inhibitory). The sign comes from the predicted neurotransmitter of the presynaptic
    neuron (acetylcholine = +1; GABA and glutamate = -1)."""
    return _read_column("Excitatory x Connectivity", np.float64)


def load_connectivity(shuffle_seed: int | None = None):
    pre, post = load_edges(shuffle_seed)
    return pre, post, load_signed_counts()


def build_network(params, exc, slnc, shuffle_seed=None):
    """Build the upstream LIF network (same as create_model() + poi() + silence() in
    upstream model.py), streaming the connectivity file in chunks to keep memory low."""
    n = len(load_ids())
    neu = b2.NeuronGroup(
        N=n, model=params["eqs"], method="linear",
        threshold=params["eq_th"], reset=params["eq_rst"], refractory="rfc",
        name="default_neurons", namespace=params,
    )
    neu.v = params["v_0"]
    neu.g = 0 * mV
    neu.rfc = params["t_rfc"]
    syn = b2.Synapses(neu, neu, "w : volt", on_pre="g += w",
                      delay=params["t_dly"], name="default_synapses")
    f = pq.ParquetFile(CONNECTIVITY, memory_map=False)
    post_all = load_edges(shuffle_seed)[1] if shuffle_seed is not None else None
    k = 0
    for g in range(f.num_row_groups):  # connect chunk by chunk (same order as upstream)
        t = f.read_row_group(g, columns=["Presynaptic_Index", "Postsynaptic_Index"])
        pre = t.column(0).to_numpy().astype(np.int32)
        post = (post_all[k:k + pre.size] if post_all is not None
                else t.column(1).to_numpy().astype(np.int32))
        del t
        syn.connect(i=pre, j=post)
        k += pre.size
    del post_all
    w_syn = float(params["w_syn"])  # 0.275 mV, in volts
    k = 0
    for g in range(f.num_row_groups):  # weight = signed synapse count x w_syn
        c = f.read_row_group(g, columns=["Excitatory x Connectivity"]).column(0).to_numpy()
        syn.w_[k:k + c.size] = c * w_syn
        k += c.size
    for i in slnc:  # same as upstream silence(): zero all synapses *from* neuron i
        syn.w[f"i == {i}"] = 0 * mV
    pois = []
    for i in exc:   # same as upstream poi()
        pois.append(b2.PoissonInput(target=neu[i], target_var="v", N=1,
                                    rate=params["r_poi"],
                                    weight=params["w_syn"] * params["f_poi"]))
        neu[i].rfc = 0 * ms
    mon = b2.SpikeMonitor(neu, record=False)  # spike counts only, saves memory
    return b2.Network(neu, syn, mon, *pois), mon, len(syn)


def _one_trial(params_r_poi_hz, exc, slnc, shuffle_seed, rand_seed, t_run_ms):
    """Build the network, run one trial, return (spike counts, n_synapses, build s, run s).
    Runs in a fresh worker process so all memory is returned after each trial
    (upstream does the same via joblib)."""
    params = dict(default_params)  # upstream constants and equations
    params["r_poi"] = params_r_poi_hz * Hz
    b2.prefs.codegen.target = pick_codegen_target()
    b2.start_scope()
    t0 = time.time()
    net, mon, n_syn = build_network(params, exc, slnc, shuffle_seed)
    t1 = time.time()
    b2.seed(rand_seed)
    net.run(t_run_ms * ms)
    return np.asarray(mon.count[:], dtype=np.int32), n_syn, t1 - t0, time.time() - t1


def run_condition(stim_ids, rate_hz=150.0, silence_ids=(), n_trials=3, t_run_ms=1000.0,
                  seed=0, shuffle_seed=None, verbose=True):
    """Run one experimental condition; return per-trial spike counts for all neurons.

    Like upstream run_trial(), the network is rebuilt for every trial (a few seconds),
    one trial at a time (peak memory about 0.6 GB). Trial k uses Brian2 random seed
    (seed * 1000 + k) for the Poisson input, so runs are repeatable.
    """
    ids = load_ids()
    id2i = {int(f): i for i, f in enumerate(ids)}
    exc = [id2i[int(f)] for f in stim_ids]          # KeyError here = ID not in v783
    slnc = [id2i[int(f)] for f in silence_ids]

    counts = np.zeros((n_trials, len(ids)), dtype=np.int32)
    build_s = sim_s = 0.0
    ctx = multiprocessing.get_context("spawn")
    for trial in range(n_trials):
        with ProcessPoolExecutor(max_workers=1, mp_context=ctx) as ex:
            c, n_syn, tb, tr = ex.submit(_one_trial, rate_hz, exc, slnc, shuffle_seed,
                                         seed * 1000 + trial, t_run_ms).result()
        counts[trial] = c
        build_s += tb
        sim_s += tr
        if verbose:
            print(f"  trial {trial + 1}/{n_trials}: {len(ids)} neurons, {n_syn} synapses, "
                  f"{int(c.sum())} spikes; build {tb:.1f} s, run {tr:.1f} s "
                  f"(target={pick_codegen_target()})", flush=True)
    return {
        "ids": ids, "counts": counts, "t_run_s": t_run_ms / 1000.0,
        "build_s": build_s, "sim_s": sim_s, "target": pick_codegen_target(),
        "stim_rate_hz": rate_hz, "n_trials": n_trials, "seed": seed,
        "shuffle_seed": shuffle_seed, "silenced": [int(s) for s in silence_ids],
    }


def rate_of(result, flywire_id):
    """Mean and s.d. firing rate (Hz) across trials for one neuron."""
    i = int(np.nonzero(result["ids"] == flywire_id)[0][0])
    r = result["counts"][:, i] / result["t_run_s"]
    return float(r.mean()), float(r.std())


def proboscis_action(mn9_rate_hz, threshold=PROBOSCIS_THRESHOLD_HZ):
    """Hand-made action mapping: extend the proboscis if MN9 fires above threshold."""
    return "EXTEND proboscis" if mn9_rate_hz > threshold else "stay still"


def top_active(result, n=10, exclude=()):
    mean = result["counts"].mean(axis=0) / result["t_run_s"]
    order = np.argsort(-mean)
    out = []
    ex = set(int(e) for e in exclude)
    for i in order:
        if mean[i] <= 0 or len(out) >= n:
            break
        if int(result["ids"][i]) in ex:
            continue
        out.append((int(result["ids"][i]), float(mean[i])))
    return out

"""Run the taste -> action experiments and write results/results.json + results/index.html.

Usage:
    python run_demo.py                      # full run: 5 trials x 1 s per condition
    python run_demo.py --trials 1 --ms 200  # quick look (numbers are noisier)
"""
from __future__ import annotations

import argparse
import html
import json
import platform
import time
from pathlib import Path

import numpy as np

import fly_taste as ft

OUT = Path(__file__).resolve().parent / "results"


def summarise(name, res, note):
    mn9_mean, mn9_sd = ft.rate_of(res, ft.MN9)
    grn = np.mean([ft.rate_of(res, g)[0] for g in ft.SUGAR_GRNS_V783])
    active = int((res["counts"].sum(axis=0) > 0).sum())
    return {
        "condition": name,
        "note": note,
        "stimulus": {"neurons": "21 right labellar sugar GRNs" if res["stim_rate_hz"] else "none",
                     "rate_hz": res["stim_rate_hz"]},
        "silenced": res["silenced"],
        "shuffle_seed": res["shuffle_seed"],
        "trials": res["n_trials"], "trial_s": res["t_run_s"], "seed": res["seed"],
        "mn9_rate_hz_mean": round(mn9_mean, 2), "mn9_rate_hz_sd": round(mn9_sd, 2),
        "mn9_rate_hz_per_trial": [float(c) / res["t_run_s"] for c in
                                  res["counts"][:, int(np.nonzero(res["ids"] == ft.MN9)[0][0])]],
        "sugar_grn_mean_rate_hz": round(float(grn), 2),
        "neurons_active": active,
        "total_spikes_per_trial": [int(x) for x in res["counts"].sum(axis=1)],
        "action": ft.proboscis_action(mn9_mean),
        "top_active_non_grn": [{"flywire_id": str(i), "rate_hz": round(r, 1)}
                               for i, r in ft.top_active(res, 10, exclude=ft.SUGAR_GRNS_V783)],
        "codegen_target": res["target"],
        "build_s": round(res["build_s"], 1), "sim_s": round(res["sim_s"], 1),
    }


def strongest_active_input_to_mn9(res):
    """Our (hand-made) choice of neuron to silence: the MN9 input partner with the largest
    (synapse count x firing rate) in the sugar run, excluding the stimulated GRNs."""
    pre, post, wsc = ft.load_connectivity()
    i_mn9 = int(np.nonzero(res["ids"] == ft.MN9)[0][0])
    sel = post == i_mn9
    p, w = pre[sel], wsc[sel]
    rate = res["counts"].mean(axis=0) / res["t_run_s"]
    grn_idx = {int(np.nonzero(res["ids"] == g)[0][0]) for g in ft.SUGAR_GRNS_V783}
    best, best_score = None, 0.0
    for k, wk in zip(p, w):
        score = wk * rate[k]
        if int(k) not in grn_idx and score > best_score:
            best, best_score = int(k), score
    return int(res["ids"][best]), float(best_score), int(sum(w[p == best]))


def write_html(data, path):
    rows = "".join(
        f"<tr><td>{html.escape(c['condition'])}</td><td>{c['stimulus']['rate_hz']:g} Hz</td>"
        f"<td>{c['trials']} x {c['trial_s']:g} s</td>"
        f"<td><b>{c['mn9_rate_hz_mean']:.1f}</b> &plusmn; {c['mn9_rate_hz_sd']:.1f}</td>"
        f"<td>{html.escape(c['action'])}</td><td>{c['neurons_active']}</td>"
        f"<td>{html.escape(c['note'])}</td></tr>"
        for c in data["conditions"])
    path.write_text(f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Taste to action: fly brain model results</title>
<style>body{{font-family:system-ui,sans-serif;max-width:60rem;margin:2rem auto;padding:0 1rem;line-height:1.5}}
table{{border-collapse:collapse;width:100%}}td,th{{border:1px solid #bbb;padding:.3rem .5rem;text-align:left;vertical-align:top}}</style>
</head><body>
<h1>Taste to action: sugar neurons &rarr; MN9 (proboscis) in a whole-brain fly model</h1>
<p>Model: leaky integrate-and-fire network from
<a href="https://github.com/philshiu/Drosophila_brain_model">philshiu/Drosophila_brain_model</a>
(Shiu et al., <i>Nature</i> 2024), FlyWire connectome v783 ({data['n_neurons']} neurons,
{data['n_synapse_rows']} connections). Run on {html.escape(data['machine'])}, {html.escape(data['date'])}.</p>
<table><tr><th>Condition</th><th>Sugar GRN input</th><th>Trials</th><th>MN9 rate (Hz, mean &plusmn; s.d.)</th>
<th>Action (rule: MN9 &gt; {ft.PROBOSCIS_THRESHOLD_HZ:g} Hz)</th><th>Neurons active</th><th>Note</th></tr>{rows}</table>
<h2>What's real here</h2>
<ul><li><b>Measured:</b> the wiring (FlyWire v783 synapse counts between neurons).</li>
<li><b>Modelled:</b> spiking dynamics (leaky integrate-and-fire), weights proportional to synapse count,
excitatory/inhibitory sign inferred from predicted neurotransmitter. No neuromodulation, no plasticity, no body.</li>
<li><b>Hand-made:</b> which neurons we stimulate and at what rate, reading out MN9, and the threshold rule
that turns an MN9 rate into "extend proboscis".</li>
<li><b>Scripted:</b> this page. Nothing here moves a real or simulated body.</li></ul>
<h2>Credits and licences</h2>
<p>Model: Shiu et al., "A Drosophila computational brain model reveals sensorimotor processing", <i>Nature</i> 634,
210&ndash;219 (2024); code <a href="https://github.com/philshiu/Drosophila_brain_model">philshiu/Drosophila_brain_model</a> (MIT).
Connectome: FlyWire public release v783 (Dorkenwald et al., <i>Nature</i> 634, 124&ndash;138, 2024; Schlegel et al.,
<i>Nature</i> 634, 139&ndash;152, 2024), licensed <b>CC BY-NC 4.0: non-commercial use only, attribution required</b>
(<a href="https://flywire.ai/guidelines">flywire.ai/guidelines</a>). These numbers are model output, not recordings from a real fly.</p>
</body></html>
""")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--trials", type=int, default=5)
    ap.add_argument("--ms", type=float, default=1000.0)
    ap.add_argument("--shuffle-seeds", type=int, nargs="*", default=[1, 2])
    ap.add_argument("--shuffle-trials", type=int, default=3)
    ap.add_argument("--skip", nargs="*", default=[], help="condition names to skip")
    a = ap.parse_args()
    OUT.mkdir(exist_ok=True)
    ids = ft.load_ids()
    data = {"model": "philshiu/Drosophila_brain_model @ 91bdd1e7 (LIF, Brian2)",
            "data_licence": "FlyWire public release data: CC BY-NC 4.0 (non-commercial, attribution required)",
            "dataset": "FlyWire FAFB v783 (Completeness_783.csv, Connectivity_783.parquet)",
            "n_neurons": int(len(ids)), "n_synapse_rows": 15091983,
            "machine": f"{platform.system()} {platform.machine()}, Python {platform.python_version()}",
            "date": time.strftime("%Y-%m-%d"), "conditions": []}

    def save():
        (OUT / "results.json").write_text(json.dumps(data, indent=1))

    def run(name, note, **kw):
        if name in a.skip:
            return None
        print(f"== {name}", flush=True)
        res = ft.run_condition(**kw)
        s = summarise(name, res, note)
        print(f"   MN9 {s['mn9_rate_hz_mean']} +/- {s['mn9_rate_hz_sd']} Hz -> {s['action']}", flush=True)
        data["conditions"].append(s)
        save()
        return res

    sugar = run("sugar_150Hz", "upstream default input rate (r_poi = 150 Hz)",
                stim_ids=ft.SUGAR_GRNS_V783, rate_hz=150.0, n_trials=a.trials, t_run_ms=a.ms)
    run("no_stimulus", "baseline: nothing stimulated", stim_ids=[], rate_hz=0.0,
        n_trials=1, t_run_ms=a.ms)
    run("sugar_50Hz", "customization 1: weaker sugar input",
        stim_ids=ft.SUGAR_GRNS_V783, rate_hz=50.0, n_trials=a.trials, t_run_ms=a.ms)
    if sugar is not None and "silence" not in a.skip:
        target, score, nsyn = strongest_active_input_to_mn9(sugar)
        data["silence_choice"] = {"flywire_id": str(target), "synapses_onto_mn9": nsyn,
                                  "rule": "largest synapse-count x rate input to MN9 in sugar_150Hz, excluding GRNs"}
        run("silence", f"customization 2: silence {target} (strongest active input to MN9)",
            stim_ids=ft.SUGAR_GRNS_V783, rate_hz=150.0, silence_ids=[target],
            n_trials=a.trials, t_run_ms=a.ms)
    for s in a.shuffle_seeds:
        run(f"shuffled_wiring_seed{s}", "control: postsynaptic targets permuted (degree-preserving)",
            stim_ids=ft.SUGAR_GRNS_V783, rate_hz=150.0, shuffle_seed=s,
            n_trials=a.shuffle_trials, t_run_ms=a.ms)
    save()
    write_html(data, OUT / "index.html")
    print(f"wrote {OUT / 'results.json'} and {OUT / 'index.html'}")


if __name__ == "__main__":
    main()

"""Minimal checks that the taste -> action project works.

Short runs (1 trial x 200 ms) so the whole file finishes in a few minutes on a 2-CPU laptop
with the pure-numpy Brian2 target. Run with:  python -m pytest -q
"""
import hashlib

import numpy as np
import pytest

import fly_taste as ft

T_MS = 200.0
EXPECTED_SHA256 = {
    "Completeness_783.csv": "bbb847a4cc2caaa7a16349722d220c087317b946d148d4d592d94d250617a311",
    "Connectivity_783.parquet": "efeb23fb99098e9c390f6869969b2a121a2ee92c833cfc45ecb2c1d8e1af0347",
}


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1 << 20), b""):
            h.update(block)
    return h.hexdigest()


@pytest.mark.parametrize("name", sorted(EXPECTED_SHA256))
def test_data_files_match_pinned_checksums(name):
    assert sha256(ft.UPSTREAM / name) == EXPECTED_SHA256[name]


def test_all_stimulus_and_readout_ids_exist_in_v783():
    ids = set(int(i) for i in ft.load_ids())
    assert len(ids) == 138639
    missing = [f for f in ft.SUGAR_GRNS_V783 + [ft.MN9] if f not in ids]
    assert missing == []


@pytest.fixture(scope="module")
def sugar():
    return ft.run_condition(ft.SUGAR_GRNS_V783, rate_hz=150.0, n_trials=1, t_run_ms=T_MS, seed=0)


def test_no_stimulus_gives_silence():
    res = ft.run_condition([], rate_hz=0.0, n_trials=1, t_run_ms=T_MS)
    assert int(res["counts"].sum()) == 0          # the model has no spontaneous activity
    assert ft.proboscis_action(ft.rate_of(res, ft.MN9)[0]) == "stay still"


def test_sugar_drives_mn9_and_proboscis(sugar):
    mn9, _ = ft.rate_of(sugar, ft.MN9)
    assert mn9 > ft.PROBOSCIS_THRESHOLD_HZ
    assert ft.proboscis_action(mn9) == "EXTEND proboscis"


def test_shuffled_wiring_changes_the_readout(sugar):
    ctrl = ft.run_condition(ft.SUGAR_GRNS_V783, rate_hz=150.0, n_trials=1, t_run_ms=T_MS,
                            seed=0, shuffle_seed=1)
    mn9_real, _ = ft.rate_of(sugar, ft.MN9)
    mn9_ctrl, _ = ft.rate_of(ctrl, ft.MN9)
    # same stimulus, same degrees and weights, different wiring -> MN9 no longer driven
    assert mn9_ctrl < ft.PROBOSCIS_THRESHOLD_HZ < mn9_real
    # and the set of active neurons is different
    act_real = set(np.nonzero(sugar["counts"].sum(axis=0))[0])
    act_ctrl = set(np.nonzero(ctrl["counts"].sum(axis=0))[0])
    assert act_real != act_ctrl

# pytest.ini
[pytest]
pythonpath = .
testpaths = tests
filterwarnings =
    ignore::DeprecationWarning

# .gitignore
.venv/
upstream/
__pycache__/
.pytest_cache/
.cython_cache/

MIT License

Copyright (c) 2026 Digital Fly Lab contributors

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

This licence covers only the files in this folder written for this project
(fly_taste.py, run_demo.py, tests/). It does not cover the upstream model code
(MIT, Copyright (c) 2023 Philip Shiu and Nico Spiller) or the FlyWire connectome
data and anything derived from it, including results/ (CC BY-NC 4.0,
non-commercial use only). See README.md.

"""OUR rate-to-drive mapping (hand-made by Digital Fly Lab, run 6, 2026-10-02; pre-registered 2026-10-02T09:35:33Z).

This is NOT part of the Shiu et al. brain model and NOT part of FlyGym. It turns descending-neuron (DN) firing rates
from a brain trial into FlyGym's two-number descending drive (one number per side, as HybridTurningController expects).
Nothing here was fitted or tuned to results. Pure numpy; no FlyGym import. Licence: MIT.

For each 100 ms bin k and side s (o = the other side), rates in Hz, R = 100 Hz fixed in advance:
    p_s = rate(DNp09_s) / R
    m_s = mean rate(MDN on side s) / R
    a_s = rate(DNa02_s) / R
    drive_s(k) = clip(p_s - m_s - 0.5 * (a_s - a_o), -1.2, 1.2)
Each bin's drive is held for 100 ms of body time.

DEVIATION 1 (logged in ../preregistration.md before the first body run): the basis check (../basis.md) found that
unilateral P9 drives IPSILATERAL turning (Bidaye et al. 2020: "P9 drives forward walking with ipsilateral turning").
With p_s on its own side, P9_L alone gives drive (p, 0): the left legs step, the right legs stand, and the fly turns
right (contralateral). The fix used for every body run:
    drive_s(k) = clip(pbar - 0.5 * (p_s - p_o) - m_s - 0.5 * (a_s - a_o), -1.2, 1.2),  pbar = (p_L + p_R) / 2
i.e. forward = mean P9 rate, plus an ipsilateral P9 turning term with the same 0.5 gain as DNa02. Algebraically
pbar - 0.5 * (p_s - p_o) = p_o, so each side's forward term is the OTHER side's P9 rate. Under bilateral P9 with
equal rates nothing changes. `drive_from_rates_prereg` keeps the pre-registered formula for comparison.
"""
from __future__ import annotations

import numpy as np

R_HZ = 100.0
CLIP = 1.2
BIN_MS = 100
FORMULA = ("OUR MAPPING (hand-made, not fitted): drive_s = clip(pbar - 0.5*(p_s - p_o) - m_s - 0.5*(a_s - a_o), "
           "-1.2, 1.2) with p = rate(DNp09)/R, pbar = mean of both sides' p, m = mean rate(MDN)/R, a = rate(DNa02)/R, "
           "R = 100 Hz, per 100 ms bin, held 100 ms; s = side (left, right), o = the other side. Pre-registered "
           "version: forward term p_s; changed to the ipsilateral-turning P9 term before the first body run "
           "(deviation 1, after the basis check)")

# FlyWire v783 root IDs (our run 6 candidate list; checked against the
# annotation table in run 6)
DN_IDS = {
    "DNp09": {"left": [720575940635872101], "right": [720575940627652358]},
    "MDN": {"left": [720575940631082808, 720575940616026939], "right": [720575940610236514, 720575940640331472]},
    "DNa02": {"left": [720575940629327659], "right": [720575940604737708]},
    # recorded and reported, NOT mapped:
    "DNa01": {"left": [720575940627787609], "right": [720575940644438551]},
    "DNp01": {"left": [720575940622838154], "right": [720575940632499757]},
    "DNg62": {"left": [720575940624319124], "right": [720575940616185531]},
}
MAPPED = ("DNp09", "MDN", "DNa02")


def drive_from_rates_prereg(p, m, a):
    """The pre-registered formula (forward term p_s). Kept for comparison; not used for the body runs."""
    p, m, a = (np.asarray(x, dtype=float) / R_HZ for x in (p, m, a))
    turn = a - a[:, ::-1]  # a_s - a_o
    return np.clip(p - m - 0.5 * turn, -CLIP, CLIP)


def drive_from_rates(p, m, a):
    """p, m, a: arrays of shape (n_bins, 2) in Hz, columns (left, right). Returns drive (n_bins, 2).
    Deviation 1: the P9 term turns ipsilaterally (pbar - 0.5*(p_s - p_o) = p_o)."""
    p, m, a = (np.asarray(x, dtype=float) / R_HZ for x in (p, m, a))
    pbar = p.mean(axis=1, keepdims=True)
    p_turn = p - p[:, ::-1]  # p_s - p_o
    a_turn = a - a[:, ::-1]  # a_s - a_o
    return np.clip(pbar - 0.5 * p_turn - m - 0.5 * a_turn, -CLIP, CLIP)


def rates_from_trial(trial: dict, n_bins: int = 10):
    """Per-bin rates (Hz) of the six candidate cell types per side from a run 6 brain trial file (field 'dn')."""
    bins = {int(x["id"]): x["bins"] for x in trial["dn"]["active"]}
    scale = 1000.0 / trial["dn"]["bin_ms"]
    out = {}
    for ct, sides in DN_IDS.items():
        out[ct] = {}
        for side, ids in sides.items():
            arr = np.array([bins.get(i, [0] * n_bins) for i in ids], dtype=float) * scale
            out[ct][side] = arr.mean(axis=0)  # mean over the neurons of this type on this side
    return out


def drive_from_trial(trial: dict):
    r = rates_from_trial(trial)
    col = lambda ct: np.stack([r[ct]["left"], r[ct]["right"]], axis=1)
    return drive_from_rates(col("DNp09"), col("MDN"), col("DNa02")), r

import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import mapping as mp  # noqa: E402

Z = np.zeros((10, 2))


def test_zero_rates_give_zero_drive():
    assert np.array_equal(mp.drive_from_rates(Z, Z, Z), np.zeros((10, 2)))


def test_p9_only_gives_equal_positive_drives():
    d = mp.drive_from_rates(np.full((10, 2), 50.0), Z, Z)
    assert np.allclose(d, 0.5) and np.all(d[:, 0] == d[:, 1])


def test_mdn_only_gives_negative_drives():
    d = mp.drive_from_rates(Z, np.full((10, 2), 40.0), Z)
    assert np.all(d < 0) and np.allclose(d, -0.4)


def test_dna02_left_slows_left():
    a = np.zeros((10, 2)); a[:, 0] = 60.0
    d = mp.drive_from_rates(np.full((10, 2), 80.0), Z, a)
    assert np.all(d[:, 0] < d[:, 1])
    assert np.allclose(d[:, 0], 0.8 - 0.3) and np.allclose(d[:, 1], 0.8 + 0.3)


def test_clipping():
    assert np.allclose(mp.drive_from_rates(np.full((10, 2), 500.0), Z, Z), 1.2)
    assert np.allclose(mp.drive_from_rates(Z, np.full((10, 2), 500.0), Z), -1.2)


def test_rates_from_trial_bins():
    tr = {"dn": {"bin_ms": 100, "active": [{"id": 720575940635872101, "bins": [10] * 10},
                                             {"id": 720575940631082808, "bins": [4] * 10}]}}
    d, r = mp.drive_from_trial(tr)
    assert np.allclose(r["DNp09"]["left"], 100.0) and np.allclose(r["DNp09"]["right"], 0.0)
    assert np.allclose(r["MDN"]["left"], 20.0)  # mean of the two left MDNs (40 Hz and 0 Hz)
    # deviation 1: P9_L alone drives the RIGHT side (the fly turns left, ipsilateral); MDN_L slows the left side
    assert np.allclose(d[:, 0], 0.0 - 0.2) and np.allclose(d[:, 1], 1.0)


def test_unilateral_p9_turns_ipsilateral():
    p = np.zeros((10, 2)); p[:, 0] = 80.0
    d = mp.drive_from_rates(p, Z, Z)
    assert np.all(d[:, 0] < d[:, 1])  # left P9 -> left side slower -> left (ipsilateral) turn
    dp = mp.drive_from_rates_prereg(p, Z, Z)
    assert np.all(dp[:, 0] > dp[:, 1])  # the pre-registered formula turned the other way


def test_fix_changes_nothing_under_equal_bilateral_p9():
    p = np.full((10, 2), 70.0); a = np.zeros((10, 2)); a[:, 1] = 30.0
    assert np.allclose(mp.drive_from_rates(p, Z, a), mp.drive_from_rates_prereg(p, Z, a))

"""Turn run 6 brain trials into two-number body drives with OUR mapping (mapping.py), plus the brainless drives.
Digital Fly Lab, run 6 (2026-10-02). Licence: MIT. numpy only.

Usage: python make_drives.py BRAIN_DIR DRIVES_DIR
Writes one DRIVES_DIR/<condition>.json per body run: {id, group, drive [10 x 2], drive_prereg, dn_rates_hz, source, note}.
Conditions (pre-registered): P9-real-0..2, P9-D3..5, SUG-real-0 (from brain trials); zero, const-1, matched, random-0..2
(brainless); fill: SUG-D3..5, P9L-real-0, P9L-D3 when their brain trials exist.
"""
import json
import os
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent))
import mapping as mp  # noqa: E402

BRAIN = {"P9-real-0": "P9-real-0", "P9-real-1": "P9-real-1", "P9-real-2": "P9-real-2", "P9-D3": "P9-D-3",
         "P9-D4": "P9-D-4", "P9-D5": "P9-D-5", "SUG-real-0": "SUG-real-0", "SUG-D3": "SUG-D-3", "SUG-D4": "SUG-D-4",
         "SUG-D5": "SUG-D-5", "P9L-real-0": "P9L-real-0", "P9L-D3": "P9L-D-3"}
def group_of(cid):
    stim, arm = cid.split("-")[0], cid.split("-")[1]
    if stim == "SUG":
        return "sugar" if arm == "real" else "sugar-scrambled"
    return "brain-real" if arm == "real" else "brain-scrambled"


def write(out_dir, rec):
    p = Path(out_dir) / f"{rec['id']}.json"
    tmp = p.with_suffix(".tmp")
    tmp.write_text(json.dumps(rec, indent=1))
    os.replace(tmp, p)


def main(brain_dir, out_dir):
    Path(out_dir).mkdir(parents=True, exist_ok=True)
    real = []
    for cid, fname in BRAIN.items():
        f = Path(brain_dir) / f"{fname}.json"
        if not f.exists():
            continue
        tr = json.loads(f.read_text())
        if tr.get("status") != "ok":
            continue
        drive, rates = mp.drive_from_trial(tr)
        col = lambda ct: np.stack([rates[ct]["left"], rates[ct]["right"]], axis=1)
        prereg = mp.drive_from_rates_prereg(col("DNp09"), col("MDN"), col("DNa02"))
        grp = group_of(cid)
        write(out_dir, {"id": cid, "group": grp, "drive": drive.round(6).tolist(), "drive_prereg": prereg.round(6).tolist(),
                        "dn_rates_hz": {ct: {s: v.round(3).tolist() for s, v in d.items()} for ct, d in rates.items()},
                        "source": f"body/results/brain/{f.name}", "mapping": mp.FORMULA})
        if cid.startswith("P9-real-"):
            real.append(drive)
    z = np.zeros((10, 2))
    write(out_dir, {"id": "zero", "group": "floor", "drive": z.tolist(), "source": "none (no command)"})
    write(out_dir, {"id": "const-1", "group": "sanity", "drive": np.ones((10, 2)).tolist(),
                    "source": "none (FlyGym's default straight walk; sanity check, not an arm)"})
    if len(real) >= 1:
        R = np.stack(real)  # runs x bins x 2
        mean_s = R.mean(axis=(0, 1))
        sd_s = R.std(axis=(0, 1), ddof=1) if R.shape[0] * R.shape[1] > 1 else np.zeros(2)
        src = ", ".join(f"body/results/brain/{BRAIN[k]}.json" for k in ("P9-real-0", "P9-real-1", "P9-real-2"))
        write(out_dir, {"id": "matched", "group": "no-brain", "drive": np.tile(mean_s, (10, 1)).round(6).tolist(),
                        "per_side_mean": mean_s.round(6).tolist(), "n_real_runs": len(real),
                        "source": f"mean over all bins of the real P9 drives ({src})"})
        for seed in range(3):
            rng = np.random.default_rng(seed)
            d = np.clip(rng.normal(mean_s, sd_s, size=(10, 2)), -mp.CLIP, mp.CLIP)
            write(out_dir, {"id": f"random-{seed}", "group": "no-brain", "drive": d.round(6).tolist(), "seed": seed,
                            "per_side_mean": mean_s.round(6).tolist(), "per_side_sd": sd_s.round(6).tolist(),
                            "source": f"numpy default_rng({seed}).normal(per-side mean, per-side sd of the real P9 drive bins), clipped to +-1.2; {src}"})
    print("drives:", sorted(p.stem for p in Path(out_dir).glob("*.json")))


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])

"""Run the FlyGym 2.1.0 body with a two-number descending drive (Digital Fly Lab, run 6, 2026-10-02). Licence: MIT.

The walking (legs, rhythm, balance, adhesion) is FlyGym's own HybridTurningController (flygym_demo.complex_terrain,
NeuroMechFly v2's two-number descending interface). This script only feeds it a drive file made by make_drives.py
(OUR rate-to-drive mapping) and records where the thorax goes. Construction follows upstream
tests/examples/test_complex_terrain_locomotion.py at v2.1.0.

Usage
  MUJOCO_GL=disable python body_run.py one DRIVE.json OUT.json [--settle 0.2] [--driven 1.0] [--smoke]
  MUJOCO_GL=disable python body_run.py queue DRIVES_DIR OUT_DIR --order a,b,c --start-cut ISO --kill ISO --log LOG
Timeline: settle s at drive (0, 0), then driven s in 10 equal bins (bin k of the drive held for driven/10 s).
Thorax x, y (mm), z and heading are saved every 10 ms. Metric: forward displacement along the heading at the start of
the driven period (mm); also heading change (deg, + = left/counter-clockwise) and path length (mm).
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import resource
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

GIB = 1024 ** 3
CG = Path("/sys/fs/cgroup/memory.current")


def now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def mem():
    try:
        return int(CG.read_text())
    except OSError:
        return None


def thorax_state(sim, mujoco, bid):
    pos = sim.mj_data.xpos[bid].copy()
    xm = sim.mj_data.xmat[bid].reshape(3, 3)
    fwd = xm[:, 0]  # body x axis in world frame
    import math
    return pos, math.atan2(fwd[1], fwd[0])


def one(drive_path, out_path, settle=0.2, driven=1.0, smoke=False):
    t0 = time.time()
    mem_launch = int(os.environ.get("BODY_MEM_AT_LAUNCH", "0")) or mem()
    import math
    import numpy as np
    import mujoco
    import flygym
    from flygym.compose import FlatGroundWorld
    from flygym.simulation import Simulation
    from flygym.utils.math import Rotation3D
    from flygym_demo.complex_terrain import (HybridControllerObservation, HybridTurningController,
                                             PreprogrammedSteps, apply_locomotion_action, make_locomotion_fly)

    spec = json.loads(Path(drive_path).read_text())
    drive = np.asarray(spec["drive"], dtype=float)
    assert drive.shape == (10, 2), drive.shape
    fly = make_locomotion_fly()
    world = FlatGroundWorld()
    world.add_fly(fly, spawn_position=[0, 0, 1.0], spawn_rotation=Rotation3D("quat", [1, 0, 0, 0]),
                  add_ground_contact_sensors=False)
    sim = Simulation(world)
    sim.reset()
    ctrl = HybridTurningController(timestep=sim.mj_model.opt.timestep, preprogrammed_steps=PreprogrammedSteps())
    seed_note = "controller has no seed argument; default construction"
    if hasattr(ctrl, "reset"):
        try:
            ctrl.reset(seed=0)
            seed_note = "controller.reset(seed=0)"
        except TypeError:
            try:
                ctrl.reset()
                seed_note = "controller.reset() (no seed argument)"
            except Exception as e:  # noqa: BLE001
                seed_note = f"controller.reset failed: {e!r}"
    names = [mujoco.mj_id2name(sim.mj_model, mujoco.mjtObj.mjOBJ_BODY, i) for i in range(sim.mj_model.nbody)]
    cand = [i for i, n in enumerate(names) if n and "thorax" in n.lower()]
    bid = cand[0]
    dt = sim.mj_model.opt.timestep
    n_settle = int(round(settle / dt))
    n_driven = 0 if smoke else int(round(driven / dt))
    per_bin = max(1, n_driven // 10)
    rec_every = max(1, int(round(0.01 / dt)))
    traj = []
    status, err = "ok", None
    start_pos = start_head = None
    try:
        for k in range(n_settle + n_driven):
            if k < n_settle:
                d = np.zeros(2)
            else:
                d = drive[min((k - n_settle) // per_bin, 9)]
                if start_pos is None:
                    start_pos, start_head = thorax_state(sim, mujoco, bid)
            obs = HybridControllerObservation.from_sim(sim, fly.name)
            action = ctrl.step(d, obs)
            apply_locomotion_action(sim, fly.name, action)
            sim.step()
            if (k + 1) % rec_every == 0:
                p, h = thorax_state(sim, mujoco, bid)
                if not np.all(np.isfinite(p)):
                    status, err = "failed-physics", f"non-finite thorax position at step {k}"
                    break
                traj.append([round((k + 1) * dt, 4), float(p[0]), float(p[1]), float(p[2]), float(h)])
    except Exception as e:  # noqa: BLE001
        status, err = "failed-physics", repr(e)
    if start_pos is None:
        start_pos, start_head = (np.array(traj[0][1:4]), traj[0][4]) if traj else (np.zeros(3), 0.0)
    T = np.asarray(traj) if traj else np.zeros((0, 5))
    res = {}
    if len(T):
        drv = T[T[:, 0] > settle + 1e-9] if not smoke else T
        last = T[T[:, 0] >= T[-1, 0] - 0.1]
        if status == "ok" and np.any(last[:, 3] < 0.3):
            status, err = "failed-physics", "thorax z below 0.3 mm in the last 0.1 s (fallen)"
        if len(drv):
            dx, dy = drv[-1, 1] - start_pos[0], drv[-1, 2] - start_pos[1]
            fwd = dx * math.cos(start_head) + dy * math.sin(start_head)
            lat = -dx * math.sin(start_head) + dy * math.cos(start_head)
            # unwrapped heading change, sample by sample (fixed after run 6: the run 6 files hold the wrapped value)
            hs = np.concatenate([[start_head], drv[:, 4]])
            dh = math.degrees(float(np.sum(np.arctan2(np.sin(np.diff(hs)), np.cos(np.diff(hs))))))
            xy = np.vstack([start_pos[:2], drv[:, 1:3]])
            res = {"forward_mm": round(float(fwd), 4), "lateral_mm": round(float(lat), 4),
                   "heading_deg": round(float(dh), 3),
                   "path_mm": round(float(np.sum(np.linalg.norm(np.diff(xy, axis=0), axis=1))), 4),
                   "min_z_mm_last_0.1s": round(float(last[:, 3].min()), 4)}
    out = {"id": spec.get("id"), "group": spec.get("group"), "status": status, "error": err,
           "settle_s": settle, "driven_s": 0 if smoke else driven, "smoke": smoke, **res,
           "start_xy_mm": [float(start_pos[0]), float(start_pos[1])], "start_heading_rad": float(start_head),
           "trajectory_cols": ["t_s", "x_mm", "y_mm", "z_mm", "heading_rad"],
           "trajectory": [[round(v, 5) for v in r] for r in T.tolist()],
           "drive": drive.tolist(), "drive_source": str(drive_path), "brain_source": spec.get("source"),
           "thorax_body": names[bid], "timestep_s": dt, "controller": "flygym_demo.complex_terrain.HybridTurningController (defaults)",
           "controller_seed": seed_note, "ground_contact_sensors": False,
           "versions": {"python": platform.python_version(), "flygym": getattr(flygym, "__version__", "?"),
                        "mujoco": mujoco.__version__, "numpy": np.__version__},
           "mujoco_gl": os.environ.get("MUJOCO_GL"), "wall_s": round(time.time() - t0, 2),
           "peak_rss_kb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, "memory_current_at_launch": mem_launch,
           "finished_at": now()}
    tmp = Path(str(out_path) + ".tmp")
    tmp.write_text(json.dumps(out))
    os.replace(tmp, out_path)
    print(f"{Path(out_path).name}: {status} fwd {res.get('forward_mm')} mm, heading {res.get('heading_deg')} deg, "
          f"path {res.get('path_mm')} mm, wall {out['wall_s']} s, rss {out['peak_rss_kb'] / 1e6:.2f} GB", flush=True)


def queue(drives_dir, out_dir, order, start_cut, kill, log):
    sc = datetime.fromisoformat(start_cut.replace("Z", "+00:00")).timestamp()
    kl = datetime.fromisoformat(kill.replace("Z", "+00:00")).timestamp()
    Path(out_dir).mkdir(parents=True, exist_ok=True)

    def L(r):
        with open(log, "a") as f:
            f.write(json.dumps(r) + "\n")
    for cid in order:
        d = Path(drives_dir) / f"{cid}.json"
        o = Path(out_dir) / f"{cid}.json"
        if o.exists() or not d.exists():
            continue
        if time.time() >= sc:
            L({"event": "start-cut", "run": cid, "at": now()})
            break
        m = mem()
        waited = 0
        while m is not None and m > 2.4 * GIB and waited < 180:
            time.sleep(30)
            waited += 30
            m = mem()
        if m is not None and m > 2.6 * GIB:
            L({"event": "skipped-memory", "run": cid, "mem": m, "at": now()})
            continue
        L({"event": "start", "run": cid, "mem": m, "waited_s": waited, "at": now()})
        t0 = time.time()
        p = subprocess.Popen([sys.executable, __file__, "one", str(d), str(o)], env=dict(os.environ, BODY_MEM_AT_LAUNCH=str(m or 0)))
        try:
            rc = p.wait(timeout=max(1, min(120, kl - t0)))
        except subprocess.TimeoutExpired:
            p.kill()
            p.wait()
            if time.time() >= kl - 1:
                L({"event": "killed-at-kill-time", "run": cid, "at": now()})
                break
            rec = {"id": cid, "status": "timeout", "wall_s": round(time.time() - t0, 1), "finished_at": now()}
            Path(str(o) + ".tmp").write_text(json.dumps(rec))
            os.replace(str(o) + ".tmp", o)
            L({"event": "timeout", "run": cid, "at": now()})
            continue
        L({"event": "end", "run": cid, "rc": rc, "wall_s": round(time.time() - t0, 1), "at": now()})
    L({"event": "queue-done", "at": now()})


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["one", "queue"])
    ap.add_argument("a")
    ap.add_argument("b")
    ap.add_argument("--settle", type=float, default=0.2)
    ap.add_argument("--driven", type=float, default=1.0)
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--order")
    ap.add_argument("--start-cut")
    ap.add_argument("--kill")
    ap.add_argument("--log")
    a = ap.parse_args()
    if a.cmd == "one":
        one(a.a, a.b, a.settle, a.driven, a.smoke)
    else:
        queue(a.a, a.b, a.order.split(","), a.start_cut, a.kill, a.log)

"""Fixed analysis of the run 6 "Add a body" test (Digital Fly Lab, 2026-10-02). Licence: MIT. Standard library only.

Usage: python3 analyze_body.py RESULTS_DIR CHART_OUT.json
Reads RESULTS_DIR/brain/*.json, RESULTS_DIR/drives/*.json, RESULTS_DIR/body/*.json; writes RESULTS_DIR/summary.json,
RESULTS_DIR/dn-table.json and CHART_OUT.json. Pre-registered readings (preregistration.md):
retention = (null - zero) / (real - zero) on the mean forward displacement; >= 0.9 "does not need the fly wiring",
<= 0.1 "the wiring shapes the walking command", between "partly"; matched-constant retention >= 0.9 "a brainless
constant command walks as well as the brain".
"""
import glob, json, math, os, re, statistics, sys

RES, CHART = sys.argv[1], sys.argv[2]
REL = "body/results/"
src = open(os.path.join(os.path.dirname(os.path.abspath(__file__)), "mapping.py")).read()
DN_IDS = eval(re.search(r"DN_IDS = (\{.*?\n\})", src, re.S).group(1))
FORMULA = eval(re.search(r"FORMULA = (\(.*?\))\n", src, re.S).group(1))
MAPPED = ("DNp09", "MDN", "DNa02")


def load(d):
    return {os.path.basename(f)[:-5]: json.load(open(f)) for f in sorted(glob.glob(os.path.join(RES, d, "*.json")))}


brain, drives, body = load("brain"), load("drives"), load("body")


def unwrapped_heading(v):
    """Heading change over the driven period, unwrapped sample by sample (10 ms). The run 6 body files store
    heading_deg wrapped to (-180, 180], which misreads spins of more than 180 deg (the P9L runs)."""
    prev, tot = v["start_heading_rad"], 0.0
    for t, x, y, z, h in v["trajectory"]:
        if t > v["settle_s"] + 1e-9:
            tot += math.atan2(math.sin(h - prev), math.cos(h - prev))
            prev = h
    return round(math.degrees(tot), 3)


for v in body.values():
    if v.get("status") == "ok" and v.get("trajectory"):
        v["heading_wrapped_deg"] = v.get("heading_deg")
        v["heading_deg"] = unwrapped_heading(v)


def ms(vals):
    vals = [v for v in vals if v is not None]
    if not vals:
        return None, None, 0
    return round(statistics.mean(vals), 4), (round(statistics.stdev(vals), 4) if len(vals) > 1 else None), len(vals)


GROUPS = {"real": ["P9-real-0", "P9-real-1", "P9-real-2"], "D": ["P9-D3", "P9-D4", "P9-D5"], "matched": ["matched"],
          "random": ["random-0", "random-1", "random-2"], "zero": ["zero"], "const-1": ["const-1"], "SUG-real": ["SUG-real-0"],
          "SUG-D": ["SUG-D3", "SUG-D4", "SUG-D5"], "P9L-real": ["P9L-real-0"], "P9L-D": ["P9L-D3"]}
ok = {k: v for k, v in body.items() if v.get("status") == "ok"}
per = {}
for k, v in body.items():
    per[k] = {f: v.get(f) for f in ("status", "error", "forward_mm", "lateral_mm", "heading_deg", "heading_wrapped_deg", "path_mm", "min_z_mm_last_0.1s", "wall_s", "peak_rss_kb")}
    per[k]["source"] = REL + f"body/{k}.json"
groups = {}
for g, ids in GROUPS.items():
    have = [i for i in ids if i in ok]
    f = ms([ok[i]["forward_mm"] for i in have])
    h = ms([ok[i]["heading_deg"] for i in have])
    p = ms([ok[i]["path_mm"] for i in have])
    groups[g] = {"runs": have, "missing_or_failed": [i for i in ids if i not in ok], "forward_mm_mean": f[0], "forward_mm_sd": f[1],
                 "heading_deg_mean": h[0], "heading_deg_sd": h[1], "abs_heading_deg_mean": ms([abs(ok[i]["heading_deg"]) for i in have])[0],
                 "path_mm_mean": p[0], "n": f[2], "status": "complete" if f[2] == len(ids) else ("incomplete" if f[2] else "not run")}
zero = groups["zero"]["forward_mm_mean"]
real = groups["real"]["forward_mm_mean"]


def ret(g):
    v = groups[g]["forward_mm_mean"]
    if v is None or real is None or zero is None or abs(real - zero) < 1e-9:
        return None
    return round((v - zero) / (real - zero), 4)


def reading(r, kind):
    if r is None:
        return "not computed"
    if kind == "D":
        return ("the walk does not need the fly wiring: it comes from the stimulated command neurons and FlyGym's controller"
                if r >= 0.9 else "the brain's wiring shapes the walking command" if r <= 0.1 else "partly")
    return "a brainless constant command walks as well as the brain" if r >= 0.9 else (
        "the brainless constant command walks less far than the brain" if r < 0.9 else "")


retention = {"D": ret("D"), "matched": ret("matched"), "random": ret("random")}
readings = {"D": reading(retention["D"], "D"), "matched": reading(retention["matched"], "matched"),
            "random": ("the brainless random drive walks as well as the brain" if (retention["random"] or 0) >= 0.9 else "the brainless random drive walks less far than the brain") if retention["random"] is not None else "not computed"}
# sugar: every mapped DN at 0 Hz with real wiring?
sug = brain.get("SUG-real-0")


def cand_rates(tr):
    b = {int(x["id"]): x["bins"] for x in tr["dn"]["active"]}
    out = {}
    for ct, sides in DN_IDS.items():
        out[ct] = {s: [round(statistics.mean([b.get(i, [0] * 10)[k] for i in ids]) * 10.0, 2) for k in range(10)] for s, ids in sides.items()}
    return out


sug_zero = None
if sug:
    r = cand_rates(sug)
    sug_zero = all(v == 0 for ct in MAPPED for s in r[ct].values() for v in s)
    readings["SUG"] = ("tasting sugar sends no walking command in this model (all mapped DNs 0 Hz; MN9 %g Hz)" % sug["mn9_hz"]) if sug_zero else "sugar drives at least one mapped DN"
# DN table
dn_table = {}
for k, tr in brain.items():
    r = cand_rates(tr)
    top = sorted(tr["dn"]["active"], key=lambda x: -x["total"])[:10]
    dn_table[k] = {"stim": tr["stim"], "arm": tr["arm"], "trial_seed": tr["trial_seed"], "shuffle_seed": tr["shuffle_seed"],
                   "mn9_hz": tr["mn9_hz"], "whole_brain_spikes": tr["whole_brain_spikes"], "neurons_active": tr["neurons_active"],
                   "dn_active": tr["dn"]["n_active"], "dn_total": tr["dn"]["n_descending"],
                   "candidates_hz_per_100ms_bin": r,
                   "candidates_hz_1s": {ct: {s: round(sum(v) / 10, 2) for s, v in d.items()} for ct, d in r.items()},
                   "top10": [{"id": x["id"], "cell_type": x["cell_type"], "side": x["side"], "hz": x["total"]} for x in top],
                   "property_checks_ok": all(tr["property_checks"].values()) if tr["property_checks"] else None,
                   "source": REL + f"brain/{k}.json"}
json.dump({"generated_by": "analyze_body.py (run 6)", "bin_ms": 100, "trials": dn_table}, open(os.path.join(RES, "dn-table.json"), "w"), indent=1)
# activity of the brain trials (display only): P9 real vs D
act = {g: ms([brain[k]["whole_brain_spikes"] for k in ks if k in brain])[0] for g, ks in
       {"P9-real": ["P9-real-0", "P9-real-1", "P9-real-2"], "P9-D": ["P9-D-3", "P9-D-4", "P9-D-5"]}.items()}
summary = {"generated_by": "analyze_body.py (run 6)", "metric": "forward displacement of the thorax along the heading at the start of the 1.0 s driven period (mm)",
           "heading_note": "heading_deg = unwrapped heading change (deg, + = left) over the driven period, from the 10 ms trajectory; the body files' own heading_deg is wrapped to (-180, 180] (kept as heading_wrapped_deg)",
           "mapping": FORMULA, "per_run": per, "groups": groups, "retention": retention, "readings": readings,
           "sugar_all_mapped_dns_zero": sug_zero, "brain_activity_whole_brain_spikes_mean": act,
           "activity_ratio_D_vs_real": round(act["P9-D"] / act["P9-real"], 3) if act["P9-D"] and act["P9-real"] else None}
json.dump(summary, open(os.path.join(RES, "summary.json"), "w"), indent=1)
# chart
LABEL = {"brain-real": "real brain, P9 stimulated", "brain-scrambled": "scrambled brain (degree-preserving), P9 stimulated",
         "no-brain": "no brain", "sugar": "real brain tasting sugar", "floor": "zero drive", "sanity": "FlyGym default drive (1, 1)",
         "sugar-scrambled": "scrambled brain tasting sugar"}
conds = []
for k, v in ok.items():
    if v.get("smoke"):
        continue
    d = drives.get(k, {})
    x0, y0 = v["start_xy_mm"]
    h0 = v["start_heading_rad"]
    pts = [[0.0, 0.0]]
    for t, x, y, z, h in v["trajectory"]:
        if t > v["settle_s"] + 1e-9:
            dx, dy = x - x0, y - y0
            pts.append([round(dx * math.cos(h0) + dy * math.sin(h0), 3), round(-dx * math.sin(h0) + dy * math.cos(h0), 3)])
    grp = {"sanity": "no-brain"}.get(v["group"], v["group"])
    lab = (LABEL.get(v["group"], v["group"]) + f" ({k})")
    conds.append({"id": k, "label": lab, "group": grp if grp in ("brain-real", "brain-scrambled", "no-brain", "sugar", "floor") else ("sugar" if "sugar" in grp else grp),
                  "path_xy_mm": pts[:121], "forward_mm": v["forward_mm"], "heading_deg": v["heading_deg"], "path_mm": v["path_mm"],
                  "drive": d.get("drive"),
                  "dn_rates_hz": ({ct: d["dn_rates_hz"][ct] for ct in MAPPED} if d.get("dn_rates_hz") else None),
                  "source": [REL + f"body/{k}.json", REL + f"drives/{k}.json"] + ([d["source"]] if d.get("source", "").startswith("body/") else [])})
chart = {"title": "Who does the walking?", "metric": summary["metric"], "path_frame": "x = forward along the starting heading, y = left (mm); 1.0 s of driven body time after a 0.2 s settle",
         "mapping": "OUR MAPPING (hand-made): " + FORMULA, "flygym_version": "2.1.0 (HybridTurningController, NeuroMechFly v2)",
         "brain_commit": "philshiu/Drosophila_brain_model@91bdd1e7 (FlyWire v783)",
         "n": {g: groups[g]["n"] for g in groups}, "retention": {"D": retention["D"], "matched": retention["matched"], "random": retention["random"]},
         "one_line": None, "caption": ["Legs, rhythm and balance come from FlyGym's walking controller.",
                                       "The brain sends two numbers every 100 ms.", "How neuron rates become those numbers is our mapping."],
         "conditions": sorted(conds, key=lambda c: c["id"])}
json.dump(chart, open(CHART, "w"), indent=1)
print(json.dumps({g: (groups[g]["forward_mm_mean"], groups[g]["forward_mm_sd"], groups[g]["heading_deg_mean"], groups[g]["n"]) for g in groups}))
print("retention", retention, "readings", readings)

"""Fair-test batch driver and trial worker (Digital Fly Lab, run 5; extended in run 6). Licence: MIT.

Run 6 (run:03aa82ff-d83b-47ec-bb36-c8edc64e696c, 2026-10-02) adds, without changing the run 5 code path:
  --stim SUG|P9|P9L   which neurons get the upstream Poisson input (SUG = the 21 sugar GRNs, as run 5; P9 = DNp09 L+R;
                      P9L = DNp09 left only). The Poisson input itself is exactly upstream's.
  --dn DN_INFO.json   record every spike (SpikeMonitor record=True; does not change the dynamics) and write, for each
                      descending neuron (FlyWire super_class "descending"), spike counts in 10 bins of 100 ms.
  --style run6        trial file name <stim>-<arm>-<seed>[-g<gain>].json
  prep-dn             DN_INFO.json (model index, root ID, cell type, side of every descending neuron) from the annotation table
  queue6              the run 6 brain queue (gate, core, G addendum with its adaptive gain rule, fill), skipping existing files

Imports fly_taste.py (the bootstrap Build-your-own project) from FLY_PROJECT, which must also hold
upstream/ (philshiu/Drosophila_brain_model at 91bdd1e7 with the FlyWire v783 files).

Commands
  python run_nulls.py prep  --annot ANNOT.tsv --sides SIDES.npz
        input/output side flags from the FlyWire annotation table (super_class, v783 root IDs)
  python run_nulls.py one   ARM SEED [--gain G] --out TRIALS_DIR [--sides SIDES.npz]
        one trial in this process (ARM: real, D, W, S, B, G, B'); writes TRIALS_DIR/<arm>-<seed>[-g<gain>].json atomically
  python run_nulls.py queue --out TRIALS_DIR [--sides SIDES.npz] [--deadline ISO] [--queue core|tier2|fill|all]
        runs the pre-registered queue, one fresh process per trial, skipping trial files that exist
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import resource
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
PROJECT = Path(os.environ.get("FLY_PROJECT", "byo"))
CGROUP_MEM = Path("/sys/fs/cgroup/memory.current")
GIB = 1024 ** 3
TIMEOUT_S = 240
UPSTREAM_COMMIT = "91bdd1e7dcf193f3e7ca5a8933497fcef63b7960"
INPUT_CLASSES = ("sensory", "ascending", "sensory_ascending")  # sensory_ascending: see report deviation 1
OUTPUT_CLASSES = ("motor", "descending", "endocrine")

ARM_SEEDS_CORE = [("B", 301), ("W", 101), ("D", 3), ("S", 201), ("B", 302), ("W", 102), ("D", 4), ("S", 202),
                  ("B", 303), ("W", 103), ("D", 5), ("S", 203), ("B", 304), ("W", 104), ("S", 204), ("B", 305),
                  ("W", 105), ("S", 205)]


def now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def mem_current():
    try:
        return int(CGROUP_MEM.read_text())
    except OSError:
        return None


def trial_name(arm, seed, gain=None, stim=None):
    a = "Bp" if arm == "B'" else arm
    pre = f"{stim}-" if stim else ""
    return pre + f"{a}-{seed}" + (f"-g{gain:g}" if gain is not None else "") + ".json"


P9_IDS = {"L": 720575940635872101, "R": 720575940627652358}  # DNp09 (FlyWire v783), from run6-dn-candidates.md


def stim_ids(stim, ft):
    if stim == "SUG":
        return list(ft.SUGAR_GRNS_V783)
    if stim == "P9":
        return [P9_IDS["L"], P9_IDS["R"]]
    if stim == "P9L":
        return [P9_IDS["L"]]
    raise SystemExit(f"unknown stim {stim}")


def prep_dn(annot, out):
    import csv
    sys.path.insert(0, str(PROJECT))
    os.environ.setdefault("FLY_UPSTREAM", str(PROJECT / "upstream"))
    import fly_taste as ft
    ids = ft.load_ids()
    id2i = {int(f): i for i, f in enumerate(ids)}
    rows = []
    with open(annot, newline="") as f:
        for row in csv.DictReader(f, delimiter="\t"):
            if row["super_class"] == "descending":
                rid = int(row["root_id"])
                if rid in id2i:
                    rows.append({"i": id2i[rid], "id": rid, "cell_type": row.get("cell_type") or "",
                                 "side": row.get("side") or ""})
    rows.sort(key=lambda r: r["i"])
    json.dump({"n": len(rows), "neurons": rows}, open(out, "w"))
    print("descending neurons in the model:", len(rows))


# ------------------------------------------------------------------ prep
def prep(annot, sides_out):
    import csv
    import numpy as np
    sys.path.insert(0, str(PROJECT))
    os.environ.setdefault("FLY_UPSTREAM", str(PROJECT / "upstream"))
    import fly_taste as ft
    ids = ft.load_ids()
    cls = {}
    with open(annot, newline="") as f:
        for row in csv.DictReader(f, delimiter="\t"):
            cls[int(row["root_id"])] = row["super_class"]
    names = sorted(set(cls.values()))
    code = np.full(ids.size, -1, dtype=np.int16)
    for i, fid in enumerate(ids):
        c = cls.get(int(fid))
        if c is not None:
            code[i] = names.index(c)
    matched = int((code >= 0).sum())
    inp = np.isin(code, [names.index(c) for c in INPUT_CLASSES if c in names])
    out = np.isin(code, [names.index(c) for c in OUTPUT_CLASSES if c in names])
    np.savez(sides_out, code=code, names=np.array(names), input_side=inp, output_side=out)
    info = {"n_model_ids": int(ids.size), "n_annot_rows": len(cls), "matched": matched,
            "match_fraction": matched / ids.size, "n_input_side": int(inp.sum()), "n_output_side": int(out.sum()),
            "classes": {n: int((code == k).sum()) for k, n in enumerate(names)}}
    print(json.dumps(info, indent=1))
    return info


# ------------------------------------------------------------------ one trial (worker)
def one(arm, seed, gain, out_dir, sides_path, stim="SUG", dn_info=None, style=None):
    t_start = time.time()
    mem_launch = int(os.environ.get("FAIRNULL_MEM_AT_LAUNCH", "0")) or mem_current()
    import numpy as np
    sys.path.insert(0, str(PROJECT))
    os.environ.setdefault("FLY_UPSTREAM", str(PROJECT / "upstream"))
    import brian2 as b2
    from brian2 import Hz, ms, mV
    import fly_taste as ft
    import nulls
    from model import default_params

    ids = ft.load_ids()
    n = ids.size
    id2i = {int(f): i for i, f in enumerate(ids)}
    exc = [id2i[int(f)] for f in stim_ids(stim, ft)]
    mn9 = id2i[ft.MN9]
    sides = np.load(sides_path) if sides_path and Path(sides_path).exists() else None

    pre = ft._read_column("Presynaptic_Index", np.int32)
    post = ft._read_column("Postsynaptic_Index", np.int32)
    cnt = ft._read_column("Excitatory x Connectivity", np.float64)
    assert np.array_equal(cnt, np.round(cnt))
    cnt = cnt.astype(np.int32)
    real = (pre, post, cnt)

    checks, stats = {}, {}
    if arm == "real":
        new = real
    elif arm in ("D", "G"):
        new = nulls.null_degree(pre, post, cnt, seed)
    elif arm == "W":
        new = nulls.null_weight(pre, post, cnt, seed)
    elif arm == "S":
        new = nulls.null_sign(pre, post, cnt, seed, n)
    elif arm == "B":
        new = nulls.null_boundary(pre, post, cnt, seed, sides["input_side"], sides["output_side"])
    elif arm == "B'":
        new = nulls.null_stimreadout(pre, post, cnt, seed, exc, mn9, n)
    else:
        raise SystemExit(f"unknown arm {arm}")
    if arm != "real":
        kw = {}
        if arm == "B":
            kw = dict(input_side=sides["input_side"], output_side=sides["output_side"])
        if arm == "B'":
            kw = dict(stim_idx=exc, readout_idx=mn9)
        checks = nulls.check(arm, real, new, n, **kw)
        if not all(checks.values()):
            print("PROPERTY CHECK FAILED", checks, file=sys.stderr)
            sys.exit(3)  # no data recorded
        s0 = nulls.structure_stats(pre, post, n)
        s1 = nulls.structure_stats(new[0], new[1], n)
        stats = {"self_loops_real": s0["self_loops"], "self_loops_null": s1["self_loops"],
                 "self_loops_created": s1["self_loops"] - s0["self_loops"],
                 "duplicate_pairs_real": s0["duplicate_pairs"], "duplicate_pairs_null": s1["duplicate_pairs"],
                 "duplicate_pairs_created": s1["duplicate_pairs"] - s0["duplicate_pairs"]}
    edge_sha = nulls.sha256_arrays(*new)
    npre, npost, ncnt = new
    del real, pre, post, cnt
    if arm != "real":
        del new
    t_null = time.time() - t_start

    # ---- build the upstream LIF network exactly as fly_taste.build_network, from arrays
    params = dict(default_params)
    params["r_poi"] = 150.0 * Hz
    b2.prefs.codegen.target = ft.pick_codegen_target()
    b2.start_scope()
    t0 = time.time()
    neu = b2.NeuronGroup(N=n, model=params["eqs"], method="linear", threshold=params["eq_th"],
                         reset=params["eq_rst"], refractory="rfc", name="default_neurons", namespace=params)
    neu.v = params["v_0"]
    neu.g = 0 * mV
    neu.rfc = params["t_rfc"]
    syn = b2.Synapses(neu, neu, "w : volt", on_pre="g += w", delay=params["t_dly"], name="default_synapses")
    step = 2_000_000
    for k in range(0, npre.size, step):
        syn.connect(i=npre[k:k + step], j=npost[k:k + step])
    w_syn = float(params["w_syn"]) * (gain if gain is not None else 1.0)  # gain scales recurrent synapses only
    for k in range(0, ncnt.size, step):
        syn.w_[k:k + step] = ncnt[k:k + step].astype(np.float64) * w_syn
    n_syn = len(syn)
    del npre, npost, ncnt
    pois = []
    for i in exc:  # Poisson input exactly as upstream poi(); the stimulus is not re-tuned
        pois.append(b2.PoissonInput(target=neu[i], target_var="v", N=1, rate=params["r_poi"],
                                    weight=params["w_syn"] * params["f_poi"]))
        neu[i].rfc = 0 * ms
    mon = b2.SpikeMonitor(neu, record=bool(dn_info))  # run 6: record=True only to bin DN spikes
    net = b2.Network(neu, syn, mon, *pois)
    t_build = time.time() - t0
    t1 = time.time()
    b2.seed(seed if arm == "real" else 0)  # trial seed: real arm = its seed; nulls = trial seed 0
    net.run(1000 * ms)
    t_sim = time.time() - t1
    c = np.asarray(mon.count[:], dtype=np.int64)

    order = np.argsort(-c, kind="stable")[:20]
    rec = {
        "stim": stim, "stim_ids": [int(ids[i]) for i in exc] if stim != "SUG" else "21 sugar GRNs (fly_taste.SUGAR_GRNS_V783)",
        "arm": arm, "null_type": {"real": None, "D": "degree-preserving", "W": "weight-shuffle", "S": "sign-shuffle",
                                  "B": "boundary-preserving", "G": "degree-preserving", "B'": "other"}[arm],
        "shuffle_seed": None if arm == "real" else seed, "trial_seed": seed if arm == "real" else 0,
        "gain": gain, "status": "ok",
        "mn9_hz": float(c[mn9]), "sugar_grn_mean_hz": float(c[exc].mean()),
        "whole_brain_spikes": int(c.sum()), "neurons_active": int((c > 0).sum()),
        "top20": [{"id": int(ids[i]), "hz": int(c[i])} for i in order if c[i] > 0],
        "n_synapses": int(n_syn), "property_checks": checks, "structure": stats, "edge_sha256": edge_sha,
        "wall_s": round(time.time() - t_start, 1), "null_s": round(t_null, 1), "build_s": round(t_build, 1),
        "sim_s": round(t_sim, 1), "peak_rss_kb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
        "memory_current_at_launch": mem_launch, "upstream_commit": UPSTREAM_COMMIT,
        "versions": {"python": platform.python_version(), "brian2": b2.__version__, "numpy": np.__version__,
                     "codegen_target": b2.prefs.codegen.target},
        "finished_at": now(),
    }
    if sides is not None:
        names = [str(x) for x in sides["names"]]
        code = sides["code"]
        act = c > 0
        for cl in ("motor", "descending"):
            if cl in names:
                m = code == names.index(cl)
                rec[f"active_{cl}"] = int((act & m).sum())
                rec[f"{cl}_total"] = int(m.sum())
    if dn_info:
        info = json.load(open(dn_info))["neurons"]
        di = np.array([r["i"] for r in info], dtype=np.int64)
        isdn = np.zeros(n, dtype=bool)
        isdn[di] = True
        sp_i = np.asarray(mon.i[:], dtype=np.int64)
        sp_t = np.asarray(mon.t[:] / ms, dtype=np.float64)
        keep = isdn[sp_i]
        b = np.minimum((sp_t[keep] // 100).astype(np.int64), 9)
        bins = {}
        for i_, k_ in zip(sp_i[keep], b):
            bins.setdefault(int(i_), [0] * 10)[int(k_)] += 1
        meta = {r["i"]: r for r in info}
        rec["dn"] = {"n_descending": int(di.size), "bin_ms": 100, "n_bins": 10, "n_active": len(bins),
                     "active": [{"id": meta[i]["id"], "cell_type": meta[i]["cell_type"], "side": meta[i]["side"],
                                 "bins": v, "total": int(sum(v))} for i, v in sorted(bins.items(), key=lambda kv: -sum(kv[1]))]}
        assert sum(x["total"] for x in rec["dn"]["active"]) == int(c[di].sum())  # bins agree with mon.count
    out = Path(out_dir) / trial_name(arm, seed, gain, stim if style == "run6" else None)
    tmp = out.with_suffix(".tmp")
    tmp.write_text(json.dumps(rec, indent=1))
    os.replace(tmp, out)
    print(f"{out.name}: MN9 {rec['mn9_hz']:.0f} Hz, spikes {rec['whole_brain_spikes']}, active {rec['neurons_active']}, "
          f"wall {rec['wall_s']} s, rss {rec['peak_rss_kb'] / 1e6:.2f} GB", flush=True)


# ------------------------------------------------------------------ queue (driver)
def build_queue(which, out_dir):
    q = []
    if which in ("core", "all"):
        q += [("real", 1, None), ("real", 2, None)] + [(a, s, None) for a, s in ARM_SEEDS_CORE]
    if which in ("tier2", "all"):
        q += [("G", 3, 1.5), ("G", 3, 2.0), ("G", 3, 3.0)]
    if which in ("fill",):
        g = os.environ.get("FAIRNULL_MATCHED_GAIN")
        if g:
            q += [("G", 4, float(g)), ("G", 5, float(g))]
        q += [("B'", 401, None), ("B'", 402, None), ("B'", 403, None)]
    return q


def queue(out_dir, sides_path, deadline, which):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    log = out_dir.parent / "queue-log.jsonl"
    dl = datetime.fromisoformat(deadline.replace("Z", "+00:00")).timestamp() if deadline else None
    timeouts = {}
    for arm, seed, gain in build_queue(which, out_dir):
        f = out_dir / trial_name(arm, seed, gain)
        if f.exists():
            continue
        if dl and time.time() >= dl:
            _log(log, {"event": "deadline", "at": now()})
            break
        if timeouts.get(arm, 0) >= 2:
            _log(log, {"event": "arm-stopped-2-timeouts", "arm": arm, "seed": seed, "at": now()})
            continue
        m = mem_current()
        waited = 0
        while m is not None and m > 2.4 * GIB and waited < 180:
            time.sleep(30)
            waited += 30
            m = mem_current()
        if m is not None and m > 2.6 * GIB:
            _log(log, {"event": "skipped-memory", "arm": arm, "seed": seed, "gain": gain, "mem": m, "at": now()})
            continue
        cmd = [sys.executable, str(HERE / "run_nulls.py"), "one", arm, str(seed), "--out", str(out_dir)]
        if gain is not None:
            cmd += ["--gain", str(gain)]
        if sides_path:
            cmd += ["--sides", str(sides_path)]
        env = dict(os.environ, FAIRNULL_MEM_AT_LAUNCH=str(m or 0))
        t0 = time.time()
        limit = TIMEOUT_S if not dl else max(1, min(TIMEOUT_S, dl - t0))
        _log(log, {"event": "start", "arm": arm, "seed": seed, "gain": gain, "mem": m, "waited_s": waited, "at": now()})
        p = subprocess.Popen(cmd, env=env)
        try:
            rc = p.wait(timeout=limit)
        except subprocess.TimeoutExpired:
            p.kill()
            p.wait()
            if dl and time.time() >= dl - 1 and limit < TIMEOUT_S:
                _log(log, {"event": "killed-at-deadline", "arm": arm, "seed": seed, "gain": gain, "at": now()})
                break
            timeouts[arm] = timeouts.get(arm, 0) + 1
            rec = {"arm": arm, "shuffle_seed": None if arm == "real" else seed, "gain": gain, "status": "timeout",
                   "wall_s": round(time.time() - t0, 1), "memory_current_at_launch": m, "finished_at": now()}
            tmp = f.with_suffix(".tmp")
            tmp.write_text(json.dumps(rec, indent=1))
            os.replace(tmp, f)
            _log(log, {"event": "timeout", "arm": arm, "seed": seed, "gain": gain, "at": now()})
            continue
        _log(log, {"event": "end", "arm": arm, "seed": seed, "gain": gain, "rc": rc,
                   "wall_s": round(time.time() - t0, 1), "at": now()})


# ------------------------------------------------------------------ run 6 queue (driver)
G_LO, G_HI, G_TARGET = 9680, 17976, 13828  # +-30% of the real whole-brain spikes (run 5 pre-registration)
G_KNOWN = {1.5: 5063, 2.0: 6865, 3.0: 24024}  # run 5 G-3 trials (fairnull/results/trials/G-3-g*.json in run 5)


def _run_one6(item, log, start_cut, kill_at):
    """item: dict(stim, arm, seed, gain, out, style). Returns 'exists'|'ran'|'cut'|'killed'|'skipped-memory'|'timeout'."""
    f = Path(item["out"]) / trial_name(item["arm"], item["seed"], item["gain"], item["stim"] if item["style"] == "run6" else None)
    if f.exists():
        return "exists"
    if time.time() >= start_cut:
        _log(log, {"event": "start-cut", "trial": f.name, "at": now()})
        return "cut"
    m = mem_current()
    waited = 0
    while m is not None and m > 2.4 * GIB and waited < 180:
        time.sleep(30)
        waited += 30
        m = mem_current()
    if m is not None and m > 2.6 * GIB:
        _log(log, {"event": "skipped-memory", "trial": f.name, "mem": m, "at": now()})
        return "skipped-memory"
    cmd = [sys.executable, str(HERE / "run_nulls.py"), "one", item["arm"], str(item["seed"]), "--out", str(item["out"]),
           "--stim", item["stim"], "--dn", item["dn"], "--sides", item["sides"]]
    if item["gain"] is not None:
        cmd += ["--gain", str(item["gain"])]
    if item["style"]:
        cmd += ["--style", item["style"]]
    env = dict(os.environ, FAIRNULL_MEM_AT_LAUNCH=str(m or 0))
    t0 = time.time()
    _log(log, {"event": "start", "trial": f.name, "mem": m, "waited_s": waited, "at": now()})
    p = subprocess.Popen(cmd, env=env)
    try:
        rc = p.wait(timeout=max(1, min(TIMEOUT_S, kill_at - t0)))
    except subprocess.TimeoutExpired:
        p.kill()
        p.wait()
        if time.time() >= kill_at - 1:
            _log(log, {"event": "killed-at-kill-time", "trial": f.name, "at": now()})
            return "killed"
        rec = {"stim": item["stim"], "arm": item["arm"], "seed": item["seed"], "gain": item["gain"], "status": "timeout",
               "wall_s": round(time.time() - t0, 1), "memory_current_at_launch": m, "finished_at": now()}
        tmp = f.with_suffix(".tmp")
        tmp.write_text(json.dumps(rec, indent=1))
        os.replace(tmp, f)
        _log(log, {"event": "timeout", "trial": f.name, "at": now()})
        return "timeout"
    _log(log, {"event": "end", "trial": f.name, "rc": rc, "wall_s": round(time.time() - t0, 1), "at": now()})
    return "ran"


def queue6(brain_out, g_out, sides, dn, start_cut_iso, kill_iso, log):
    sc = datetime.fromisoformat(start_cut_iso.replace("Z", "+00:00")).timestamp()
    kl = datetime.fromisoformat(kill_iso.replace("Z", "+00:00")).timestamp()
    Path(brain_out).mkdir(parents=True, exist_ok=True)
    Path(g_out).mkdir(parents=True, exist_ok=True)
    B = dict(out=brain_out, style="run6", sides=sides, dn=dn, gain=None)
    gate = [dict(B, stim="SUG", arm="real", seed=0)]
    core = [dict(B, stim="P9", arm=a, seed=s) for a, s in (("real", 0), ("D", 3), ("D", 4), ("D", 5), ("real", 1), ("real", 2))]
    for it in gate + core:
        if _run_one6(it, log, sc, kl) in ("cut", "killed"):
            return
    # G addendum (pre-registered): D shuffle 3 at 2.5; below the band -> 2.75, above -> 2.25; then bisect inside the
    # tightest bracket (run 5 points included); at most 3 new gains; matched = closest to 13,828 within the band.
    G = dict(out=g_out, style=None, sides=sides, dn=dn, stim="SUG", arm="G")
    known = dict(G_KNOWN)
    gain = 2.5
    for k in range(3):
        st = _run_one6(dict(G, seed=3, gain=gain), log, sc, kl)
        if st in ("cut", "killed"):
            return
        f = Path(g_out) / trial_name("G", 3, gain)
        r = json.loads(f.read_text())
        if r.get("status") != "ok":
            break
        known[gain] = r["whole_brain_spikes"]
        if G_LO <= r["whole_brain_spikes"] <= G_HI:
            break
        lo = max([g for g, v in known.items() if v < G_LO], default=None)
        hi = min([g for g, v in known.items() if v > G_HI], default=None)
        if k == 0:
            gain = 2.75 if r["whole_brain_spikes"] < G_LO else 2.25
        else:
            gain = round((lo + hi) / 2, 4) if lo is not None and hi is not None else None
        if gain is None or gain in known:
            break
    inband = {g: v for g, v in known.items() if G_LO <= v <= G_HI and g not in G_KNOWN}
    if inband:
        mg = min(inband, key=lambda g: abs(inband[g] - G_TARGET))
        _log(log, {"event": "matched-gain", "gain": mg, "spikes": inband[mg], "at": now()})
        for s in (4, 5):
            if _run_one6(dict(G, seed=s, gain=mg), log, sc, kl) in ("cut", "killed"):
                return
    else:
        _log(log, {"event": "no-matched-gain", "known": {str(k): v for k, v in known.items()}, "at": now()})
    fill = [dict(B, stim="SUG", arm="D", seed=s) for s in (3, 4, 5)] + \
           [dict(B, stim="P9L", arm="real", seed=0), dict(B, stim="P9L", arm="D", seed=3)]
    for it in fill:
        if _run_one6(it, log, sc, kl) in ("cut", "killed"):
            return
    _log(log, {"event": "queue-done", "at": now()})


def _log(path, rec):
    with open(path, "a") as f:
        f.write(json.dumps(rec) + "\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["prep", "one", "queue", "prep-dn", "queue6"])
    ap.add_argument("--stim", default="SUG", choices=["SUG", "P9", "P9L"])
    ap.add_argument("--dn")
    ap.add_argument("--style")
    ap.add_argument("--g-out")
    ap.add_argument("--start-cut")
    ap.add_argument("--kill")
    ap.add_argument("--log")
    ap.add_argument("arm", nargs="?")
    ap.add_argument("seed", nargs="?", type=int)
    ap.add_argument("--gain", type=float)
    ap.add_argument("--out")
    ap.add_argument("--sides")
    ap.add_argument("--annot")
    ap.add_argument("--deadline")
    ap.add_argument("--queue", default="all")
    a = ap.parse_args()
    if a.cmd == "prep":
        prep(a.annot, a.sides)
    elif a.cmd == "one":
        one(a.arm, a.seed, a.gain, a.out, a.sides, a.stim, a.dn, a.style)
    elif a.cmd == "prep-dn":
        prep_dn(a.annot, a.out)
    elif a.cmd == "queue6":
        queue6(a.out, a.g_out, a.sides, a.dn, a.start_cut, a.kill, a.log)
    else:
        queue(a.out, a.sides, a.deadline, a.queue)


if __name__ == "__main__":
    main()

"""Null (scrambled-wiring) generators and property checks for the fair test (Digital Fly Lab, run 5).

Licence: MIT (same as the bootstrap project). Pre-registration: ../preregistration.md.

Every generator takes the real edge list as numpy arrays
    pre   int32 [E]  presynaptic neuron index (Brian2 index = row of Completeness_783.csv)
    post  int32 [E]  postsynaptic neuron index
    cnt   int32 [E]  signed synapse count (upstream "Excitatory x Connectivity": count x +1/-1)
and returns new (pre, post, cnt) arrays. Edge order is kept, so the number of edges and the
total synapse count (sum |cnt|) are always unchanged.

Arms (null_type in the controls ledger):
  D  degree-preserving   permute the post column over all edges (identical to the bootstrap
                         fly_taste.load_edges(seed): np.random.default_rng(seed).shuffle(post))
  W  weight-shuffle      permute |cnt| over all edges; graph and each edge's sign kept
  S  sign-shuffle        permute the per-neuron transmitter sign over presynaptic neurons
                         (Dale's law: one sign per neuron); graph and |cnt| kept
  B  boundary-preserving permute the post column over interior edges only (pre not on the input
                         side and post not on the output side); boundary edges unchanged
  B' other               permute the post column over all edges except those out of the stimulus
                         neurons and those into the readout neuron
"""
from __future__ import annotations

import hashlib

import numpy as np


# ---------------------------------------------------------------- generators
def null_degree(pre, post, cnt, seed):
    p = post.copy()
    np.random.default_rng(seed).shuffle(p)
    return pre, p, cnt


def null_weight(pre, post, cnt, seed):
    a = np.abs(cnt)
    perm = np.random.default_rng(seed).permutation(a.size)
    return pre, post, (np.sign(cnt) * a[perm]).astype(cnt.dtype)


def neuron_signs(pre, cnt, n):
    """Per-neuron sign (+1/-1) for every neuron with outgoing edges, 0 otherwise.
    Raises if a presynaptic neuron has edges of both signs (Dale's law violated in the input)."""
    s = np.zeros(n, dtype=np.int8)
    sg = np.sign(cnt).astype(np.int8)
    s[pre] = sg
    if not np.array_equal(s[pre], sg):
        raise ValueError("input data has a presynaptic neuron with mixed signs")
    return s


def null_sign(pre, post, cnt, seed, n):
    s = neuron_signs(pre, cnt, n)
    has = np.nonzero(s != 0)[0]
    new = np.zeros_like(s)
    new[has] = np.random.default_rng(seed).permutation(s[has])
    return pre, post, (np.abs(cnt) * new[pre]).astype(cnt.dtype)


def interior_mask_boundary(pre, post, input_side, output_side):
    """Interior edge = presynaptic neuron not on the input side AND postsynaptic neuron not on the output side."""
    return ~input_side[pre] & ~output_side[post]


def null_masked(pre, post, cnt, seed, interior):
    """Permute post only among edges where interior is True."""
    p = post.copy()
    idx = np.nonzero(interior)[0]
    vals = p[idx]
    np.random.default_rng(seed).shuffle(vals)
    p[idx] = vals
    return pre, p, cnt


def null_boundary(pre, post, cnt, seed, input_side, output_side):
    m = interior_mask_boundary(pre, post, input_side, output_side)
    return null_masked(pre, post, cnt, seed, m)


def null_stimreadout(pre, post, cnt, seed, stim_idx, readout_idx, n):
    keep = np.zeros(n, dtype=bool)
    keep[np.asarray(stim_idx)] = True
    m = ~keep[pre] & (post != readout_idx)
    return null_masked(pre, post, cnt, seed, m)


# ---------------------------------------------------------------- property checks
def _deg(a, n):
    return np.bincount(a, minlength=n)


def check(arm, real, null, n, input_side=None, output_side=None, stim_idx=None, readout_idx=None):
    """Return a dict of named boolean checks for the arm; all must be True."""
    pre0, post0, cnt0 = real
    pre1, post1, cnt1 = null
    c = {
        "n_edges_kept": pre1.size == pre0.size == post1.size == cnt1.size,
        "total_synapses_kept": int(np.abs(cnt1).sum(dtype=np.int64)) == int(np.abs(cnt0).sum(dtype=np.int64)),
    }
    if arm in ("D", "G", "B", "B'"):
        c["out_degree_kept"] = np.array_equal(pre1, pre0)
        c["in_degree_kept"] = np.array_equal(_deg(post1, n), _deg(post0, n))
        c["edge_counts_kept_with_pre"] = np.array_equal(cnt1, cnt0)
    if arm == "B":
        m = interior_mask_boundary(pre0, post0, input_side, output_side)
        c["boundary_edges_unchanged"] = np.array_equal(post1[~m], post0[~m])
        c["interior_stays_interior"] = bool((~output_side[post1[m]]).all())
    if arm == "B'":
        keep = np.zeros(n, dtype=bool)
        keep[np.asarray(stim_idx)] = True
        fixed = keep[pre0] | (post0 == readout_idx)
        c["stim_out_and_readout_in_unchanged"] = np.array_equal(post1[fixed], post0[fixed])
        c["no_new_edges_into_readout"] = int((post1 == readout_idx).sum()) == int((post0 == readout_idx).sum())
    if arm == "W":
        c["graph_kept"] = np.array_equal(pre1, pre0) and np.array_equal(post1, post0)
        c["edge_signs_kept"] = np.array_equal(np.sign(cnt1), np.sign(cnt0))
        c["abs_count_multiset_kept"] = np.array_equal(np.sort(np.abs(cnt1)), np.sort(np.abs(cnt0)))
    if arm == "S":
        c["graph_kept"] = np.array_equal(pre1, pre0) and np.array_equal(post1, post0)
        c["abs_counts_kept"] = np.array_equal(np.abs(cnt1), np.abs(cnt0))
        try:
            s1 = neuron_signs(pre1, cnt1, n)
            s0 = neuron_signs(pre0, cnt0, n)
            c["one_sign_per_presynaptic_neuron"] = True
            c["n_excitatory_neurons_kept"] = int((s1 > 0).sum()) == int((s0 > 0).sum())
            c["n_inhibitory_neurons_kept"] = int((s1 < 0).sum()) == int((s0 < 0).sum())
        except ValueError:
            c["one_sign_per_presynaptic_neuron"] = False
    return {k: bool(v) for k, v in c.items()}


def structure_stats(pre, post, n):
    """Self-loops and duplicate (pre, post) pairs (extra copies beyond the first)."""
    key = pre.astype(np.int64) * n + post.astype(np.int64)
    key.sort()
    dup = int((key[1:] == key[:-1]).sum())
    del key
    return {"self_loops": int((pre == post).sum()), "duplicate_pairs": dup}


def sha256_arrays(*arrays):
    h = hashlib.sha256()
    for a in arrays:
        h.update(np.ascontiguousarray(a).tobytes())
    return h.hexdigest()

absl-py==2.5.0
asttokens==3.0.2
contourpy==1.4.0
cycler==0.12.1
etils==1.14.0
executing==2.2.1
flygym==2.1.0
fonttools==4.66.1
fsspec==2026.9.0
glfw==2.10.2
imageio==2.38.0
imageio-ffmpeg==0.6.0
ipython==9.17.1
ipython-pygments-lexers==1.1.1
jaxtyping==0.3.7
jedi==0.20.0
kiwisolver==1.5.1
llvmlite==0.50.0
loguru==0.7.3
matplotlib==3.11.2
matplotlib-inline==0.2.2
mediapy==1.2.7
mujoco==3.9.0
numba==0.68.0
numpy==2.5.3
packaging==26.3
parso==0.8.7
pexpect==4.9.0
pillow==12.3.0
prompt-toolkit==3.0.53
psutil==7.2.2
ptyprocess==0.7.0
pure-eval==0.2.4
pygments==2.21.0
pyopengl==3.1.10
pyparsing==3.3.3
python-dateutil==2.9.0.post0
pyyaml==6.0.3
scipy==1.18.1
six==1.17.0
stack-data==0.6.3
tabulate==0.9.0
traitlets==5.16.1
typing-extensions==4.16.0
wadler-lindig==0.1.7
wcwidth==0.9.1
zipp==4.1.0

import csv, json, sys, hashlib
annot, compl, out = sys.argv[1:4]
ids = set()
with open(compl) as f:
    r = csv.reader(f); h = next(r)
    for row in r: ids.add(int(row[0]))
sets = {"L": {"LC4": [], "LPLC2": []}, "R": {"LC4": [], "LPLC2": []}}
other_side = {"LC4": {}, "LPLC2": {}}
read = {}
RO = {"DNp01", "DNa02", "DNa01", "DNp09", "DNg62", "CB0701", "MDN"}
hemi_alias = []
with open(annot, newline="") as f:
    for row in csv.DictReader(f, delimiter="\t"):
        ct = row.get("cell_type") or ""
        side = row.get("side") or ""
        rid = int(row["root_id"])
        if ct in ("LC4", "LPLC2"):
            if side == "left": sets["L"][ct].append(rid)
            elif side == "right": sets["R"][ct].append(rid)
            else: other_side[ct][side] = other_side[ct].get(side, 0) + 1
        if ct in RO:
            read.setdefault(ct, []).append({"id": rid, "side": side, "in_v783": rid in ids})
        ht = row.get("hemibrain_type") or ""
        if ht in RO and ct not in RO:
            hemi_alias.append({"id": rid, "cell_type": ct, "hemibrain_type": ht, "side": side})
res = {"source": "flyconnectome/flywire_annotations supplemental_files/Supplemental_file1_neuron_annotations.tsv",
       "annotation_commit": "a83b2776d60d5764cef36b927f5f9679c16c47a2",
       "annotation_sha256": hashlib.sha256(open(annot,'rb').read()).hexdigest(),
       "completeness": "philshiu/Drosophila_brain_model Completeness_783.csv at 91bdd1e7 (%d neurons)" % len(ids),
       "selector": "cell_type exactly LC4 or LPLC2; side == left/right (soma side), as fly67 src/circuits.js",
       "counts": {s: {t: len(v) for t, v in d.items()} | {"total": sum(len(v) for v in d.values())} for s, d in sets.items()},
       "expected": {"L": {"LC4": 54, "LPLC2": 108, "total": 162}, "R": {"LC4": 50, "LPLC2": 102, "total": 152}},
       "missing_from_v783": [i for s in sets.values() for v in s.values() for i in v if i not in ids],
       "dropped_other_side": other_side,
       "readouts": read, "hemibrain_alias_not_used": hemi_alias,
       "ids": {s: sorted(i for v in d.values() for i in v) for s, d in sets.items()}}
res["matches_expected"] = all(res["counts"][s][k] == res["expected"][s][k] for s in "LR" for k in ("LC4", "LPLC2"))
json.dump(res, open(out, "w"), indent=1)
print(json.dumps({k: res[k] for k in ("counts", "matches_expected", "missing_from_v783", "dropped_other_side")}))
print({k: [(r["id"], r["side"], r["in_v783"]) for r in v] for k, v in read.items()})
print(hemi_alias)

"""Run 8 looming worker (Digital Fly Lab, 2026-10-05). Licence: MIT.

A copy of run 6's run_nulls.py (the Add-a-body worker on this page) with three additions that do not
change the run 5/6 code path for SUG/P9 trials:
  --stim LOOML|LOOMR   every v783 neuron with cell_type LC4 or LPLC2 and soma side left/right (IDs from ../step0.json,
                       built by step0.py from the Schlegel et al. annotation table, fly67's selector)
  --rate HZ            Poisson rate per stimulated neuron (default 150 Hz = upstream; looming uses 80 Hz as fly67)
  --silence GF         zero every synapse FROM both DNp01 (giant fibre), as upstream silence()
  queue8               the pre-registered run 8 brain queue (gate, 12 core trials, GF-off and G fill), skipping
                       existing files; refuses to start a LOOM trial unless ../preregistration.md exists
Trial files: results/brain/<STIM>-<arm>-<seed>[-g<gain>][-gfoff].json, written atomically.
Original docstring follows.
Fair-test batch driver and trial worker (Digital Fly Lab, run 5; extended in run 6). Licence: MIT.

Run 6 (run:03aa82ff-d83b-47ec-bb36-c8edc64e696c, 2026-10-02) adds, without changing the run 5 code path:
  --stim SUG|P9|P9L   which neurons get the upstream Poisson input (SUG = the 21 sugar GRNs, as run 5; P9 = DNp09 L+R;
                      P9L = DNp09 left only). The Poisson input itself is exactly upstream's.
  --dn DN_INFO.json   record every spike (SpikeMonitor record=True; does not change the dynamics) and write, for each
                      descending neuron (FlyWire super_class "descending"), spike counts in 10 bins of 100 ms.
  --style run6        trial file name <stim>-<arm>-<seed>[-g<gain>].json
  prep-dn             DN_INFO.json (model index, root ID, cell type, side of every descending neuron) from the annotation table
  queue6              the run 6 brain queue (gate, core, G addendum with its adaptive gain rule, fill), skipping existing files

Imports fly_taste.py (the bootstrap Build-your-own project) from FLY_PROJECT, which must also hold
upstream/ (philshiu/Drosophila_brain_model at 91bdd1e7 with the FlyWire v783 files).

Commands
  python run_nulls.py prep  --annot ANNOT.tsv --sides SIDES.npz
        input/output side flags from the FlyWire annotation table (super_class, v783 root IDs)
  python run_nulls.py one   ARM SEED [--gain G] --out TRIALS_DIR [--sides SIDES.npz]
        one trial in this process (ARM: real, D, W, S, B, G, B'); writes TRIALS_DIR/<arm>-<seed>[-g<gain>].json atomically
  python run_nulls.py queue --out TRIALS_DIR [--sides SIDES.npz] [--deadline ISO] [--queue core|tier2|fill|all]
        runs the pre-registered queue, one fresh process per trial, skipping trial files that exist
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import resource
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
PROJECT = Path(os.environ.get("FLY_PROJECT", "byo"))
CGROUP_MEM = Path("/sys/fs/cgroup/memory.current")
GIB = 1024 ** 3
TIMEOUT_S = 240
UPSTREAM_COMMIT = "91bdd1e7dcf193f3e7ca5a8933497fcef63b7960"
INPUT_CLASSES = ("sensory", "ascending", "sensory_ascending")  # sensory_ascending: see report deviation 1
OUTPUT_CLASSES = ("motor", "descending", "endocrine")

ARM_SEEDS_CORE = [("B", 301), ("W", 101), ("D", 3), ("S", 201), ("B", 302), ("W", 102), ("D", 4), ("S", 202),
                  ("B", 303), ("W", 103), ("D", 5), ("S", 203), ("B", 304), ("W", 104), ("S", 204), ("B", 305),
                  ("W", 105), ("S", 205)]


def now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def mem_current():
    try:
        return int(CGROUP_MEM.read_text())
    except OSError:
        return None


def trial_name(arm, seed, gain=None, stim=None, silence=None):
    a = "Bp" if arm == "B'" else arm
    pre = f"{stim}-" if stim else ""
    return pre + f"{a}-{seed}" + (f"-g{gain:g}" if gain is not None else "") + ("-gfoff" if silence else "") + ".json"


P9_IDS = {"L": 720575940635872101, "R": 720575940627652358}  # DNp09 (FlyWire v783), from run6-dn-candidates.md


STEP0 = HERE.parent / "step0.json"
GF_IDS = [720575940622838154, 720575940632499757]  # DNp01 left, right (step0.json readouts)


def silence_counts(npre, ncnt, slnc_idx):
    """As upstream silence(): zero every synapse FROM the silenced neurons. Returns a copy if anything is silenced."""
    import numpy as np
    if not len(slnc_idx):
        return ncnt
    out = ncnt.copy()
    out[np.isin(npre, slnc_idx)] = 0
    return out


def stim_ids(stim, ft):
    if stim in ("LOOML", "LOOMR"):
        return list(json.load(open(STEP0))["ids"][stim[-1]])
    if stim == "SUG":
        return list(ft.SUGAR_GRNS_V783)
    if stim == "P9":
        return [P9_IDS["L"], P9_IDS["R"]]
    if stim == "P9L":
        return [P9_IDS["L"]]
    raise SystemExit(f"unknown stim {stim}")


def prep_dn(annot, out):
    import csv
    sys.path.insert(0, str(PROJECT))
    os.environ.setdefault("FLY_UPSTREAM", str(PROJECT / "upstream"))
    import fly_taste as ft
    ids = ft.load_ids()
    id2i = {int(f): i for i, f in enumerate(ids)}
    rows = []
    with open(annot, newline="") as f:
        for row in csv.DictReader(f, delimiter="\t"):
            if row["super_class"] == "descending":
                rid = int(row["root_id"])
                if rid in id2i:
                    rows.append({"i": id2i[rid], "id": rid, "cell_type": row.get("cell_type") or "",
                                 "side": row.get("side") or ""})
    rows.sort(key=lambda r: r["i"])
    json.dump({"n": len(rows), "neurons": rows}, open(out, "w"))
    print("descending neurons in the model:", len(rows))


# ------------------------------------------------------------------ prep
def prep(annot, sides_out):
    import csv
    import numpy as np
    sys.path.insert(0, str(PROJECT))
    os.environ.setdefault("FLY_UPSTREAM", str(PROJECT / "upstream"))
    import fly_taste as ft
    ids = ft.load_ids()
    cls = {}
    with open(annot, newline="") as f:
        for row in csv.DictReader(f, delimiter="\t"):
            cls[int(row["root_id"])] = row["super_class"]
    names = sorted(set(cls.values()))
    code = np.full(ids.size, -1, dtype=np.int16)
    for i, fid in enumerate(ids):
        c = cls.get(int(fid))
        if c is not None:
            code[i] = names.index(c)
    matched = int((code >= 0).sum())
    inp = np.isin(code, [names.index(c) for c in INPUT_CLASSES if c in names])
    out = np.isin(code, [names.index(c) for c in OUTPUT_CLASSES if c in names])
    np.savez(sides_out, code=code, names=np.array(names), input_side=inp, output_side=out)
    info = {"n_model_ids": int(ids.size), "n_annot_rows": len(cls), "matched": matched,
            "match_fraction": matched / ids.size, "n_input_side": int(inp.sum()), "n_output_side": int(out.sum()),
            "classes": {n: int((code == k).sum()) for k, n in enumerate(names)}}
    print(json.dumps(info, indent=1))
    return info


# ------------------------------------------------------------------ one trial (worker)
def one(arm, seed, gain, out_dir, sides_path, stim="SUG", dn_info=None, style=None, rate=150.0, silence=None):
    t_start = time.time()
    mem_launch = int(os.environ.get("FAIRNULL_MEM_AT_LAUNCH", "0")) or mem_current()
    import numpy as np
    sys.path.insert(0, str(PROJECT))
    os.environ.setdefault("FLY_UPSTREAM", str(PROJECT / "upstream"))
    import brian2 as b2
    from brian2 import Hz, ms, mV
    import fly_taste as ft
    import nulls
    from model import default_params

    ids = ft.load_ids()
    n = ids.size
    id2i = {int(f): i for i, f in enumerate(ids)}
    exc = [id2i[int(f)] for f in stim_ids(stim, ft)]
    mn9 = id2i[ft.MN9]
    sides = np.load(sides_path) if sides_path and Path(sides_path).exists() else None

    pre = ft._read_column("Presynaptic_Index", np.int32)
    post = ft._read_column("Postsynaptic_Index", np.int32)
    cnt = ft._read_column("Excitatory x Connectivity", np.float64)
    assert np.array_equal(cnt, np.round(cnt))
    cnt = cnt.astype(np.int32)
    real = (pre, post, cnt)

    checks, stats = {}, {}
    if arm == "real":
        new = real
    elif arm in ("D", "G"):
        new = nulls.null_degree(pre, post, cnt, seed)
    elif arm == "W":
        new = nulls.null_weight(pre, post, cnt, seed)
    elif arm == "S":
        new = nulls.null_sign(pre, post, cnt, seed, n)
    elif arm == "B":
        new = nulls.null_boundary(pre, post, cnt, seed, sides["input_side"], sides["output_side"])
    elif arm == "B'":
        new = nulls.null_stimreadout(pre, post, cnt, seed, exc, mn9, n)
    else:
        raise SystemExit(f"unknown arm {arm}")
    if arm != "real":
        kw = {}
        if arm == "B":
            kw = dict(input_side=sides["input_side"], output_side=sides["output_side"])
        if arm == "B'":
            kw = dict(stim_idx=exc, readout_idx=mn9)
        checks = nulls.check(arm, real, new, n, **kw)
        if not all(checks.values()):
            print("PROPERTY CHECK FAILED", checks, file=sys.stderr)
            sys.exit(3)  # no data recorded
        s0 = nulls.structure_stats(pre, post, n)
        s1 = nulls.structure_stats(new[0], new[1], n)
        stats = {"self_loops_real": s0["self_loops"], "self_loops_null": s1["self_loops"],
                 "self_loops_created": s1["self_loops"] - s0["self_loops"],
                 "duplicate_pairs_real": s0["duplicate_pairs"], "duplicate_pairs_null": s1["duplicate_pairs"],
                 "duplicate_pairs_created": s1["duplicate_pairs"] - s0["duplicate_pairs"]}
    edge_sha = nulls.sha256_arrays(*new)
    npre, npost, ncnt = new
    slnc_idx = [id2i[i] for i in GF_IDS] if silence == "GF" else []
    ncnt = silence_counts(npre, ncnt, slnc_idx)
    del real, pre, post, cnt
    if arm != "real":
        del new
    t_null = time.time() - t_start

    # ---- build the upstream LIF network exactly as fly_taste.build_network, from arrays
    params = dict(default_params)
    params["r_poi"] = float(rate) * Hz
    b2.prefs.codegen.target = ft.pick_codegen_target()
    b2.start_scope()
    t0 = time.time()
    neu = b2.NeuronGroup(N=n, model=params["eqs"], method="linear", threshold=params["eq_th"],
                         reset=params["eq_rst"], refractory="rfc", name="default_neurons", namespace=params)
    neu.v = params["v_0"]
    neu.g = 0 * mV
    neu.rfc = params["t_rfc"]
    syn = b2.Synapses(neu, neu, "w : volt", on_pre="g += w", delay=params["t_dly"], name="default_synapses")
    step = 2_000_000
    for k in range(0, npre.size, step):
        syn.connect(i=npre[k:k + step], j=npost[k:k + step])
    w_syn = float(params["w_syn"]) * (gain if gain is not None else 1.0)  # gain scales recurrent synapses only
    for k in range(0, ncnt.size, step):
        syn.w_[k:k + step] = ncnt[k:k + step].astype(np.float64) * w_syn
    n_syn = len(syn)
    del npre, npost, ncnt
    pois = []
    for i in exc:  # Poisson input exactly as upstream poi(); the stimulus is not re-tuned
        pois.append(b2.PoissonInput(target=neu[i], target_var="v", N=1, rate=params["r_poi"],
                                    weight=params["w_syn"] * params["f_poi"]))
        neu[i].rfc = 0 * ms
    mon = b2.SpikeMonitor(neu, record=bool(dn_info))  # run 6: record=True only to bin DN spikes
    net = b2.Network(neu, syn, mon, *pois)
    t_build = time.time() - t0
    t1 = time.time()
    b2.seed(seed if arm == "real" else 0)  # trial seed: real arm = its seed; nulls = trial seed 0
    net.run(1000 * ms)
    t_sim = time.time() - t1
    c = np.asarray(mon.count[:], dtype=np.int64)

    order = np.argsort(-c, kind="stable")[:20]
    rec = {
        "stim": stim, "stim_ids": [int(ids[i]) for i in exc] if stim != "SUG" else "21 sugar GRNs (fly_taste.SUGAR_GRNS_V783)",
        "arm": arm, "null_type": {"real": None, "D": "degree-preserving", "W": "weight-shuffle", "S": "sign-shuffle",
                                  "B": "boundary-preserving", "G": "degree-preserving", "B'": "other"}[arm],
        "shuffle_seed": None if arm == "real" else seed, "trial_seed": seed if arm == "real" else 0,
        "gain": gain, "status": "ok", "rate_hz": float(rate), "silenced": [int(ids[i]) for i in slnc_idx],
        "stim_n": len(exc), "stim_spikes": int(c[exc].sum()),
        "mn9_hz": float(c[mn9]), "sugar_grn_mean_hz": float(c[exc].mean()),
        "whole_brain_spikes": int(c.sum()), "neurons_active": int((c > 0).sum()),
        "top20": [{"id": int(ids[i]), "hz": int(c[i])} for i in order if c[i] > 0],
        "n_synapses": int(n_syn), "property_checks": checks, "structure": stats, "edge_sha256": edge_sha,
        "wall_s": round(time.time() - t_start, 1), "null_s": round(t_null, 1), "build_s": round(t_build, 1),
        "sim_s": round(t_sim, 1), "peak_rss_kb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
        "memory_current_at_launch": mem_launch, "upstream_commit": UPSTREAM_COMMIT,
        "versions": {"python": platform.python_version(), "brian2": b2.__version__, "numpy": np.__version__,
                     "codegen_target": b2.prefs.codegen.target},
        "finished_at": now(),
    }
    if sides is not None:
        names = [str(x) for x in sides["names"]]
        code = sides["code"]
        act = c > 0
        for cl in ("motor", "descending"):
            if cl in names:
                m = code == names.index(cl)
                rec[f"active_{cl}"] = int((act & m).sum())
                rec[f"{cl}_total"] = int(m.sum())
    if dn_info:
        info = json.load(open(dn_info))["neurons"]
        di = np.array([r["i"] for r in info], dtype=np.int64)
        isdn = np.zeros(n, dtype=bool)
        isdn[di] = True
        sp_i = np.asarray(mon.i[:], dtype=np.int64)
        sp_t = np.asarray(mon.t[:] / ms, dtype=np.float64)
        keep = isdn[sp_i]
        b = np.minimum((sp_t[keep] // 100).astype(np.int64), 9)
        bins = {}
        for i_, k_ in zip(sp_i[keep], b):
            bins.setdefault(int(i_), [0] * 10)[int(k_)] += 1
        meta = {r["i"]: r for r in info}
        rec["dn"] = {"n_descending": int(di.size), "bin_ms": 100, "n_bins": 10, "n_active": len(bins),
                     "active": [{"id": meta[i]["id"], "cell_type": meta[i]["cell_type"], "side": meta[i]["side"],
                                 "bins": v, "total": int(sum(v))} for i, v in sorted(bins.items(), key=lambda kv: -sum(kv[1]))]}
        assert sum(x["total"] for x in rec["dn"]["active"]) == int(c[di].sum())  # bins agree with mon.count
    out = Path(out_dir) / trial_name(arm, seed, gain, stim if style in ("run6", "run8") else None, silence)
    tmp = out.with_suffix(".tmp")
    tmp.write_text(json.dumps(rec, indent=1))
    os.replace(tmp, out)
    print(f"{out.name}: MN9 {rec['mn9_hz']:.0f} Hz, spikes {rec['whole_brain_spikes']}, active {rec['neurons_active']}, "
          f"wall {rec['wall_s']} s, rss {rec['peak_rss_kb'] / 1e6:.2f} GB", flush=True)


# ------------------------------------------------------------------ queue (driver)
def build_queue(which, out_dir):
    q = []
    if which in ("core", "all"):
        q += [("real", 1, None), ("real", 2, None)] + [(a, s, None) for a, s in ARM_SEEDS_CORE]
    if which in ("tier2", "all"):
        q += [("G", 3, 1.5), ("G", 3, 2.0), ("G", 3, 3.0)]
    if which in ("fill",):
        g = os.environ.get("FAIRNULL_MATCHED_GAIN")
        if g:
            q += [("G", 4, float(g)), ("G", 5, float(g))]
        q += [("B'", 401, None), ("B'", 402, None), ("B'", 403, None)]
    return q


def queue(out_dir, sides_path, deadline, which):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    log = out_dir.parent / "queue-log.jsonl"
    dl = datetime.fromisoformat(deadline.replace("Z", "+00:00")).timestamp() if deadline else None
    timeouts = {}
    for arm, seed, gain in build_queue(which, out_dir):
        f = out_dir / trial_name(arm, seed, gain)
        if f.exists():
            continue
        if dl and time.time() >= dl:
            _log(log, {"event": "deadline", "at": now()})
            break
        if timeouts.get(arm, 0) >= 2:
            _log(log, {"event": "arm-stopped-2-timeouts", "arm": arm, "seed": seed, "at": now()})
            continue
        m = mem_current()
        waited = 0
        while m is not None and m > 2.4 * GIB and waited < 180:
            time.sleep(30)
            waited += 30
            m = mem_current()
        if m is not None and m > 2.6 * GIB:
            _log(log, {"event": "skipped-memory", "arm": arm, "seed": seed, "gain": gain, "mem": m, "at": now()})
            continue
        cmd = [sys.executable, str(HERE / "run_nulls.py"), "one", arm, str(seed), "--out", str(out_dir)]
        if gain is not None:
            cmd += ["--gain", str(gain)]
        if sides_path:
            cmd += ["--sides", str(sides_path)]
        env = dict(os.environ, FAIRNULL_MEM_AT_LAUNCH=str(m or 0))
        t0 = time.time()
        limit = TIMEOUT_S if not dl else max(1, min(TIMEOUT_S, dl - t0))
        _log(log, {"event": "start", "arm": arm, "seed": seed, "gain": gain, "mem": m, "waited_s": waited, "at": now()})
        p = subprocess.Popen(cmd, env=env)
        try:
            rc = p.wait(timeout=limit)
        except subprocess.TimeoutExpired:
            p.kill()
            p.wait()
            if dl and time.time() >= dl - 1 and limit < TIMEOUT_S:
                _log(log, {"event": "killed-at-deadline", "arm": arm, "seed": seed, "gain": gain, "at": now()})
                break
            timeouts[arm] = timeouts.get(arm, 0) + 1
            rec = {"arm": arm, "shuffle_seed": None if arm == "real" else seed, "gain": gain, "status": "timeout",
                   "wall_s": round(time.time() - t0, 1), "memory_current_at_launch": m, "finished_at": now()}
            tmp = f.with_suffix(".tmp")
            tmp.write_text(json.dumps(rec, indent=1))
            os.replace(tmp, f)
            _log(log, {"event": "timeout", "arm": arm, "seed": seed, "gain": gain, "at": now()})
            continue
        _log(log, {"event": "end", "arm": arm, "seed": seed, "gain": gain, "rc": rc,
                   "wall_s": round(time.time() - t0, 1), "at": now()})


# ------------------------------------------------------------------ run 6 queue (driver)
G_LO, G_HI, G_TARGET = 9680, 17976, 13828  # +-30% of the real whole-brain spikes (run 5 pre-registration)
G_KNOWN = {1.5: 5063, 2.0: 6865, 3.0: 24024}  # run 5 G-3 trials (fairnull/results/trials/G-3-g*.json in run 5)


def _run_one6(item, log, start_cut, kill_at):
    """item: dict(stim, arm, seed, gain, out, style). Returns 'exists'|'ran'|'cut'|'killed'|'skipped-memory'|'timeout'."""
    f = Path(item["out"]) / trial_name(item["arm"], item["seed"], item["gain"], item["stim"] if item["style"] == "run6" else None)
    if f.exists():
        return "exists"
    if time.time() >= start_cut:
        _log(log, {"event": "start-cut", "trial": f.name, "at": now()})
        return "cut"
    m = mem_current()
    waited = 0
    while m is not None and m > 2.4 * GIB and waited < 180:
        time.sleep(30)
        waited += 30
        m = mem_current()
    if m is not None and m > 2.6 * GIB:
        _log(log, {"event": "skipped-memory", "trial": f.name, "mem": m, "at": now()})
        return "skipped-memory"
    cmd = [sys.executable, str(HERE / "run_nulls.py"), "one", item["arm"], str(item["seed"]), "--out", str(item["out"]),
           "--stim", item["stim"], "--dn", item["dn"], "--sides", item["sides"]]
    if item["gain"] is not None:
        cmd += ["--gain", str(item["gain"])]
    if item["style"]:
        cmd += ["--style", item["style"]]
    env = dict(os.environ, FAIRNULL_MEM_AT_LAUNCH=str(m or 0))
    t0 = time.time()
    _log(log, {"event": "start", "trial": f.name, "mem": m, "waited_s": waited, "at": now()})
    p = subprocess.Popen(cmd, env=env)
    try:
        rc = p.wait(timeout=max(1, min(TIMEOUT_S, kill_at - t0)))
    except subprocess.TimeoutExpired:
        p.kill()
        p.wait()
        if time.time() >= kill_at - 1:
            _log(log, {"event": "killed-at-kill-time", "trial": f.name, "at": now()})
            return "killed"
        rec = {"stim": item["stim"], "arm": item["arm"], "seed": item["seed"], "gain": item["gain"], "status": "timeout",
               "wall_s": round(time.time() - t0, 1), "memory_current_at_launch": m, "finished_at": now()}
        tmp = f.with_suffix(".tmp")
        tmp.write_text(json.dumps(rec, indent=1))
        os.replace(tmp, f)
        _log(log, {"event": "timeout", "trial": f.name, "at": now()})
        return "timeout"
    _log(log, {"event": "end", "trial": f.name, "rc": rc, "wall_s": round(time.time() - t0, 1), "at": now()})
    return "ran"


def queue6(brain_out, g_out, sides, dn, start_cut_iso, kill_iso, log):
    sc = datetime.fromisoformat(start_cut_iso.replace("Z", "+00:00")).timestamp()
    kl = datetime.fromisoformat(kill_iso.replace("Z", "+00:00")).timestamp()
    Path(brain_out).mkdir(parents=True, exist_ok=True)
    Path(g_out).mkdir(parents=True, exist_ok=True)
    B = dict(out=brain_out, style="run6", sides=sides, dn=dn, gain=None)
    gate = [dict(B, stim="SUG", arm="real", seed=0)]
    core = [dict(B, stim="P9", arm=a, seed=s) for a, s in (("real", 0), ("D", 3), ("D", 4), ("D", 5), ("real", 1), ("real", 2))]
    for it in gate + core:
        if _run_one6(it, log, sc, kl) in ("cut", "killed"):
            return
    # G addendum (pre-registered): D shuffle 3 at 2.5; below the band -> 2.75, above -> 2.25; then bisect inside the
    # tightest bracket (run 5 points included); at most 3 new gains; matched = closest to 13,828 within the band.
    G = dict(out=g_out, style=None, sides=sides, dn=dn, stim="SUG", arm="G")
    known = dict(G_KNOWN)
    gain = 2.5
    for k in range(3):
        st = _run_one6(dict(G, seed=3, gain=gain), log, sc, kl)
        if st in ("cut", "killed"):
            return
        f = Path(g_out) / trial_name("G", 3, gain)
        r = json.loads(f.read_text())
        if r.get("status") != "ok":
            break
        known[gain] = r["whole_brain_spikes"]
        if G_LO <= r["whole_brain_spikes"] <= G_HI:
            break
        lo = max([g for g, v in known.items() if v < G_LO], default=None)
        hi = min([g for g, v in known.items() if v > G_HI], default=None)
        if k == 0:
            gain = 2.75 if r["whole_brain_spikes"] < G_LO else 2.25
        else:
            gain = round((lo + hi) / 2, 4) if lo is not None and hi is not None else None
        if gain is None or gain in known:
            break
    inband = {g: v for g, v in known.items() if G_LO <= v <= G_HI and g not in G_KNOWN}
    if inband:
        mg = min(inband, key=lambda g: abs(inband[g] - G_TARGET))
        _log(log, {"event": "matched-gain", "gain": mg, "spikes": inband[mg], "at": now()})
        for s in (4, 5):
            if _run_one6(dict(G, seed=s, gain=mg), log, sc, kl) in ("cut", "killed"):
                return
    else:
        _log(log, {"event": "no-matched-gain", "known": {str(k): v for k, v in known.items()}, "at": now()})
    fill = [dict(B, stim="SUG", arm="D", seed=s) for s in (3, 4, 5)] + \
           [dict(B, stim="P9L", arm="real", seed=0), dict(B, stim="P9L", arm="D", seed=3)]
    for it in fill:
        if _run_one6(it, log, sc, kl) in ("cut", "killed"):
            return
    _log(log, {"event": "queue-done", "at": now()})


# ------------------------------------------------------------------ run 8 queue (driver)
MEM_START_MAX = 2.2 * GIB   # run 8 memory guard: wait (30 s steps, up to 2 min) until below 2.2 GiB
MEM_STOP = 2.7 * GIB        # during a trial: stop it above 2.7 GiB and log stopped-memory


def _run_one8(item, log, start_cut, kill_at, prereg):
    f = Path(item["out"]) / trial_name(item["arm"], item["seed"], item["gain"], item["stim"], item.get("silence"))
    if f.exists():
        return "exists"
    if time.time() >= start_cut:
        _log(log, {"event": "not-run (clock)", "trial": f.name, "at": now()})
        return "cut"
    if item["stim"].startswith("LOOM") and not Path(prereg).exists():
        _log(log, {"event": "refused-no-preregistration", "trial": f.name, "at": now()})
        raise SystemExit("pre-registration file missing")
    m = mem_current(); waited = 0
    while m is not None and m > MEM_START_MAX and waited < 120:
        time.sleep(30); waited += 30; m = mem_current()
    if m is not None and m > MEM_START_MAX:
        _log(log, {"event": "skipped-memory", "trial": f.name, "mem": m, "at": now()})
        return "skipped-memory"
    cmd = [sys.executable, str(HERE / "run_loom.py"), "one", item["arm"], str(item["seed"]), "--out", str(item["out"]),
           "--stim", item["stim"], "--dn", item["dn"], "--sides", item["sides"], "--style", "run8",
           "--rate", str(item["rate"])]
    if item["gain"] is not None:
        cmd += ["--gain", str(item["gain"])]
    if item.get("silence"):
        cmd += ["--silence", item["silence"]]
    env = dict(os.environ, FAIRNULL_MEM_AT_LAUNCH=str(m or 0))
    t0 = time.time()
    _log(log, {"event": "start", "trial": f.name, "mem": m, "waited_s": waited, "at": now()})
    p = subprocess.Popen(cmd, env=env)
    peak = m or 0
    while True:
        try:
            rc = p.wait(timeout=2)
            break
        except subprocess.TimeoutExpired:
            mm = mem_current() or 0
            peak = max(peak, mm)
            if mm > MEM_STOP:
                p.kill(); p.wait()
                _log(log, {"event": "stopped-memory", "trial": f.name, "mem": mm, "at": now()})
                return "stopped-memory"
            if time.time() >= kill_at or time.time() - t0 > TIMEOUT_S:
                p.kill(); p.wait()
                ev = "killed-at-kill-time" if time.time() >= kill_at else "timeout"
                _log(log, {"event": ev, "trial": f.name, "at": now()})
                return "killed" if ev.startswith("killed") else "timeout"
    _log(log, {"event": "end", "trial": f.name, "rc": rc, "wall_s": round(time.time() - t0, 1),
               "cgroup_peak_during": peak, "at": now()})
    return "ran"


def queue8(out, sides, dn, start_cut_iso, fill_cut_iso, kill_iso, log, prereg):
    ts = lambda x: datetime.fromisoformat(x.replace("Z", "+00:00")).timestamp()
    sc, fc, kl = ts(start_cut_iso), ts(fill_cut_iso), ts(kill_iso)
    Path(out).mkdir(parents=True, exist_ok=True)
    B = dict(out=out, sides=sides, dn=dn, gain=None, rate=80.0)
    gate = [dict(B, stim="SUG", arm="real", seed=0, rate=150.0)]
    core = []
    for k, (arm, seed) in enumerate([("real", 0), ("D", 3), ("real", 1), ("D", 4), ("real", 2), ("D", 5)]):
        for side in ("L", "R"):
            core.append(dict(B, stim="LOOM" + side, arm=arm, seed=seed))
    for it in gate + core:
        if _run_one8(it, log, sc, kl, prereg) in ("cut", "killed"):
            return
    # fill 1: GF-off (both DNp01 silenced), L-real-0 and R-real-0
    for side in ("L", "R"):
        if _run_one8(dict(B, stim="LOOM" + side, arm="real", seed=0, silence="GF"), log, fc, kl, prereg) in ("cut", "killed"):
            return
    # fill 2: per-shuffle G gain search on LOOM-L: gains 1.5, 2, 3 until whole-brain spikes within +-30% of L-real-0
    ref = json.loads((Path(out) / trial_name("real", 0, None, "LOOML")).read_text())["whole_brain_spikes"]
    lo, hi = 0.7 * ref, 1.3 * ref
    for s in (3, 4, 5):
        for g in (1.5, 2.0, 3.0):
            st = _run_one8(dict(B, stim="LOOML", arm="G", seed=s, gain=g), log, fc, kl, prereg)
            if st in ("cut", "killed"):
                return
            f = Path(out) / trial_name("G", s, g, "LOOML")
            if f.exists():
                r = json.loads(f.read_text())
                if r.get("status") == "ok" and lo <= r["whole_brain_spikes"] <= hi:
                    _log(log, {"event": "matched-gain", "seed": s, "gain": g, "spikes": r["whole_brain_spikes"],
                               "band": [lo, hi], "at": now()})
                    break
    _log(log, {"event": "queue-done", "at": now()})


def _log(path, rec):
    with open(path, "a") as f:
        f.write(json.dumps(rec) + "\n")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["prep", "one", "queue", "prep-dn", "queue6", "queue8"])
    ap.add_argument("--stim", default="SUG", choices=["SUG", "P9", "P9L", "LOOML", "LOOMR"])
    ap.add_argument("--rate", type=float, default=150.0)
    ap.add_argument("--silence", choices=["GF"])
    ap.add_argument("--prereg")
    ap.add_argument("--fill-cut")
    ap.add_argument("--dn")
    ap.add_argument("--style")
    ap.add_argument("--g-out")
    ap.add_argument("--start-cut")
    ap.add_argument("--kill")
    ap.add_argument("--log")
    ap.add_argument("arm", nargs="?")
    ap.add_argument("seed", nargs="?", type=int)
    ap.add_argument("--gain", type=float)
    ap.add_argument("--out")
    ap.add_argument("--sides")
    ap.add_argument("--annot")
    ap.add_argument("--deadline")
    ap.add_argument("--queue", default="all")
    a = ap.parse_args()
    if a.cmd == "prep":
        prep(a.annot, a.sides)
    elif a.cmd == "one":
        one(a.arm, a.seed, a.gain, a.out, a.sides, a.stim, a.dn, a.style, a.rate, a.silence)
    elif a.cmd == "prep-dn":
        prep_dn(a.annot, a.out)
    elif a.cmd == "queue8":
        queue8(a.out, a.sides, a.dn, a.start_cut, a.fill_cut, a.kill, a.log, a.prereg)
    elif a.cmd == "queue6":
        queue6(a.out, a.g_out, a.sides, a.dn, a.start_cut, a.kill, a.log)
    else:
        queue(a.out, a.sides, a.deadline, a.queue)


if __name__ == "__main__":
    main()

"""Per-trial looming readouts from a run 8 brain trial file (Digital Fly Lab, 2026-10-05). Licence: MIT.
Rates in Hz over 1 s (trial = 1,000 ms, so spike count = Hz). fly67 set means: gf = mean of both DNp01,
back = mean of the 4 MDN, mn9 = mean of both CB0701, groom = mean of both DNg62.
Gate (fly67's own rule, validation tolerances in tools/probe.mjs): gf > 50 and DNa02_contra > 3 * DNa02_ipsi + 5.
LI = (DNa02_contra - DNa02_ipsi) / (DNa02_contra + DNa02_ipsi + 1 Hz). Selectivity = DNa02_contra / whole-brain
spikes per s x 1000 (display only)."""
import json
import sys

IDS = {"DNp01": {"L": 720575940622838154, "R": 720575940632499757},
       "DNa02": {"L": 720575940629327659, "R": 720575940604737708},
       "DNa01": {"L": 720575940627787609, "R": 720575940644438551},
       "DNp09": {"L": 720575940635872101, "R": 720575940627652358},
       "DNg62": {"L": 720575940624319124, "R": 720575940616185531},
       "CB0701": {"L": 720575940618238523, "R": 720575940660219265}}
MDN = {"L": [720575940631082808, 720575940616026939], "R": [720575940610236514, 720575940640331472]}


def readout(t):
    tot = {int(x["id"]): x["total"] for x in t["dn"]["active"]}
    top = {int(x["id"]): x["hz"] for x in t.get("top20", [])}
    g = lambda i: float(tot.get(i, top.get(i, 0)))
    side = t["stim"][-1] if t["stim"].startswith("LOOM") else None
    r = {f"{k.lower()}_{s.lower()}": g(v[s]) for k, v in IDS.items() for s in "LR"}
    r["mdn_l"] = sum(g(i) for i in MDN["L"]) / 2
    r["mdn_r"] = sum(g(i) for i in MDN["R"]) / 2
    r["gf_mean"] = (r["dnp01_l"] + r["dnp01_r"]) / 2
    r["mdn_mean"] = (r["mdn_l"] + r["mdn_r"]) / 2
    r["mn9_mean"] = (r["cb0701_l"] + r["cb0701_r"]) / 2
    if side:
        o = "r" if side == "L" else "l"
        ip = "l" if side == "L" else "r"
        c, i_ = r[f"dna02_{o}"], r[f"dna02_{ip}"]
        r["dna02_contra"], r["dna02_ipsi"] = c, i_
        r["li"] = (c - i_) / (c + i_ + 1.0)
        r["gate_gf"] = r["gf_mean"] > 50
        r["gate_turn"] = c > 3 * i_ + 5
        r["gate_pass"] = r["gate_gf"] and r["gate_turn"]
        r["selectivity"] = c / t["whole_brain_spikes"] * 1000 if t["whole_brain_spikes"] else None
    r.update(side=side, whole_brain_spikes=t["whole_brain_spikes"], active_neurons=t["neurons_active"],
             stim_spikes=t.get("stim_spikes"), stim_n=t.get("stim_n"))
    return r


if __name__ == "__main__":
    for f in sys.argv[1:]:
        t = json.load(open(f))
        print(f.split("/")[-1], json.dumps({k: (round(v, 3) if isinstance(v, float) else v) for k, v in readout(t).items()}))

"""Run 8 drive builder: brain trial -> FlyGym two-number drive (Digital Fly Lab, 2026-10-05). Licence: MIT.

OUR MAPPING (hand-made, not fitted; pre-registered 2026-10-05T09:42:47Z, loom/preregistration.md):
    drive_s = clip(0.8 - m_s - 0.5 * (a_s - a_o), -1.2, 1.2)
    a = rate(DNa02 on side s) / R, m = mean rate(MDN on side s) / R, R = 100 Hz, per 100 ms brain bin.
    0.8 = constant EXTERNAL walking drive (not the brain), as in fly67 and optogenetic experiments. No P9 term.
Timeline per drive file: 12 bins of 100 ms = 2 pre-onset bins at (0.8, 0.8) + the 10 brain bins; onset_s = 0.2.
Conditions:
    F                floor: (0.8, 0.8) for all 12 bins
    <S>-real-k, <S>-D<k>, GFoff-<S>   from the brain trial files (S = L or R, the looming side)
    rand-<S><k>      brainless: turning term t ~ N(0, s) per brain bin, s = RMS of the real runs' turning terms
                     0.5*(a_s - a_o) over all bins and both sides; drive = (0.8 - t, 0.8 + t); no MDN term.
                     Generator: numpy default_rng(seed + 100 * (side == 'R')) (implementation detail fixed before
                     the first body run; logged as a deviation in preregistration.md)
Identical drives (sha256 over the float64 drive array) are marked dedup_of = the first condition with that drive.
Usage: python make_loom_drives.py BRAIN_DIR DRIVES_DIR
"""
from __future__ import annotations

import hashlib
import json
import sys
from pathlib import Path

import numpy as np

R_HZ, CLIP, FWD, PRE_BINS = 100.0, 1.2, 0.8, 2
IDS = {"DNa02": {"left": [720575940629327659], "right": [720575940604737708]},
       "MDN": {"left": [720575940631082808, 720575940616026939], "right": [720575940610236514, 720575940640331472]}}
FORMULA = ("OUR MAPPING (hand-made, not fitted): drive_s = clip(0.8 - m_s - 0.5*(a_s - a_o), -1.2, 1.2); "
           "a = rate(DNa02_s)/R, m = mean rate(MDN_s)/R, R = 100 Hz, per 100 ms bin; 0.8 = constant external walking "
           "drive, not the brain; 2 pre-onset bins at (0.8, 0.8) then 10 brain bins")


def rates(trial):
    bins = {int(x["id"]): x["bins"] for x in trial["dn"]["active"]}
    sc = 1000.0 / trial["dn"]["bin_ms"]
    out = {}
    for ct, sides in IDS.items():
        out[ct] = np.stack([np.mean([np.array(bins.get(i, [0] * 10), float) for i in sides[s]], axis=0) * sc
                            for s in ("left", "right")], axis=1)  # (10, 2)
    return out


def drive_from_trial(trial):
    r = rates(trial)
    a, m = r["DNa02"] / R_HZ, r["MDN"] / R_HZ
    turn = 0.5 * (a - a[:, ::-1])
    raw = FWD - m - turn
    d = np.clip(raw, -CLIP, CLIP)
    return d, turn, raw


def sha(d):
    return hashlib.sha256(np.ascontiguousarray(np.asarray(d, dtype=np.float64)).tobytes()).hexdigest()


def full(d10):
    return np.vstack([np.full((PRE_BINS, 2), FWD), d10])


def main(brain_dir, drives_dir):
    B, D = Path(brain_dir), Path(drives_dir)
    D.mkdir(parents=True, exist_ok=True)
    conds = []  # (id, group, side, wiring, seed, brain file)
    for k in (0, 1, 2):
        for s in "LR":
            conds.append((f"{s}-real-{k}", "real", s, "real", k, f"LOOM{s}-real-{k}.json"))
    for k in (3, 4, 5):
        for s in "LR":
            conds.append((f"{s}-D{k}", "scrambled", s, "D", k, f"LOOM{s}-D-{k}.json"))
    for s in "LR":
        conds.append((f"GFoff-{s}", "gf-off", s, "real-gf-silenced", 0, f"LOOM{s}-real-0-gfoff.json"))
    seen, turns_real, written = {}, [], []

    def write(cid, rec):
        drive = np.asarray(rec["drive"], float)
        h = sha(drive)
        rec["drive_sha256"] = h
        rec["dedup_of"] = seen.get(h)
        seen.setdefault(h, cid)
        rec["clipped_share"] = float(np.mean(np.abs(drive[PRE_BINS:]) >= CLIP - 1e-12))
        rec.update(id=cid, onset_s=0.2, bin_s=0.1, formula=FORMULA)
        p = D / f"{cid}.json"
        tmp = p.with_suffix(".tmp")
        tmp.write_text(json.dumps(rec, indent=1))
        tmp.replace(p)
        written.append(cid)

    write("F", {"group": "floor", "side": None, "wiring": None, "seed": None, "source": None,
                "drive": np.full((12, 2), FWD).tolist()})
    for cid, grp, s, wir, k, fn in conds:
        f = B / fn
        if not f.exists():
            continue
        t = json.loads(f.read_text())
        if t.get("status") != "ok":
            continue
        d, turn, raw = drive_from_trial(t)
        if grp == "real":
            turns_real.append(turn)
        write(cid, {"group": grp, "side": s, "wiring": wir, "seed": k, "source": f"results/brain/{fn}",
                    "drive": full(d).tolist(), "turn_term": turn.tolist(), "unclipped_brain_bins": raw.tolist()})
    if turns_real:
        sd = float(np.sqrt(np.mean(np.concatenate([x.ravel() for x in turns_real]) ** 2)))
        for s in "LR":
            for k in (0, 1, 2):
                rng = np.random.default_rng(k + 100 * (s == "R"))
                t = rng.normal(0.0, sd, size=10) if sd > 0 else np.zeros(10)
                d = np.clip(np.stack([FWD - t, FWD + t], axis=1), -CLIP, CLIP)
                write(f"rand-{s}{k}", {"group": "random", "side": s, "wiring": None, "seed": k, "source": None,
                                       "random_sd": sd, "n_real_runs_for_sd": len(turns_real),
                                       "rng": f"numpy.default_rng({k + 100 * (s == 'R')})", "turn_term": t.tolist(),
                                       "drive": full(d).tolist()})
    print("drives:", ", ".join(written))


if __name__ == "__main__":
    main(*sys.argv[1:3])

"""Run 8 copy (2026-10-05): the drive may have any number of 100 ms bins (run 8: 12 = two
pre-onset bins at the constant external drive (0.8, 0.8) plus ten brain bins); `onset_s` in the drive file marks the
looming onset in driven time, and the result adds `heading_after_onset_deg` (unwrapped, + = left), the thorax path
every 50 ms after onset (`path_xy_mm`, relative to the onset position) and `forward_after_onset_mm`. Usage as below,
with --driven = number of bins x 0.1 s (set automatically when --driven is omitted).

Run the FlyGym 2.1.0 body with a two-number descending drive (Digital Fly Lab, run 6, 2026-10-02). Licence: MIT.

The walking (legs, rhythm, balance, adhesion) is FlyGym's own HybridTurningController (flygym_demo.complex_terrain,
NeuroMechFly v2's two-number descending interface). This script only feeds it a drive file made by make_drives.py
(OUR rate-to-drive mapping) and records where the thorax goes. Construction follows upstream
tests/examples/test_complex_terrain_locomotion.py at v2.1.0.

Usage
  MUJOCO_GL=disable python body_run.py one DRIVE.json OUT.json [--settle 0.2] [--driven 1.0] [--smoke]
  MUJOCO_GL=disable python body_run.py queue DRIVES_DIR OUT_DIR --order a,b,c --start-cut ISO --kill ISO --log LOG
Timeline: settle s at drive (0, 0), then driven s in 10 equal bins (bin k of the drive held for driven/10 s).
Thorax x, y (mm), z and heading are saved every 10 ms. Metric: forward displacement along the heading at the start of
the driven period (mm); also heading change (deg, + = left/counter-clockwise) and path length (mm).
"""
from __future__ import annotations

import argparse
import json
import os
import platform
import resource
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

GIB = 1024 ** 3
CG = Path("/sys/fs/cgroup/memory.current")


def now():
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def mem():
    try:
        return int(CG.read_text())
    except OSError:
        return None


def thorax_state(sim, mujoco, bid):
    pos = sim.mj_data.xpos[bid].copy()
    xm = sim.mj_data.xmat[bid].reshape(3, 3)
    fwd = xm[:, 0]  # body x axis in world frame
    import math
    return pos, math.atan2(fwd[1], fwd[0])


def one(drive_path, out_path, settle=0.2, driven=None, smoke=False):
    t0 = time.time()
    mem_launch = int(os.environ.get("BODY_MEM_AT_LAUNCH", "0")) or mem()
    import math
    import numpy as np
    import mujoco
    import flygym
    from flygym.compose import FlatGroundWorld
    from flygym.simulation import Simulation
    from flygym.utils.math import Rotation3D
    from flygym_demo.complex_terrain import (HybridControllerObservation, HybridTurningController,
                                             PreprogrammedSteps, apply_locomotion_action, make_locomotion_fly)

    spec = json.loads(Path(drive_path).read_text())
    drive = np.asarray(spec["drive"], dtype=float)
    assert drive.ndim == 2 and drive.shape[1] == 2, drive.shape
    nb = drive.shape[0]
    if driven is None:
        driven = round(nb * 0.1, 6)
    onset = float(spec.get("onset_s", 0.0))
    fly = make_locomotion_fly()
    world = FlatGroundWorld()
    world.add_fly(fly, spawn_position=[0, 0, 1.0], spawn_rotation=Rotation3D("quat", [1, 0, 0, 0]),
                  add_ground_contact_sensors=False)
    sim = Simulation(world)
    sim.reset()
    ctrl = HybridTurningController(timestep=sim.mj_model.opt.timestep, preprogrammed_steps=PreprogrammedSteps())
    seed_note = "controller has no seed argument; default construction"
    if hasattr(ctrl, "reset"):
        try:
            ctrl.reset(seed=0)
            seed_note = "controller.reset(seed=0)"
        except TypeError:
            try:
                ctrl.reset()
                seed_note = "controller.reset() (no seed argument)"
            except Exception as e:  # noqa: BLE001
                seed_note = f"controller.reset failed: {e!r}"
    names = [mujoco.mj_id2name(sim.mj_model, mujoco.mjtObj.mjOBJ_BODY, i) for i in range(sim.mj_model.nbody)]
    cand = [i for i, n in enumerate(names) if n and "thorax" in n.lower()]
    bid = cand[0]
    dt = sim.mj_model.opt.timestep
    n_settle = int(round(settle / dt))
    n_driven = 0 if smoke else int(round(driven / dt))
    per_bin = max(1, n_driven // nb)
    rec_every = max(1, int(round(0.01 / dt)))
    traj = []
    status, err = "ok", None
    start_pos = start_head = None
    try:
        for k in range(n_settle + n_driven):
            if k < n_settle:
                d = np.zeros(2)
            else:
                d = drive[min((k - n_settle) // per_bin, nb - 1)]
                if start_pos is None:
                    start_pos, start_head = thorax_state(sim, mujoco, bid)
            obs = HybridControllerObservation.from_sim(sim, fly.name)
            action = ctrl.step(d, obs)
            apply_locomotion_action(sim, fly.name, action)
            sim.step()
            if (k + 1) % rec_every == 0:
                p, h = thorax_state(sim, mujoco, bid)
                if not np.all(np.isfinite(p)):
                    status, err = "failed-physics", f"non-finite thorax position at step {k}"
                    break
                traj.append([round((k + 1) * dt, 4), float(p[0]), float(p[1]), float(p[2]), float(h)])
    except Exception as e:  # noqa: BLE001
        status, err = "failed-physics", repr(e)
    if start_pos is None:
        start_pos, start_head = (np.array(traj[0][1:4]), traj[0][4]) if traj else (np.zeros(3), 0.0)
    T = np.asarray(traj) if traj else np.zeros((0, 5))
    res = {}
    if len(T):
        drv = T[T[:, 0] > settle + 1e-9] if not smoke else T
        last = T[T[:, 0] >= T[-1, 0] - 0.1]
        if status == "ok" and np.any(last[:, 3] < 0.3):
            status, err = "failed-physics", "thorax z below 0.3 mm in the last 0.1 s (fallen)"
        if len(drv):
            dx, dy = drv[-1, 1] - start_pos[0], drv[-1, 2] - start_pos[1]
            fwd = dx * math.cos(start_head) + dy * math.sin(start_head)
            lat = -dx * math.sin(start_head) + dy * math.cos(start_head)
            # unwrapped heading change, sample by sample (fixed after run 6: the run 6 files hold the wrapped value)
            hs = np.concatenate([[start_head], drv[:, 4]])
            dh = math.degrees(float(np.sum(np.arctan2(np.sin(np.diff(hs)), np.cos(np.diff(hs))))))
            xy = np.vstack([start_pos[:2], drv[:, 1:3]])
            res = {"forward_mm": round(float(fwd), 4), "lateral_mm": round(float(lat), 4),
                   "heading_deg": round(float(dh), 3),
                   "path_mm": round(float(np.sum(np.linalg.norm(np.diff(xy, axis=0), axis=1))), 4),
                   "min_z_mm_last_0.1s": round(float(last[:, 3].min()), 4)}
            # run 8: after onset
            t_on = settle + onset
            on = T[T[:, 0] >= t_on - 1e-9]
            if len(on) > 1:
                hs2 = on[:, 4]
                dh2 = math.degrees(float(np.sum(np.arctan2(np.sin(np.diff(hs2)), np.cos(np.diff(hs2))))))
                x0, y0, h0 = on[0, 1], on[0, 2], on[0, 4]
                ddx, ddy = on[-1, 1] - x0, on[-1, 2] - y0
                res["heading_after_onset_deg"] = round(dh2, 3)
                res["forward_after_onset_mm"] = round(float(ddx * math.cos(h0) + ddy * math.sin(h0)), 4)
                res["onset_t_s"] = round(float(on[0, 0]), 4)
                sel = on[::5]
                res["path_xy_mm"] = [[round(float(r[1] - x0), 4), round(float(r[2] - y0), 4)] for r in sel]
    out = {"id": spec.get("id"), "group": spec.get("group"), "status": status, "error": err,
           "settle_s": settle, "driven_s": 0 if smoke else driven, "smoke": smoke, **res,
           "start_xy_mm": [float(start_pos[0]), float(start_pos[1])], "start_heading_rad": float(start_head),
           "trajectory_cols": ["t_s", "x_mm", "y_mm", "z_mm", "heading_rad"],
           "trajectory": [[round(v, 5) for v in r] for r in T.tolist()],
           "drive": drive.tolist(), "drive_source": str(drive_path), "brain_source": spec.get("source"),
           "thorax_body": names[bid], "timestep_s": dt, "controller": "flygym_demo.complex_terrain.HybridTurningController (defaults)",
           "controller_seed": seed_note, "ground_contact_sensors": False,
           "versions": {"python": platform.python_version(), "flygym": getattr(flygym, "__version__", "?"),
                        "mujoco": mujoco.__version__, "numpy": np.__version__},
           "mujoco_gl": os.environ.get("MUJOCO_GL"), "wall_s": round(time.time() - t0, 2),
           "peak_rss_kb": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, "memory_current_at_launch": mem_launch,
           "finished_at": now()}
    tmp = Path(str(out_path) + ".tmp")
    tmp.write_text(json.dumps(out))
    os.replace(tmp, out_path)
    print(f"{Path(out_path).name}: {status} fwd {res.get('forward_mm')} mm, heading {res.get('heading_deg')} deg, "
          f"path {res.get('path_mm')} mm, wall {out['wall_s']} s, rss {out['peak_rss_kb'] / 1e6:.2f} GB", flush=True)


def queue(drives_dir, out_dir, order, start_cut, kill, log):
    sc = datetime.fromisoformat(start_cut.replace("Z", "+00:00")).timestamp()
    kl = datetime.fromisoformat(kill.replace("Z", "+00:00")).timestamp()
    Path(out_dir).mkdir(parents=True, exist_ok=True)

    def L(r):
        with open(log, "a") as f:
            f.write(json.dumps(r) + "\n")
    for cid in order:
        d = Path(drives_dir) / f"{cid}.json"
        o = Path(out_dir) / f"{cid}.json"
        if o.exists():
            continue
        if not d.exists():
            L({"event": "no-drive-file", "run": cid, "at": now()})
            continue
        dd = json.loads(d.read_text())
        if dd.get("dedup_of"):  # identical drive (same sha256) already run: physics is deterministic, take that result
            src = Path(out_dir) / f"{dd['dedup_of']}.json"
            if src.exists():
                r = json.loads(src.read_text()); r["id"] = cid; r["dedup_of"] = dd["dedup_of"]; r["group"] = dd.get("group")
                Path(str(o) + ".tmp").write_text(json.dumps(r)); os.replace(str(o) + ".tmp", o)
                L({"event": "dedup", "run": cid, "of": dd["dedup_of"], "at": now()})
                continue
        if time.time() >= sc:
            L({"event": "not-run (clock)", "run": cid, "at": now()})
            continue
        m = mem()
        waited = 0
        while m is not None and m > 2.2 * GIB and waited < 120:
            time.sleep(30)
            waited += 30
            m = mem()
        if m is not None and m > 2.2 * GIB:
            L({"event": "skipped-memory", "run": cid, "mem": m, "at": now()})
            continue
        L({"event": "start", "run": cid, "mem": m, "waited_s": waited, "at": now()})
        t0 = time.time()
        p = subprocess.Popen([sys.executable, __file__, "one", str(d), str(o)], env=dict(os.environ, BODY_MEM_AT_LAUNCH=str(m or 0)))
        try:
            rc = p.wait(timeout=max(1, min(180, kl - t0)))
        except subprocess.TimeoutExpired:
            p.kill()
            p.wait()
            if time.time() >= kl - 1:
                L({"event": "killed-at-kill-time", "run": cid, "at": now()})
                break
            rec = {"id": cid, "status": "timeout", "wall_s": round(time.time() - t0, 1), "finished_at": now()}
            Path(str(o) + ".tmp").write_text(json.dumps(rec))
            os.replace(str(o) + ".tmp", o)
            L({"event": "timeout", "run": cid, "at": now()})
            continue
        L({"event": "end", "run": cid, "rc": rc, "wall_s": round(time.time() - t0, 1), "at": now()})
    L({"event": "queue-done", "at": now()})


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["one", "queue"])
    ap.add_argument("a")
    ap.add_argument("b")
    ap.add_argument("--settle", type=float, default=0.2)
    ap.add_argument("--driven", type=float, default=None)
    ap.add_argument("--smoke", action="store_true")
    ap.add_argument("--order")
    ap.add_argument("--start-cut")
    ap.add_argument("--kill")
    ap.add_argument("--log")
    a = ap.parse_args()
    if a.cmd == "one":
        one(a.a, a.b, a.settle, a.driven, a.smoke)
    else:
        queue(a.a, a.b, a.order.split(","), a.start_cut, a.kill, a.log)

"""Run 8 looming analysis (Digital Fly Lab, 2026-10-05). Licence: MIT. Applies the pre-registered rules in
../preregistration.md (2026-10-05T09:42:47Z, with deviations 3-5 logged before any body run).
Usage: python3 analyze_loom.py RESULTS_DIR CHART_OUT
Reads RESULTS_DIR/brain/*.json, drives/*.json, body/*.json; writes RESULTS_DIR/summary.json, dn-table.json and the
chart data file."""
import glob, json, math, os, sys
from pathlib import Path
from statistics import mean, pstdev

sys.path.insert(0, str(Path(__file__).resolve().parent))
from brain_readout import readout  # noqa: E402

RD = Path(sys.argv[1]); CHART = Path(sys.argv[2])
REL = "loom/results/"
FLY67 = {"url": "https://github.com/znatgost/fly67/blob/45f62fd9b867632141f84cc6daa77d871a2259bb/validation/circuits.json",
         "commit": "45f62fd9b867632141f84cc6daa77d871a2259bb", "engine": "the author's own JavaScript engine (not ours)",
         "loom_left": {"gf": 114.8, "turnR": 25, "turnL": 0, "walkL": 0.7, "back": 0.8},
         "loom_right": {"gf": 120.3, "turnL": 43, "turnR": 0, "back": 3.3}, "no_input": "every readout 0"}
CAVEAT = ("One model (Shiu et al. LIF on FlyWire v783), open loop, our hand-made mapping, a constant external walking "
          "drive, 1 s; not a real fly.")


def atomic(path, obj):
    tmp = Path(str(path) + ".tmp"); tmp.write_text(json.dumps(obj, indent=1)); os.replace(tmp, path)


def tid(t):
    s = t["stim"][-1]
    base = f"{s}-real-{t['trial_seed']}" if t["arm"] == "real" else f"{s}-{t['arm']}{t['shuffle_seed']}"
    if t.get("gain"):
        base += f"-g{t['gain']:g}"
    if t.get("silenced"):
        base = f"GFoff-{s}"
    return base


brain_rows, dn_rows, checks = [], [], []
for f in sorted(glob.glob(str(RD / "brain" / "LOOM*.json"))):
    t = json.load(open(f))
    if t.get("status") != "ok":
        brain_rows.append({"trial": Path(f).stem, "status": t.get("status"), "source": REL + "brain/" + Path(f).name}); continue
    r = readout(t)
    row = {"trial": tid(t), "file": Path(f).name, "side": r["side"],
           "wiring": "real" if t["arm"] == "real" else ("D" if t["arm"] == "D" else t["arm"]),
           "seed": t["trial_seed"] if t["arm"] == "real" else t["shuffle_seed"], "gain": t.get("gain"),
           "silenced": t.get("silenced") or [],
           "gf_l": r["dnp01_l"], "gf_r": r["dnp01_r"], "gf_mean": r["gf_mean"], "dna02_l": r["dna02_l"],
           "dna02_r": r["dna02_r"], "dna02_contra": r["dna02_contra"], "dna02_ipsi": r["dna02_ipsi"],
           "li": round(r["li"], 4), "dna01_l": r["dna01_l"], "dna01_r": r["dna01_r"], "mdn_l": r["mdn_l"], "mdn_r": r["mdn_r"],
           "mdn_mean": r["mdn_mean"], "dnp09_l": r["dnp09_l"], "dnp09_r": r["dnp09_r"], "dng62_mean": (r["dng62_l"] + r["dng62_r"]) / 2,
           "mn9_l": r["cb0701_l"], "mn9_r": r["cb0701_r"], "whole_brain_spikes": r["whole_brain_spikes"],
           "active_neurons": r["active_neurons"], "stim_n": r["stim_n"], "stim_spikes": r["stim_spikes"],
           "selectivity": round(r["selectivity"], 4) if r["selectivity"] is not None else None,
           "gate_gf": r["gate_gf"], "gate_turn": r["gate_turn"], "gate_pass": r["gate_pass"],
           "wall_s": t["wall_s"], "sim_s": t["sim_s"], "peak_rss_gb": round(t["peak_rss_kb"] / 1e6, 3),
           "source": REL + "brain/" + Path(f).name}
    brain_rows.append(row)
    if t["arm"] != "real":
        checks.append({"trial": row["trial"], "property_checks": t["property_checks"], "all_pass": all(t["property_checks"].values()),
                       "edge_sha256": t["edge_sha256"], "structure": t["structure"]})
    dn_rows.append({"trial": row["trial"], "n_active_dn": t["dn"]["n_active"],
                    "top_dn": [{k: x[k] for k in ("id", "cell_type", "side", "total")} for x in t["dn"]["active"][:25]]})

ok = [r for r in brain_rows if "gf_mean" in r]
core = [r for r in ok if not r["silenced"] and not r["gain"]]
real = [r for r in core if r["wiring"] == "real"]
D = [r for r in core if r["wiring"] == "D"]
gate = {s: next((r for r in real if r["side"] == s and r["seed"] == 0), None) for s in "LR"}
gate_pass = all(g and g["gate_pass"] for g in gate.values())
brain = {
    "gate": {s: ({k: g[k] for k in ("gf_mean", "gf_l", "gf_r", "dna02_contra", "dna02_ipsi", "gate_gf", "gate_turn", "gate_pass")} if g else None) for s, g in gate.items()},
    "gate_result": "reproduced in our Brian2 run" if gate_pass else "not reproduced in our Brian2 run",
    "n": {"real": len(real), "D": len(D), "gfoff": len([r for r in ok if r["silenced"]]), "G": len([r for r in ok if r["gain"]])},
    "real_mean": {k: round(mean(r[k] for r in real), 3) for k in ("li", "gf_mean", "dna02_contra", "dna02_ipsi", "whole_brain_spikes", "active_neurons", "selectivity", "mdn_mean")} if real else None,
    "D_mean": {k: round(mean(r[k] for r in D), 3) for k in ("li", "gf_mean", "dna02_contra", "dna02_ipsi", "whole_brain_spikes", "active_neurons", "selectivity", "mdn_mean")} if D else None,
    "fly67_rule_pass": {"real": sum(r["gate_pass"] for r in real), "D": sum(r["gate_pass"] for r in D)},
    "shuffle_checks_all_pass": all(c["all_pass"] for c in checks) if checks else None,
}
if real and D:
    rm, dm = brain["real_mean"], brain["D_mean"]
    brain["retention_display"] = {"LI_D_over_real": round(dm["li"] / rm["li"], 4) if rm["li"] else None,
                                  "GF_D_over_real": round(dm["gf_mean"] / rm["gf_mean"], 4) if rm["gf_mean"] else None,
                                  "selectivity_D_over_real": round(dm["selectivity"] / rm["selectivity"], 4) if rm["selectivity"] else None,
                                  "activity_ratio_spikes_D_over_real": round(dm["whole_brain_spikes"] / rm["whole_brain_spikes"], 4)}

# ------------------------------------------------------------------ body
drives = {Path(f).stem: json.load(open(f)) for f in glob.glob(str(RD / "drives" / "*.json"))}
bodies = {Path(f).stem: json.load(open(f)) for f in glob.glob(str(RD / "body" / "*.json"))}
body_runs, body = [], {}
F = bodies.get("F")
f_dh = F.get("heading_after_onset_deg") if F and F.get("status") == "ok" else None
f_away = {"L": -f_dh, "R": f_dh} if f_dh is not None else None
for cid, b in sorted(bodies.items()):
    d = drives.get(cid, {})
    side = d.get("side")
    dh = b.get("heading_after_onset_deg")
    away = None if dh is None or side is None else (-dh if side == "L" else dh)
    rel = None if away is None or f_away is None else away - f_away[side]
    dr = d.get("drive")
    clip = None
    if dr:
        vals = [abs(v) for row in dr[2:] for v in row]
        clip = sum(v >= 1.2 - 1e-12 for v in vals) / len(vals)
    body_runs.append({"id": cid, "group": d.get("group") or b.get("group"), "side": side, "wiring": d.get("wiring"),
                      "seed": d.get("seed"), "status": b.get("status"), "heading_after_onset_deg": dh,
                      "heading_away_deg": None if away is None else round(away, 3),
                      "heading_away_minus_floor_deg": None if rel is None else round(rel, 3),
                      "turns_away": None if rel is None else rel > 10, "forward_mm": b.get("forward_after_onset_mm"),
                      "forward_driven_mm": b.get("forward_mm"), "clipped_share": clip,
                      "drive_sha256": d.get("drive_sha256"), "dedup_of": d.get("dedup_of"),
                      "path_xy_mm": b.get("path_xy_mm"), "source": REL + "body/" + cid + ".json",
                      "drive_source": REL + "drives/" + cid + ".json", "brain_source": d.get("source") and "loom/" + d["source"],
                      "wall_s": b.get("wall_s")})
grp = lambda g: [r for r in body_runs if r["group"] == g and r["heading_away_deg"] is not None]
R_, D_, RND, GF = grp("real"), grp("scrambled"), grp("random"), grp("gf-off")
if f_away is not None and R_:
    def fl(rs): return mean(f_away[r["side"]] for r in rs)
    real_away, real_f = mean(r["heading_away_deg"] for r in R_), fl(R_)
    ret = None
    if D_ and abs(real_away - real_f) > 1e-9:
        ret = (mean(r["heading_away_deg"] for r in D_) - fl(D_)) / (real_away - real_f)
    rnd_ret = ((mean(r["heading_away_deg"] for r in RND) - fl(RND)) / (real_away - real_f)) if RND else None
    n_real, k_real = len(R_), sum(r["turns_away"] for r in R_)
    need = math.ceil(5 / 6 * n_real)
    clip_vals = [r["clipped_share"] for r in R_ + D_ if r["clipped_share"] is not None]
    clip_share = mean(clip_vals) if clip_vals else None
    rand_k = sum(r["turns_away"] for r in RND)
    if k_real < need:
        reading = "no reliable turn in our body test"
    elif ret is not None and ret <= 0.1:
        reading = "the turn away needs the wiring"
    elif ret is not None and ret >= 0.9:
        reading = "it does not"
    elif ret is not None:
        reading = "partly"
    else:
        reading = "no D runs"
    flags = []
    if clip_share is not None and clip_share > 0.10:
        flags.append("partly by construction")
    if len(RND) >= 6 and rand_k >= 5:
        flags.append("indistinguishable from random steering")
    body = {"floor": {"heading_after_onset_deg": f_dh, "away_L": f_away["L"], "away_R": f_away["R"],
                      "forward_mm": F.get("forward_after_onset_mm"), "forward_driven_mm": F.get("forward_mm"),
                      "walks_5mm": (F.get("forward_mm") or 0) >= 5},
            "n": {"real": n_real, "D": len(D_), "D_unique_drives": len({r["drive_sha256"] for r in D_}), "random": len(RND), "gfoff": len(GF)},
            "real_turns_away": k_real, "required": need, "D_turns_away": sum(r["turns_away"] for r in D_),
            "random_turns_away": rand_k, "gfoff_turns_away": sum(r["turns_away"] for r in GF),
            "side_swap": {"L_real_away": sum(r["turns_away"] for r in R_ if r["side"] == "L"), "L_real_n": sum(r["side"] == "L" for r in R_),
                          "R_real_away": sum(r["turns_away"] for r in R_ if r["side"] == "R"), "R_real_n": sum(r["side"] == "R" for r in R_)},
            "mean_away_deg": {"real": round(real_away, 3), "D": round(mean(r["heading_away_deg"] for r in D_), 3) if D_ else None,
                              "random": round(mean(r["heading_away_deg"] for r in RND), 3) if RND else None,
                              "floor_for_real": round(real_f, 3)},
            "sd_away_deg": {"real": round(pstdev([r["heading_away_deg"] for r in R_]), 3) if len(R_) > 1 else None,
                            "D": round(pstdev([r["heading_away_deg"] for r in D_]), 3) if len(D_) > 1 else None,
                            "random": round(pstdev([r["heading_away_deg"] for r in RND]), 3) if len(RND) > 1 else None},
            "retention_D": None if ret is None else round(ret, 4), "retention_random": None if rnd_ret is None else round(rnd_ret, 4),
            "clip_share_real_and_D": clip_share, "dedup": [{"id": r["id"], "of": r["dedup_of"]} for r in body_runs if r["dedup_of"]],
            "reading": reading, "flags": flags}
summary = {"checked_at": "2026-10-05", "preregistered_utc": "2026-10-05T09:42:47Z", "brain": brain, "body": body,
           "shuffle_checks": checks, "brain_rows": brain_rows, "body_runs": [{k: v for k, v in r.items() if k != "path_xy_mm"} for r in body_runs]}
atomic(RD / "summary.json", summary)
atomic(RD / "dn-table.json", {"checked_at": "2026-10-05", "note": "top 25 descending neurons by spikes per trial (1 s, so spikes = Hz)", "trials": dn_rows})
SENT = {
    "the turn away needs the wiring": "We ran the map. A looming shadow on one eye turned our simulated fly away in {X} of {N} runs. With the same map scrambled (same number of connections per neuron) it turned away in {Y} of {M}. Here the map's wiring, not the stimulus or our code, makes the turn.",
    "it does not": "We ran the map. A looming shadow turned our simulated fly away, but scrambled maps turned it away just as much. The turn comes from the stimulus and our mapping, not from the map's wiring.",
    "partly": "We ran the map. Scrambled maps kept about {Z}% of the turn away that the real map produced.",
    "no reliable turn in our body test": "We ran the map, and it did not reliably turn the fly away from a looming shadow in our test. {WHY}",
}
reading = body.get("reading") if body else None
sentence = None
if reading in SENT:
    sentence = SENT[reading].format(X=body["real_turns_away"], N=body["n"]["real"], Y=body["D_turns_away"], M=body["n"]["D"],
                                    Z=None if body["retention_D"] is None else round(100 * body["retention_D"]),
                                    WHY="The body turns were below 10°." if gate_pass else "The brain gate failed.")
chart = {"checked_at": "2026-10-05", "preregistered_utc": "2026-10-05T09:42:47Z",
         "model": "Shiu et al. 2024 LIF (philshiu/Drosophila_brain_model@91bdd1e7), FlyWire v783, Brian2 2.9.0 numpy, dt 0.1 ms, 1 s per trial; LOOM-L = 54 LC4 + 108 LPLC2, LOOM-R = 50 LC4 + 102 LPLC2, Poisson 80 Hz each",
         "body": "FlyGym 2.1.0 HybridTurningController, flat terrain, 0.2 s settle, 0.2 s at (0.8, 0.8) before onset, then 1.0 s of brain bins",
         "mapping": "OUR MAPPING (hand-made, not fitted): drive_s = clip(0.8 - m_s - 0.5*(a_s - a_o), -1.2, 1.2); a = DNa02 rate / 100 Hz, m = mean MDN rate on side s / 100 Hz; 0.8 = external walking drive, not the brain",
         "caveat": CAVEAT, "reading": reading, "flags": body.get("flags") if body else None, "sentence": sentence,
         "brain_gate": brain["gate_result"],
         "clip_share": body.get("clip_share_real_and_D") if body else None, "retention": body.get("retention_D") if body else None,
         "heading_sign": "heading_away_deg > 0 = turned away from the looming side; FlyGym heading is + to the left",
         "brain_rows": [{k: r.get(k) for k in ("trial", "side", "wiring", "seed", "gf_l", "gf_r", "gf_mean", "dna02_l", "dna02_r", "li", "mdn_mean", "dnp09_l", "dnp09_r", "whole_brain_spikes", "active_neurons", "selectivity", "gate_pass", "source")} for r in ok],
         "body_runs": [{k: r.get(k) for k in ("id", "group", "side", "wiring", "seed", "heading_away_deg", "heading_away_minus_floor_deg", "turns_away", "forward_mm", "clipped_share", "drive_sha256", "dedup_of", "path_xy_mm", "source")} for r in body_runs],
         "fly67_reference": FLY67, "summary_source": REL + "summary.json"}
atomic(CHART, chart)
print(json.dumps({"brain": {k: brain[k] for k in ("gate_result", "n", "real_mean", "D_mean", "fly67_rule_pass", "shuffle_checks_all_pass")} | {"ret": brain.get("retention_display")}, "body": {k: v for k, v in body.items() if k not in ("dedup",)}, "sentence": sentence}, indent=1))

# tests/test_loom.py
"""Run 8 looming worker checks (no FlyWire data needed except ../step0.json). Licence: MIT."""
import json
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent))
import run_loom  # noqa: E402


def test_stim_sets_match_step0():
    s0 = json.load(open(HERE.parent.parent / "step0.json"))
    L, R = run_loom.stim_ids("LOOML", None), run_loom.stim_ids("LOOMR", None)
    assert len(L) == s0["counts"]["L"]["total"] == 162
    assert len(R) == s0["counts"]["R"]["total"] == 152
    assert not set(L) & set(R)


def test_trial_names():
    assert run_loom.trial_name("real", 0, None, "LOOML") == "LOOML-real-0.json"
    assert run_loom.trial_name("D", 3, None, "LOOMR") == "LOOMR-D-3.json"
    assert run_loom.trial_name("G", 4, 1.5, "LOOML") == "LOOML-G-4-g1.5.json"
    assert run_loom.trial_name("real", 0, None, "LOOML", "GF") == "LOOML-real-0-gfoff.json"


def test_silence_zeroes_only_outgoing_edges():
    pre = np.array([0, 1, 2, 1, 3], dtype=np.int32)
    post = np.array([1, 2, 1, 0, 1], dtype=np.int32)
    cnt = np.array([5, -3, 2, 4, 7], dtype=np.int32)
    out = run_loom.silence_counts(pre, cnt, [1])
    assert out.tolist() == [5, 0, 2, 0, 7]
    assert cnt.tolist() == [5, -3, 2, 4, 7]  # input unchanged
    assert run_loom.silence_counts(pre, cnt, []) is cnt


def test_gf_ids_are_dnp01():
    s0 = json.load(open(HERE.parent.parent / "step0.json"))
    assert sorted(r["id"] for r in s0["readouts"]["DNp01"]) == sorted(run_loom.GF_IDS)

# tests/test_nulls.py
"""Synthetic-graph tests for nulls.py (no FlyWire data; seconds). Run: python -m pytest -q tests/test_nulls.py"""
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import nulls  # noqa: E402

N = 200


def graph(seed=0, e=5000):
    rng = np.random.default_rng(seed)
    pre = rng.integers(0, N, e).astype(np.int32)
    post = rng.integers(0, N, e).astype(np.int32)
    sign = np.where(rng.random(N) < 0.6, 1, -1)  # one sign per neuron (Dale's law)
    cnt = (rng.integers(1, 50, e) * sign[pre]).astype(np.int32)
    return pre, post, cnt


def sides():
    inp = np.zeros(N, bool)
    inp[:20] = True
    out = np.zeros(N, bool)
    out[-15:] = True
    return inp, out


def test_degree():
    r = graph()
    n = nulls.null_degree(*r, seed=3)
    c = nulls.check("D", r, n, N)
    assert all(c.values()), c
    assert not np.array_equal(n[1], r[1])


def test_degree_matches_bootstrap_load_edges():
    r = graph()
    p = r[1].copy()
    np.random.default_rng(1).shuffle(p)
    assert np.array_equal(nulls.null_degree(*r, seed=1)[1], p)


def test_weight():
    r = graph()
    n = nulls.null_weight(*r, seed=101)
    c = nulls.check("W", r, n, N)
    assert all(c.values()), c
    assert not np.array_equal(n[2], r[2])


def test_sign():
    r = graph()
    n = nulls.null_sign(*r, seed=201, n=N)
    c = nulls.check("S", r, n, N)
    assert all(c.values()), c
    assert not np.array_equal(n[2], r[2])


def test_boundary():
    r = graph()
    inp, out = sides()
    n = nulls.null_boundary(*r, seed=301, input_side=inp, output_side=out)
    c = nulls.check("B", r, n, N, input_side=inp, output_side=out)
    assert all(c.values()), c
    assert not np.array_equal(n[1], r[1])


def test_stimreadout():
    r = graph()
    n = nulls.null_stimreadout(*r, seed=401, stim_idx=[0, 1, 2], readout_idx=N - 1, n=N)
    c = nulls.check("B'", r, n, N, stim_idx=[0, 1, 2], readout_idx=N - 1)
    assert all(c.values()), c


def test_checks_catch_violations():
    r = graph()
    inp, out = sides()
    bad = (r[0], np.random.default_rng(9).permutation(r[1]), r[2])  # full shuffle moves boundary edges
    assert not all(nulls.check("B", r, bad, N, input_side=inp, output_side=out).values())
    badw = (r[0], r[1], -r[2])  # flips every sign
    assert not nulls.check("W", r, badw, N)["edge_signs_kept"]
    mixed = r[2].copy()
    mixed[0] = -mixed[0]  # one edge of a neuron gets the other sign
    assert not nulls.check("S", r, (r[0], r[1], mixed), N)["one_sign_per_presynaptic_neuron"]
    shorter = (r[0][:-1], r[1][:-1], r[2][:-1])
    assert not nulls.check("D", r, shorter, N)["n_edges_kept"]


def test_structure_stats():
    pre = np.array([0, 0, 1, 2], np.int32)
    post = np.array([1, 1, 1, 2], np.int32)
    s = nulls.structure_stats(pre, post, 3)
    assert s == {"self_loops": 2, "duplicate_pairs": 1}