"""
The pause test (hold_vs_pulse_revised.md, section 6, approved by Colson on October 13, 2026; prelock_amendment.md,
"The cause"). Does picking up again from kept progress, as the nudge test's run code does, make a run part from the
same run made in one go, on this machine?

What runs (set before the run). One old start, start 7 (seed 6). Its saved run (coherence-growth-2026-10-11/
sim_fresh4000.npz) parts earliest from today's unbroken runs: first gap difference at t = 320, first coherence
difference at t = 1,834 (machine_check.json). The nudge test's run code (nudge_test.py's lyapunov), imported unchanged,
runs it twice side by side, with the saved run's settings (N = 4,000, K = 1.90, settling 300, then 6,000 time units,
gap 1e-7 rms put back every 5, step 0.025):
  A. in one go (no deadline);
  B. with one pause: a deadline 15 s after it begins, so it stops at the first put-back after that and keeps its
     progress; then it picks up again from the kept progress with its own code, as on October 11, and runs to the end.
Its progress is kept in generated-artifacts/cause-prelock-2026-10-13/pause_test/ rather than the nudge test's folder,
as coherence_growth.py kept it in its own folder. Nothing else is changed.

Compared (set before the run): A with B, A with the saved run, B with the saved run. For each pair: the settling trace,
the coherence every time unit (complex64, as the code returns it) and every gap log growth (float64), all exactly;
where each pair first differs; and how large the differences are afterwards, by 1,000-unit stretch.

How it will be read (set before the run):
- If B parts from A: on this machine, picking up again changes the numbers by itself, so the pauses are enough to
  explain why the two paused batches part from today's unbroken runs. It would not show that their pauses did cause
  it: their pause points and their machine were not recorded.
- If B does not part from A over the 6,000 units: picking up again changes nothing here, so on this machine the
  pauses do not explain the departures, and the cause stays open (a different or busier machine remains possible).
- A is expected to part from the saved run at t = 320 (gap) and t = 1,834 (coherence), as the cause test's base-run
  code did. If not, that is reported as found.
- Noted before the run, from the saved run's own record: start 7 took 1,289 s in all, about 4.9 time units per
  second. At that speed its first pause, at most 230 s into the first piece, would have come near t = 830 (settling
  included), after its first gap difference at t = 320. Unless its first piece ran much slower than its average, its
  departure did not begin at a pause. This test cannot settle that.

Usage: python3 -B pause_test.py   (about 3 minutes in the foreground; two processes; nothing left running)
A sandbox script kept with the project. Nothing here is served by the site.
"""
import hashlib
import json
import sys
import time
from multiprocessing import Pool
from pathlib import Path

import prelock_common as pc  # noqa: E402  (one thread per worker; the cause test's modules, unchanged)

np, nt = pc.np, pc.nt
HERE, ROOT = pc.HERE, pc.ROOT
NUDGE = HERE.parent / "nudge-test-2026-10-10"
GROWTH = HERE.parent / "coherence-growth-2026-10-11"
SAVED = GROWTH / "sim_fresh4000.npz"
WORK = pc.WORK / "pause_test"
PARTS = WORK / "parts"
nt.LYAP_PARTS = PARTS          # after every import: coherence_growth.py, imported by the cause test, sets its own folder
START = 6                      # seed 6: old start 7
RUN = 6000.0
PAUSE_AFTER = 15.0             # seconds; the pause comes at the first put-back after this
TASK = (pc.N, pc.K, (START, 1), nt.SETTLE, RUN, nt.LYAP_EVERY, nt.LYAP_GAP, nt.DT)
STRETCH = 1000
OUT_JSON, OUT_TXT, OUT_NPZ = HERE / "pause_test_results.json", HERE / "pause_test_results.txt", WORK / "runs.npz"


def guarded():
    if not str(nt.LYAP_PARTS).startswith(str(pc.WORK)):
        raise RuntimeError(f"progress would go to {nt.LYAP_PARTS}, outside {pc.WORK}")


