"""Check 1 of the pilot review: does the timing gap between groups A and B last into the growth window?

After the fact and exploratory. This describes the pilot's saved results; it is not a reading fixed in plan.md and not a
new experiment. It runs no model. It reads only the pilot's 120 saved moment files (starts 103 to 114, moments 1 to 10)
in generated-artifacts/medium-test/pilot, checks each against the hash recorded in pilot_results.json, and writes
check1_timing.txt and check1_timing.json beside this file. Nothing else is changed.

Question (from pilot_review.md): at the restoration point R the medium had shifted the timing between groups A and B
(closer when raised, farther apart when lowered), while the matched pushes left it alone. Does that shift last into the
growth window (R + 200, R + 500], or fade before it?

Measure: the A-B timing gap at each time unit, |wrap(psi_A - psi_B)| in radians, where psi_g is the angle of group g's
saved coherence (all its oscillators). It is the report's |A-B timing| at R, followed over time; the script first checks
that it reproduces the report's values at R.

Fixed before looking at the curves:
- The split: half of (the lowered arm's gap minus the raised arm's gap). Positive means the raised arm's groups sit
  closer together in timing than the lowered arm's, the direction seen at R.
- The common part: the average of the two arms' gaps minus the untouched gap.
- Both are read at R and over (R + 200, R + 500], and in 100-unit blocks (the medium and untouched from T, the pushes
  and shams from R), like the report's coherence curves.
- The share of the medium's split at R still there over (R + 200, R + 500]: a ratio of averages, like the report's
  persistence.
- Medium against matched push, like for like: only moments with both pushes built.
- "Lasts" if the medium's split over (R + 200, R + 500] is clear with the sign it had at R; "reverses" if it is clear
  with the other sign; "fades before" if it is clear at R and not clear over (R + 200, R + 500].
- Averages and error bars as in the report: each start's moments are averaged first, then the 12 starts. The error bar
  is the larger of the resampled one (starts drawn with replacement, 400 times, seed 20261016) and SD / sqrt(12).
  Clear means more than 3 bars.
"""

import hashlib
import json
import sys
import time
import warnings
from datetime import datetime
from pathlib import Path
from zoneinfo import ZoneInfo

import numpy as np

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[2]
PARTS = ROOT / "generated-artifacts" / "medium-test" / "pilot"
OUT_TXT, OUT_JSON = HERE / "check1_timing.txt", HERE / "check1_timing.json"
STARTS, MOMENTS = tuple(range(103, 115)), tuple(range(1, 11))
W, AFTER, BLOCK, LATER = 200, 500, 100, (200, 500)        # window, followed after R, block, growth window after R
RESAMPLES, SEED, CLEAR = 400, 20261016, 3.0                # as in the report
DIRECTIONS = ("raised", "lowered")


def sha256(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def gap(z):
    """|A-B timing| at every time unit, from saved coherences with the groups on the second-to-last axis."""
    return np.abs(np.angle(z[..., 0, :] * np.conj(z[..., 1, :])))


def by_start(values):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=RuntimeWarning)
        return np.nanmean(np.asarray(values, float), axis=1)


def summary(per_start):
    """The report's average over the starts and its error bar."""
    v = np.asarray(per_start, float)
    v = v[np.isfinite(v)]
    n = len(v)
    if n < 2:
        return dict(mean=float(v.mean()) if n else float("nan"), bar=float("nan"), n=n, above=int((v > 0).sum()),
                    clear=False)
    boot = v[np.random.default_rng(SEED).integers(0, n, (RESAMPLES, n))].mean(-1)
    mean = float(v.mean())
    bar = max(float(boot.std(ddof=1)), float(v.std(ddof=1) / np.sqrt(n)))
    return dict(mean=mean, bar=bar, n=n, above=int((v > 0).sum()), clear=bool(abs(mean) > CLEAR * bar))


