"""T0.1 — validate the QR/Benettin FTLE estimator core against systems with KNOWN spectra.

Reimplements the IDENTICAL accumulation used in diagnose_{trm,hrm}_joint.py:
  Q in R^{n x k} init random-orthonormal; each step apply the (known) Jacobian to Q's columns;
  every t_ons steps QR-decompose, accumulate sum of log|diag(R)|; LE_i = sum / n_qr_steps.

Test systems (known answers):
  (a) diagonal linear map     LE_i = log|d_i|           (exact at all T)
  (b) symmetric linear map     LE_i = log|eig_i|          (exact; eig=singular values)
  (c) non-normal (shear) map   LE_i = log|eig_i| asympt.  (finite-time transient from singular values)
  (d) Henon map (a=1.4,b=0.3)  LE = {+0.41922, -1.62319}  (nonlinear chaotic; literature value)

A passing result = recovered exponents match known to within tolerance, confirming the QR core
(orthonormalization cadence, log|diag R| bookkeeping, ordering, averaging) is correct.
No GPU, no model — this isolates the numerical estimator.
"""
from __future__ import annotations
import numpy as np

RNG = np.random.default_rng(0)


def qr_le(jac_fn, x0, n_steps, k, t_ons=1, warmup=0):
    """Benettin/QR LE estimate. jac_fn(x)->(x_next, J) gives next state and Jacobian at x.
    Mirrors diagnose_*_joint.py: QR every t_ons steps, accumulate log|diag R|, average over QR steps."""
    x = np.asarray(x0, float)
    d = x.shape[0]
    Q, _ = np.linalg.qr(RNG.standard_normal((d, k)))
    log_R_sum = np.zeros(k)
    n_qr = 0
    for t in range(n_steps):
        x, J = jac_fn(x)
        Q = J @ Q
        if (t + 1) % t_ons == 0:
            Q, R = np.linalg.qr(Q)
            if t >= warmup:
                log_R_sum += np.log(np.clip(np.abs(np.diag(R)), 1e-30, None))
                n_qr += 1
    return np.sort(log_R_sum / max(n_qr, 1))[::-1]


def run():
    out = ["# T0.1 estimator validation (QR/Benettin core vs known spectra)", ""]
    tol = 5e-3

    # (a) diagonal
    d_vals = np.array([1.5, 0.8, 0.3, 0.05])
    M = np.diag(d_vals)
    known = np.sort(np.log(np.abs(d_vals)))[::-1]
    est = qr_le(lambda x: (x, M), np.ones(4), 4000, k=4)  # linear: J state-independent, don't grow x
    out += [f"(a) diagonal linear: known {np.round(known,4)}",
            f"    recovered        {np.round(est,4)}   max|err|={np.max(np.abs(est-known)):.2e}  "
            f"{'PASS' if np.max(np.abs(est-known))<tol else 'FAIL'}"]

    # (b) symmetric
    A = RNG.standard_normal((5, 5)); S = (A + A.T) / 2
    # scale so spectral radius < ~1.3 (keep magnitudes spread, finite)
    S = 0.9 * S / np.max(np.abs(np.linalg.eigvalsh(S)))
    eig = np.linalg.eigvalsh(S)
    known = np.sort(np.log(np.abs(eig)))[::-1]
    est = qr_le(lambda x: (x, S), np.ones(5), 8000, k=5)
    out += [f"(b) symmetric linear: known {np.round(known,4)}",
            f"    recovered         {np.round(est,4)}   max|err|={np.max(np.abs(est-known)):.2e}  "
            f"{'PASS' if np.max(np.abs(est-known))<tol else 'FAIL'}"]

    # (c) non-normal shear: LE -> log|eig| asymptotically; finite-time transient from singular values
    N = np.array([[1.1, 5.0], [0.0, 0.6]])  # eigenvalues 1.1, 0.6 (triangular); highly non-normal
    known = np.sort(np.log(np.abs(np.linalg.eigvals(N))))[::-1]
    est_long = qr_le(lambda x: (x, N), np.ones(2), 40000, k=2)
    sv = np.sort(np.log(np.linalg.svd(N, compute_uv=False)))[::-1]
    est_short = qr_le(lambda x: (x, N), np.ones(2), 5, k=2)
    out += [f"(c) non-normal shear: known asymptotic log|eig| {np.round(known,4)}",
            f"    recovered (T=40000)                        {np.round(est_long,4)}   "
            f"max|err|={np.max(np.abs(est_long-known)):.2e}  "
            f"{'PASS' if np.max(np.abs(est_long-known))<1e-2 else 'FAIL'}",
            f"    single-step log singular values            {np.round(sv,4)}  (finite-time transient ref)",
            f"    recovered (T=5, finite-time)               {np.round(est_short,4)}  "
            f"(should sit between sv and asymptotic -> confirms finite-time != asymptotic)"]

    # (d) Henon map
    a, b = 1.4, 0.3
    def henon(x):
        xn = np.array([1 - a * x[0] ** 2 + x[1], b * x[0]])
        J = np.array([[-2 * a * x[0], 1.0], [b, 0.0]])
        return xn, J
    # settle onto attractor first
    x = np.array([0.1, 0.1])
    for _ in range(1000):
        x, _ = henon(x)
    known = np.array([0.41922, -1.62319])  # literature (Sprott)
    est = qr_le(henon, x, 200000, k=2, warmup=1000)
    out += [f"(d) Henon (a=1.4,b=0.3): literature {np.round(known,4)}  (sum={known.sum():.4f})",
            f"    recovered                       {np.round(est,4)}  (sum={est.sum():.4f})   "
            f"|err λ1|={abs(est[0]-known[0]):.2e}  "
            f"{'PASS' if abs(est[0]-known[0])<5e-3 else 'FAIL'}"]

    out += ["", "Interpretation: (a)(b) confirm exact recovery for normal maps; (c) confirms the",
            "estimator converges to log|eig| asymptotically while finite-time windows reflect",
            "singular-value growth (the regime our paper operates in); (d) confirms correct",
            "recovery on a known chaotic nonlinear system. The QR core is validated."]
    print("\n".join(out))
    from pathlib import Path
    Path(__file__).resolve().parent.joinpath("validation_results.md").write_text("\n".join(out))


if __name__ == "__main__":
    run()