#!/usr/bin/env python3
"""
Earlier-growth check (October 11, 2026, evening): does coherence say anything about later growth that the earlier
growth does not already say?

Exploratory. Chosen after the coherence-growth test (coherence_growth.py) and run on the same saved runs, so it is not
a test, whatever it shows. A sandbox run kept with the project; nothing here is served by the site.

What led here
    The test found that more coherence now went with slower growth 200 to 500 time units later: -0.2061 (0.0186),
    below 0 in all six fresh starts. That could happen without coherence carrying anything of its own: growth may slow
    by itself after a burst, while coherence, which rises after growth, only marks the burst. Holding the earlier growth
    fixed tells the two apart. If coherence only marks the burst, the link should fade to about 0; if coherence carries
    more, the link should stay.

What it does (all of this was set before it was run)
    1. The test's second measure exactly as before (the coherence in each 5-unit interval against the growth rate L
       time units later, start by start over t = 2,000 to 6,000, averaged over L = 200 to 500), now with the earlier
       growth held fixed: the coherence and the later growth are both freed of what the earlier growth predicts (least
       squares, start by start and lag by lag), then correlated (a partial correlation).
       Earlier growth, main version: the growth rate in each of the five 100-unit windows that end with the coherence
       interval (the 500 time units up to it). For the first intervals this reaches back before t = 2,000 (to 1,500 at
       most); the coherence and the later growth are exactly the test's.
    2. Also shown: holding only the 200 time units before fixed; a check on the check, holding fixed another start's
       earlier growth (start 7's coherence with start 8's growth, and so on), which knows nothing about the start and
       so shows how much holding five smooth series fixed moves the link by itself; the growth's own link, the growth
       rate over the 200 before against the growth rate 200 to 500 later (below 0: growth slows by itself after a
       burst); the link at every lag from 5 to 500, with and without the main version held fixed; and the same question
       in the test's windows of 200, one regression per start: growth in a later window = a + b x coherence + c x growth,
       both in the window (the next window, and the one after it).
    Error bars: from the starts, resampled 400 times with the test's seed. Read with the test's rule, on the fresh
    starts: still there if the main version is more than 3 error bars from 0; otherwise the earlier growth could account
    for the link. Both sets are shown: the fresh starts (7 to 12, where the test was run) and the earlier run (1 to 6,
    where the pattern was first noticed). Neither is fresh for this check.

Usage
    python3 earlier_growth_check.py     # earlier_growth_check.txt, earlier_growth_check.json and the figure

Needs Python 3, numpy and matplotlib, coherence_growth.py and its saved runs (unchanged) and ../nudge-test-2026-10-10
(unchanged). The figure goes to screenshots/nudge-test-earlier-growth-check.png.
"""
import json
import sys
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve().parent
sys.dont_write_bytecode = True                             # leave both folders as they are
sys.path.insert(0, str(HERE))
import coherence_growth as cg  # noqa: E402  (the test, its settings and its saved runs, unchanged)

nt = cg.nt
EVERY = nt.LYAP_EVERY                                      # 5 time units per interval
MAIN = (5, 100.0)                                          # main version: the growth in five windows of 100 before
SHORT = (1, 200.0)                                         # also shown: the growth over the 200 before
OWN = 200.0                                                # the growth's own link: the growth over the 200 before
CURVE = np.arange(5, 501, 5)                               # the link at every lag, with and without the main version
STEPS = (1, 2)                                             # windows of 200: the next window and the one after it
OUT_TXT, OUT_JSON = HERE / "earlier_growth_check.txt", HERE / "earlier_growth_check.json"
FIGURE = HERE.parent.parent / "screenshots" / "nudge-test-earlier-growth-check.png"


def corr(x, y):
    x = x - x.mean()
    y = y - y.mean()
    return float((x @ y) / np.sqrt((x @ x) * (y @ y)))


def partial(x, y, held):
    """The correlation of x and y after both are freed of what the columns of `held` predict (least squares)."""
    design = np.column_stack([np.ones(len(x)), held])
    free_x = x - design @ np.linalg.lstsq(design, x, rcond=None)[0]
    free_y = y - design @ np.linalg.lstsq(design, y, rcond=None)[0]
    return corr(free_x, free_y)


def standardized(y, columns):
    """Least-squares weights of y on the columns, everything scaled to mean 0 and spread 1."""
    scale = lambda v: (v - v.mean()) / v.std()  # noqa: E731
    design = np.column_stack([scale(c) for c in columns])
    return np.linalg.lstsq(design, scale(y), rcond=None)[0].tolist()


