import json
from pathlib import Path
import numpy as np

SEED = 1632
J = np.array([[0.0, -1.0], [1.0, 0.0]])
I2 = np.eye(2)


def transition(gamma, omega, s):
    # Euler step for dx/dt=(-gamma I+s*omega J)x, dt=1.
    return (1.0 - gamma) * I2 + s * omega * J


def antisym(x):
    return 0.5 * (x - x.T)


def response(A, q):
    return np.linalg.matrix_power(A, q)


def stationary_cov(A, noise=1.0, n=200000, burn=100):
    rng = np.random.default_rng(1000 + int(abs(A[0, 1]) * 10000))
    x = np.zeros(2)
    xs = []
    for t in range(n + burn):
        x = A @ x + noise * rng.normal(size=2)
        if t >= burn:
            xs.append(x.copy())
    x = np.asarray(xs)
    return np.cov(x, rowvar=False, bias=True), x


def lag_cov(xs, lag):
    a, b = xs[lag:], xs[:-lag]
    a = a - a.mean(0)
    b = b - b.mean(0)
    return a.T @ b / len(a)


def fdr_residual_empirical(xs, A, tau=1.0, lag=1):
    # For discrete dynamics, dC/dt is approximated by C(lag)-C(0).
    c0 = np.cov(xs, rowvar=False, bias=True)
    cl = lag_cov(xs, lag)
    # impulse response at lag and current time; current impulse is identity.
    chi_l = np.linalg.matrix_power(A, lag)
    chi_0 = I2
    lhs = cl - c0
    rhs = -tau * (chi_l - chi_0)
    return np.linalg.norm(lhs - rhs) / (np.linalg.norm(lhs) + np.linalg.norm(rhs) + 1e-12)


def spectral_boundary_sweep():
    # Predicted stability boundary for Euler A=(1-gamma)I+s omega J:
    # rho(A)<1 iff (1-gamma)^2+omega^2<1.
    gamma = 0.25
    omegas = np.linspace(0.0, 1.15, 24)
    rows = []
    for w in omegas:
        A = transition(gamma, w, 1)
        rho = max(abs(np.linalg.eigvals(A)))
        rows.append((float(w), float(rho), bool(rho < 1.0)))
    predicted = float(np.sqrt(1.0 - (1.0 - gamma) ** 2))
    observed = omegas[np.where(np.array([r[2] for r in rows]))[0][-1]]
    return {'gamma': gamma, 'predicted_omega_boundary': predicted,
            'observed_grid_boundary': float(observed), 'rows': rows}


def oc_sweep():
    # Exact prediction: chi_q^A(+omega)+chi_q^A(-omega)=0; magnitude is
    # approximately q*omega for small omega and vanishes at omega=0.
    gamma = 0.2
    qs = [1, 2, 3]
    omegas = np.linspace(0.0, 0.03, 7)
    data = []
    for q in qs:
        vals = []
        mags = []
        for w in omegas:
            Ap, Am = transition(gamma, w, 1), transition(gamma, w, -1)
            vals.append(np.linalg.norm(antisym(response(Ap, q)) + antisym(response(Am, q))))
            mags.append(np.linalg.norm(antisym(response(Ap, q))))
        slope = np.polyfit(omegas[1:], mags[1:], 1)[0]
        data.append({'q': q, 'max_reversal_residual': float(max(vals)),
                     'zero_omega_magnitude': float(mags[0]),
                     'small_omega_slope': float(slope),
                     'predicted_small_omega_slope': float(np.sqrt(2) * q * (1-gamma)**(q-1))})
    return {'gamma': gamma, 'data': data}


def fdr_sweep():
    # Test that the empirical residual decreases with sample size and that
    # the correct effective temperature is the injected noise scale squared.
    gamma, omega, noise = 0.35, 0.12, 0.7
    A = transition(gamma, omega, 1)
    # Stationary covariance solves C=A C A^T+noise^2 I; tau=noise^2.
    c, xs = stationary_cov(A, noise=noise, n=80000)
    taus = np.linspace(0.25, 1.0, 16)
    residuals = []
    for tau in taus:
        # Use exact covariance and exact discrete lag covariance, removing MC error.
        cl = A @ c
        exact = np.linalg.norm((cl-c) + tau*(A-I2))
        residuals.append(float(exact))
    best_tau = float(taus[int(np.argmin(residuals))])
    sample_sizes = [500, 2000, 8000, 32000]
    empirical = []
    for n in sample_sizes:
        empirical.append(float(fdr_residual_empirical(xs[:n], A, tau=noise**2)))
    return {'gamma': gamma, 'omega': omega, 'noise_scale': noise,
            'predicted_tau': noise**2, 'best_grid_tau': best_tau,
            'exact_residual_at_predicted_tau': float(residuals[np.argmin(abs(taus-noise**2))]),
            'sample_sizes': sample_sizes, 'empirical_normalized_residuals': empirical,
            'tau_grid': taus.tolist(), 'exact_residuals': residuals}


def fit_comparison():
    # Tiny equal-compute linear transition fitting. Regularizer is the OC
    # penalty on the learned antisymmetric component under orientation reversal.
    rng = np.random.default_rng(SEED)
    gamma, omega = 0.2, 0.18
    Aplus, Aminus = transition(gamma, omega, 1), transition(gamma, omega, -1)
    n = 5000
    x = rng.normal(size=(n,2))
    y = x @ Aplus.T + 0.08*rng.normal(size=(n,2))
    # Baseline least squares on the same plus-orientation data.
    Ab = np.linalg.lstsq(x, y, rcond=None)[0].T
    vals = []
    for lam in [0.0, 0.001, 0.01, 0.1, 1.0]:
        # Objective ||AX-Y||^2 + lambda ||Anti(A)+Anti(Aminus)||^2.
        # Solve by deterministic gradient descent, same 300 steps for all.
        A = np.zeros((2,2))
        lr = 0.08
        for _ in range(300):
            grad = 2*(A @ (x.T@x)/n - y.T@x/n)
            R = antisym(A) + antisym(Aminus)
            grad += 2*lam*R
            A -= lr*grad
        testx = rng.normal(size=(3000,2))
        testy = testx @ Aplus.T
        mse = np.mean((testx @ A.T-testy)**2)
        oc = np.linalg.norm(antisym(A)+antisym(Aminus))
        vals.append({'lambda': lam, 'test_mse': float(mse), 'oc_residual': float(oc)})
    testx = rng.normal(size=(3000,2))
    testy = testx @ Aplus.T
    return {'baseline_test_mse': float(np.mean((testx @ Ab.T-testy)**2)), 'baseline_oc_residual': float(np.linalg.norm(antisym(Ab)+antisym(Aminus))), 'regularized_runs': vals}


def main():
    out = {'seed': SEED, 'oc_sweep': oc_sweep(), 'stability_sweep': spectral_boundary_sweep(),
           'fdr_sweep': fdr_sweep(), 'fit_comparison': fit_comparison()}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()