import json, math
from pathlib import Path
import numpy as np
from scipy.linalg import eigvals
from scipy.optimize import minimize
from scipy.special import lambertw

# Scalar delayed recurrent core:
# xdot(t) = -a*x(t) - b*x(t-h) + r(t), p=x, q=r.
# The stated LMI then certifies ||p||_2^2 <= gamma ||q||_2^2.

def pade_system(a, b, h):
    if h == 0:
        return np.array([[-a-b]]), np.array([[1.]]), np.array([[1.]]), np.array([[0.]])
    # y_delay = -x + 2 z, z_dot = -2/h*z + 2/h*x
    A = np.array([[-a+b, -2*b], [2/h, -2/h]], float)
    B = np.array([[1.], [0.]])
    C = np.array([[1., 0.]])
    D = np.array([[0.]])
    return A, B, C, D

def exact_roots(a,b,h,branches=range(-20,21)):
    if h == 0: return np.array([-a-b+0j])
    z = -b*h*np.exp(a*h)
    return np.array([-a + lambertw(z,k)/h for k in branches])

def exact_stable(a,b,h):
    return np.max(exact_roots(a,b,h).real) < 0

def critical_delay(a,b):
    # for b>a, xdot=-a x-b x(t-h), first imaginary crossing
    if b <= a: return float('inf')
    return math.acos(-a/b)/math.sqrt(b*b-a*a)

def hinf(A,B,C,D, n=30000):
    if np.max(eigvals(A).real) >= -1e-8: return float('inf')
    poles = np.abs(eigvals(A))
    wmax = max(100., 100.*float(np.max(poles)))
    ws = np.concatenate(([0.], np.logspace(-5, math.log10(wmax), n-1)))
    best=0.
    I=np.eye(A.shape[0])
    for w in ws:
        H = (C @ np.linalg.solve(1j*w*I-A,B) + D)[0,0]
        best=max(best,abs(H))
    return float(best)

def lmi_margin(A,B,C,D,P,gamma):
    """Standard bounded-real LMI: dV + ||p||^2 - gamma ||r||^2 <= 0.
    Here gamma is the squared gain bound and P is constrained positive definite.
    """
    n=A.shape[0]
    M=np.block([[A.T@P+P@A + C.T@C, P@B + C.T@D],
                [B.T@P + D.T@C, D.T@D - gamma*np.eye(D.shape[1])]])
    return float(np.max(np.linalg.eigvalsh(M)))

def lmi_feasible(A,B,C,D,gamma):
    n=A.shape[0]
    # Cholesky-like parametrization P=L L' + eps I, optimized against largest eigenvalue.
    def unpack(v):
        L=np.zeros((n,n)); k=0
        for i in range(n):
            for j in range(i+1):
                L[i,j]=v[k]; k+=1
        return L@L.T + 1e-7*np.eye(n)
    x0=np.zeros(n*(n+1)//2)
    for i in range(n): x0[i*(i+1)//2+i]=1.
    fun=lambda v: max(0., lmi_margin(A,B,C,D,unpack(v),gamma))**2 + 1e-8*np.dot(v,v)
    res=minimize(fun,x0,method='Nelder-Mead',options={'maxiter':1000,'xatol':1e-9,'fatol':1e-12})
    P=unpack(res.x); m=lmi_margin(A,B,C,D,P,gamma)
    return bool(m <= -1e-6 and np.min(np.linalg.eigvalsh(P)) >= 1e-6), float(m), P

def lmi_gamma(A,B,C,D):
    lo=max(1e-8,hinf(A,B,C,D)**2*0.95); hi=max(1.,hinf(A,B,C,D)**2*2+1e-3)
    if not np.isfinite(lo): return float('inf'), float('inf')
    # Feasibility is monotone in gamma for this bounded-real form.
    for _ in range(22):
        mid=(lo+hi)/2
        ok,_,_=lmi_feasible(A,B,C,D,mid)
        if ok: hi=mid
        else: lo=mid
    ok,m,_=lmi_feasible(A,B,C,D,hi*1.001)
    return float(hi),float(m)

def run():
    np.random.seed(0); a=1.; b=1.4
    hc=critical_delay(a,b)
    # Prediction 1: direct characteristic roots cross at hc.
    delays=np.array([.5*hc,.8*hc,.95*hc,.99*hc,1.01*hc,1.2*hc,1.5*hc])
    root_rows=[{'h':float(h),'stable':bool(exact_stable(a,b,float(h))),'max_real':float(np.max(exact_roots(a,b,float(h)).real))} for h in delays]
    below=min(delays,key=lambda h: abs(h-.95*hc)); above=min(delays,key=lambda h: abs(h-1.05*hc))
    # Prediction 2: h=0 gain is exactly 1/(a+b), and certified gamma is its square.
    zero=[]
    for bb in [.2,.5,1.,2.,4.]:
        A,B,C,D=pade_system(a,bb,0.)
        g=hinf(A,B,C,D); zero.append({'b':bb,'measured_gain':g,'predicted_gain':1/(a+bb),'relative_error':abs(g-1/(a+bb))/(1/(a+bb))})
    # Prediction 3: gain increases toward the boundary and Padé remains stable below it.
    near=[]
    for h in [0.,.2,.4,.6,.8,.95*hc,1.02*hc]:
        A,B,C,D=pade_system(a,b,h); stable=bool(np.max(eigvals(A).real)<0)
        g=hinf(A,B,C,D); gam= float('inf') if not np.isfinite(g) else g*g
        near.append({'h':h,'exact_stable':exact_stable(a,b,h),'pade_stable':stable,'gain':g,'gamma_hinf':gam,'lmi_gamma':(None if not stable else lmi_gamma(A,B,C,D)[0])})
    # Secondary same-system comparison: unconstrained near-boundary versus certified lower-delay block.
    A1,B1,C1,D1=pade_system(a,b,.95*hc); A2,B2,C2,D2=pade_system(a,b,.35*hc)
    comparison={'unconstrained_delay':.95*hc,'certified_delay':.35*hc,
                'unconstrained_gain':hinf(A1,B1,C1,D1),'certified_gain':hinf(A2,B2,C2,D2),
                'unconstrained_stable':bool(np.max(eigvals(A1).real)<0),'certified_stable':bool(np.max(eigvals(A2).real)<0)}
    out={'parameters':{'a':a,'b':b,'predicted_critical_delay':hc},'prediction_1_boundary':{'rows':root_rows,'below':below,'above':above,'below_stable':exact_stable(a,b,below),'above_stable':exact_stable(a,b,above)},'prediction_2_zero_delay_scaling':zero,'prediction_3_delay_gain':near,'secondary_comparison':comparison}
    Path('results.json').write_text(json.dumps(out,indent=2, default=lambda x: x.item() if hasattr(x, 'item') else str(x)))
    print(json.dumps(out,indent=2, default=lambda x: x.item() if hasattr(x, 'item') else str(x)))

if __name__=='__main__': run()