def before(logs, js, version):
    """The growth rate in `windows` windows of `length` time units, the newest ending with interval j (included)."""
    windows, length = version
    size, total = int(round(length / EVERY)), np.concatenate([[0.0], np.cumsum(logs)])
    return np.column_stack([(total[js + 1 - k * size] - total[js + 1 - (k + 1) * size]) / length
                            for k in range(windows)])


def checked(z, logs, rng):
    """Every number for one set of starts. z: the untouched copy's coherence, (start, group, time); logs: the growth of
    the log of the gap in each 5-unit interval, (start, interval)."""
    r = np.abs(z.astype(np.complex128)).mean(1)
    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 cg.SPAN)
    picks = rng.integers(0, count, (400, count))                           # as in the test: the same seed, the same draws

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

    def link(start, lag, version=None, held_from=None):
        s = int(round(lag / EVERY))
        js = np.arange(a, b - s)
        x, y = r5[start, js], logs[start, js + s]
        if version is None:
            return corr(x, y)
        return partial(x, y, before(logs[start if held_from is None else held_from], js, version))

    def own(start, lag):
        s = int(round(lag / EVERY))
        js = np.arange(a, b - s)
        return corr(before(logs[start], js, (1, OWN))[:, 0], logs[start, js + s])

    later = np.arange(cg.LAGS[0], cg.LAGS[1] + EVERY / 2, EVERY)           # 200, 205, ..., 500, as in the test
    over = lambda f: [float(np.mean([f(i, lag) for lag in later])) for i in range(count)]  # noqa: E731

    size = int(round(cg.WINDOW / EVERY))
    n = (b - a) // size
    coherence = r5[:, a:a + n * size].reshape(count, n, size).mean(-1)      # the test's windows of 200
    growth = logs[:, a:a + n * size].reshape(count, n, size).sum(-1) / cg.WINDOW
    windows = {}
    for m in STEPS:
        x, y, g = coherence[:, :-m], growth[:, m:], growth[:, :-m]
        windows[str(m)] = dict(
            plain=summary([corr(x[i], y[i]) for i in range(count)]),
            held=summary([partial(x[i], y[i], g[i][:, None]) for i in range(count)]),
            weights=[standardized(y[i], (x[i], g[i])) for i in range(count)])   # [b, c] per start, scaled

    curve_plain = np.array([[link(i, lag) for lag in CURVE] for i in range(count)])
    curve_held = np.array([[link(i, lag, MAIN) for lag in CURVE] for i in range(count)])
    return dict(
        starts=count,
        test=summary(over(lambda i, lag: link(i, lag))),
        main=summary(over(lambda i, lag: link(i, lag, MAIN))),
        short=summary(over(lambda i, lag: link(i, lag, SHORT))),
        other=summary(over(lambda i, lag: link(i, lag, MAIN, held_from=(i + 1) % count))),
        own=summary(over(own)),
        windows=windows,
        curve=dict(lags=CURVE.tolist(), plain=curve_plain.mean(0).tolist(),
                   plain_se=curve_plain[picks].mean(1).std(0).tolist(), held=curve_held.mean(0).tolist(),
                   held_se=curve_held[picks].mean(1).std(0).tolist()))


def reading(s):
    if abs(s["mean"]) <= cg.RULE * s["se"]:
        return f"within {cg.RULE:g} error bars of 0: the earlier growth could account for the link"
    if s["mean"] < 0:
        return "still there: coherence says something about later growth that the earlier growth does not"
    return f"more than {cg.RULE:g} error bars above 0: the link turns around once the earlier growth is held fixed"


