#!/usr/bin/env python3
"""
Check of the slow-drift hold (H2) before any lock (hold_check_plan.md), October 13, 2026. On old starts 13 to 18 with
new moments: after a first kick to target at T - H, each group is corrected every 5 units by a set fraction lambda of
the gap between its own 25-unit average and its target, (1 +- 0.25) x the untouched run's 25-unit average; fast swings
run freely. lambda is chosen on starts 13 to 15 by the plan's rule and confirmed on 16 to 18. Only the untouched first
copy is used: no gap is carried and no growth is measured. Nothing here is locked.

A sandbox script kept with the project. Nothing here is served by the site.

Usage (each command runs in the foreground and ends by itself; nothing is left running)
    python3 -B hold_check.py dryrun   # the machinery on old start 1 for 700 units; writes only to /tmp
    python3 -B hold_check.py run      # one piece of whatever the plan's rule says comes next, 7 tasks at a time
    python3 -B hold_check.py report   # once the rule has nothing left to run: the criteria and the outputs

Each is run with a hard limit: timeout -k 10 285 python3 -B hold_check.py <command>
"""
import json
import shutil
import sys
import time
from pathlib import Path

sys.dont_write_bytecode = True
import numpy as np  # noqa: E402

from prelock_common import (HERE, WORK, SEED, N, K, EPS, DT, EVERY, ct, nt, stamp, shown,  # noqa: E402
                            sha256, write_json, save_npz, run_piece, guard, lock_changes, folder_fingerprints)
from prelock_calibration import calibrate, reproduced  # noqa: E402  (the base run in pieces; V1 as preregistered)
from prelock_feasibility import (one_sided, summarise, natural_spreads,  # noqa: E402
                                 moment_table as first_check_moments)

PARTS = WORK / "hold_check"
nt.LYAP_PARTS = PARTS / "unused"       # importing the cause test points this at another folder; nothing here uses it

STARTS = tuple(range(12, 18))          # old starts 13 to 18 (seeds 12 to 17)
TUNE, CONFIRM = STARTS[:3], STARTS[3:]
MOMENTS, SPACING, FIRST = 10, 350, 2000
STEP = int(EVERY)                      # a kick every 5 units
HOLDS = (500, 100)
SIZE = 0.25
WINDOW = 25                            # the target base and the branch's own average: the past 25 units
BLOCK = 100                            # the registered blocks
LAMBDAS = (0.05, 0.10, 0.20)           # tried one at a time, in this order
AFTER = 500                            # untouched after T, in the confirmation
HALF_BAND, SHAM_BAND, REACHED = 0.025, (0.95, 1.05), 1e-5
ARMS = ("up", "down", "sham")
CRITERIA = dict(most_missing=1, fewest_kept=5, achieved=(0.75, 1.25), achieved_mean=(0.85, 1.15), share=0.9,
                floor=0.25, block_mean=(0.75, 1.25), sham=(0.8, 1.2), sham_history=(-0.25, 0.25),
                draft_block=(0.5, 1.5), first_check_block=(0.5, 1.5))
WINDOWS = dict(early=(0, 200), later=(200, 500))
BUDGET, START_BY = 190.0, 120.0
RUNLOG = HERE / "hold_check_runlog.json"
PLAN = HERE / "hold_check_plan.md"
OUT_TXT, OUT_JSON = HERE / "hold_check_results.txt", HERE / "hold_check_results.json"
FIRST_CHECK = HERE / "feasibility_results.json"
DRY = Path("/tmp/hold-check-dryrun")


def offsets():
    """o for each start: 5 x an integer from 0 to 69 drawn from (SEED, 5, seed); if it equals the first check's o for
    that start, the next draw from the same generator."""
    old = {s: T - FIRST for s, m, T in first_check_moments() if m == 0}
    out = {}
    for s in STARTS:
        draw = np.random.default_rng((SEED, 5, s))
        o = STEP * int(draw.integers(0, 70))
        while o == old[s]:
            o = STEP * int(draw.integers(0, 70))
        out[s] = o
    return out


def moment_table(starts=STARTS):
    o = offsets()
    return [(s, m, FIRST + o[s] + SPACING * m) for s in starts for m in range(MOMENTS)]


def trailing(R, width=WINDOW):
    """The mean of R over t - width + 1, ..., t for every t (nan before there are `width` values), (group, time)."""
    cum = np.concatenate([np.zeros((R.shape[0], 1)), np.cumsum(R, -1)], -1)
    out = np.full(R.shape, np.nan)
    out[:, width - 1:] = (cum[:, width:] - cum[:, :-width]) / width
    return out


# ---------------------------------------------------------------------------------------------------------------------
# the sham and the hold

