import json
import numpy as np
from scipy.linalg import expm
from pathlib import Path


def block(gf, gs, k):
    return np.array([[-gf, k], [k, -gs]], dtype=float)


def euler(A, h, n, x0=None, forcing=None):
    x = np.zeros((n + 1, A.shape[0]))
    if x0 is not None:
        x[0] = x0
    for t in range(n):
        u = 0.0 if forcing is None else forcing[t]
        b = np.ones(A.shape[0]) * u
        x[t + 1] = x[t] + h * (A @ x[t] + b)
    return x


def frequency_transfer(A, w, b=None, c=None):
    b = np.ones(A.shape[0]) if b is None else np.asarray(b)
    c = np.ones(A.shape[0]) if c is None else np.asarray(c)
    z = 1j * w
    return c @ np.linalg.solve(z * np.eye(A.shape[0]) - A, b)


def decay_rates_from_impulse(A, dt=0.002, duration=12.0):
    # Independently simulate exact sampled propagation, then fit each modal
    # amplitude. This tests decay rather than simply reporting eig(A).
    vals, vecs = np.linalg.eigh(A)
    x = np.ones(2)
    n = int(duration / dt)
    modal = np.empty((n + 1, 2))
    modal[0] = np.linalg.solve(vecs, x)
    P = expm(A * dt)
    for i in range(n):
        x = P @ x
        modal[i + 1] = np.linalg.solve(vecs, x)
    t = np.arange(n + 1) * dt
    fitted = []
    for j in range(2):
        # Avoid the final numerical underflow tail and fit log magnitude.
        keep = (np.abs(modal[:, j]) > 1e-10) & (t > 0.1) & (t < duration * .8)
        slope = np.polyfit(t[keep], np.log(np.abs(modal[keep, j])), 1)[0]
        fitted.append(-slope)
    return np.sort(-vals), np.sort(np.asarray(fitted))


def least_squares_reconstruction(gf=10.0, gs=1.0, k=0.15, seed=0):
    rng = np.random.default_rng(seed)
    n_train, n_test = 3000, 1500
    t = np.arange(n_train + n_test)
    y = (np.sin(2 * np.pi * 0.025 * t) + 0.7 * np.sin(2 * np.pi * 0.22 * t + .4))
    x = y + 0.03 * rng.standard_normal(len(t))
    # Discrete stable dynamics; each state is a causal filtered view of x.
    h = 0.03
    Ad = np.eye(2) + h * block(gf, gs, k)
    states = np.zeros((len(x) + 1, 2))
    for i, u in enumerate(x):
        states[i + 1] = Ad @ states[i] + h * np.array([u, u])
    # Baseline: ordinary one-timescale slow state, matched parameter count as
    # closely as possible by fitting its readout on the same training segment.
    sb = np.zeros(len(x) + 1)
    ad = 1 - h * gs
    for i, u in enumerate(x):
        sb[i + 1] = ad * sb[i] + h * u
    def fit_predict(features):
        F = np.column_stack([features, np.ones(len(features))])
        coef = np.linalg.lstsq(F[:n_train], y[:n_train], rcond=None)[0]
        return F[n_train:] @ coef
    pred_dual = fit_predict(states[1:])
    pred_single = fit_predict(sb[1:, None])
    return float(np.mean((pred_single - y[n_train:]) ** 2)), float(np.mean((pred_dual - y[n_train:]) ** 2))


def main():
    out = {}
    # Prediction 1: symmetric-part stability changes sign at k=sqrt(gf*gs).
    gf, gs = 10.0, 1.0
    kcrit = np.sqrt(gf * gs)
    ks = np.linspace(0, 1.5 * kcrit, 301)
    lmax = np.array([np.linalg.eigvalsh((block(gf, gs, k) + block(gf, gs, k).T) / 2).max() for k in ks])
    crossing = ks[np.argmin(np.abs(lmax))]
    out['stability_prediction'] = {
        'predicted_kcrit': float(kcrit), 'observed_kcrit_grid': float(crossing),
        'relative_error': float(abs(crossing-kcrit)/kcrit),
        'stable_at_0.8kcrit': bool(np.linalg.eigvalsh(block(gf, gs, .8*kcrit)).max() <= 1e-12),
        'unstable_at_1.2kcrit': bool(np.linalg.eigvalsh(block(gf, gs, 1.2*kcrit)).max() > 1e-12),
    }
    # Prediction 2: a gap yields two decay rates, with timescale ratio tracking
    # the imposed gap (small exchange perturbation).
    gap_rows = []
    for ratio in [2, 5, 10, 20, 40]:
        gs0, gf0 = 1.0, float(ratio)
        kk = .05 * np.sqrt(gf0 * gs0)
        rates, fitted = decay_rates_from_impulse(block(gf0, gs0, kk))
        gap_rows.append({'ratio': ratio, 'predicted_rates': rates.tolist(), 'measured_rates': fitted.tolist(),
                         'measured_rate_ratio': float(fitted[1]/fitted[0])})
    out['gap_prediction'] = gap_rows
    # Prediction 3: slow-only reduction loses the fast response around w~gf.
    A = block(10., 1., .15)
    freqs = np.logspace(-2, 2, 300)
    full = np.array([abs(frequency_transfer(A, w)) for w in freqs])
    slow = np.array([abs(1 / (1j*w + 1)) for w in freqs])
    relerr = np.abs(full-slow) / np.maximum(full, 1e-12)
    low = float(np.mean(relerr[freqs < 0.3]))
    high = float(np.mean(relerr[freqs > 10.]))
    near = float(np.mean(relerr[(freqs > 7) & (freqs < 14)]))
    out['frequency_prediction'] = {'fast_rate': 10.0, 'observed_peak_error_frequency': float(freqs[np.argmax(relerr)]),
                                   'mean_relative_error_low_w': low, 'mean_relative_error_near_fast_rate': near,
                                   'mean_relative_error_high_w': high, 'error_ratio_near_vs_low': near/max(low,1e-12)}
    single_mse, dual_mse = least_squares_reconstruction()
    out['toy_reconstruction'] = {'single_timescale_mse': single_mse, 'dual_band_mse': dual_mse,
                                 'relative_mse_change': (single_mse-dual_mse)/single_mse}
    Path('results.json').write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()