def main():
    fresh_data, old = np.load(cg.SAVED), np.load(cg.EARLIER)
    assert tuple(int(s) for s in fresh_data["starts"]) == cg.FRESH and float(fresh_data["run"]) == cg.RUN
    i = [round(float(x), 2) for x in old["ks"]].index(cg.K)
    fresh = checked(fresh_data["z"], fresh_data["logs"], np.random.default_rng(cg.SEED))
    earlier = checked(old["z"][i], old["logs"][i], np.random.default_rng(cg.SEED))
    test = json.loads((HERE / "coherence_growth.json").read_text())
    for mine, theirs in ((fresh, test["fresh"]), (earlier, test["earlier"])):   # the same code gives the test's numbers
        assert abs(mine["test"]["mean"] - theirs["later"]["mean"]) < 1e-12
        assert abs(mine["test"]["se"] - theirs["later"]["se"]) < 1e-12
    result = dict(
        exploratory="chosen after the coherence-growth test and run on the same saved runs: not a test",
        model=dict(n=cg.N, k=cg.K, span=cg.SPAN, later=cg.LAGS, rule=cg.RULE, main=dict(windows=MAIN[0], length=MAIN[1]),
                   short=dict(windows=SHORT[0], length=SHORT[1]), own=OWN, window=cg.WINDOW, steps=STEPS),
        fresh=fresh, earlier=earlier, reading=reading(fresh["main"]), reading_earlier=reading(earlier["main"]))
    lines = []

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

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

    def row(name, key, group=None):
        cells = []
        for m in (fresh, earlier):
            s = m[key] if group is None else m["windows"][group][key]
            cells.append(f"{cell(s)}, below 0 in {s['below']} of {m['starts']}")
        say(f"  {name:<52}{cells[0]:<40}{cells[1]}")

    say(f"Earlier-growth check, October 11, 2026: N = {cg.N:,} per group, K = {cg.K:.2f}, eps = {nt.EPS:g}")
    say("Exploratory: chosen after the coherence-growth test and run on the same saved runs, so not a test.")
    say(f"Question: does the coherence say anything about the growth {cg.LAGS[0]:.0f} to {cg.LAGS[1]:.0f} time units later"
        " that the earlier growth does not? Start by start over t = 2,000 to 6,000; error bars from the starts.")
    say(f"Main version (set before it was run): the earlier growth = the growth rate in each of the {MAIN[0]} windows of"
        f" {MAIN[1]:.0f} time units up to the coherence, held fixed by least squares (a partial correlation).")
    say()
    say(f"Coherence now against the growth rate {cg.LAGS[0]:.0f} to {cg.LAGS[1]:.0f} later (the test's second measure)")
    say(f"  {'':<52}{'fresh starts (7 to 12)':<40}earlier run (1 to 6)")
    row("as in the test, nothing held fixed", "test")
    row(f"the earlier growth held fixed: the {MAIN[0] * MAIN[1]:.0f} before (main)", "main")
    row(f"the earlier growth held fixed: the {SHORT[0] * SHORT[1]:.0f} before", "short")
    row("check on the check: another start's earlier growth", "other")
    say(f"  starts, main version: fresh {', '.join(f'{v:+.2f}' for v in fresh['main']['starts'])};"
        f" earlier {', '.join(f'{v:+.2f}' for v in earlier['main']['starts'])}")
    say(f"Reading (main version, fresh starts): {result['reading']}")
    say(f"Earlier run, main version: {result['reading_earlier']}")
    say()
    say(f"The growth's own link: the growth rate over the {OWN:.0f} before against the growth rate {cg.LAGS[0]:.0f} to"
        f" {cg.LAGS[1]:.0f} later (below 0: growth slows by itself after a burst)")
    row("", "own")
    say()
    say(f"In the test's windows of {cg.WINDOW:.0f}, one regression per start: growth in a later window = a + b x coherence"
        " + c x growth, both in the window")
    for m in STEPS:
        label = "the next window" if m == 1 else f"{m} windows on"
        row(f"{label}: coherence and later growth, plain", "plain", str(m))
        row(f"{label}: the window's own growth held fixed (b)", "held", str(m))
    for name, m in (("fresh", fresh), ("earlier", earlier)):
        weights = np.array(m["windows"]["1"]["weights"])
        say(f"  scaled weights, the next window, {name}: b (coherence) {weights[:, 0].mean():+.2f},"
            f" c (growth) {weights[:, 1].mean():+.2f} (averages over the starts)")
    say()
    say("The link at chosen lags L (time units), averaged over the starts: nothing held fixed / the main version held"
        " fixed")
    shown = (5, 50, 100, 200, 300, 400, 500)
    say("            " + "".join(f"{lag:>16d}" for lag in shown))
    for name, m in (("fresh", fresh), ("earlier", earlier)):
        lags = m["curve"]["lags"]
        say(f"  {name:<10}" + "".join(f"{m['curve']['plain'][lags.index(lag)]:>+8.3f}/{m['curve']['held'][lags.index(lag)]:<+7.3f}"
                                      for lag in shown))
    OUT_JSON.write_text(json.dumps(result, indent=1))
    OUT_TXT.write_text("\n".join(lines) + "\n")
    figure(result)
    print(f"\nWrote {OUT_TXT.name}, {OUT_JSON.name} and {FIGURE.relative_to(HERE.parent.parent)}")


