#!/usr/bin/env python3
"""
Coherence and growth (October 11, 2026): when a group is more coherent, do tiny differences grow more slowly?

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

What led here
    The long run at the threshold (Test 16 of Memory in Motion; nudge_test.py, run earlier on October 11: N = 4,000
    per group, K = 1.90, 6 starts, 6,000 time units) printed the coherence |z| and the growth rate of a tiny gap in
    windows of 1,000 time units. Averaged over the starts, the windows with higher coherence had slower growth. That
    was noticed after the run, not tested.

Looked at again before this test (the same run, so not a test either)
    Start by start, over t = 2,000 to 6,000 (after the slow settling):
    - In windows of 25 to 500 time units, higher coherence went with faster growth in the same window: a correlation
      of about +0.3 at every window length, above 0 in every start at 25 to 200.
    - In windows of 1,000 time units the sign depends on where the windows are cut: -0.42 with the windows as printed
      (2,000-3,000 to 5,000-6,000), then +0.07, +0.49 and +0.34 with the windows moved by 250, 500 and 750 time
      units (3 or 4 windows per start).
    - Over time: higher coherence went with slower growth 200 to 500 time units later (about -0.2 to -0.3, below 0
      in every start), while higher growth went with higher coherence later.
    So the noticing may rest on where the windows were cut, and two other patterns turned up. Patterns found by
    looking at one run need fresh runs to test them.

The test (everything in this section was set before the run)
    Six fresh starts at the same setting: starts 7 to 12 (random phases and gap directions never used at this size),
    N = 4,000, K = 1.90, the same model, settling, gap (rms 1e-7 radians) and resets (every 5 time units), for 6,000
    time units, run by nudge_test.py's own code, imported unchanged. Measured start by start over t = 2,000 to 6,000,
    from the untouched copy's coherence (average of A and B) and the growth of the log of the gap:
    1. Same time: the correlation, window by window, between the coherence and the growth rate in windows of 200 time
       units (20 per start).
    2. Later: the correlation between the coherence in each 5-unit interval and the growth rate L time units later,
       averaged over L = 200 to 500.
    Each is averaged over the six starts, with error bars from the starts (resampled 400 times, as in every earlier
    part). Reading: a link if the average is more than 3 error bars from 0; below 0, more coherence goes with slower
    growth; above 0, with faster growth. Otherwise: no clear link.
    Shown but not read as tests: the same-time correlation in windows of 25 to 500 time units, windows of 1,000 at the
    four placements, the correlation at lags from -500 to +500, and all of it for the earlier run beside the fresh one.
    Also, beside Test 16: the fresh starts' long-run rate over t = 3,000 to 6,000 (Test 16's span) against Test 16's
    +0.0147 (0.0006); they agree if they differ by no more than 3 combined error bars.
    Check: start 1, run by this script for 50 time units, must repeat the long run exactly (the same code gives the
    same numbers).

Usage
    python3 coherence_growth.py check      # start 1 for 50 time units against the long run
    python3 coherence_growth.py run        # again and again until it saves sim_fresh4000.npz; each piece stops
                                           # after about 4 minutes, and unfinished starts keep their progress in
                                           # parts/ and carry on the next time
    python3 coherence_growth.py report     # coherence_growth.txt, coherence_growth.json and the figure

Needs Python 3, numpy and matplotlib, and ../nudge-test-2026-10-10/nudge_test.py (unchanged). The figure goes to
screenshots/nudge-test-coherence-growth.png.
"""
import json
import sys
import time
from multiprocessing import Pool
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
NUDGE = HERE.parent / "nudge-test-2026-10-10"
sys.dont_write_bytecode = True                             # leave the nudge test's folder as it is
sys.path.insert(0, str(NUDGE))
import nudge_test as nt  # noqa: E402  (the model and the gap method, unchanged)

nt.LYAP_PARTS = HERE / "parts"                             # progress of unfinished starts stays in this folder

