"""Check 2 of the pilot review: do moments with bigger slower-half or timing differences show bigger medium-versus-push
growth differences?

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
check2_moments.txt and check2_moments.json beside this file. Nothing else is changed.

Question (from pilot_review.md): at R the medium and its matched push had the same coherence but held it differently.
The medium put more of it in the slower-turning half of each group, and it had shifted the timing between groups A and
B. Across the 119 moments with both pushes built, do the moments where these differences were bigger also show a bigger
medium-versus-push difference in later growth? Check 1 found that the push reaches most of the timing shift within 100
units, so timing goes in as the weaker candidate.

Per moment (moments with both pushes built), each in the form of the report's S difference, so that each is positive
in the direction seen on average:
- Growth difference Y: half of [(lowered medium - lowered push) - (raised medium - raised push)] in later growth over
  (R + 200, R + 500]. It is the report's S difference, moment by moment; the script first checks that it reproduces
  the report's value.
- Slower-half difference X_slow: half of [(raised medium - raised push) - (lowered medium - lowered push)] in the
  coherence of each group's slower-turning half at R, averaged over the two groups. Positive: the medium holds more of
  its coherence there than its push does, relative to the other arm.
- Timing difference X_time: half of [(lowered medium - lowered push) - (raised medium - raised push)] in the A-B timing
  gap at R. Positive: the medium's raised groups sit closer in timing than its push's, relative to the other arm. It is
  check 1's like-for-like difference at R, moment by moment.
- The size of the coherence change at R, dC: half of (raised medium - lowered medium) in coherence. The pushes share it
  by construction.

Fixed before looking at the numbers:
- Main measure, for each of X_slow and X_time: the correlation with Y across each start's moments (Pearson), averaged
  over the 12 starts, with the report's error bar (the larger of the resampled one, starts drawn with replacement 400
  times with seed 20261016, and SD / sqrt(12)). Clear means more than 3 bars.
- "Goes with" if clear and positive, the direction the candidate explanation expects; "goes against" if clear and
  negative; "not clear" otherwise.
- Descriptions, not readings:
  - the same correlation with dC held fixed (a partial correlation within each start);
  - the raised and lowered arms separately;
  - the growth difference in each start's half of moments with the biggest X against its half with the smallest (the
    middle moment left out when a start has an odd number);
  - the share of the average growth difference that the link would account for if it held all the way down to no
    difference: the pooled within-start slope times the average X, over the average Y, with its error bar from the same
    resampling of starts;
  - the correlations of X_slow and X_time with each other and with dC, within starts, to show how entangled they are.

Limits stated before looking:
- Each moment's growth difference is noisy; the average difference itself is 3.6 bars over 12 starts. With about ten
  moments per start, a correlation's error bar should be roughly 0.1, so only a correlation of about 0.3 or more can be
  clear. "Not clear" would not show that a feature doesn't matter.
- A link across moments is not cause. Moments where the medium acted more strongly could show bigger differences of
  every kind; holding dC fixed addresses only the size of the coherence change.
"""

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 / "check2_moments.txt", HERE / "check2_moments.json"
STARTS, MOMENTS = tuple(range(103, 115)), tuple(range(1, 11))
W, LATER = 200, (200, 500)                                 # window; growth window after R
RESAMPLES, SEED, CLEAR = 400, 20261016, 3.0                # as in the report
STATES = ("untouched", "medium raised", "medium lowered", "push raised", "push lowered", "sham raised", "sham lowered")
TOL = 1e-12


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


