"""Pullback random-attractor monitor for a linear noisy damped oscillator.

The shared additive noise cancels between pullback replicas, so their diameter
measures contraction of the stochastic one-step map rather than noise variance.
"""
import json
from pathlib import Path
import numpy as np


def oscillator_matrix(h, omega, damping, method):
    """Return the 2x2 deterministic state transition matrix."""
    # x' = v, v' = -omega^2*x - damping*v + noise
    if method == "euler":
        return np.array([[1.0, h], [-h * omega**2, 1.0 - h * damping]])
    if method == "semi_implicit":
        # Backward Euler only on velocity damping/force, then symplectic x update.
        den = 1.0 + h * damping
        v_x = -h * omega**2 / den
        v_v = 1.0 / den
        return np.array([[1.0 + h * v_x, h * v_v], [v_x, v_v]])
    raise ValueError(method)


def predicted_rho(M):
    return float(np.max(np.abs(np.linalg.eigvals(M))))


def monitor(h, method, omega, damping, K=32, steps=128, seed=0, sigma=0.15):
    rng = np.random.default_rng(seed)
    # Fixed bounded ball: random directions with radius in [0.8, 1.2].
    z = rng.normal(size=(K, 2))
    z /= np.linalg.norm(z, axis=1, keepdims=True)
    z *= rng.uniform(0.8, 1.2, size=(K, 1))
    M = oscillator_matrix(h, omega, damping, method)
    diam = [pairwise_rms(z)]
    # One identical noise vector is replayed for every replica at every step.
    noises = rng.normal(size=(steps, 2)) * sigma * np.sqrt(h)
    for noise in noises:
        z = z @ M.T + noise[None, :]
        diam.append(pairwise_rms(z))
    return np.asarray(diam), M


def pairwise_rms(z):
    # RMS of all pair distances, a stable diameter proxy.
    d = z[:, None, :] - z[None, :, :]
    return float(np.sqrt(np.mean(np.sum(d * d, axis=-1))))


def fit_slope(diam, h, start=8, end=None):
    end = len(diam) if end is None else end
    y = np.log(np.maximum(diam[start:end], 1e-30))
    x = np.arange(start, end) * h
    return float(np.polyfit(x, y, 1)[0])


def independent_growth_rate(M, h, steps=300, seed=991):
    rng = np.random.default_rng(seed)
    v = rng.normal(size=2); v /= np.linalg.norm(v)
    logs = []
    for _ in range(steps):
        v = M @ v
        n = np.linalg.norm(v)
        logs.append(np.log(n))
        v /= n
    # This is the accumulated Lyapunov sum for repeatedly renormalized vectors.
    return float(np.mean(np.asarray(logs)[50:]) / h)


def stability_sweep(omega=1.0, damping=0.2):
    # Prediction: Euler is stable iff rho(M)<1, with boundary h= damping/omega^2.
    hs = [0.10, 0.18, 0.20, 0.22, 0.40]
    rows = []
    for method in ("euler", "semi_implicit"):
        for h in hs:
            d, M = monitor(h, method, omega, damping, steps=128, seed=10)
            rho = predicted_rho(M)
            slope = fit_slope(d, h, start=10, end=100)
            rows.append({"method": method, "h": h, "predicted_rho": rho,
                         "measured_slope": slope, "D0": d[0], "D128": d[-1],
                         "contracting_predicted": bool(rho < 1.0),
                         "contracting_observed": bool(d[-1] < d[0])})
    return rows


def rate_sweep(omega=1.0, h=0.05):
    # Prediction: as damping varies, lambda_h=log(rho(M))/h and measured log D
    # slope track it; near the boundary contraction time is ~1/|lambda_h|.
    rows = []
    for damping in [0.05, 0.10, 0.20, 0.40, 0.80]:
        d, M = monitor(h, "euler", omega, damping, steps=500, seed=20)
        rho = predicted_rho(M)
        theory = np.log(rho) / h
        measured = fit_slope(d, h, start=12, end=150)
        independent = independent_growth_rate(M, h)
        target = d[0] * 0.1
        hit = np.where(d <= target)[0]
        hit_time = float(hit[0] * h) if len(hit) else None
        predicted_time = float(np.log(10.0) / (-theory)) if theory < 0 else None
        rows.append({"damping": damping, "predicted_lambda": theory,
                     "measured_pullback_slope": measured,
                     "independent_vector_slope": independent,
                     "predicted_time_to_0.1": predicted_time,
                     "observed_time_to_0.1": hit_time})
    return rows


def noise_control(h=0.05, damping=0.4, omega=1.0, K=32, steps=600):
    """Compare the monitor's shared-noise replay with independent noise.
    Prediction: shared noise contracts; independent noise reaches a noise floor.
    """
    rng = np.random.default_rng(123)
    z = rng.normal(size=(K, 2)); z /= np.linalg.norm(z, axis=1, keepdims=True)
    z *= rng.uniform(0.8, 1.2, size=(K, 1))
    M = oscillator_matrix(h, omega, damping, "euler")
    shared = z.copy(); independent = z.copy()
    shared_d, independent_d = [pairwise_rms(shared)], [pairwise_rms(independent)]
    for _ in range(steps):
        eps = rng.normal(size=(K, 2)) * 0.15 * np.sqrt(h)
        shared = shared @ M.T + eps[0][None, :]
        independent = independent @ M.T + eps
        shared_d.append(pairwise_rms(shared)); independent_d.append(pairwise_rms(independent))
    return {"shared_initial": shared_d[0], "shared_final": shared_d[-1],
            "independent_initial": independent_d[0], "independent_final": independent_d[-1],
            "shared_slope": fit_slope(np.asarray(shared_d), h, 10, 90),
            "independent_slope": fit_slope(np.asarray(independent_d), h, 10, 90)}

def main():
    out = {"stability_boundary_prediction": "Euler rho=1 at h=damping/omega^2=0.20; semi-implicit should remain stable here",
           "rate_scaling_prediction": "pullback log-diameter slope and independent vector growth equal log(rho)/h; time scales as 1/abs(lambda)",
           "noise_replay_prediction": "shared replay contracts while independent noise produces a nonzero diameter floor",
           "noise_control": noise_control(),
           "stability_sweep": stability_sweep(), "rate_sweep": rate_sweep()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()