N, K = 4000, 1.9
FRESH = tuple(range(6, 12))                                # starts 7 to 12: never run at this size
RUN = 6000.0                                               # time units after settling, as in the long run
SPAN = (2000.0, 6000.0)                                    # measured here, after the slow settling (set before the run)
WINDOW = 200.0                                             # measure 1: windows of this length
LAGS = (200.0, 500.0)                                      # measure 2: the growth this much later, averaged over the range
RULE = 3.0                                                 # a link if the average is more than this many error bars from 0
LONG_FROM = 3000.0                                         # the long-run rate is read from here on (Test 16's span)
SIZES = (25.0, 50.0, 100.0, 200.0, 500.0)                  # shown: the same-time link at other window lengths
PLACEMENTS = (0.0, 250.0, 500.0, 750.0)                    # shown: windows of 1,000 moved by these amounts
CURVE_LAGS = np.arange(-500, 501, 5)                       # shown: the link at every lag (negative: the growth first)
TABLE_LAGS = (-500, -300, -200, -100, -50, 0, 50, 100, 200, 300, 500)
CHECK_RUN = 50.0
SAVED, EARLIER = HERE / "sim_fresh4000.npz", NUDGE / "sim_lyapunov4000_long.npz"
CHECK_FILE = HERE / "check.json"
FIGURE = HERE.parent.parent / "screenshots" / "nudge-test-coherence-growth.png"
SEED = 20261015                                            # resampling the starts for the error bars


def check():
    """Start 1, run by this script for 50 time units, against the long run's own numbers."""
    res = nt.lyapunov((N, K, (0, 1), nt.SETTLE, CHECK_RUN, nt.LYAP_EVERY, nt.LYAP_GAP, nt.DT, None))
    old = np.load(EARLIER)
    samples, resets = int(CHECK_RUN) + 1, int(round(CHECK_RUN / nt.LYAP_EVERY))
    same = dict(coherence=bool(np.array_equal(res["z"], old["z"][0, :1, :, :samples])),
                gap=bool(np.array_equal(res["logs"], old["logs"][0, :1, :resets])),
                settling=bool(np.array_equal(res["trace"], old["trace"][0, :1])))
    CHECK_FILE.write_text(json.dumps(dict(start=1, run=CHECK_RUN, **same, seconds=round(res["seconds"], 1)), indent=1))
    print(f"Check: start 1 for {CHECK_RUN:.0f} time units against the long run: settling {same['settling']}, coherence"
          f" {same['coherence']}, gap {same['gap']} ({'repeats exactly' if all(same.values()) else 'did NOT repeat'};"
          f" {res['seconds']:.0f} s)")


def run():
    """The six fresh starts, in pieces of at most nudge_test's LYAP_BUDGET seconds."""
    if SAVED.exists():
        print(f"  nothing left to run: {SAVED.name} is saved")
        return
    deadline = time.time() + nt.LYAP_BUDGET
    tasks = [(N, K, (s, 1), nt.SETTLE, RUN, nt.LYAP_EVERY, nt.LYAP_GAP, nt.DT, deadline) for s in FRESH]
    clock, results = time.time(), []
    with Pool(len(tasks)) as pool:
        for res in pool.imap_unordered(nt.lyapunov, tasks):
            results.append(res)
            where = "finished" if res["finished"] else f"paused at t = {res['step'] * nt.DT:,.0f}"
            print(f"  {where}: start {res['first'] + 1} ({res['seconds']:.0f} s so far)", flush=True)
    done = sum(res["finished"] for res in results)
    if done < len(results):
        print(f"  {done} of {len(results)} starts finished; run the same command again to carry on"
              f" (this piece {time.time() - clock:.0f} s)")
        return
    results.sort(key=lambda res: res["first"])
    np.savez_compressed(SAVED, starts=np.array(FRESH), n=N, k=K, dt=nt.DT, settle=nt.SETTLE, run=RUN,
                        every=nt.LYAP_EVERY, gap=nt.LYAP_GAP, seed=nt.LYAP_SEED,
                        **{key: np.concatenate([res[key] for res in results]) for key in ("z", "logs", "trace")},
                        seconds=np.array([res["seconds"] for res in results]))
    for res in results:
        nt.lyap_progress(N, K, res["first"], RUN).unlink()
    if nt.LYAP_PARTS.exists() and not any(nt.LYAP_PARTS.iterdir()):
        nt.LYAP_PARTS.rmdir()
    print(f"  saved {SAVED.name} (this piece {time.time() - clock:.0f} s)")