def figure(result):
    """The link at every lag, with and without the earlier growth held fixed, and the four versions side by side."""
    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, 7.4))
    grid = fig.add_gridspec(1, 2, left=0.06, right=0.985, top=0.83, bottom=0.2, wspace=0.2, width_ratios=(1.15, 1))
    fig.suptitle(f"Earlier-growth check (October 11, 2026, exploratory): {cg.N:,} per group, K = {cg.K:.2f}\nDoes"
                 " coherence say anything about later growth that the earlier growth does not?", fontsize=14)

    ax = fig.add_subplot(grid[0, 0])
    lags = np.array(fresh["curve"]["lags"])
    ax.axvspan(*cg.LAGS, color="#fff3c4", zorder=0)
    ax.text(np.mean(cg.LAGS), 0.97, "the test's range", transform=ax.get_xaxis_transform(), ha="center", va="top",
            fontsize=9, color="0.3")
    for m, color, name, width in ((earlier, grey, "earlier run", 1.4), (fresh, blue, "fresh starts", 2.2)):
        curve = m["curve"]
        if m is fresh:
            held, se = np.array(curve["held"]), np.array(curve["held_se"])
            ax.fill_between(lags, held - se, held + se, color=blue, alpha=0.2, lw=0)
        ax.plot(lags, curve["plain"], color=color, lw=width, ls="--", label=f"{name}: nothing held fixed")
        ax.plot(lags, curve["held"], color=color, lw=width, label=f"{name}: earlier growth ({MAIN[0] * MAIN[1]:.0f}"
                                                                  " before) held fixed")
    ax.axhline(0, color="0.5", lw=0.9)
    ax.set(xlim=(0, lags[-1]), xlabel="L, time units", ylabel="correlation (average over the starts)",
           title="Coherence now against the growth rate L later\n(band: error bar of the fresh starts, held fixed)")
    ax.legend(fontsize=8.5, loc="lower left")

    ax = fig.add_subplot(grid[0, 1])
    keys = ("test", "main", "short", "other")
    names = ("as in the test", f"earlier growth\n({MAIN[0] * MAIN[1]:.0f} before)\nheld fixed",
             f"earlier growth\n({SHORT[0] * SHORT[1]:.0f} before)\nheld fixed", "another start's\nearlier growth\n(check)")
    for x, key in enumerate(keys):
        s, e = fresh[key], earlier[key]
        ax.add_patch(plt.Rectangle((x - 0.2, -cg.RULE * s["se"]), 0.4, 2 * cg.RULE * s["se"], color="0.87", lw=0))
        ax.plot(np.full(len(s["starts"]), x + 0.07), s["starts"], "o", ms=5, color=blue, alpha=0.55)
        ax.errorbar(x - 0.07, s["mean"], yerr=s["se"], fmt="D", ms=9, color=blue, capsize=4, mew=2)
        ax.plot(x + 0.24, 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")
    ax.add_patch(plt.Rectangle((0, 0), 0, 0, color="0.87", label=f"within {cg.RULE:g} error bars of 0"))
    ax.axhline(0, color="0.5", lw=0.9)
    ax.set(xlim=(-0.5, len(keys) - 0.5), xticks=range(len(keys)), ylabel="correlation",
           title=f"Coherence now against the growth {cg.LAGS[0]:.0f} to {cg.LAGS[1]:.0f} later")
    ax.set_xticklabels(names, fontsize=9)
    ax.legend(fontsize=8.5, loc="lower right")

    for y, text in ((0.075, f"Exploratory: chosen after the coherence-growth test and run on the same saved runs (six"
                            f" fresh starts, 7 to 12, and the earlier run, 1 to 6; {cg.RUN:,.0f} time units each),"
                            f" so not a test."),
                    (0.05, f"Held fixed: the coherence and the later growth are both freed of what the earlier growth"
                           f" predicts (least squares), then correlated. Measured start by start over"
                           f" t = {cg.SPAN[0]:,.0f} to {cg.SPAN[1]:,.0f}."),
                    (0.025, f"Main reading (fresh starts): {result['reading']}. 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__":
    main()