def growth(logs, lo, hi):
    """The report's growth over (lo, hi]: the gap's log growth per time unit, from logs kept every 5 units."""
    return float(np.sum(logs[lo // 5:hi // 5]) / (hi - lo))


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 pearson(x, y):
    return float(np.corrcoef(x, y)[0, 1])


def partial(x, y, c):
    """The correlation of x and y with c held fixed: both regressed on c (with a constant), residuals correlated."""
    A = np.column_stack([np.ones_like(c), c])
    rx = x - A @ np.linalg.lstsq(A, x, rcond=None)[0]
    ry = y - A @ np.linalg.lstsq(A, y, rcond=None)[0]
    return pearson(rx, ry)


def halves(x, y):
    """Average y over the half of moments with the biggest x, minus the half with the smallest."""
    order, h = np.argsort(x), len(x) // 2
    return float(y[order[-h:]].mean() - y[order[:h]].mean())


def per_start(fn, *grids):
    """fn applied within each start, to the moments where every grid has a value."""
    out = []
    for i in range(len(STARTS)):
        keep = np.all([np.isfinite(g[i]) for g in grids], axis=0)
        out.append(fn(*(g[i][keep] for g in grids)) if keep.sum() >= 4 else np.nan)
    return np.array(out, float)


def share(X, Y, idx):
    """Pooled within-start slope of Y on X, times the average X, over the average Y, for the starts in idx."""
    sxy = sxx = 0.0
    xm, ym = [], []
    for i in idx:
        keep = np.isfinite(X[i]) & np.isfinite(Y[i])
        x, y = X[i][keep], Y[i][keep]
        sxy += float(np.sum((x - x.mean()) * (y - y.mean())))
        sxx += float(np.sum((x - x.mean()) ** 2))
        xm.append(x.mean())
        ym.append(y.mean())
    slope = sxy / sxx
    return slope * float(np.mean(xm)) / float(np.mean(ym)), slope


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


def main():
    began = time.time()
    report = json.loads((HERE / "pilot_results.json").read_text())
    hashes = report["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:
                rows.append({key: p[key] for key in ("z_medium", "logs_medium", "logs_push", "built", "desc_inner",
                                                     "desc_ab")})
    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)
    lo, hi = LATER

    # later growth over (R + 200, R + 500]: medium logs count from T, push logs from R
    g = dict(medium_raised=grid(lambda r: growth(r["logs_medium"][0], W + lo, W + hi)),
             medium_lowered=grid(lambda r: growth(r["logs_medium"][1], W + lo, W + hi)),
             push_raised=grid(lambda r: growth(r["logs_push"][0], lo, hi)) * mask,
             push_lowered=grid(lambda r: growth(r["logs_push"][1], lo, hi)) * mask)
    Y_raised = -(g["medium_raised"] - g["push_raised"])      # positive: the medium's raised arm grew more slowly
    Y_lowered = g["medium_lowered"] - g["push_lowered"]      # positive: the medium's lowered arm grew faster
    Y = (Y_raised + Y_lowered) / 2

    # the states at R, averaged over the two groups as in the report
    inner = {state: grid(lambda r, i=i: np.mean(np.abs(r["desc_inner"][i]))) for i, state in enumerate(STATES)}
    ab = {state: grid(lambda r, i=i: abs(float(r["desc_ab"][i]))) for i, state in enumerate(STATES)}
    X_slow_raised = (inner["medium raised"] - inner["push raised"]) * mask
    X_slow_lowered = -(inner["medium lowered"] - inner["push lowered"]) * mask
    X_slow = (X_slow_raised + X_slow_lowered) / 2
    X_time_raised = -(ab["medium raised"] - ab["push raised"]) * mask
    X_time_lowered = (ab["medium lowered"] - ab["push lowered"]) * mask
    X_time = (X_time_raised + X_time_lowered) / 2

    def C(z):
        return np.abs(z).mean(-2)

    dC = grid(lambda r: (C(r["z_medium"][0])[W] - C(r["z_medium"][1])[W]) / 2)

    # the numbers reproduce the report (and check 1) before anything else is read
    checks = {}
    mine = summary(by_start(Y))
    theirs = report["later"]["S_difference"]
    checks["growth difference (S difference)"] = max(abs(mine["mean"] - theirs["mean"]), abs(mine["bar"] - theirs["bar"]))
    checks["size of the coherence change (dC)"] = abs(float(np.mean(by_start(dC))) - report["dC"])
    worst_state = 0.0
    for i, state in enumerate(STATES):
        m_ = mask if i >= 3 else 1.0
        for key, values in (("inner", inner[state]), ("ab", ab[state])):
            worst_state = max(worst_state, abs(float(np.mean(by_start(values * m_)))
                                               - report["descriptions"]["states"][state][key]))
    checks["slower-half and timing descriptions at R"] = worst_state
    check1 = HERE / "check1_timing.json"
    if check1.exists():
        c1 = json.loads(check1.read_text())["result"]["at_R"]["difference"]["mean"]
        checks["timing difference at R (check 1)"] = abs(float(np.mean(by_start(X_time))) - c1)
    bad = {key: value for key, value in checks.items() if not value <= (1e-9 if "check 1" in key else TOL)}
    if bad:
        sys.exit(f"These differ from the report: {bad}. Stopping.")

    # main measure and descriptions
    predictors = dict(slow=(X_slow, X_slow_raised, X_slow_lowered), time=(X_time, X_time_raised, X_time_lowered))
    result = dict(n_moments=int(np.isfinite(mask).sum()), averages=dict(
        growth_difference=summary(by_start(Y)), growth_difference_raised=summary(by_start(Y_raised)),
        growth_difference_lowered=summary(by_start(Y_lowered)), size_of_change=summary(by_start(dC * mask))))
    rng_idx = np.random.default_rng(SEED).integers(0, S_, (RESAMPLES, S_))
    verdicts = {}
    for name, (X, Xr, Xl) in predictors.items():
        main_ = summary(per_start(pearson, X, Y))
        verdicts[name] = ("goes with" if main_["clear"] and main_["mean"] > 0
                          else "goes against" if main_["clear"] else "not clear")
        value, slope = share(X, Y, range(S_))
        boot = np.array([share(X, Y, idx)[0] for idx in rng_idx])
        result[name] = dict(
            average=summary(by_start(X)), correlation=main_, verdict=verdicts[name],
            correlation_dC_fixed=summary(per_start(partial, X, Y, dC * mask)),
            correlation_raised=summary(per_start(pearson, Xr, Y_raised)),
            correlation_lowered=summary(per_start(pearson, Xl, Y_lowered)),
            halves=summary(per_start(halves, X, Y)),
            share=dict(mean=value, bar=float(boot.std(ddof=1)), slope=slope,
                       clear=bool(abs(value) > CLEAR * float(boot.std(ddof=1)))))
    result["entangled"] = dict(
        slow_with_time=summary(per_start(pearson, X_slow, X_time)),
        slow_with_dC=summary(per_start(pearson, X_slow, dC * mask)),
        time_with_dC=summary(per_start(pearson, X_time, dC * mask)),
        growth_with_dC=summary(per_start(pearson, Y, dC * mask)))

    made = datetime.now(ZoneInfo("America/Halifax")).strftime("%Y-%m-%d %H:%M")
    label = dict(slow="Slower half (X_slow)", time="Timing (X_time)")
    lines = [
        f"Check 2 of the pilot review: do moments with bigger slower-half or timing differences show bigger "
        f"medium-versus-push growth differences? (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; "
        f"{result['n_moments']} moments with both pushes built. Reproduced before reading on: "
        + "; ".join(f"{key} (largest difference {value:.1e})" for key, value in checks.items()) + ".",
        "Each difference is medium minus matched push, in the form of the report's S difference: positive is the "
        "direction seen on average. Correlations are within each start, averaged over the 12 starts. Error bars and "
        "'clear' (more than 3 bars) as in the report.",
        "",
        f"Answer: slower half: {verdicts['slow']}; timing: {verdicts['time']}.",
        "",
        "Averages (like for like)",
        f"  growth difference Y: {shown(result['averages']['growth_difference'], 5)}",
        f"    raised arm: {shown(result['averages']['growth_difference_raised'], 5)}; lowered arm: "
        f"{shown(result['averages']['growth_difference_lowered'], 5)}",
        f"  slower-half difference X_slow: {shown(result['slow']['average'])}",
        f"  timing difference X_time (radians): {shown(result['time']['average'])}",
        f"  size of the coherence change dC: {shown(result['averages']['size_of_change'])}",
        "",
    ]
    for name in ("slow", "time"):
        r = result[name]
        lines += [
            label[name],
            f"  correlation with the growth difference (the main measure): {shown(r['correlation'], 3)}",
            f"    with the size of the coherence change held fixed: {shown(r['correlation_dC_fixed'], 3)}",
            f"    raised arm alone: {shown(r['correlation_raised'], 3)}",
            f"    lowered arm alone: {shown(r['correlation_lowered'], 3)}",
            f"  growth difference, biggest-difference half of moments minus smallest: {shown(r['halves'], 5)}",
            f"  share of the average growth difference the link would account for: {r['share']['mean']:+.2f} +- "
            f"{r['share']['bar']:.2f} (pooled within-start slope {r['share']['slope']:+.4g}; an extrapolation, "
            f"meaningful only if the link is clear)",
            "",
        ]
    e = result["entangled"]
    lines += [
        "How entangled the measures are (correlations within starts)",
        f"  slower half with timing: {shown(e['slow_with_time'], 3)}",
        f"  slower half with the size of the coherence change: {shown(e['slow_with_dC'], 3)}",
        f"  timing with the size of the coherence change: {shown(e['time_with_dC'], 3)}",
        f"  growth difference with the size of the coherence change: {shown(e['growth_with_dC'], 3)}",
        "",
        "Limits",
        "  Each moment's growth difference is noisy, so only fairly strong links can be clear here; 'not clear' does not "
        "show that a feature doesn't matter.",
        "  A link across moments is not cause: moments where the medium acted more strongly could show bigger "
        "differences of every kind. Holding the size of the coherence change fixed addresses only that one.",
        "  The slower-half and timing differences are read at R only; the individual oscillators after R weren't saved.",
        "",
        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, verdicts=verdicts, reproduced=checks, result=result),
                                   indent=1))
    print("\n".join(lines))


if __name__ == "__main__":
    main()