def mixed_sham(along, draw, sizes, first):
    """The sham's move for one group at one kick: the cause test's paired signs (ct.paired_signs, its draws
    unchanged), with the pairs (neighbours in the ranking by sin^2) taking up's and down's sizes in turn down the
    ranking, starting with sizes[first]. A leftover oscillator (an odd count) takes their root mean square."""
    signs = ct.paired_signs(along, draw)
    order = np.argsort(along * along, kind="stable")                  # the ranking paired_signs uses
    count = len(order)
    pairs = count // 2
    size = np.empty(count)
    turn = (np.arange(2 * pairs) // 2 + first) % 2
    size[order[:2 * pairs]] = np.where(turn == 0, sizes[0], sizes[1])
    if count % 2:
        size[order[-1]] = np.sqrt((sizes[0] ** 2 + sizes[1] ** 2) / 2)
    return signs * size * along


def h2_hold(u0, f0, T, H, R_none, base, key, lam, size=SIZE, after=0, first_change=True):
    """One hold from the untouched first copy's state at T - H (u0: (group, pulled), f0: (group, free)): up, down and
    sham side by side (rows 0, 1, 2), a kick at every multiple of 5 in [T - H, T], then untouched to T + after.
    R_none: the untouched run's R_g every time unit from t = 0; base: its trailing 25-unit mean, both (group, time).
    The first kick sets up's and down's R_g to (1 +- size) base_g(T - H); each later kick aims at R_g x (target_g /
    b_g)^lam, b_g being the branch's own mean of R_g over the past 25 units since T - H (the value just before the kick
    last). The sham: mixed_sham() with up's and down's sizes, signs drawn from `key`, group A before B. If a kick
    cannot be solved or is outside its band, the hold stops there and is not carried out; nothing is clipped or
    replaced. first_change=False (the dry run's check only): the first kick asks for no change."""
    slow, stepper = nt.model(N, K, EPS)
    count, kicks = u0.shape[-1], H // STEP + 1
    U, F = np.repeat(u0[None], 3, 0), np.repeat(f0[None], 3, 0)
    step, none = stepper(DT), np.zeros(3)
    draw = np.random.default_rng(key)
    per_unit = int(round(1.0 / DT))
    nan = np.nan
    rec = dict(R=np.full((3, 2, H + after + 1), nan),                 # R_g every unit from T - H (after the kick at kicks)
               before=np.full((3, 2, kicks), nan),                    # R_g just before each kick
               own=np.full((2, 2, kicks), nan), goal=np.full((2, 2, kicks), nan), amp=np.full((2, 2, kicks), nan),
               status=np.zeros((2, 2, kicks), int), ratio=np.full((3, 2, kicks), nan),
               inband=np.zeros((3, 2, kicks), bool), sumsq=np.full((3, kicks), nan), sumabs=np.full((3, kicks), nan),
               largest=np.full((3, kicks), nan), kick_t=T - H + STEP * np.arange(kicks), kicks=kicks, done=0,
               carried=False, reason="", lam=float(lam))
    factor = np.array([1 + size, 1 - size])
    for k in range(kicks):
        t = T - H + STEP * k
        z = ct.harmonic(U, F)
        R, psi = np.abs(z), np.angle(z)
        rec["before"][:, :, k] = R
        along = np.sin(psi[..., None] - np.angle(U))                  # sin(psi_g - theta_j), (row, group, pulled)
        target = factor[:, None] * base[None, :, t]                   # (arm, group)
        if k == 0:
            goal = target if first_change else R[:2].copy()
        else:
            low = max(0, STEP * k - (WINDOW - 1))
            own = (rec["R"][:2, :, low:STEP * k].sum(-1) + R[:2]) / (STEP * k - low + 1)
            rec["own"][:, :, k] = own
            goal = R[:2] * (target / own) ** lam
        rec["goal"][:, :, k] = goal
        for a in range(2):
            for g in range(2):
                rec["amp"][a, g, k], rec["status"][a, g, k] = one_sided(U[a, g], along[a, g], z[a, g], goal[a, g])
        if rec["status"][:, :, k].any():
            bad = [f"{ARMS[a]} {'AB'[g]} {'infeasible' if rec['status'][a, g, k] == 1 else 'unsolved'}"
                   for a in range(2) for g in range(2) if rec["status"][a, g, k]]
            rec["reason"] = f"kick {k + 1} at t = {t}: " + ", ".join(bad)
            return rec
        dtheta = np.zeros((3, 2, count))
        for g in range(2):
            dtheta[0, g] = rec["amp"][0, g, k] * along[0, g]
            dtheta[1, g] = rec["amp"][1, g, k] * along[1, g]
            dtheta[2, g] = mixed_sham(along[2, g], draw, (abs(rec["amp"][0, g, k]), abs(rec["amp"][1, g, k])), k % 2)
        U *= np.exp(1j * dtheta)
        now = np.abs(ct.harmonic(U, F))
        rec["ratio"][:2, :, k] = now[:2] / goal
        rec["ratio"][2, :, k] = now[2] / R[2]
        if k == 0 and first_change:
            for a in range(2):
                landed = now[a] / base[:, t]
                rec["inband"][a, :, k] = (factor[a] - HALF_BAND <= landed) & (landed <= factor[a] + HALF_BAND)
        else:
            rec["inband"][:2, :, k] = np.abs(rec["ratio"][:2, :, k] - 1) <= REACHED
        rec["inband"][2, :, k] = (SHAM_BAND[0] <= rec["ratio"][2, :, k]) & (rec["ratio"][2, :, k] <= SHAM_BAND[1])
        rec["sumsq"][:, k] = np.square(dtheta).sum((-2, -1))
        rec["sumabs"][:, k] = np.abs(dtheta).sum((-2, -1))
        rec["largest"][:, k] = np.abs(dtheta).max((-2, -1))
        rec["R"][:, :, STEP * k] = now
        rec["done"] = k + 1
        if not rec["inband"][:, :, k].all():
            bad = [f"{arm} {'AB'[g]} at {rec['ratio'][a, g, k]:.6f}" for a, arm in enumerate(ARMS) for g in range(2)
                   if not rec["inband"][a, g, k]]
            rec["reason"] = f"kick {k + 1} at t = {t} outside its band: " + ", ".join(bad)
            return rec
        units = STEP - 1 if k < kicks - 1 else after                  # the 5th unit is the next kick's "before"
        for unit in range(1, units + 1):
            for _ in range(per_unit):
                step(U, F, none)
            rec["R"][:, :, STEP * k + unit] = np.abs(ct.harmonic(U, F))
        if k < kicks - 1:
            for _ in range(per_unit):
                step(U, F, none)
    rec["carried"] = True
    return rec


# ---------------------------------------------------------------------------------------------------------------------
# the stages and the pieces

def base_times(s):
    return sorted({t for ss, m, T in moment_table() if ss == s for t in (T - 500, T - 100)})


def base_tasks():
    return [dict(name=f"base{s + 1:02d}", kind="base", seed=s) for s in STARTS]


def tune_tasks(i):
    return [dict(name=f"t{i}_{s + 1:02d}_{m}", kind="tune", seed=s, moment=m, T=T, lam=LAMBDAS[i], after=0)
            for s, m, T in moment_table(TUNE)]


def confirm_tasks(i):
    rows = moment_table(CONFIRM)
    return ([dict(name=f"c500_{s + 1:02d}_{m}", kind="c500", seed=s, moment=m, T=T, lam=LAMBDAS[i], after=AFTER)
             for s, m, T in rows] +
            [dict(name=f"c100_{s + 1:02d}_{m}", kind="c100", seed=s, moment=m, T=T, lam=LAMBDAS[i], after=AFTER)
             for s, m, T in rows])                                    # the longer tasks first


def base_work(task, deadline):
    """The base run of one start (ct.base_run step for step, with no fresh copies), in pieces if it must, keeping the
    untouched first copy's state at T - 500 and T - 100 of every moment."""
    s = task["seed"]
    times = base_times(s)
    state, part = PARTS / f"{task['name']}.state.npz", PARTS / f"{task['name']}.npz"
    out = calibrate(s, deadline, state, moments=np.zeros(0, int), keep=times)
    if out is None:
        return dict(status="saved")
    if len(out["kept_u"]) != len(times):
        raise RuntimeError("a kept state is missing")
    same = reproduced(out, s)
    save_npz(part, start=s + 1, times=np.array(times), z=out["z"], same=np.array([same["coherence"], same["gap"],
             same["settling"]]), us=out["kept_u"], uf=out["kept_f"], spent=out["spent"], pieces=out["pieces"])
    state.unlink(missing_ok=True)
    return dict(status="done", sha256=sha256(part), error="" if all(same.values()) else f"reproduced {same}")


def moment_work(task):
    s, m, T, lam, after = task["seed"], task["moment"], task["T"], task["lam"], task["after"]
    with np.load(PARTS / f"base{s + 1:02d}.npz") as b:
        times = [int(t) for t in b["times"]]
        R_none, us, uf = np.abs(b["z"]), b["us"], b["uf"]
    base = trailing(R_none)
    holds = HOLDS if task["kind"] == "tune" else (int(task["kind"][1:]),)
    out = {}
    for H in holds:
        j = times.index(T - H)
        rec = h2_hold(us[j].copy(), uf[j].copy(), T, H, R_none, base, (SEED, 6, s, m, H), lam, after=after)
        out.update({f"h{H}_{key}": value for key, value in rec.items()})
    part = PARTS / f"{task['name']}.npz"
    save_npz(part, start=s + 1, moment=m, T=T, lam=lam, after=after, **out)
    return dict(status="done", sha256=sha256(part))


def work(task, deadline):
    return base_work(task, deadline) if task["kind"] == "base" else moment_work(task)


def job(args):
    return guard(work, args)


def done(task):
    return (PARTS / f"{task['name']}.npz").exists()


def log_decision(text, **fields):
    log = json.loads(RUNLOG.read_text())
    log.setdefault("decisions", [])
    if not any(d["text"] == text for d in log["decisions"]):
        log["decisions"].append(dict(made=stamp(), text=text, **fields))
        write_json(RUNLOG, log)
        print(text)


def next_stage():
    """The plan's rule: the base runs; then the candidates for lambda in order on starts 13 to 15, each run whole and
    judged on the 500-unit hold (F1, F2, F3a, F3b, F4a, F4c), stopping at the first that meets them all; then the
    confirmation on starts 16 to 18 with it. Returns (stage, unfinished tasks, chosen index or None)."""
    todo = [t for t in base_tasks() if not done(t)]
    if todo:
        return "base runs", todo, None
    for i, lam in enumerate(LAMBDAS):
        todo = [t for t in tune_tasks(i) if not done(t)]
        if todo:
            return f"choosing lambda: candidate {lam}", todo, None
        res = judge(i, TUNE, 500, 0)
        if res["met"]:
            log_decision(f"lambda = {lam} chosen: on starts 13 to 15 it meets every criterion for the 500-unit hold",
                         lam=lam, criteria={key: res[key] for key in CRITERION_KEYS})
            todo = [t for t in confirm_tasks(i) if not done(t)]
            return ("confirmation", todo, i) if todo else ("report", [], i)
        log_decision(f"lambda = {lam} not chosen: on starts 13 to 15 it misses "
                     + ", ".join(key for key in CRITERION_KEYS if not res[key]) + " for the 500-unit hold",
                     lam=lam, criteria={key: res[key] for key in CRITERION_KEYS})
    log_decision("no candidate meets the criteria: H2 is not feasible at +-25% as specified; no confirmation")
    return "report", [], None


def run():
    PARTS.mkdir(parents=True, exist_ok=True)
    stage, todo, chosen = next_stage()
    if not todo:
        print("Nothing left to run: run report.")
        return
    print(f"stage: {stage}")
    run_piece(RUNLOG, todo, job, BUDGET, START_BY)


# ---------------------------------------------------------------------------------------------------------------------
# the numbers of one hold, and the criteria over a set of moments

def numbers(rec, R_none, base, T, H, size=SIZE, after=0):
    c = rec["R"].mean(1)                                               # (arm, time from T - H)
    none = R_none[:, T - H:T + after + 1]
    c0 = none.mean(0)
    m = base[:, T - H:T].mean(0)                                       # the target base, (time)
    M, M0 = c[:, :H].mean(-1), c0[:H].mean()
    A = np.array([M[0] - M0, M0 - M[1], M[2] - M0]) / (size * M0)
    nb = H // BLOCK
    cb, mb = c[:, :H].reshape(3, nb, BLOCK).mean(-1), m.reshape(nb, BLOCK).mean(-1)
    u, d = (cb[0] / mb - 1) / size, (1 - cb[1] / mb) / size
    q, q0 = c[:, :H].reshape(3, -1, 25).mean(-1), c0[:H].reshape(-1, 25).mean(-1)
    first_check = np.stack([q[0] - q0, q0 - q[1]]) / (size * q0)      # the first check's 25-unit block fractions
    kick_t = rec["kick_t"]
    factor = np.array([1 + size, 1 - size])[:, None, None]
    pre = rec["before"][:2, :, 1:] / (factor * base[None, :, kick_t[1:]]) - 1
    sumsq, sumabs = rec["sumsq"], rec["sumabs"]
    arms_sq, arms_abs = sumsq[:2].mean(0), sumabs[:2].mean(0)
    moved = arms_sq > 0
    left = {}
    for name, (a, b) in WINDOWS.items():
        if b > after:
            continue
        span = slice(H + a + 1, H + b + 1)
        level = c0[span].mean()
        left[name] = (c[:, span].mean(-1) - level) / (size * level) * np.array([1.0, -1.0, 1.0])
    return dict(M=M, M0=M0, A=A, dM=(M[0] - M[1]) / 2, integral=H * (M - M0), u=u, d=d, s=(u + d) / 2,
                first_check=first_check, drift=pre.mean((1, 2)), spread=pre.std((1, 2)),
                p5=np.percentile(pre, 5, axis=(1, 2)), p95=np.percentile(pre, 95, axis=(1, 2)),
                sumsq=sumsq.sum(-1), sumabs=sumabs.sum(-1), largest=rec["largest"].max(-1),
                first_share=sumsq[:, 0] / sumsq.sum(-1), later=np.median(sumsq[:, 1:], -1),
                sham_sq=sumsq[2].sum() / arms_sq.sum(), sham_abs=sumabs[2].sum() / arms_abs.sum(),
                kick_sq=sumsq[2][moved] / arms_sq[moved], kick_abs=sumabs[2][moved] / arms_abs[moved],
                same_times=bool(np.array_equal(sumsq[2] > 0, moved)), left=left)


CRITERION_KEYS = ("F1", "F2", "F3a", "F3b", "F4a", "F4c")


def part_path(stage, s, m, H):
    return PARTS / (f"c{H}_{s + 1:02d}_{m}.npz" if stage == "confirm" else f"t{stage}_{s + 1:02d}_{m}.npz")


def load_base(s):
    with np.load(PARTS / f"base{s + 1:02d}.npz") as b:
        R = np.abs(b["z"])
        return dict(R=R, base=trailing(R), same=b["same"].copy(), seconds=float(b["spent"]))


def load_rec(stage, s, m, H):
    with np.load(part_path(stage, s, m, H)) as p:
        prefix = f"h{H}_"
        return {key[len(prefix):]: p[key][()] if p[key].ndim == 0 else p[key].copy()
                for key in p.files if key.startswith(prefix)}


def judge(stage, starts, H, after, bases=None, keep=False):
    """The plan's criteria for one hold length over the moments of `starts` in one stage (a candidate's index, or
    "confirm"). keep: also return the per-moment numbers."""
    bases = bases or {s: load_base(s) for s in starts}
    table = moment_table(starts)
    recs = {(s, m): load_rec(stage, s, m, H) for s, m, T in table}
    carried = np.array([[bool(recs[(s, m)]["carried"]) for m in range(MOMENTS)] for s in starts])
    missing, fewest = int(carried.size - carried.sum()), int(carried.sum(1).min())
    res = dict(moments=len(table), carried=int(carried.sum()), missing=missing, fewest_per_start=fewest,
               reasons=[f"start {s + 1}, moment {m} (T = {T}): {recs[(s, m)]['reason']}" for s, m, T in table
                        if not recs[(s, m)]["carried"]],
               F1=bool(missing <= CRITERIA["most_missing"] and fewest >= CRITERIA["fewest_kept"]))
    rows = [(s, m, T) for s, m, T in table if recs[(s, m)]["carried"]]
    nums = [numbers(recs[(s, m)], bases[s]["R"], bases[s]["base"], T, H, after=after) for s, m, T in rows]
    if not nums:
        res.update(F2=False, F3a=False, F3b=False, F4a=False, F4c=False, F4b=False, met=False)
        return (res, nums, rows) if keep else res
    A = np.array([x["A"] for x in nums])
    lo, hi = CRITERIA["achieved"]
    mlo, mhi = CRITERIA["achieved_mean"]
    inside = (A[:, 0] >= lo) & (A[:, 0] <= hi) & (A[:, 1] >= lo) & (A[:, 1] <= hi)
    F2 = bool(inside.mean() >= CRITERIA["share"] and mlo <= A[:, 0].mean() <= mhi and mlo <= A[:, 1].mean() <= mhi)
    u, d, sep = (np.array([x[key] for x in nums]) for key in ("u", "d", "s"))   # (moment, block)
    sustained = (sep >= CRITERIA["floor"]).all(1)
    F3a = bool(sustained.mean() >= CRITERIA["share"])
    blo, bhi = CRITERIA["block_mean"]
    F3b = bool(((u.mean(0) >= blo) & (u.mean(0) <= bhi) & (d.mean(0) >= blo) & (d.mean(0) <= bhi)).all())
    slo, shi = CRITERIA["sham"]
    sham_sq, sham_abs = np.array([x["sham_sq"] for x in nums]), np.array([x["sham_abs"] for x in nums])
    same_times = all(x["same_times"] for x in nums)
    F4a = bool(((sham_sq >= slo) & (sham_sq <= shi)).all())
    F4c = bool(((sham_abs >= slo) & (sham_abs <= shi)).all() and same_times)
    hlo, hhi = CRITERIA["sham_history"]
    F4b = bool(hlo <= A[:, 2].mean() <= hhi)
    dlo, dhi = CRITERIA["draft_block"]
    draft = ((u >= dlo) & (u <= dhi) & (d >= dlo) & (d <= dhi)).all(1)
    flo, fhi = CRITERIA["first_check_block"]
    first = np.array([bool(((x["first_check"] >= flo) & (x["first_check"] <= fhi)).all()) for x in nums])
    kick_sq = np.concatenate([x["kick_sq"] for x in nums])
    kick_abs = np.concatenate([x["kick_abs"] for x in nums])
    res.update(
        F2=F2, F3a=F3a, F3b=F3b, F4a=F4a, F4c=F4c, F4b=F4b, met=bool(res["F1"] and F2 and F3a and F3b and F4a and F4c),
        achieved={arm: summarise(A[:, a]) for a, arm in enumerate(ARMS)}, inside_share=float(inside.mean()),
        sustained_share=float(sustained.mean()), lowest_difference=float(sep.min()),
        block_mean_up=u.mean(0).tolist(), block_mean_down=d.mean(0).tolist(), block_mean_difference=sep.mean(0).tolist(),
        block_sd_up=u.std(0, ddof=1).tolist() if len(u) > 1 else [], block_sd_down=d.std(0, ddof=1).tolist() if len(d) > 1 else [],
        block_sd_difference=sep.std(0, ddof=1).tolist() if len(sep) > 1 else [],
        draft_share=float(draft.mean()), first_check_share=float(first.mean()),
        dM=summarise([x["dM"] for x in nums]), integral={arm: summarise([x["integral"][a] for x in nums])
                                                         for a, arm in enumerate(ARMS)},
        drift={arm: summarise([x["drift"][a] for x in nums]) for a, arm in enumerate(ARMS[:2])},
        spread={arm: summarise([x["spread"][a] for x in nums]) for a, arm in enumerate(ARMS[:2])},
        pre_5_95={arm: [float(np.mean([x["p5"][a] for x in nums])), float(np.mean([x["p95"][a] for x in nums]))]
                  for a, arm in enumerate(ARMS[:2])},
        kicks=H // STEP + 1, same_times=bool(same_times),
        sumsq={arm: summarise([x["sumsq"][a] for x in nums]) for a, arm in enumerate(ARMS)},
        sumabs={arm: summarise([x["sumabs"][a] for x in nums]) for a, arm in enumerate(ARMS)},
        largest={arm: float(max(x["largest"][a] for x in nums)) for a, arm in enumerate(ARMS)},
        first_share={arm: summarise([x["first_share"][a] for x in nums]) for a, arm in enumerate(ARMS)},
        later={arm: summarise([x["later"][a] for x in nums]) for a, arm in enumerate(ARMS)},
        sham_sq=summarise(sham_sq), sham_abs=summarise(sham_abs),
        kick_sq=dict(median=float(np.median(kick_sq)), low=float(kick_sq.min()), high=float(kick_sq.max())),
        kick_abs=dict(median=float(np.median(kick_abs)), low=float(kick_abs.min()), high=float(kick_abs.max())),
        left={name: {arm: summarise([x["left"][name][a] for x in nums]) for a, arm in enumerate(ARMS)}
              for name in WINDOWS if all(name in x["left"] for x in nums)})
    return (res, nums, rows) if keep else res


# ---------------------------------------------------------------------------------------------------------------------
# the dry run

def dryrun():
    """Old start 1 for 700 units, one moment at T = 600 (holds from t = 100 and 500), 50 units after release: the base
    run in pieces against ct.base_run and the saved run (exactly); a hold that asks for no change (lambda = 0, no
    first change), which must follow the untouched run; the target base and the branch's own average against a direct
    computation; then both holds at +-25% with lambda = 0.10, and the sham's matching."""
    if DRY.exists():
        shutil.rmtree(DRY)
    DRY.mkdir(parents=True)
    s, length, T, after, lam = 0, 700.0, 600, 50, 0.10
    times = [T - 500, T - 100]
    clock = time.time()
    res = ct.base_run(s, keep=times, run=length)
    base_seconds = time.time() - clock
    stops, out = 0, None
    while out is None:
        out = calibrate(s, time.time() + 5.0, DRY / "base.state.npz", moments=np.zeros(0, int), run=length, keep=times)
        stops += out is None
    same = reproduced(out, s, length)
    kept = bool(len(out["kept_u"]) == len(times) and all(
        np.array_equal(out["kept_u"][j], res["states"][j][0][0]) and np.array_equal(out["kept_f"][j], res["states"][j][1][0])
        for j in range(len(times))))
    agree = bool(np.array_equal(out["z"], res["z"][0]) and np.array_equal(out["logs"], res["logs"][0]))
    R_none = np.abs(out["z"])
    base = trailing(R_none)
    direct = np.array([[R_none[g, t - WINDOW + 1:t + 1].mean() for t in range(WINDOW - 1, R_none.shape[1])]
                       for g in range(2)])
    base_error = float(np.abs(base[:, WINDOW - 1:] - direct).max())
    state = {t: (out["kept_u"][j].copy(), out["kept_f"][j].copy()) for j, t in enumerate(times)}
    lines = [f"Dry run, {shown(stamp())}: old start 1 (seed 0), 700 units, one moment at T = 600, 50 units after, "
             f"lambda = {lam}",
             f"  base run in pieces: stopped {stops} times, {out['spent']:.1f} s in all (ct.base_run in one go "
             f"{base_seconds:.1f} s); reproduces the saved run exactly: {same}; the same coherence, gap and kept "
             f"states as ct.base_run to the last digit: {agree and kept}",
             f"  target base (trailing 25-unit mean) against a direct mean at every t: largest difference "
             f"{base_error:.2g} (needs at most 1e-15)"]
    zero = h2_hold(*state[T - 100], T, 100, R_none, base, (SEED, 99, 0), 0.0, size=0.0, after=after,
                   first_change=False)
    follow = float(np.abs(zero["R"] - R_none[None, :, T - 100:T + after + 1]).max())
    lines.append(f"  no change asked for (lambda = 0, no first change), 100-unit hold: carried {zero['carried']}; "
                 f"largest |R - R_none| over the hold and after {follow:.3g}; largest kick {np.nanmax(zero['largest']):.3g}")
    ok = all(same.values()) and agree and kept and base_error <= 1e-15 and zero["carried"] and follow <= 1e-9
    for H in HOLDS:
        clock = time.time()
        rec = h2_hold(*state[T - H], T, H, R_none, base, (SEED, 99, H), lam, after=after)
        seconds = time.time() - clock
        first = float(np.abs(rec["before"][:, :, 0] - R_none[None, :, T - H]).max())
        line = (f"  {H}-unit hold at +-25%: {seconds:.1f} s; carried {rec['carried']} {rec['reason']}; all three rows "
                f"start on the untouched run (largest difference {first:.2g}, needs at most 1e-15)")
        ok = ok and rec["carried"] and first <= 1e-15
        if rec["carried"]:
            k = rec["kicks"] - 1                                       # the branch's own average at the last kick
            t = T - H + STEP * k
            direct_own = np.concatenate([rec["R"][:2, :, STEP * k - WINDOW + 1:STEP * k],
                                         rec["before"][:2, :, k][..., None]], -1).mean(-1)
            own_error = float(np.abs(direct_own - rec["own"][:, :, k]).max())
            goal_error = float(np.abs(rec["goal"][:, :, k] - rec["before"][:2, :, k] * (
                np.array([1.25, 0.75])[:, None] * base[None, :, t] / direct_own) ** lam).max())
            ok = ok and own_error <= 1e-15 and goal_error <= 1e-15
            x = numbers(rec, R_none, base, T, H, after=after)
            line += (f"\n    the branch's own average and the goal at the last kick against a direct computation: "
                     f"largest differences {own_error:.2g} and {goal_error:.2g} (need at most 1e-15)"
                     f"\n    achieved fraction up {x['A'][0]:.3f}, down {x['A'][1]:.3f}, sham {x['A'][2]:+.3f}; "
                     f"100-unit blocks against the target: up {np.round(x['u'], 2).tolist()}, down "
                     f"{np.round(x['d'], 2).tolist()}, difference {np.round(x['s'], 2).tolist()}"
                     f"\n    ratio to target just before a kick, drift: up {x['drift'][0]:+.4f}, down {x['drift'][1]:+.4f}; "
                     f"spread up {x['spread'][0]:.4f}, down {x['spread'][1]:.4f}"
                     f"\n    sum of squared moves: up {x['sumsq'][0]:.3f}, down {x['sumsq'][1]:.3f}, sham "
                     f"{x['sumsq'][2]:.3f}; first kick's share {np.round(x['first_share'], 3).tolist()}; largest move "
                     f"{np.round(x['largest'], 3).tolist()}"
                     f"\n    sham against the mean of up and down: sum of squared moves {x['sham_sq']:.3f}, sum of "
                     f"|moves| {x['sham_abs']:.3f}; per kick {x['kick_sq'].min():.3f} to {x['kick_sq'].max():.3f} and "
                     f"{x['kick_abs'].min():.3f} to {x['kick_abs'].max():.3f}; the same kick times: {x['same_times']}")
        lines.append(line)
    lines.append(f"  the cause test's locked files unchanged: {not lock_changes()}")
    lines.append(f"  dry run {'passed' if ok else 'FAILED'}")
    (DRY / "dryrun.txt").write_text("\n".join(lines) + "\n")
    print("\n".join(lines))


# ---------------------------------------------------------------------------------------------------------------------
# the report

def fmt(x, places=3):
    return f"{x['mean']:.{places}f} (SD {x['sd']:.{places}f}, {x['low']:.{places}f} to {x['high']:.{places}f})"


def describe(res, H):
    yes = {True: "met", False: "NOT met"}
    lines = [f"    carried out {res['carried']} of {res['moments']} (fewest per start {res['fewest_per_start']}); "
             f"F1 {yes[res['F1']]}"]
    lines += [f"      not carried out: {r}" for r in res["reasons"]]
    if res["carried"]:
        lines += [
            f"    F2 {yes[res['F2']]}: both fractions within 0.75 to 1.25 in {res['inside_share']:.0%} of moments; "
            f"A_up {fmt(res['achieved']['up'])}, A_down {fmt(res['achieved']['down'])}",
            f"    F3a {yes[res['F3a']]}: every block keeps at least a quarter of the planned difference in "
            f"{res['sustained_share']:.0%} of moments (lowest block difference {res['lowest_difference']:.2f})",
            f"    F3b {yes[res['F3b']]}: block means over the moments, up {np.round(res['block_mean_up'], 3).tolist()}, "
            f"down {np.round(res['block_mean_down'], 3).tolist()} (difference "
            f"{np.round(res['block_mean_difference'], 3).tolist()})",
            f"    F4a {yes[res['F4a']]}: sham's sum of squared moves / the mean of up's and down's {fmt(res['sham_sq'])}",
            f"    F4c {yes[res['F4c']]}: sham's sum of |moves| / the mean of up's and down's {fmt(res['sham_abs'])}; "
            f"the same kick times in every moment: {res['same_times']}",
            f"    F4b (reported) {'within' if res['F4b'] else 'OUTSIDE'} -0.25 to +0.25: A_sham {fmt(res['achieved']['sham'])}",
            f"    the draft's standard (every 100-unit block fraction of up and down within 0.5 to 1.5): "
            f"{res['draft_share']:.0%} of moments; the first check's F3 (25-unit blocks): {res['first_check_share']:.0%}"]
    return lines


def report():
    stage, todo, chosen = next_stage()
    if todo:
        sys.exit(f"Not finished: {stage} has {len(todo)} tasks left. Run run.")
    log = json.loads(RUNLOG.read_text())
    bases = {s: load_base(s) for s in STARTS}
    tried = [i for i in range(len(LAMBDAS)) if all(done(t) for t in tune_tasks(i))]
    tuning = {LAMBDAS[i]: {H: judge(i, TUNE, H, 0, bases) for H in HOLDS} for i in tried}
    confirmation = {}
    nums = {}
    if chosen is not None:
        for H in HOLDS:
            res, x, rows = judge("confirm", CONFIRM, H, AFTER, bases, keep=True)
            confirmation[H], nums[H] = res, (x, rows)
    natural = natural_spreads(bases)
    first_check = json.loads(FIRST_CHECK.read_text())
    folders_now = folder_fingerprints()
    unchanged = dict(lock=not lock_changes() and not log["lock_changes_before"],
                     folders=folders_now == log["folders_before"],
                     plan_sha256=sha256(PLAN))
    base_ok = {s + 1: bool(bases[s]["same"].all()) for s in STARTS}
    verdict500 = bool(chosen is not None and confirmation[500]["met"])
    verdict100 = bool(chosen is not None and confirmation[100]["met"])
    out = dict(made=stamp(), plan_sha256=unchanged["plan_sha256"], unchanged=unchanged, base_reproduced=base_ok,
               offsets={s + 1: o for s, o in offsets().items()}, decisions=log.get("decisions", []),
               lam=LAMBDAS[chosen] if chosen is not None else None,
               tuning={str(lam): {str(H): r for H, r in v.items()} for lam, v in tuning.items()},
               confirmation={str(H): r for H, r in confirmation.items()},
               feasible={"500": verdict500, "100": verdict100},
               natural={str(H): v for H, v in natural.items()}, criteria=CRITERIA, lambdas=LAMBDAS)
    write_json(OUT_JSON, out)

    lines = [f"Check of the slow-drift hold (H2), {shown(out['made'])}",
             f"Plan: hold_check_plan.md (sha256 {unchanged['plan_sha256'][:16]}...), not edited after it was written",
             f"The cause test's locked files unchanged: {unchanged['lock']}; the protected folders unchanged: "
             f"{unchanged['folders']}",
             f"Base runs reproduce their saved runs exactly: {base_ok}",
             f"Moments: T = 2,000 + o + 350 m; o by start {out['offsets']}", "",
             "Choosing lambda (old starts 13 to 15, 30 moments, nothing after T)"]
    for d in out["decisions"]:
        lines.append(f"  {shown(d['made'])}: {d['text']}")
    for lam, v in tuning.items():
        for H in HOLDS:
            lines.append(f"  lambda = {lam}, {H}-unit hold{' (judged)' if H == 500 else ' (reported)'}")
            lines += describe(v[H], H)
    lines.append("")
    if chosen is None:
        lines.append("No lambda was chosen: H2 is NOT FEASIBLE at +-25% as specified. The confirmation was not run.")
    else:
        lines.append(f"Confirmation (old starts 16 to 18, 30 new moments, lambda = {LAMBDAS[chosen]}, untouched to T + 500)")
        for H in HOLDS:
            r = confirmation[H]
            lines.append(f"  {H}-unit hold: {'FEASIBLE' if r['met'] else 'NOT FEASIBLE'} at +-25% as specified")
            lines += describe(r, H)
        lines.append("")
    # the plan's "reported separately": for the confirmation if it ran; if not, for every candidate tried (nothing
    # after T was run then, so there is nothing after release)
    if chosen is not None:
        sets = [("confirmation", "", confirmation)]
    else:
        sets = [("choosing lambda, old starts 13 to 15", f"lambda = {lam}, ", tuning[lam]) for lam in tuning]
    for title, prefix, results in sets:
        lines.append(f"Reported separately ({title})")
        for H in HOLDS:
            r = results[H]
            if not r["carried"]:
                continue
            planned = SIZE * natural[H]["level"]
            fc = first_check["holds"][str(H)]
            lines += [
                f"  {prefix}{H}-unit hold",
                f"    history: dM = {fmt(r['dM'], 5)} against the planned {planned:.5f}; dM / the natural SD of "
                f"{H}-unit history means ({natural[H]['sd']:.5f}) = {r['dM']['mean'] / natural[H]['sd']:.2f}; block SDs "
                f"over the moments: up {np.round(r['block_sd_up'], 2).tolist()}, down "
                f"{np.round(r['block_sd_down'], 2).tolist()}, difference {np.round(r['block_sd_difference'], 2).tolist()}",
                f"      ratio to target just before a kick: drift up {fmt(r['drift']['up'], 4)}, down "
                f"{fmt(r['drift']['down'], 4)}; spread up {r['spread']['up']['mean']:.3f}, down "
                f"{r['spread']['down']['mean']:.3f}; 5th to 95th percentile up {np.round(r['pre_5_95']['up'], 3).tolist()}, "
                f"down {np.round(r['pre_5_95']['down'], 3).tolist()}",
                f"    effort: {r['kicks']} kicks; sum of squared moves up {fmt(r['sumsq']['up'], 2)}, down "
                f"{fmt(r['sumsq']['down'], 2)}, sham {fmt(r['sumsq']['sham'], 2)} (the first check's pinned hold: up "
                f"{fc['sumsq']['up']['mean']:.1f}, down {fc['sumsq']['down']['mean']:.1f}, sham "
                f"{fc['sumsq']['sham']['mean']:.1f}; one pulse at +-25%: {first_check['pulse']['sumsq']['up']['mean']:.2f})",
                f"      sum of |moves| up {fmt(r['sumabs']['up'], 1)}, down {fmt(r['sumabs']['down'], 1)}, sham "
                f"{fmt(r['sumabs']['sham'], 1)}; largest move up {r['largest']['up']:.3f}, down {r['largest']['down']:.3f}, "
                f"sham {r['largest']['sham']:.3f}",
                f"      first kick's share up {r['first_share']['up']['mean']:.3f}, down {r['first_share']['down']['mean']:.3f}, "
                f"sham {r['first_share']['sham']['mean']:.3f}; median later kick (sum of squared moves) up "
                f"{r['later']['up']['mean']:.4f}, down {r['later']['down']['mean']:.4f}",
                f"      each kick's sham-to-arms ratio: sum of squared moves median {r['kick_sq']['median']:.3f} "
                f"({r['kick_sq']['low']:.3f} to {r['kick_sq']['high']:.3f}); sum of |moves| median "
                f"{r['kick_abs']['median']:.3f} ({r['kick_abs']['low']:.3f} to {r['kick_abs']['high']:.3f})"]
            for name, (a, b) in WINDOWS.items():
                if name in r["left"]:
                    lines.append(f"    after release, mean c over (T + {a}, T + {b}] minus untouched, per unit of the "
                                 f"planned change (down: positive while below): up {fmt(r['left'][name]['up'])}, down "
                                 f"{fmt(r['left'][name]['down'])}, sham {fmt(r['left'][name]['sham'])}")
    matching = sum(1 for name, v in log["parts"].items()
                   if (PARTS / f"{name}.npz").exists() and sha256(PARTS / f"{name}.npz") == v["sha256"])
    environments = {json.dumps(p["machine"], sort_keys=True) for p in log["pieces"]}
    first, last = log["pieces"][0], log["pieces"][-1]
    running = sum(p["seconds"] for p in log["pieces"])
    lines += ["",
              f"Pieces: {len(log['pieces'])}, {first['started']['halifax'][11:16]} to {last['ended']['halifax'][11:16]} "
              f"Halifax, {running / 60:.0f} minutes of running (the plan said about 30 to 45 minutes in all, with a "
              f"confirmation); tasks took up to about 190 s with seven running at once",
              f"Parts matching the run log: {matching} of {len(log['parts'])}; one numerical environment in every piece: "
              f"{len(environments) == 1} ({log['machine']['cpu']}, Python {log['machine']['python']}, "
              f"{log['machine']['libc']}, numpy {log['machine']['numpy']}, BLAS threads: "
              f"{log['machine']['blas_threads']})",
              "Seen before the check ran (dryruns/hold_check_dryrun.txt; old start 1, one moment at T = 600, "
              "lambda = 0.10; counts toward nothing): 500-unit hold A_up 0.759, A_down 0.853, A_sham +0.025, up's "
              "100-unit blocks [0.43, 0.84, 0.94, 0.86, 0.55]; 100-unit hold A_up 1.255, A_down 0.369. Both were "
              "carried out.",
              "Changes to the report code after the run, reporting only (no criterion, candidate or number changed): "
              "before the report was first made, the plan's 'reported separately' items were added for the candidates "
              "tried, since no confirmation ran, and the 100-unit hold's verdict was worded as not judged; after it "
              "was first made, these last lines (pieces, parts, environment and the dry run's numbers) were added."]
    hundred = (f"the 100-unit hold is {'FEASIBLE' if verdict100 else 'NOT FEASIBLE'}" if chosen is not None else
               "the 100-unit hold was not judged, since the plan judges it on the confirmation, which did not run "
               "(its numbers on starts 13 to 15 are above, reported only)")
    lines += ["", f"Verdict: the 500-unit hold under H2 is {'FEASIBLE' if verdict500 else 'NOT FEASIBLE'} at +-25% "
                  f"as specified; {hundred}."]
    OUT_TXT.write_text("\n".join(lines) + "\n")
    print("\n".join(lines))


if __name__ == "__main__":
    command = sys.argv[1] if len(sys.argv) > 1 else ""
    if command == "dryrun":
        dryrun()
    elif command == "run":
        run()
    elif command == "report":
        report()
    else:
        sys.exit(__doc__)
