"""
What the two checks before any lock share (prelock_checks_plan.md): the cause test's code, imported unchanged; the
paths and seeds; the checks that nothing fixed has changed; and the runner of foreground pieces.

A sandbox script kept with the project. Nothing here is served by the site.
"""
import json
import os
import sys
import time
from multiprocessing import Pool
from pathlib import Path

NUMPY_FIRST = "numpy" in sys.modules                      # numpy loaded before the setting below takes effect
os.environ["OPENBLAS_NUM_THREADS"] = os.environ["OMP_NUM_THREADS"] = "1"   # one thread per worker (set before numpy)
import numpy as np  # noqa: E402

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent.parent
CAUSE = HERE.parent / "cause-test-2026-10-12"
EXPLORE = HERE.parent / "cause-explore-2026-10-12"
sys.dont_write_bytecode = True                             # leave every folder as it is
sys.path.insert(0, str(CAUSE))
import cause_test as ct  # noqa: E402  (the model runs, the amplitude solve, the paired signs; unchanged)

nt, rule, egc = ct.nt, ct.rule, ct.egc
N, K, EPS, DT, GAP, EVERY, RUN = ct.N, ct.K, ct.EPS, ct.DT, ct.GAP, ct.EVERY, ct.RUN
SEED = 20261013          # every draw of these checks is (SEED, use, ...): 1 the calibration's fresh directions,
                         # 2 the feasibility moments, 3 a hold's sham signs, 4 a pulse's sham signs
WORKERS = 7
WORK = ROOT / "generated-artifacts" / "cause-prelock-2026-10-13"   # checkpoints, saved states, parts (git ignores it)
stamp, shown, sha256, write_json, plain = ct.stamp, ct.shown, ct.sha256, ct.write_json, ct.plain


SAME_FIELDS = ("cpu", "python", "libc", "numpy", "numpy_dispatch")   # what must match between runs that are compared


def machine():
    """The numerical environment, recorded with every piece and every reference run (prelock_amendment.md): the
    processor, Python, the C library, numpy and the routines numpy picked for this processor (SAME_FIELDS, which must
    match); and, for information only, the BLAS thread setting, the kernel's boot id and when this container started."""
    import platform
    from numpy._core._multiarray_umath import __cpu_dispatch__, __cpu_features__
    model = next((line.split(":", 1)[1].strip() for line in open("/proc/cpuinfo") if line.startswith("model name")), "")
    up = float(Path("/proc/uptime").read_text().split()[0])
    return dict(cpu=model, python=platform.python_version(), libc=" ".join(platform.libc_ver()), numpy=np.__version__,
                numpy_dispatch=[name for name in __cpu_dispatch__ if __cpu_features__.get(name)],
                blas_threads="the library default (numpy loaded first)" if NUMPY_FIRST else "1 per worker",
                boot_id=Path("/proc/sys/kernel/random/boot_id").read_text().strip(),
                container_started=stamp(round(time.time() - up)))


def same_environment(a, b):
    """True if the two environments agree on every SAME_FIELDS entry that both recorded."""
    return all(a[key] == b[key] for key in SAME_FIELDS if key in a and key in b)


def reference_checks(kind):
    """From machine_check.json, for each old start (numbered from 1), the latest comparison of the given kind
    ("calibration" or "feasibility") with the cause test's unchanged base-run code run in one go, with that check's
    environment and where the start first differs from its saved run."""
    path = HERE / "machine_check.json"
    out = {}
    if path.exists():
        for record in json.loads(path.read_text())["checks"]:
            for r in record["results"]:
                if kind in r:
                    out[r["start"]] = dict(same=r[kind], saved=r["saved"], machine=record["machine"],
                                           made=record["made"])
    return out


def lock_changes():
    """The names of the cause test's locked files that no longer match cause_lock.json (none expected)."""
    lock = json.loads((CAUSE / "cause_lock.json").read_text())
    now = ct.fingerprints()
    return [name for name in ct.LOCKED if lock["sha256"].get(name) != now[name]]


