import json
import math
import os
import numpy as np
from scipy.linalg import expm
from scipy.signal import savgol_filter
from scipy.optimize import least_squares

SEED = 2112
rng = np.random.default_rng(SEED)


def exact_trajectory(r, alpha, beta, a0, q0, t):
    """Exact linear dynamics [a,q]' = M[a,q]."""
    M = np.array([[r-alpha, beta], [alpha, -beta]], dtype=float)
    x0 = np.array([a0, q0], dtype=float)
    return np.array([(expm(M * ti) @ x0) for ti in t])


def true_coefficients(r, alpha, beta):
    return beta * r, r - alpha - beta


def finite_derivatives(y, dt):
    """Centered second-order finite differences, excluding endpoints."""
    v = (y[2:] - y[:-2]) / (2.0 * dt)
    acc = (y[2:] - 2.0 * y[1:-1] + y[:-2]) / (dt * dt)
    return y[1:-1], v, acc


def fit_reduced_from_samples(y, dt):
    ym, v, acc = finite_derivatives(y, dt)
    X = np.column_stack([ym, v])
    coef, *_ = np.linalg.lstsq(X, acc, rcond=None)
    pred = X @ coef
    return coef, float(np.sqrt(np.mean((pred - acc) ** 2)))


def fit_first_order_from_samples(y, dt):
    v = (y[2:] - y[:-2]) / (2.0 * dt)
    ym = y[1:-1]
    k = float(np.dot(ym, v) / (np.dot(ym, ym) + 1e-12))
    return k


def rollout_reduced(y0, v0, c, dt, n):
    y, v = float(y0), float(v0)
    out = [y]
    cy, cv = c
    for _ in range(n):
        # RK4 for y'=v, v'=cy*y+cv*v
        def f(z): return np.array([z[1], cy*z[0] + cv*z[1]])
        z = np.array([y, v])
        k1 = f(z); k2 = f(z + dt*k1/2); k3 = f(z + dt*k2/2); k4 = f(z + dt*k3)
        z = z + dt*(k1 + 2*k2 + 2*k3 + k4)/6
        y, v = z
        out.append(y)
    return np.asarray(out)


def rollout_first_order(y0, k, dt, n):
    return y0 * np.exp(k * dt * np.arange(n+1))


def invariant_sweep():
    """Prediction 1: fitted cy,cv obey exact rate relations over rate sweep."""
    rows = []
    r = 0.20
    t = np.arange(0, 20.0001, 0.01)
    for alpha in [0.05, 0.15, 0.40, 0.80]:
        for beta in [0.03, 0.12, 0.35, 0.70]:
            x = exact_trajectory(r, alpha, beta, 1.0, 0.4, t)
            c, rmse = fit_reduced_from_samples(x[:, 0] + x[:, 1], t[1]-t[0])
            truth = np.array(true_coefficients(r, alpha, beta))
            rel = np.abs((c-truth) / np.maximum(np.abs(truth), 1e-8))
            rows.append({'alpha':alpha, 'beta':beta, 'cy_hat':c[0], 'cv_hat':c[1],
                         'cy_true':truth[0], 'cv_true':truth[1], 'max_relative_error':float(np.max(rel)),
                         'accel_fit_rmse':rmse})
    return rows


def sampling_sweep():
    """Prediction 2: derivative/reduced accuracy degrades as dt approaches 1/(alpha+beta)."""
    r, alpha, beta = 0.20, 0.60, 0.40
    tau = 1.0/(alpha+beta)
    truth = np.array(true_coefficients(r, alpha, beta))
    rows = []
    # Includes well below, near, and above the claimed fastest switching scale.
    for ratio in [0.05, 0.10, 0.20, 0.50, 1.0, 1.5, 2.0]:
        dt = ratio * tau
        t = np.arange(0, 30 + 0.5*dt, dt)
        x = exact_trajectory(r, alpha, beta, 1.0, 0.25, t)
        y = x[:, 0] + x[:, 1]
        c, rmse = fit_reduced_from_samples(y, dt)
        rel = float(np.linalg.norm(c-truth)/(np.linalg.norm(truth)+1e-12))
        rows.append({'dt_over_tau':ratio, 'dt':dt, 'coefficient_relative_error':rel,
                     'accel_rmse':rmse, 'cy_hat':c[0], 'cv_hat':c[1]})
    return rows


def noisy_rollout_comparison():
    """Prediction 3: structural second-order model extrapolates better than y'=k y."""
    r, alpha, beta = 0.18, 0.45, 0.25
    dt = 0.08
    t_train = np.arange(0, 8+1e-9, dt)
    t_test = np.arange(0, 40+1e-9, dt)
    train = exact_trajectory(r, alpha, beta, 1.0, 0.8, t_train)
    test = exact_trajectory(r, alpha, beta, 1.0, 0.8, t_test)
    ytrain = train.sum(axis=1)
    ytrue = test.sum(axis=1)
    noise = 0.002 * np.std(ytrain) * rng.normal(size=ytrain.shape)
    c, fit_rmse = fit_reduced_from_samples(ytrain + noise, dt)
    k = fit_first_order_from_samples(ytrain + noise, dt)
    v0 = (ytrain[1]-ytrain[0])/dt
    pred2 = rollout_reduced(ytrain[0], v0, c, dt, len(t_test)-1)
    pred1 = rollout_first_order(ytrain[0], k, dt, len(t_test)-1)
    horizon = {'reduced_rmse':float(np.sqrt(np.mean((pred2-ytrue)**2))),
               'first_order_rmse':float(np.sqrt(np.mean((pred1-ytrue)**2))),
               'reduced_tail_rmse':float(np.sqrt(np.mean((pred2[len(pred2)//2:]-ytrue[len(ytrue)//2:])**2))),
               'first_order_tail_rmse':float(np.sqrt(np.mean((pred1[len(pred1)//2:]-ytrue[len(ytrue)//2:])**2))),
               'reduced_coefficients':c.tolist(), 'first_order_k':k, 'fit_accel_rmse':fit_rmse}
    return horizon


def main():
    inv = invariant_sweep()
    samp = sampling_sweep()
    comp = noisy_rollout_comparison()
    inv_err = [x['max_relative_error'] for x in inv]
    # Quantitative criteria: exact math on fine samples, and clear sampling transition.
    low = np.mean([x['coefficient_relative_error'] for x in samp if x['dt_over_tau'] <= .2])
    high = np.mean([x['coefficient_relative_error'] for x in samp if x['dt_over_tau'] >= 1.0])
    result = {
      'seed': SEED,
      'prediction_1_invariant_rate_sweep': {
        'predicted': 'cy=beta*r and cv=r-alpha-beta for every alpha,beta',
        'observed_max_relative_error': float(max(inv_err)),
        'observed_median_relative_error': float(np.median(inv_err)),
        'n_cases':len(inv)},
      'prediction_2_sampling_transition': {
        'predicted': 'error small for dt/tau << 1 and degrades around dt/tau >= 1',
        'tau':1.0, 'low_ratio_mean_error':float(low), 'high_ratio_mean_error':float(high),
        'sweep':samp},
      'prediction_3_rollout': {
        'predicted':'second-order reduced dynamics has lower long-horizon error than first-order y_dot=k*y',
        **comp},
      'invariant_rows':inv}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()