def ratio_summary(numerator, denominator):
    """The report's ratio of averages over the starts, with the resampled error bar."""
    a, b = np.asarray(numerator, float), np.asarray(denominator, float)
    keep = np.isfinite(a) & np.isfinite(b)
    a, b = a[keep], b[keep]
    idx = np.random.default_rng(SEED).integers(0, len(a), (RESAMPLES, len(a)))
    mean, bar = float(a.mean() / b.mean()), float((a[idx].mean(-1) / b[idx].mean(-1)).std(ddof=1))
    return dict(mean=mean, bar=bar, n=len(a), clear=bool(abs(mean) > CLEAR * bar))


def shown(s):
    return (f"{s['mean']:+.4f} +- {s['bar']:.4f} ({s['mean'] / s['bar']:+.1f} bars; "
            f"{s['above']} of {s['n']} starts above zero)")


def main():
    began = time.time()
    hashes = json.loads((HERE / "pilot_results.json").read_text())["moment_sha256"]
    rows = []
    for s in STARTS:
        for m in MOMENTS:
            path = PARTS / f"moment_{s:03d}_{m:02d}.npz"
            if hashes.get(path.name) != sha256(path):
                sys.exit(f"{path.name} differs from the hash recorded in pilot_results.json. Stopping.")
            with np.load(path) as p:
                r = {key: p[key] for key in ("z_none", "z_medium", "z_push", "z_sham", "built", "desc_ab")}
            r.update(g_none=gap(r["z_none"]), g_medium=gap(r["z_medium"]), g_push=gap(r["z_push"]),
                     g_sham=gap(r["z_sham"]))       # untouched and medium from T (701 units); push and sham from R (501)
            rows.append(r)
    S_, M_ = len(STARTS), len(MOMENTS)

    def grid(fn):
        return np.array([[fn(rows[i * M_ + m]) for m in range(M_)] for i in range(S_)], float)

    built = np.array([[rows[i * M_ + m]["built"] for m in range(M_)] for i in range(S_)], bool)
    mask = np.where(built.all(-1), 1.0, np.nan)

    # the measure reproduces the report's |A-B timing| at R (pushes and shams only where built)
    worst = 0.0
    for r in rows:
        at_r = (r["g_none"][W], r["g_medium"][0][W], r["g_medium"][1][W], r["g_push"][0][0], r["g_push"][1][0],
                r["g_sham"][0][0], r["g_sham"][1][0])
        for i, value in enumerate(at_r):
            if i < 3 or r["built"][(i - 3) % 2]:
                worst = max(worst, abs(value - abs(float(r["desc_ab"][i]))))
    if worst > 1e-9:
        sys.exit(f"The gap at R differs from the report's description by up to {worst:.2e}. Stopping.")

    lo, hi = LATER
    at = dict(none=grid(lambda r: r["g_none"][W]))
    over = dict(none=grid(lambda r: r["g_none"][W + lo + 1:W + hi + 1].mean()))
    for d, name in enumerate(DIRECTIONS):
        at[name] = grid(lambda r, d=d: r["g_medium"][d][W])
        over[name] = grid(lambda r, d=d: r["g_medium"][d][W + lo + 1:W + hi + 1].mean())
        for arm in ("push", "sham"):
            at[f"{arm} {name}"] = grid(lambda r, d=d, arm=arm: r[f"g_{arm}"][d][0]) * mask
            over[f"{arm} {name}"] = grid(lambda r, d=d, arm=arm: r[f"g_{arm}"][d][lo + 1:hi + 1].mean()) * mask

    def parts(g, arm=""):
        """Split and common part for the medium (arm "") or for the pushes or shams (arm "push " or "sham ")."""
        raised, lowered = g[f"{arm}raised"], g[f"{arm}lowered"]
        none = g["none"] * (mask if arm else 1.0)
        return dict(split=(lowered - raised) / 2, common=(raised + lowered) / 2 - none, raised=raised - none,
                    lowered=lowered - none)

    medium_at, medium_over = parts(at), parts(over)
    push_at, push_over = parts(at, "push "), parts(over, "push ")
    sham_at, sham_over = parts(at, "sham "), parts(over, "sham ")
    result = dict(
        at_R=dict(medium={k: summary(by_start(v)) for k, v in medium_at.items()},
                  push={k: summary(by_start(v)) for k, v in push_at.items()},
                  sham={k: summary(by_start(v)) for k, v in sham_at.items()},
                  difference=summary(by_start(medium_at["split"] * mask - push_at["split"]))),
        growth_window=dict(medium={k: summary(by_start(v)) for k, v in medium_over.items()},
                           push={k: summary(by_start(v)) for k, v in push_over.items()},
                           sham={k: summary(by_start(v)) for k, v in sham_over.items()},
                           difference=summary(by_start(medium_over["split"] * mask - push_over["split"]))),
        share=dict(medium=ratio_summary(by_start(medium_over["split"]), by_start(medium_at["split"])),
                   push=ratio_summary(by_start(push_over["split"]), by_start(push_at["split"]))))

    # 100-unit blocks: the medium and untouched from T, the pushes and shams from R
    blocks = dict(medium_split=[], medium_common=[], push_split=[], sham_split=[], difference=[])
    levels = {key: [] for key in ("untouched", "medium raised", "medium lowered", "push raised", "push lowered",
                                  "sham raised", "sham lowered")}
    for b in range(0, W + AFTER, BLOCK):
        g = dict(none=grid(lambda r, b=b: r["g_none"][b + 1:b + BLOCK + 1].mean()))
        for d, name in enumerate(DIRECTIONS):
            g[name] = grid(lambda r, b=b, d=d: r["g_medium"][d][b + 1:b + BLOCK + 1].mean())
        p = parts(g)
        blocks["medium_split"].append(summary(by_start(p["split"])))
        blocks["medium_common"].append(summary(by_start(p["common"])))
        levels["untouched"].append(float(np.mean(by_start(g["none"]))))
        for name in DIRECTIONS:
            levels[f"medium {name}"].append(float(np.mean(by_start(g[name]))))
    for b in range(0, AFTER, BLOCK):
        g = dict(none=grid(lambda r, b=b: r["g_none"][W + b + 1:W + b + BLOCK + 1].mean()))
        for d, name in enumerate(DIRECTIONS):
            g[name] = grid(lambda r, b=b, d=d: r["g_medium"][d][W + b + 1:W + b + BLOCK + 1].mean())
            for arm in ("push", "sham"):
                g[f"{arm} {name}"] = grid(lambda r, b=b, d=d, arm=arm: r[f"g_{arm}"][d][b + 1:b + BLOCK + 1].mean()) \
                    * mask
                levels[f"{arm} {name}"].append(float(np.mean(by_start(g[f"{arm} {name}"]))))
        pm, pp, ps = parts(g), parts(g, "push "), parts(g, "sham ")
        blocks["push_split"].append(summary(by_start(pp["split"])))
        blocks["sham_split"].append(summary(by_start(ps["split"])))
        blocks["difference"].append(summary(by_start(pm["split"] * mask - pp["split"])))

    s_r, s_l = result["at_R"]["medium"]["split"], result["growth_window"]["medium"]["split"]
    if s_l["clear"] and np.sign(s_l["mean"]) == np.sign(s_r["mean"]):
        verdict = "lasts"
    elif s_l["clear"]:
        verdict = "reverses"
    elif s_r["clear"]:
        verdict = "fades before"
    else:
        verdict = "not clear at R"
    clear_blocks = [i for i, s in enumerate(blocks["medium_split"]) if s["clear"] and s["mean"] > 0]

    names_from_t = ["(T, T+100]", "(T+100, R]"] + [f"(R+{b}, R+{b + BLOCK}]" for b in range(0, AFTER, BLOCK)]
    names_from_r = [f"(R+{b}, R+{b + BLOCK}]" for b in range(0, AFTER, BLOCK)]
    made = datetime.now(ZoneInfo("America/Halifax")).strftime("%Y-%m-%d %H:%M")
    lines = [
        f"Check 1 of the pilot review: the timing gap between groups A and B over time (run {made} Halifax).",
        "After the fact and exploratory: a description of the pilot's saved results, not a reading fixed in plan.md and "
        "not a new experiment. No model was run.",
        "A toy model: structure inside the model, not evidence about cells. Histories, not memory: differences are "
        "between intervention histories and the states they leave.",
        f"Read: the pilot's 120 moment files (starts 103 to 114), each matching the hash in pilot_results.json. The gap "
        f"reproduces the report's |A-B timing| at R (largest difference {worst:.1e} rad).",
        "Gap: |psi_A - psi_B| in radians; 0 means the two groups' rhythms are in step, about 1.571 (pi/2) is what "
        "unrelated rhythms average. Split: half of (lowered gap minus raised gap); positive means the raised arm's "
        "groups sit closer in timing. Error bars and 'clear' (more than 3 bars) as in the report.",
        "",
        f"Answer: {verdict}.",
        "",
        "The medium (all 120 moments)",
        f"  split at R: {shown(s_r)}",
        f"  split over (R+200, R+500]: {shown(s_l)}",
        f"  share of the split at R still there over (R+200, R+500]: {result['share']['medium']['mean']:+.3f} +- "
        f"{result['share']['medium']['bar']:.3f}",
        f"  raised minus untouched: at R {shown(result['at_R']['medium']['raised'])}",
        f"                          over (R+200, R+500] {shown(result['growth_window']['medium']['raised'])}",
        f"  lowered minus untouched: at R {shown(result['at_R']['medium']['lowered'])}",
        f"                           over (R+200, R+500] {shown(result['growth_window']['medium']['lowered'])}",
        f"  common part: at R {shown(result['at_R']['medium']['common'])}",
        f"               over (R+200, R+500] {shown(result['growth_window']['medium']['common'])}",
        "",
        "Pushes and shams (moments with both pushes built)",
        f"  push split: at R {shown(result['at_R']['push']['split'])}",
        f"              over (R+200, R+500] {shown(result['growth_window']['push']['split'])}",
        f"  sham split: at R {shown(result['at_R']['sham']['split'])}",
        f"              over (R+200, R+500] {shown(result['growth_window']['sham']['split'])}",
        "",
        "Medium against matched push, like for like (split difference)",
        f"  at R: {shown(result['at_R']['difference'])}",
        f"  over (R+200, R+500]: {shown(result['growth_window']['difference'])}",
        "",
        "The medium's split by 100-unit block (from T; R = T+200)",
        *(f"  {name}: {shown(s)}" for name, s in zip(names_from_t, blocks["medium_split"])),
        f"  blocks with a clear positive split: {', '.join(names_from_t[i] for i in clear_blocks) or 'none'}",
        "The medium's common part by block (from T)",
        *(f"  {name}: {shown(s)}" for name, s in zip(names_from_t, blocks["medium_common"])),
        "Push split, sham split and medium-minus-push split difference by block (from R)",
        *(f"  {name}: push {shown(a)}; sham {shown(b_)}; difference {shown(c)}"
          for name, a, b_, c in zip(names_from_r, blocks["push_split"], blocks["sham_split"], blocks["difference"])),
        "",
        "Average gap by block (radians; untouched and medium from T, pushes and shams from R)",
        *(f"  {key}: " + "  ".join(f"{v:.4f}" for v in values) for key, values in levels.items()),
        "",
        "Limits",
        "  Lower coherence makes each group's rhythm noisier, which pulls the average gap toward about 1.571. The lowered "
        "arm stays less coherent after R, so part of its wider gap can come from that. The matched push, at nearly the "
        "same coherence, is the comparison that separates timing from coherence level.",
        "  Whether the gap carries the later-growth difference is not tested here; this only says whether it is still "
        "there when that difference appears.",
        "",
        f"Compute: {time.time() - began:.1f} seconds, in the foreground.",
    ]
    OUT_TXT.write_text("\n".join(lines) + "\n")
    OUT_JSON.write_text(json.dumps(dict(made_halifax=made, verdict=verdict, worst_at_R=worst, result=result,
                                        blocks=blocks, levels=levels, block_names_from_T=names_from_t,
                                        block_names_from_R=names_from_r), indent=1))
    print("\n".join(lines))


if __name__ == "__main__":
    main()
