import json
from pathlib import Path
import numpy as np

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


def certificate(a, gain, gamma, eps=None, kappa=None):
    # Scalar A=-a, C=1, P=1, phi_gamma(x)=gamma*tanh(x).
    # OSL rho=gamma and QIB alpha=gamma^2 are globally valid.
    if eps is None:
        eps = gamma if gamma > 0 else 1.0
    M = -2.0 * (a + gain) + 2.0 * gamma + gamma * gamma / eps + eps
    # largest certifiable margin for this fixed P and selected epsilon
    return float(M), float(-M)


def dynamics_error(e, a, gain, gamma):
    # true x and estimate differ by e; nonlinear incremental term is exact.
    return -(a + gain) * e + gamma * np.tanh(e)


def rk4_step(e, h, a, gain, gamma):
    f = lambda z: dynamics_error(z, a, gain, gamma)
    k1 = f(e); k2 = f(e + .5*h*k1); k3 = f(e + .5*h*k2); k4 = f(e + h*k3)
    return e + h*(k1 + 2*k2 + 2*k3 + k4)/6


def estimate_rate(times, vals, start=0.2):
    mask = (times >= start) & (vals > 1e-14)
    if mask.sum() < 3: return float('nan')
    return float(np.polyfit(times[mask], np.log(vals[mask]), 1)[0])


def verify_bounds():
    # Numerical check of OSL/QIB over a broad pair grid, including unequal signs.
    gamma = 1.7
    xs = np.linspace(-8, 8, 401)
    worst_os = -np.inf
    worst_q = -np.inf
    for x in xs[::4]:
        for xhat in xs[::4]:
            e = xhat-x
            if abs(e) < 1e-9: continue
            dp = gamma*(np.tanh(xhat)-np.tanh(x))
            worst_os = max(worst_os, e*dp/e**2)
            worst_q = max(worst_q, dp**2/e**2)
    return {'rho_theory': gamma, 'rho_sample_max': float(worst_os),
            'alpha_theory': gamma**2, 'alpha_sample_max': float(worst_q),
            'bounds_hold': bool(worst_os <= gamma+1e-10 and worst_q <= gamma**2+1e-10)}


def boundary_sweep(a=1.0, gain=0.5):
    # With epsilon=gamma, M=-2(a+L)+4 gamma; sufficient certificate boundary is (a+L)/2.
    rows=[]
    for gamma in np.linspace(0, 2.2, 23):
        M, kap = certificate(a, gain, gamma)
        e=2.0; h=.002; n=int(8/h); vals=[]
        for _ in range(n):
            vals.append(abs(e)); e=rk4_step(e,h,a,gain,gamma)
        # observed contraction if final error is below initial by a meaningful amount
        observed = bool(vals[-1] < vals[0]*1e-3)
        rows.append((gamma, M, observed, vals[-1]))
    # Certificate sign transition (not finite-time threshold): M first becomes nonnegative.
    cert_boundary = next((float(g) for g,m,o,v in rows if m >= 0), float('nan'))
    # Exact local tanh stability boundary is gamma=a+L because tanh'(0)=1.
    exact_boundary = a + gain
    return rows, cert_boundary, exact_boundary


def rate_sweep(a=1.0, gain=0.5):
    # Prediction: log V slope is bounded above by M (the certificate matrix scalar).
    rows=[]
    for gamma in [0.0, .2, .5, .8, 1.0, 1.2, 1.4]:
        M, kap = certificate(a,gain,gamma)
        h=.001; times=np.arange(0, 3+h/2, h); e=0.02; vals=[]
        for t in times:
            vals.append(e*e); e=rk4_step(e,h,a,gain,gamma)
        obs=estimate_rate(times,np.asarray(vals),.3)
        rows.append({'gamma':gamma,'certificate_upper_bound_logV_slope':M,'observed_logV_slope':obs,'bound_holds':bool(obs <= M + 1e-3)})
    return rows


def gain_sweep(a=1.0, gamma=1.2, target_kappa=0.8):
    # From the optimized Young certificate: kappa=2(a+L)-4 gamma, so L_required=2 gamma-a+kappa/2.
    predicted = 2*gamma-a+target_kappa/2
    rows=[]
    for gain in [0,.1,.2,.3,.4,.5,.7,1.0,1.5,1.8,2.0]:
        M,kap=certificate(a,gain,gamma)
        rows.append({'gain':gain,'certificate_kappa':kap,'meets_target':bool(kap>=target_kappa-1e-9)})
    observed=min((r['gain'] for r in rows if r['meets_target']), default=float('nan'))
    return predicted, rows, observed


def noisy_comparison(a=1.0, gamma=0.5):
    # Same Euler observer setup; baseline has no correction, idea uses certified L.
    # State is driven by a known bounded input, while only noisy state output is observed.
    h=.01; steps=1000; noise=.08; trials=40
    out={}
    for name,gain in [('baseline_no_observer',0.0),('certified_observer',1.0)]:
        errs=[]
        for tr in range(trials):
            x=1.5; xhat=0.0; se=0.
            for k in range(steps):
                u=.7*np.sin(.025*k)
                # nonlinear plant: stable linear part plus bounded nonlinear residual and input
                x += h*(-a*x + gamma*np.tanh(x) + u)
                y=x + rng.normal(0,noise)
                xhat += h*(-a*xhat + gamma*np.tanh(xhat) + u + gain*(y-xhat))
                if k >= steps//2: se += (xhat-x)**2
            errs.append(se/(steps//2))
        out[name]={'mean_mse':float(np.mean(errs)),'std_mse':float(np.std(errs))}
    return out


def main():
    a=1.0; gain=.5
    bounds=verify_bounds()
    boundary_rows, observed_boundary, exact_boundary=boundary_sweep(a,gain)
    rates=rate_sweep(a,gain)
    pred_gain,gain_rows,observed_gain=gain_sweep()
    noisy=noisy_comparison()
    result={
      'seed':SEED,
      'model':'scalar stable nonlinear plant with phi_gamma(x)=gamma*tanh(x), P=1, C=1',
      'math':{'a':a,'observer_gain':gain,'rho': 'gamma','alpha':'gamma^2',
              'certificate_at_eps_gamma':'M=-2*(a+L)+4*gamma',
              'predicted_certificate_boundary_gamma':(a+gain)/2,
              'predicted_exact_local_boundary_gamma':a+gain},
      'bound_check':bounds,
      'boundary_sweep':{'predicted_certificate_boundary':(a+gain)/2,'observed_certificate_sign_grid_gamma':observed_boundary,
                        'predicted_exact_local_boundary':a+gain,'observed_exact_local_boundary':exact_boundary,
                        'rows':[{'gamma':g,'M':m,'contracted':o,'final_abs_error':v} for g,m,o,v in boundary_rows]},
      'rate_sweep':rates,
      'gain_sweep':{'predicted_min_gain_for_kappa_0.8':pred_gain,'observed_grid_min_gain':observed_gain,'rows':gain_rows},
      'noisy_comparison':noisy}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps({'bound_check':bounds,'predicted_certificate_boundary':(a+gain)/2,
      'observed_certificate_sign_grid':observed_boundary,
      'predicted_exact_local_boundary':a+gain,'rate_sweep':rates,
      'predicted_gain':pred_gain,'observed_gain_grid':observed_gain,'noisy':noisy},indent=2))

if __name__=='__main__': main()