import json, math
import numpy as np
from scipy.linalg import eigvalsh, solve_discrete_lyapunov

SEED = 2075

def hessians():
    a=.42; R=np.array([[math.cos(a),-math.sin(a)],[math.sin(a),math.cos(a)]])
    return [np.diag([1.,8.]), R@np.diag([1.5,5.5])@R.T]

def lmi_matrices(Hs,K,alpha,sigma,lam=1.):
    n=K.shape[0]; Abar=-K@sum(Hs)/len(Hs)
    P=solve_discrete_lyapunov(np.eye(n)+.02*Abar,np.eye(n)); P=P/np.trace(P)*n
    out=[]
    for H in Hs:
        A=-K@H; B=A
        M=np.block([[A.T@P+P@A+2*alpha*P+lam*sigma*sigma*np.eye(n),P@B],
                    [B.T@P,-lam*np.eye(n)]])
        out.append((M+M.T)/2)
    return P,out

def margin(Hs,K,alpha,sigma,lam=1.):
    return max(float(eigvalsh(M).max()) for M in lmi_matrices(Hs,K,alpha,sigma,lam)[1])

def max_certified_sigma(Hs,K,alpha=.02):
    lo,hi=0.,1.
    while margin(Hs,K,alpha,hi)<0 and hi<100: hi*=2
    for _ in range(55):
        mid=(lo+hi)/2
        if margin(Hs,K,alpha,mid)<0: lo=mid
        else: hi=mid
    return lo

def simulate(Hs,K,h,sigma,steps=1200,max_interval=10000):
    x=np.array([2.,-1.5]); last=x.copy(); events=0; since=0; intervals=[]; norms=[]; losses=[]
    for t in range(steps):
        H=Hs[t%len(Hs)]; x=x-h*K@(H@last); since+=1
        d=last-x
        if sigma==0 or d@d>sigma*sigma*(x@x+1e-30) or since>=max_interval:
            last=x.copy(); events+=1; intervals.append(since); since=0
        norms.append(float(np.linalg.norm(x))); losses.append(float(.5*x@H@x))
        if not np.all(np.isfinite(x)) or norms[-1]>1e12: break
    return {'stable':len(norms)==steps and max(norms)<1e12,'events':events,'fraction':events/steps,'mean_interval':float(np.mean(intervals)) if intervals else steps,'final_norm':norms[-1],'tail_loss':float(np.mean(losses[-max(20,steps//10):])),'max_norm':max(norms)}

def main():
    Hs=hessians(); Hb=sum(Hs)/2
    Kd=np.diag(.72/np.diag(Hb)); Kf=.72*np.linalg.inv(Hb)
    print('HESSIANS',json.dumps([H.tolist() for H in Hs]))
    print('LMI_AND_TRIGGER')
    for name,K in [('diagonal',Kd),('dense',Kf)]:
        sig=max_certified_sigma(Hs,K)
        vals=[]
        for s in [0,.5*sig,.9*sig,1.1*sig]: vals.append({'sigma':s,'margin':margin(Hs,K,.02,s),'run':simulate(Hs,K,.16,s)})
        print(json.dumps({'method':name,'certified_sigma':sig,'sweep':vals}))
    print('DISCRETE_BOUNDARY')
    for name,K in [('diagonal',Kd),('dense',Kf)]:
        rho=max(max(abs(np.linalg.eigvals(K@H))) for H in Hs); boundary=2/rho
        rows=[]
        for mult in [.9,1.02]:
            r=simulate(Hs,K,mult*boundary,0,steps=200,max_interval=1)
            rows.append({'mult':mult,'h':mult*boundary,'stable':r['stable'],'max_norm':r['max_norm']})
        print(json.dumps({'method':name,'rho':float(rho),'boundary':boundary,'rows':rows}))
    print('LMI_RANDOM_IMPLICATION')
    rng=np.random.default_rng(SEED); K=Kf; sig=.8*max_certified_sigma(Hs,K); P,Ms=lmi_matrices(Hs,K,.02,sig)
    worst=-1e9
    for M in Ms:
        for _ in range(20000):
            e=rng.normal(size=2); d=rng.normal(size=2); d*=sig*np.linalg.norm(e)/(np.linalg.norm(d)+1e-30)
            z=np.r_[e,d]; worst=max(worst,float(z@M@z))
    print(json.dumps({'sigma':sig,'max_sampled_quadratic_form':worst,'all_negative':worst<0}))

if __name__=='__main__': main()