def run(kind):
    """A: in one go. B: one pause, then picked up again from the kept progress. Both with lyapunov unchanged."""
    nt.LYAP_PARTS = PARTS
    guarded()
    clock = time.time()
    if kind == "A":
        res = nt.lyapunov(TASK + (None,))
        return dict(kind=kind, z=res["z"], logs=res["logs"], trace=res["trace"], wall=time.time() - clock)
    first = nt.lyapunov(TASK + (time.time() + PAUSE_AFTER,))
    if first["finished"]:
        raise RuntimeError("B finished without a pause")
    progress = nt.lyap_progress(pc.N, pc.K, START, RUN)
    kept = {}
    with np.load(progress) as saved:
        kept["us_c_order"], kept["uf_c_order"] = bool(saved["us"].flags.c_contiguous), bool(saved["uf"].flags.c_contiguous)
    second = nt.lyapunov(TASK + (time.time() + 1e6,))
    if not second["finished"]:
        raise RuntimeError("B did not finish after picking up again")
    progress.unlink()
    return dict(kind=kind, z=second["z"], logs=second["logs"], trace=second["trace"], wall=time.time() - clock,
                paused_at=first["step"] * nt.DT, read_back_c_order=kept)


def layouts():
    """The memory layout of the arrays as the run builds them, and as they come back from a kept file."""
    slow, _ = nt.model(pc.N, pc.K, nt.EPS)
    us, uf = nt.starting_phases(pc.N, slow, START, 1)
    us, uf = np.concatenate([us, us]), np.concatenate([uf, uf])
    return dict(built=dict(us=us.strides, uf=uf.strides), read_back_c_order_strides=dict(
        us=np.ascontiguousarray(us).strides, uf=np.ascontiguousarray(uf).strides))


def tree(folder):
    return {str(p.relative_to(ROOT)): hashlib.sha256(p.read_bytes()).hexdigest()
            for p in sorted(folder.rglob("*")) if p.is_file()}