def folder_fingerprints():
    """Every file in the cause test's folder and the exploratory folder, which these checks only read."""
    out = {}
    for folder in (CAUSE, EXPLORE):
        for path in sorted(folder.rglob("*")):
            if path.is_file():
                out[str(path.relative_to(ROOT))] = sha256(path)
    return out


def save_npz(path, **arrays):
    """Written whole or not at all (uncompressed, so every number comes back exactly)."""
    partial = path.with_name(path.name + ".partial.npz")
    np.savez(partial, **arrays)
    os.replace(partial, path)


def open_log(path):
    if path.exists():
        return json.loads(path.read_text())
    return dict(created=stamp(), lock_changes_before=lock_changes(), folders_before=folder_fingerprints(),
                pieces=[], parts={})


def run_piece(log_path, tasks, work, budget, start_by):
    """One foreground piece. tasks: the unfinished tasks in order (dicts with a "name"); work((task, deadline,
    start_by)), a top-level function that calls guard(), runs one and returns dict(name, status, seconds, error):
    status "done" (its part is written), "saved" (stopped at the deadline with its exact state kept, to go on in the
    next piece), "skipped" (not begun: too late in the piece) or "failed". WORKERS at a time; the piece ends when every
    task has returned."""
    log = open_log(log_path)
    if log["lock_changes_before"] or lock_changes():
        sys.exit("A file fixed by the cause test's lock has changed. Stopping.")
    env = machine()
    if "machine" in log and not same_environment(log["machine"], env):
        sys.exit(f"The numerical environment differs from the earlier pieces' ({log['machine']}; now {env}). "
                 "Stopping: a start must run wholly in one environment (prelock_amendment.md).")
    if "machine" not in log:
        log["machine"], log["machine_recorded_from_piece"] = env, len(log["pieces"]) + 1
    clock = time.time()
    entry = dict(piece=len(log["pieces"]) + 1, started=stamp(clock), tasks=len(tasks), machine=env, results={})
    log["pieces"].append(entry)
    write_json(log_path, log)
    print(f"piece {entry['piece']}: {len(tasks)} unfinished tasks, {WORKERS} at a time", flush=True)
    jobs = [(task, clock + budget, clock + start_by) for task in tasks]
    with Pool(WORKERS) as pool:
        for res in pool.imap_unordered(work, jobs):
            entry["results"][res["name"]] = dict(status=res["status"], seconds=round(res["seconds"], 1),
                                                 error=res.get("error", ""))
            if res["status"] == "done":
                log["parts"][res["name"]] = dict(sha256=res["sha256"], piece=entry["piece"], finished=stamp())
            if res["status"] != "skipped":
                print(f"  {res['status']}: {res['name']} ({res['seconds']:.0f} s) {res.get('error', '')}", flush=True)
            write_json(log_path, log)
    entry["ended"], entry["seconds"] = stamp(), round(time.time() - clock, 1)
    write_json(log_path, log)
    counts = {status: sum(r["status"] == status for r in entry["results"].values())
              for status in ("done", "saved", "skipped", "failed")}
    print(f"piece {entry['piece']} ended after {entry['seconds']:.0f} s: {counts}", flush=True)
    return log


def guard(work, job):
    """work(task, deadline) inside a guard: a task that is begun too late returns "skipped"; an error returns
    "failed" with its message instead of ending the piece. (Each script's worker is a top-level function calling
    this, so the pool can hand it to its workers.)"""
    task, deadline, start_by = job
    clock = time.time()
    if clock > start_by:
        return dict(name=task["name"], status="skipped", seconds=0.0)
    try:
        out = work(task, deadline)
    except Exception as error:  # noqa: BLE001  (reported in the run log, never hidden)
        return dict(name=task["name"], status="failed", seconds=time.time() - clock, error=repr(error))
    out.setdefault("name", task["name"])
    out["seconds"] = time.time() - clock
    return out