"""Minimal analytically-grounded toy reproducing 'failure trajectories are more chaotic'.

Framing: recursive reasoning = chaotic SEARCH of latent space until the trajectory lands in the
solution basin, then it converges (the answer is a stable fixed point). At a fixed readout time T:
  - captured by T  => SUCCESS, finite-time Lyapunov exponent (FTLE) is low (chaotic search for a
                      while, then contraction).
  - not captured   => FAILURE, FTLE stays at the chaotic-saddle value.

Map on [0,1] (input s = solution location, drawn per 'puzzle'):
  search phase:  x <- 4 x (1-x)            (fully chaotic, lambda = ln 2 ~ +0.693)
  capture:       if |x - s| < eps  -> converging phase: x <- s + 0.5 (x - s)   (lambda = ln 0.5 < 0)

This is a textbook chaotic-saddle / transient-chaos system: escape (capture) times are ~geometric,
and FTLE over [0,T] = (ln2 * t_search + ln0.5 * (T - t_search)) / T. So FTLE separates outcome
PURELY as a finite-time artifact of capture time -- matching our real-model findings:
  (i) failure more chaotic, (ii) separation is concurrent/finite-time (vanishes as T->inf since
  everyone eventually captures), (iii) it is convergence detection, not a deep correctness signal.

Run: python toy_transient_chaos.py   (CPU, seconds). Produces toy_transient_chaos.png + stats.
"""
from __future__ import annotations
from pathlib import Path
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = Path(__file__).resolve().parent
rng = np.random.default_rng(0)


def run(n=20000, T=16, eps=0.04, seed=0):
    rg = np.random.default_rng(seed)
    s = rg.uniform(0.15, 0.85, n)           # 'puzzle': solution location
    x = rg.uniform(0, 1, n)                  # initial latent
    captured = np.zeros(n, bool)
    cap_time = np.full(n, T, int)
    log_deriv_sum = np.zeros(n)
    drift = np.zeros((n, T))
    for t in range(T):
        x_prev = x.copy()
        search = ~captured
        # log|f'| this step
        ld = np.empty(n)
        ld[search] = np.log(np.abs(4 * (1 - 2 * x[search])) + 1e-12)
        ld[~search] = np.log(0.5)
        # update
        x_new = x.copy()
        x_new[search] = 4 * x[search] * (1 - x[search])
        x_new[~search] = s[~search] + 0.5 * (x[~search] - s[~search])
        # capture check (enter basin during search)
        newly = search & (np.abs(x_new - s) < eps)
        captured |= newly
        cap_time[newly] = t + 1
        x = x_new
        log_deriv_sum += ld
        drift[:, t] = np.abs(x - x_prev)
    ftle = log_deriv_sum / T
    success = captured & (np.abs(x - s) < 0.05)   # near solution at readout
    return dict(ftle=ftle, success=success, cap_time=cap_time, drift=drift, T=T, eps=eps)


def auc(score, y):
    p, n = score[y == 1], score[y == 0]
    if len(p) == 0 or len(n) == 0:
        return float("nan")
    a = np.concatenate([p, n]); o = np.argsort(a); r = np.empty(len(a)); r[o] = np.arange(1, len(a) + 1)
    return float((r[:len(p)].sum() - len(p) * (len(p) + 1) / 2) / (len(p) * len(n)))


def main():
    d = run(T=16)
    y = d["success"].astype(int); ftle = d["ftle"]
    print(f"=== Toy transient-chaos (T=16, eps={d['eps']}) ===")
    print(f"success rate = {y.mean():.3f}")
    print(f"FTLE: success median={np.median(ftle[y==1]):+.3f}  failure median={np.median(ftle[y==0]):+.3f}")
    print(f"AUC(-FTLE -> success) = {auc(-ftle, y):.3f}   (>0.9 => reproduces 'failure more chaotic')")
    late_drift = d["drift"][:, -4:].mean(1)
    print(f"late drift: success median={np.median(late_drift[y==1]):.3f} failure median={np.median(late_drift[y==0]):.3f}")
    print(f"AUC(-late_drift -> success) = {auc(-late_drift, y):.3f}")

    # FTLE separation vs readout time T (the finite-time / concurrent-not-antecedent prediction)
    Ts = [4, 8, 16, 32, 64, 128, 256]
    seps, accs = [], []
    for T in Ts:
        dd = run(T=T)
        yy = dd["success"].astype(int)
        seps.append(auc(-dd["ftle"], yy)); accs.append(yy.mean())
    print("\nFTLE-AUC and success-rate vs readout time T (separation is a finite-time effect):")
    for T, a, ac in zip(Ts, seps, accs):
        print(f"  T={T:>3}: AUC(-FTLE->success)={a:.3f}  success_rate={ac:.3f}")

    fig, ax = plt.subplots(1, 3, figsize=(15, 4.2))
    bins = np.linspace(-0.5, 0.75, 50)
    ax[0].hist(ftle[y == 1], bins=bins, alpha=0.6, color="tab:green", density=True, label=f"success (n={int(y.sum())})")
    ax[0].hist(ftle[y == 0], bins=bins, alpha=0.6, color="tab:red", density=True, label=f"failure (n={int((1-y).sum())})")
    ax[0].axvline(0, color="gray", lw=0.6); ax[0].set_xlabel("finite-time Lyapunov exponent (T=16)")
    ax[0].set_title(f"Toy: failure more chaotic\nAUC={auc(-ftle,y):.3f}"); ax[0].legend(fontsize=8)
    ax[1].hist(d["cap_time"][d["cap_time"] < 16], bins=range(0, 17), color="tab:blue", alpha=0.8)
    ax[1].set_xlabel("capture (escape) time"); ax[1].set_ylabel("count")
    ax[1].set_title("escape-time distribution (~geometric:\nchaotic-saddle signature)")
    ax[2].plot(Ts, seps, "o-", label="AUC(-FTLE->success)")
    ax[2].plot(Ts, accs, "s--", color="tab:purple", label="success rate")
    ax[2].set_xscale("log"); ax[2].axhline(0.5, color="gray", lw=0.5)
    ax[2].set_xlabel("readout time T"); ax[2].set_title("separation is FINITE-TIME:\nvanishes as T->inf (all escape)")
    ax[2].legend(fontsize=8)
    fig.tight_layout(); fig.savefig(HERE / "toy_transient_chaos.png", dpi=140)
    print("\nsaved toy_transient_chaos.png")


if __name__ == "__main__":
    main()