def measures(z, logs, rng):
    """Every measure for one set of starts. z: the untouched copy's coherence, (start, group, time), every time unit
    from t = 0; logs: the growth of the log of the gap in each 5-unit interval, (start, interval)."""
    every = nt.LYAP_EVERY
    r = np.abs(z.astype(np.complex128)).mean(1)                            # (start, time): A and B averaged
    count = len(r)
    r5 = r[:, 1:].reshape(count, -1, int(round(every))).mean(-1)           # interval j: t in (5j, 5j + 5], as logs[:, j]
    assert r5.shape == logs.shape
    a, b = (int(round(x / every)) for x in SPAN)
    picks = rng.integers(0, count, (400, count))

    def corr(x, y):                                                         # start by start
        x = x - x.mean(-1, keepdims=True)
        y = y - y.mean(-1, keepdims=True)
        return (x * y).sum(-1) / np.sqrt((x * x).sum(-1) * (y * y).sum(-1))

    def windows(length, shift=0.0):                                         # coherence and growth rate, window by window
        size, low = int(round(length / every)), a + int(round(shift / every))
        cut = slice(low, low + (b - low) // size * size)
        return (r5[:, cut].reshape(count, -1, size).mean(-1), logs[:, cut].reshape(count, -1, size).sum(-1) / length)

    def lagged(lag):                                                        # coherence against the growth `lag` later
        steps = int(round(lag / every))
        if steps >= 0:
            return corr(r5[:, a:b - steps], logs[:, a + steps:b])
        return corr(r5[:, a - steps:b], logs[:, a:b + steps])

    def summary(values):
        return dict(mean=float(values.mean()), se=float(values[picks].mean(1).std()), starts=values.tolist(),
                    below=int((values < 0).sum()))

    later_lags = np.arange(LAGS[0], LAGS[1] + every / 2, every)
    curve = np.stack([lagged(lag) for lag in CURVE_LAGS], 1)               # (start, lag)
    x, y = windows(WINDOW)
    long = logs[:, int(round(LONG_FROM / every)):].sum(1) / (RUN - LONG_FROM)
    return dict(
        starts=count,
        same=summary(corr(x, y)),
        later=summary(np.mean([lagged(lag) for lag in later_lags], 0)),
        sizes={f"{w:g}": summary(corr(*windows(w))) for w in SIZES},
        placements={f"{p:g}": summary(corr(*windows(1000.0, p))) for p in PLACEMENTS},
        curve=curve.mean(0).tolist(), curve_se=curve[picks].mean(1).std(0).tolist(),
        long_rate=float(long.mean()), long_rate_se=float(long[picks].mean(1).std()), long_rate_starts=long.tolist(),
        coherence=float(r[:, int(SPAN[0]):].mean()),
        scatter=dict(coherence=(x - x.mean(1, keepdims=True)).tolist(), growth=(y - y.mean(1, keepdims=True)).tolist()))


def reading(s, later=False):
    if abs(s["mean"]) <= RULE * s["se"]:
        return "no clear link"
    if later:
        return "more coherence now, slower growth later" if s["mean"] < 0 else "more coherence now, faster growth later"
    return "more coherence, slower growth" if s["mean"] < 0 else "more coherence, faster growth"


def report():
    fresh_data, old = np.load(SAVED), np.load(EARLIER)
    assert int(fresh_data["n"]) == N and float(fresh_data["run"]) == RUN == float(old["run"])
    assert all(float(fresh_data[key]) == float(old[key]) for key in ("dt", "settle", "every", "gap", "seed"))
    assert tuple(int(s) for s in fresh_data["starts"]) == FRESH
    i = [round(float(x), 2) for x in old["ks"]].index(K)
    fresh = measures(fresh_data["z"], fresh_data["logs"], np.random.default_rng(SEED))
    earlier = measures(old["z"][i], old["logs"][i], np.random.default_rng(SEED))
    test16 = json.loads((NUDGE / "lyapunov_long.json").read_text())
    diff = fresh["long_rate"] - test16["rate"]
    combined = float(np.hypot(fresh["long_rate_se"], test16["rate_se"]))
    agree = bool(abs(diff) <= RULE * combined)
    checked = json.loads(CHECK_FILE.read_text()) if CHECK_FILE.exists() else None
    result = dict(
        model=dict(n=N, k=K, starts=[s + 1 for s in FRESH], run=RUN, span=SPAN, window=WINDOW, later=LAGS, rule=RULE,
                   gap=nt.LYAP_GAP, put_back_every=nt.LYAP_EVERY, delta=nt.DELTA, eps=nt.EPS, dt=nt.DT,
                   settle=nt.SETTLE),
        check=checked, fresh=fresh, earlier=earlier,
        readings=dict(same=reading(fresh["same"]), later=reading(fresh["later"], later=True)),
        test16=dict(rate=test16["rate"], rate_se=test16["rate_se"], difference=diff, combined_se=combined,
                    agree=agree))
    lines = []

    def say(text=""):
        print(text)
        lines.append(text)

    def cell(s):
        return f"{s['mean']:+.4f} ({s['se']:.4f})"

    def starts(s):
        return ", ".join(f"{v:+.2f}" for v in s["starts"])

    say(f"Coherence and growth, October 11, 2026: N = {N:,} per group, K = {K:.2f} (the threshold), eps = {nt.EPS:g}")
    say(f"Six fresh starts (starts {FRESH[0] + 1} to {FRESH[-1] + 1}) of the long run at the threshold: the same model,"
        f" settling, gap (rms {nt.LYAP_GAP:g} radians, put back every {nt.LYAP_EVERY:g} time units) and length"
        f" ({RUN:,.0f} time units), run by nudge_test.py's own code.")
    say(f"Set before the run: measured start by start over t = {SPAN[0]:,.0f} to {SPAN[1]:,.0f}. 1. Same time: the"
        f" correlation between coherence and growth rate in windows of {WINDOW:.0f} time units. 2. Later: the correlation"
        f" between coherence and the growth rate {LAGS[0]:.0f} to {LAGS[1]:.0f} time units later. A link if the average"
        f" over the starts is more than {RULE:g} error bars from 0. Error bars: from the starts.")
    if checked:
        how = "exactly" if all(checked[key] for key in ("settling", "coherence", "gap")) else "NOT (check failed)"
        say(f"Check: start 1, run by this script for {checked['run']:.0f} time units, repeats the long run {how}.")
    say(f"\n1. Same time (windows of {WINDOW:.0f}): {cell(fresh['same'])}; starts {starts(fresh['same'])}; below 0 in"
        f" {fresh['same']['below']} of {fresh['starts']}: {result['readings']['same']}")
    say(f"2. Later ({LAGS[0]:.0f} to {LAGS[1]:.0f} later): {cell(fresh['later'])}; starts {starts(fresh['later'])}; below 0"
        f" in {fresh['later']['below']} of {fresh['starts']}: {result['readings']['later']}")
    say(f"Beside Test 16: the fresh starts' long-run rate, t = {LONG_FROM:,.0f} to {RUN:,.0f}:"
        f" {fresh['long_rate']:+.4f} ({fresh['long_rate_se']:.4f}); Test 16: {test16['rate']:+.4f} ({test16['rate_se']:.4f});"
        f" difference {diff:+.4f} ({combined:.4f}): {'agree' if agree else 'differ'}")
    say(f"\nShown, not read as tests: the fresh starts, then the earlier run (where the patterns were noticed)")
    say("Same time, correlation by window length (time units)")
    heads = [f"{w:g}" for w in SIZES] + [f"1,000 (+{p:g})" for p in PLACEMENTS]
    say("            " + "".join(f"{h:>18}" for h in heads))
    for name, m in (("fresh", fresh), ("earlier", earlier)):
        cells = [m["sizes"][f"{w:g}"] for w in SIZES] + [m["placements"][f"{p:g}"] for p in PLACEMENTS]
        say(f"  {name:<10}" + "".join(f"{cell(c):>18}" for c in cells))
    say("Coherence against the growth rate L time units later (negative L: the growth comes first)")
    say("            " + "".join(f"{lag:>+9d}" for lag in TABLE_LAGS))
    for name, m in (("fresh", fresh), ("earlier", earlier)):
        say(f"  {name:<10}" + "".join(f"{m['curve'][list(CURVE_LAGS).index(lag)]:>+9.3f}" for lag in TABLE_LAGS))
    say(f"Coherence |z| over t = {SPAN[0]:,.0f} to {SPAN[1]:,.0f} (average of A and B over the starts): fresh"
        f" {fresh['coherence']:.4f}, earlier {earlier['coherence']:.4f}")
    say(f"The two measures for the earlier run: same time {cell(earlier['same'])}; later {cell(earlier['later'])}")
    (HERE / "coherence_growth.json").write_text(json.dumps(result, indent=1))
    (HERE / "coherence_growth.txt").write_text("\n".join(lines) + "\n")
    figure(result)
    print(f"\nWrote coherence_growth.json, coherence_growth.txt and {FIGURE.relative_to(HERE.parent.parent)}")


def figure(result):
    """The two measures set before the run, the link at every lag and at every window length."""
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    blue, grey = "#377eb8", "0.55"
    fresh, earlier = result["fresh"], result["earlier"]
    fig = plt.figure(figsize=(16, 11.6))
    grid = fig.add_gridspec(2, 2, left=0.06, right=0.985, top=0.885, bottom=0.12, hspace=0.33, wspace=0.2)
    fig.suptitle(f"Coherence and growth (October 11, 2026): {N:,} per group, K = {K:.2f}, six fresh starts\nWhen the"
                 " group is more coherent, do tiny differences grow more slowly?", fontsize=14)

    ax = fig.add_subplot(grid[0, 0])
    for x, y in zip(earlier["scatter"]["coherence"], earlier["scatter"]["growth"]):
        ax.plot(x, y, "o", ms=3, color=grey, alpha=0.5, label=None)
    colors = plt.cm.viridis(np.linspace(0, 0.9, len(fresh["scatter"]["coherence"])))
    for j, (x, y) in enumerate(zip(fresh["scatter"]["coherence"], fresh["scatter"]["growth"])):
        ax.plot(x, y, "o", ms=5, color=colors[j], label=f"start {result['model']['starts'][j]}")
    ax.plot([], [], "o", ms=3, color=grey, label="earlier run (starts 1-6)")
    ax.axhline(0, color="0.6", lw=0.8)
    ax.axvline(0, color="0.6", lw=0.8)
    ax.set(xlabel="coherence in the window, minus the start's average",
           ylabel="growth rate in the window, minus the start's average",
           title=f"1. Same time: each dot is a window of {WINDOW:.0f} time units\n(set before the run; reading:"
                 f" {result['readings']['same']})")
    ax.legend(fontsize=8, loc="upper left", ncol=2)

    ax = fig.add_subplot(grid[0, 1])
    lags = np.array(CURVE_LAGS)
    ax.axvspan(*LAGS, color="#fff3c4", zorder=0)
    ax.text(np.mean(LAGS), 0.97, "set before the run", transform=ax.get_xaxis_transform(), ha="center", va="top",
            fontsize=9, color="0.3")
    ax.plot(lags, earlier["curve"], color=grey, lw=1.6, label="earlier run (where it was noticed)")
    mean, se = np.array(fresh["curve"]), np.array(fresh["curve_se"])
    ax.fill_between(lags, mean - se, mean + se, color=blue, alpha=0.2, lw=0)
    ax.plot(lags, mean, color=blue, lw=2.2, label="fresh starts (band: error bar)")
    ax.axhline(0, color="0.5", lw=0.9)
    ax.axvline(0, color="0.7", lw=0.8, ls=":")
    ax.set(xlim=(lags[0], lags[-1]), xlabel="L, time units (negative: the growth comes first)",
           ylabel="correlation", title=f"2. Coherence now against the growth rate L later\n(shaded: set before the run;"
                                       f" reading: {result['readings']['later']})")
    ax.legend(fontsize=8.5, loc="lower left")

    ax = fig.add_subplot(grid[1, 0])
    xs = list(SIZES)
    for m, color, dx, name in ((earlier, grey, 0.97, "earlier run"), (fresh, blue, 1.03, "fresh starts")):
        ys = [m["sizes"][f"{w:g}"] for w in SIZES]
        ax.errorbar([x * dx for x in xs], [c["mean"] for c in ys], yerr=[c["se"] for c in ys], color=color,
                    marker="o", capsize=3, lw=1.4, label=name)
        for p in PLACEMENTS:
            c = m["placements"][f"{p:g}"]
            ax.errorbar(1000 * dx * (1 + p / 4000), c["mean"], yerr=c["se"], color=color, marker="s", ms=4,
                        capsize=2, lw=0, elinewidth=1)
    ax.axhline(0, color="0.5", lw=0.9)
    ax.set_xscale("log")
    ax.set(xticks=[25, 50, 100, 200, 500, 1000], xticklabels=["25", "50", "100", "200", "500", "1,000\n(4 placements)"],
           xlabel="window length (time units)", ylabel="same-time correlation (average over the starts)",
           title=f"3. Same time, by window length (only {WINDOW:.0f} was set as the test)")
    ax.legend(fontsize=8.5, loc="lower left")

    ax = fig.add_subplot(grid[1, 1])
    for x, key in ((0, "same"), (1, "later")):
        s, e = fresh[key], earlier[key]
        ax.add_patch(plt.Rectangle((x - 0.18, -RULE * s["se"]), 0.36, 2 * RULE * s["se"], color="0.85", lw=0))
        ax.plot(np.full(len(s["starts"]), x + 0.06), s["starts"], "o", ms=5, color=blue, alpha=0.55)
        ax.errorbar(x - 0.06, s["mean"], yerr=s["se"], fmt="D", ms=9, color=blue, capsize=4, mew=2)
        ax.plot(x + 0.22, e["mean"], "D", ms=7, mfc="white", mec=grey, mew=2)
    ax.plot([], [], "D", ms=8, color=blue, label="fresh starts: average and error bar (dots: single starts)")
    ax.plot([], [], "D", ms=7, mfc="white", mec=grey, mew=2, label="earlier run, for reference")
    ax.add_patch(plt.Rectangle((0, 0), 0, 0, color="0.85", label=f"within {RULE:g} error bars of 0: no clear link"))
    ax.axhline(0, color="0.5", lw=0.9)
    ax.set(xlim=(-0.5, 1.5), xticks=[0, 1], ylabel="correlation", title="4. The two measures set before the run")
    ax.set_xticklabels([f"1. same time\n(windows of {WINDOW:.0f})", f"2. later\n({LAGS[0]:.0f} to {LAGS[1]:.0f})"])
    ax.legend(fontsize=8.5, loc="upper right")

    test16 = result["test16"]
    for y, text in ((0.05, f"Six fresh starts (starts {result['model']['starts'][0]} to {result['model']['starts'][-1]}),"
                           f" the same model, gap (rms {nt.LYAP_GAP:g} radians, put back every {nt.LYAP_EVERY:g} time"
                           f" units) and length ({RUN:,.0f}) as the long run at the threshold. Measured over"
                           f" t = {SPAN[0]:,.0f} to {SPAN[1]:,.0f}, start by start."),
                    (0.031, f"Coherence: the untouched copy's |z|, average of A and B. Growth: of the log of the gap."
                            f" Fresh long-run rate (t = {LONG_FROM:,.0f}-{RUN:,.0f}): {result['fresh']['long_rate']:+.4f}"
                            f" ± {result['fresh']['long_rate_se']:.4f}; Test 16: {test16['rate']:+.4f} ±"
                            f" {test16['rate_se']:.4f} ({'agree' if test16['agree'] else 'differ'})."),
                    (0.012, f"A link if the average over the starts is more than {RULE:g} error bars from 0. Error"
                            " bars: from the starts. A toy model: this shows structure, not proof about anything"
                            " beyond it.")):
        fig.text(0.06, y, text, fontsize=9.5, color="0.25")
    FIGURE.parent.mkdir(exist_ok=True)
    fig.savefig(FIGURE, dpi=115)


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