# jam3_spiral_ar2.py — JAM 3, line 2 instrument (founder).
# If the settled lexicon relaxes through a complex pair r·e^{±iθ},
# every scalar census projection x_t obeys, asymptotically,
#   x_{t+2} = 2r·cosθ · x_{t+1} − r² · x_t
# once the fixed point is removed. Differencing consecutive censuses
# kills the fixed point without estimating it, so we fit that AR(2)
# per census basis and read the pair straight from the roots:
# identical angle across bases = the walker's prediction; scattered
# angles = the complex-pair reading dies. TV's absolute value hides
# the phase; the recurrence does not.
from itertools import groupby
import numpy as np

LO, HI = 38, 58
BASES = (1, 2, 4, 6)


def las_step(s: str) -> str:
    return "".join(str(len(list(g))) + k for k, g in groupby(s))


def census(s: str, k: int) -> dict:
    counts, total = {}, len(s) - k + 1
    for i in range(total):
        b = s[i:i + k]
        counts[b] = counts.get(b, 0) + 1
    return {b: c / total for b, c in counts.items()}


def fit_pair(rows):
    keys = sorted(set().union(*[set(r) for r in rows]))
    M = np.array([[r.get(k, 0.0) for k in keys] for r in rows])
    Y = np.diff(M, axis=0)          # removes the fixed point exactly
    A, b = [], []
    for j in range(Y.shape[1]):
        y = Y[:, j]
        m = np.max(np.abs(y))
        if m < 1e-9:
            continue
        y = y / m                   # equal weight per surviving column
        for t in range(len(y) - 2):
            A.append([y[t + 1], y[t]])
            b.append(y[t + 2])
    A, b = np.array(A), np.array(b)
    (a1, a2), res, *_ = np.linalg.lstsq(A, b, rcond=None)
    lam = np.roots([1.0, -a1, -a2])
    r = float(np.abs(lam[0]))
    th = float(np.abs(np.angle(lam[0])))
    rel = float(np.sqrt(res[0] / np.sum(b * b))) if len(res) else float("nan")
    complex_pair = bool(np.iscomplex(lam[0]))
    return r, th, complex_pair, rel, len(b)


def fit_order(rows, order):
    """AR(order) on differenced censuses; returns roots and residual."""
    keys = sorted(set().union(*[set(r) for r in rows]))
    M = np.array([[r.get(k, 0.0) for k in keys] for r in rows])
    Y = np.diff(M, axis=0)
    A, b = [], []
    for j in range(Y.shape[1]):
        y = Y[:, j]
        m = np.max(np.abs(y))
        if m < 1e-9:
            continue
        y = y / m
        for t in range(len(y) - order):
            A.append([y[t + order - 1 - i] for i in range(order)])
            b.append(y[t + order])
    A, b = np.array(A), np.array(b)
    coef, res, *_ = np.linalg.lstsq(A, b, rcond=None)
    lam = np.roots([1.0] + list(-coef))
    rel = float(np.sqrt(res[0] / np.sum(b * b))) if len(res) else float("nan")
    order_desc = sorted(lam, key=lambda z: -abs(z))
    return order_desc, rel, len(b)


def main():
    s = "1"
    stored = {}
    for n in range(1, HI + 1):
        s = las_step(s)
        if n >= LO:
            stored[n] = s
    print("basis | modulus r | angle theta (rad, deg) | period 2pi/theta"
          " | complex? | rel-resid | eqs")
    for k in BASES:
        rows = [census(stored[n], k) for n in range(LO, HI + 1)]
        r, th, cp, rel, neq = fit_pair(rows)
        period = 2 * np.pi / th if th > 1e-9 else float("inf")
        print("  %d-blocks: r=%.4f  theta=%.4f rad (%.1f deg)  period=%.2f"
              "  complex=%s  resid=%.3f  eqs=%d"
              % (k, r, th, np.degrees(th), period, cp, rel, neq))
    print("\nAR(4) — room for two modes (roots by descending modulus):")
    for k in BASES:
        rows = [census(stored[n], k) for n in range(LO, HI + 1)]
        roots, rel, neq = fit_order(rows, 4)
        desc = "  ".join("r=%.3f th=%.2f" % (abs(z), abs(np.angle(z)))
                         for z in roots)
        print("  %d-blocks: %s  resid=%.3f  eqs=%d" % (k, desc, rel, neq))


if __name__ == "__main__":
    main()