import numpy as np
from scipy.linalg import expm, solve_continuous_lyapunov


def stationary_cov(A, D):
    # A C + C A^T = 2D for dX=-A X dt + sqrt(2D)dW
    return -solve_continuous_lyapunov(A, -2.0 * D)


def exact_step(A, D, dt):
    C = stationary_cov(A, D)
    F = expm(-A * dt)
    Q = (C - F @ C @ F.T)
    Q = (Q + Q.T) / 2
    vals, vecs = np.linalg.eigh(Q)
    Q = vecs @ np.diag(np.maximum(vals, 1e-12)) @ vecs.T
    return C, F, Q


def simulate(A, D, dt, n, rng):
    C, F, Q = exact_step(A, D, dt)
    out = np.empty((n, 2))
    out[0] = rng.multivariate_normal(np.zeros(2), C)
    noise = rng.multivariate_normal(np.zeros(2), Q, size=n - 1)
    for k in range(n - 1):
        out[k + 1] = F @ out[k] + noise[k]
    return out, C, F, Q


def gaussian_logpdf(v, cov):
    sign, logdet = np.linalg.slogdet(cov)
    inv = np.linalg.inv(cov)
    return -.5 * (logdet + v @ inv @ v + 2*np.log(2*np.pi))


def forward_reverse_action_gap(path, F, Q, dt):
    """Discretized path-action gap using the forward transition both ways.
    This is the standard forward path versus its time-reversed ordering;
    unlike a fitted conditional-entropy difference it detects probability currents.
    """
    vals = []
    for x, y in zip(path[:-1], path[1:]):
        vals.append(gaussian_logpdf(y - F @ x, Q) -
                    gaussian_logpdf(x - F @ y, Q))
    return float(np.mean(vals) / dt)


def analytic_spectrum(omega, A, D):
    out = np.empty_like(omega, dtype=float)
    for i, w in enumerate(omega):
        H = np.linalg.inv(A + 1j * w * np.eye(2))
        out[i] = np.real(2 * (H @ D @ H.conj().T)[0, 0])
    return out


def paper_sigma(a, b, Dx, Dm, Dc):
    return (a * (Dc + Dm) + b * Dx) ** 2 / (a * (Dx * Dm - Dc * Dc))


def autocorr(path, maxlag):
    z = path[:, 0] - path[:, 0].mean()
    den = np.dot(z, z)
    return np.array([np.dot(z[:-k] if k else z, z[k:] if k else z) /
                     (den if k == 0 else np.sqrt(np.dot(z[:-k], z[:-k]) * np.dot(z[k:], z[k:])))
                     for k in range(maxlag + 1)])


def main():
    rng = np.random.default_rng(20250308)
    # Exact adaptation: A_yy=0 and nontrivial coupling. Stable since det(A)>0.
    a, b, g = 1.0, 1.0, -1.0
    A = np.array([[a, b], [g, 0.0]])
    Dx = Dm = 1.0
    dcs = [-0.75, 0.0, 0.75]
    omega = np.logspace(-2, 2, 300)
    reference = analytic_spectrum(omega, A, np.diag([Dx, Dm]))
    print('A=', A.tolist())
    print('Analytic blind-direction check (max relative PSD error):')
    rows = []
    for dc in dcs:
        D = np.array([[Dx, dc], [dc, Dm]])
        spec = analytic_spectrum(omega, A, D)
        rel = np.max(np.abs(spec - reference) / np.maximum(reference, 1e-12))
        predicted = paper_sigma(a, b, Dx, Dm, dc)
        path, C, F, Q = simulate(A, D, dt=.03, n=180000, rng=rng)
        gap = forward_reverse_action_gap(path, F, Q, .03)
        # Compare observed autocorrelation to its exact stationary prediction.
        lags = np.arange(1, 101) * .03
        exact_ac = np.array([(expm(-A*t) @ C)[0, 0] / C[0, 0] for t in lags])
        sample_ac = autocorr(path, 100)[1:]
        ac_err = float(np.max(np.abs(sample_ac - exact_ac)))
        rows.append((dc, rel, predicted, gap, ac_err))
    sigmas = np.array([r[2] for r in rows]); gaps = np.array([r[3] for r in rows])
    print(' dc     analytic_rel_err    sigma_formula    action_gap       max_AC_error')
    for r in rows:
        print('% .2f   % .3e          % .6f       % .6f        % .4f' % r)
    print('formula sigma range:', float(sigmas.min()), float(sigmas.max()))
    print('action-gap range:', float(gaps.min()), float(gaps.max()))
    print('spectrum remains blind:', max(r[1] for r in rows) < 1e-10)
    print('autocorrelation agreement:', max(r[4] for r in rows) < .03)
    print('gap spread / formula spread:', float(np.ptp(gaps)), float(np.ptp(sigmas)))

if __name__ == '__main__':
    main()