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

SEED = 1418
np.random.seed(SEED)


def spectral_abscissa(M):
    return float(np.max(np.real(eigvals(M))))


def linear_sweep():
    # x' = A x, y = [1,0]x; z' = A z + G(y-Cz)
    # The kernel-projected H is always [a22], independent of G.
    A = np.array([[0., 1.], [1., .5]])
    C = np.array([[1., 0.]])
    g1 = 2.0
    gs = np.linspace(0.0, 4.0, 81)
    rows = []
    for g2 in gs:
        G = np.array([[g1], [g2]])
        M = A - G @ C
        # U_K = [0,1]^T, hence H = .5.
        H = np.array([[0.5]])
        rows.append((float(g2), spectral_abscissa(M), float(H[0, 0])))
    stable = [g for g, ab, _ in rows if ab < 0]
    observed_boundary = min(stable) if stable else None
    # exact characteristic polynomial gives det(M)=g2-2 and trace=-1.5.
    predicted_boundary = 2.0

    # Check exponential error decay at a stable and unstable gain with exact flow.
    x0 = np.array([1., -1.])
    times = np.linspace(0, 8, 161)
    checks = {}
    for g2 in (1.5, 3.0):
        M = A - np.array([[g1], [g2]]) @ C
        norms = np.array([np.linalg.norm(expm(M*t) @ x0) for t in times])
        # fit after transient, avoiding zero/roundoff
        slope = float(np.polyfit(times[20:], np.log(norms[20:] + 1e-30), 1)[0])
        checks[str(g2)] = {"spectral_abscissa": spectral_abscissa(M), "log_norm_slope": slope,
                           "norm_t0": float(norms[0]), "norm_t8": float(norms[-1])}
    return {"predicted_stability_boundary_g2": predicted_boundary,
            "observed_first_grid_stable_g2": observed_boundary,
            "grid_step": float(gs[1]-gs[0]),
            "projected_H": 0.5,
            "checks": checks}


def gain_sweep():
    A = np.array([[0., 1.], [1., .5]])
    C = np.array([[1., 0.]])
    U = np.array([[0.], [1.]])
    rows = []
    for g2 in np.linspace(0., 4., 17):
        G = np.array([[2.], [g2]])
        M = A - G @ C
        H = U.T @ M @ U
        rows.append({"g2": float(g2), "full_abscissa": spectral_abscissa(M),
                     "projected_H": float(H[0, 0])})
    return rows


def euler_boundary_sweep():
    # Prediction: explicit Euler is stable iff max |1+dt*lambda_i|<1.
    A = np.array([[0., 1.], [1., .5]])
    C = np.array([[1., 0.]])
    g1 = 2.0
    g2 = 3.0
    M = A - np.array([[g1], [g2]]) @ C
    dts = np.linspace(.001, 3.5, 3500)
    stable = []
    for dt in dts:
        rho = max(abs(eigvals(np.eye(2) + dt*M)))
        if rho < 1: stable.append(dt)
    # solve numerically at fine grid; report observed and direct formula check.
    observed = max(stable)
    eig = eigvals(M)
    # Exact disk condition for each (possibly complex) eigenvalue:
    # |1 + dt*lambda| < 1 => dt < -2 Re(lambda)/|lambda|^2.
    predicted = min(float(-2*l.real/(abs(l)**2)) for l in eig)
    return {"g2": g2, "eigenvalues": [[float(l.real), float(l.imag)] for l in eig],
            "predicted_dt_max": predicted, "observed_dt_max_grid": observed,
            "dt_grid": float(dts[1]-dts[0])}


def lorenz_rhs(x):
    sig, rho, beta = 10., 28., 8./3.
    return np.array([sig*(x[1]-x[0]), x[0]*(rho-x[2])-x[1], x[0]*x[1]-beta*x[2]])


def rk4_step(fun, x, dt):
    k1=fun(x); k2=fun(x+dt*k1/2); k3=fun(x+dt*k2/2); k4=fun(x+dt*k3)
    return x + dt*(k1+2*k2+2*k3+k4)/6


def lorenz_observer():
    # Known vector field, partial continuous observations x,y. This is a practical
    # nudging check, not a trained neural ODE benchmark.
    dt=.005; T=10.; n=int(T/dt)
    x=np.array([1.,1.,1.]); z0=np.array([-8.,7.,18.])
    C=np.array([[1.,0.,0.],[0.,1.,0.]])
    gains={"open_loop":np.zeros((3,2)), "nudging":np.array([[8.,0.],[0.,8.],[12.,-4.]])}
    out={}
    for name,G in gains.items():
        truth=x.copy(); z=z0.copy(); errors=[]
        for _ in range(n):
            y=C@truth
            def obs(q): return lorenz_rhs(q)+G@(y-C@q)
            truth=rk4_step(lorenz_rhs,truth,dt)
            z=rk4_step(obs,z,dt)
            errors.append(np.linalg.norm(truth-z))
        out[name]={"initial_error":float(np.linalg.norm(x-z0)),
                   "final_error":float(errors[-1]),
                   "median_last_20pct":float(np.median(errors[int(.8*n):]))}
    return out


def main():
    result={"seed":SEED, "linear_contraction_sweep":linear_sweep(),
            "gain_sweep":gain_sweep(),
            "euler_stability_sweep":euler_boundary_sweep(),
            "lorenz_partial_observation":lorenz_observer(),
            "interpretation":{
              "prediction_1":"Full 2D observer stability begins at g2=2 from det(A-GC)>0; observed grid boundary should agree.",
              "prediction_2":"At stable g2=3, Euler stability has a finite dt boundary predicted by -2 Re(lambda)/|lambda|^2; observed sweep should agree.",
              "prediction_gain":"Projected H remains +0.5 across all gains, while full abscissa crosses zero at g2=2; this falsifies H-only sufficiency here.",
              "prediction_3":"The proposed H on ker(C) remains +0.5 for every gain, despite full observer becoming stable; therefore the stated projected certificate is not sufficient in this coupled example."}}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()