def compare(a, b):
    """Where two runs first differ, exactly, and how far apart they are afterwards."""
    gap_diff = np.flatnonzero((a["logs"] != b["logs"]).any(0))
    coh_diff = np.flatnonzero((a["z"] != b["z"]).any((0, 1)))
    out = dict(settling_same=bool(np.array_equal(a["trace"], b["trace"])),
               gap_same=bool(gap_diff.size == 0), coherence_same=bool(coh_diff.size == 0),
               first_gap_t=float(nt.LYAP_EVERY * (gap_diff[0] + 1)) if gap_diff.size else None,
               first_coherence_t=int(coh_diff[0]) if coh_diff.size else None, stretches=[])
    Ra, Rb = np.abs(a["z"][0].astype(np.complex128)), np.abs(b["z"][0].astype(np.complex128))
    per = int(STRETCH / nt.LYAP_EVERY)
    for s in range(int(RUN // STRETCH)):
        la, lb = a["logs"][0, s * per:(s + 1) * per], b["logs"][0, s * per:(s + 1) * per]
        span = slice(s * STRETCH + 1, (s + 1) * STRETCH + 1)
        out["stretches"].append(dict(t=f"{s * STRETCH:,} to {(s + 1) * STRETCH:,}",
                                     largest_log_difference=float(np.abs(la - lb).max()),
                                     largest_coherence_difference=float(np.abs(Ra[:, span] - Rb[:, span]).max()),
                                     typical_coherence=float(Ra[:, span].mean()),
                                     coherence_r=float(np.corrcoef(Ra[:, span].ravel(), Rb[:, span].ravel())[0, 1])))
    return out


def main():
    if PARTS.exists() and any(PARTS.iterdir()):
        sys.exit(f"progress left from an earlier attempt in {PARTS}; look at it before running again")
    PARTS.mkdir(parents=True, exist_ok=True)
    guarded()
    before = dict(lock=pc.lock_changes(), cause_explore=pc.folder_fingerprints(), nudge=tree(NUDGE), growth=tree(GROWTH))
    env, started = pc.machine(), pc.stamp()
    print(f"pause test: start 7, A in one go and B with one pause after {PAUSE_AFTER:.0f} s, side by side", flush=True)
    with Pool(2) as pool:
        runs = {res["kind"]: res for res in pool.map(run, ["A", "B"])}
    ended = pc.stamp()
    with np.load(SAVED) as d:
        i = list(d["starts"]).index(START)
        saved = dict(z=d["z"][i:i + 1], logs=d["logs"][i:i + 1], trace=d["trace"][i:i + 1])
    A, B = runs["A"], runs["B"]
    pc.save_npz(OUT_NPZ, zA=A["z"], logsA=A["logs"], traceA=A["trace"], zB=B["z"], logsB=B["logs"], traceB=B["trace"])
    pairs = {"A with B": compare(A, B), "A with the saved run": compare(A, saved), "B with the saved run": compare(B, saved)}
    after = dict(lock=pc.lock_changes(), cause_explore=pc.folder_fingerprints(), nudge=tree(NUDGE), growth=tree(GROWTH))
    untouched = {name: before[name] == after[name] for name in before}
    if PARTS.exists() and not any(PARTS.iterdir()):
        PARTS.rmdir()
    ab = pairs["A with B"]
    parted = not (ab["gap_same"] and ab["coherence_same"] and ab["settling_same"])
    result = dict(title="Pause test: start 7 with the nudge test's run code, in one go (A) and with one pause (B)",
                  started=started, ended=ended, machine=env, task=dict(n=pc.N, k=pc.K, seed=START, start=START + 1,
                  settle=nt.SETTLE, run=RUN, every=nt.LYAP_EVERY, gap=nt.LYAP_GAP, dt=nt.DT),
                  pause_after_seconds=PAUSE_AFTER, paused_at=B["paused_at"], read_back=B["read_back_c_order"],
                  layouts=layouts(), wall_seconds=dict(A=round(A["wall"], 1), B=round(B["wall"], 1)),
                  pairs=pairs, B_parts_from_A=parted, folders_unchanged=untouched,
                  lock_changes=dict(before=before["lock"], after=after["lock"]))
    pc.write_json(OUT_JSON, result)
    lines = [result["title"], f"Run {started['halifax'][:16].replace('T', ' ')} to {ended['halifax'][11:16]} Halifax"
             f" (A {A['wall']:.0f} s, B {B['wall']:.0f} s). B paused at t = {B['paused_at']:,.0f} and picked up again"
             " from its kept progress.", ""]
    for name, p in pairs.items():
        if p["gap_same"] and p["coherence_same"] and p["settling_same"]:
            lines.append(f"{name}: the same to the last digit over all 6,000 units (settling, coherence, every gap).")
        else:
            lines.append(f"{name}: settling {'same' if p['settling_same'] else 'DIFFERS'}; first gap difference at t ="
                         f" {p['first_gap_t']:,.0f}; first coherence difference at t = "
                         + (f"{p['first_coherence_t']:,}" if p['first_coherence_t'] is not None else "none") + ".")
            for s in p["stretches"]:
                lines.append(f"    t = {s['t']}: largest gap log difference {s['largest_log_difference']:.2e}; largest"
                             f" coherence difference {s['largest_coherence_difference']:.2e} (typical coherence"
                             f" {s['typical_coherence']:.4f}; r = {s['coherence_r']:.3f})")
    lines += ["", "Reading, as set before the run: " + (
        "B parts from A, so on this machine picking up again changes the numbers by itself. The pauses are enough to"
        " explain the paused batches' departures; that their pauses caused them is not shown." if parted else
        "B does not part from A, so on this machine picking up again changes nothing, and the pauses do not explain"
        " the departures. The cause stays open."),
        f"Layout: built uf strides {result['layouts']['built']['uf']}; read back in C order "
        f"{result['layouts']['read_back_c_order_strides']['uf']}.",
        "Folders and locked files unchanged: " + ", ".join(f"{k} {'yes' if v else 'NO'}" for k, v in untouched.items())]
    OUT_TXT.write_text("\n".join(lines) + "\n")
    print("\n".join(lines))


if __name__ == "__main__":
    main()
