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

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


def make_hessians():
    # Two non-commuting positive definite curvature vertices.
    r = np.array([[math.cos(.42), -math.sin(.42)], [math.sin(.42), math.cos(.42)]])
    h1 = np.diag([1.0, 8.0])
    h2 = r @ np.diag([1.5, 5.5]) @ r.T
    return [h1, h2]


def lmi_margin(Hs, K, alpha, sigma, P=None, lam=1.0):
    """Return largest eigenvalue of supplied S-procedure matrices.
    Negative means the fixed-P certificate is strict. P is normalized if absent.
    """
    n = Hs[0].shape[0]
    if P is None:
        # A valid Lyapunov candidate for the no-staleness closed loop.
        Abar = -K @ sum(Hs) / len(Hs)
        P = solve_discrete_lyapunov(np.eye(n) + .02*Abar, np.eye(n))
        P = P / np.trace(P) * n
    worst = -np.inf
    for H in Hs:
        A = -K @ H
        B = A.copy()
        tl = A.T@P + P@A + 2*alpha*P + lam*sigma*sigma*np.eye(n)
        M = np.block([[tl, P@B], [B.T@P, -lam*np.eye(n)]])
        worst = max(worst, eigvalsh((M+M.T)/2).max())
    return float(worst), P


def run(h, K, sigma, steps=2500, seed=0, max_interval=10):
    rng = np.random.default_rng(seed)
    Hs = make_hessians()
    x = np.array([2.0, -1.5])
    last = x.copy()
    total = 0.0
    events = 0
    max_norm = 0.0
    losses = []
    intervals = []
    since = 0
    # Alternating vertex is a deliberately adversarial but bounded curvature sequence.
    for t in range(steps):
        H = Hs[t % len(Hs)]
        g_stale = H @ last
        x = x - h * (K @ g_stale)
        since += 1
        # e proxy is x itself because x*=0, matching the proposed trigger proxy.
        d = last - x
        if (d @ d > sigma*sigma*(x @ x + 1e-12)) or since >= max_interval:
            last = x.copy()
            events += 1
            intervals.append(since)
            since = 0
        q = .5*x @ H @ x
        total += q
        losses.append(q)
        max_norm = max(max_norm, float(np.linalg.norm(x)))
        if not np.isfinite(q) or max_norm > 1e12:
            return dict(stable=False, final_loss=float('inf'), events=events,
                        rate=events/steps, mean_interval=np.mean(intervals) if intervals else steps,
                        max_norm=max_norm, losses=losses)
    return dict(stable=True, final_loss=float(losses[-1]), events=events,
                rate=events/steps, mean_interval=np.mean(intervals) if intervals else steps,
                max_norm=max_norm, losses=losses)


def main():
    Hs = make_hessians()
    # Dense K is aligned to average Hessian; diagonal baseline has same trace/average scale.
    Hbar = sum(Hs)/len(Hs)
    K_dense = np.linalg.inv(Hbar) * 0.72
    K_diag = np.diag(0.72/np.diag(Hbar))
    # Prediction 1: no-stale discrete quadratic stability boundary h*rho(KH)<2.
    print('HESSIANS', [h.tolist() for h in Hs])
    for name,K in [('diag',K_diag),('dense',K_dense)]:
        rho=max(abs(np.linalg.eigvals(K@H)) .max() for H in Hs)
        predicted=2/rho
        rows=[]
        for mult in [0.90, 1.02]:
            h=mult*predicted
            rr=run(h,K,0.0,steps=160,seed=1,max_interval=1)
            rows.append({'mult':mult,'h':h,'stable':rr['stable'],'max_norm':rr['max_norm']})
        print('STABILITY',name,json.dumps({'rho':float(rho),'predicted_h_boundary':predicted,'observed':rows}))
    # Prediction 2: trigger threshold zero gives stale-free updates; larger sigma permits longer intervals
    # and should increase error relative to a synchronized reference.
    h=.16
    print('TRIGGER_SWEEP')
    for name,K in [('diag',K_diag),('dense',K_dense)]:
        ref=run(h,K,0.0,steps=800,seed=3,max_interval=1)
        for sigma in [0.0,.01,.03,.08,.16,.30]:
            rr=run(h,K,sigma,steps=800,seed=3,max_interval=50)
            # average loss over final 20%, normalized to stale-free final loss floor
            tail=float(np.mean(rr['losses'][-160:]))
            ref_tail=float(np.mean(ref['losses'][-160:]))
            print(json.dumps({'kind':'trigger','method':name,'sigma':sigma,
                'events':rr['events'],'event_fraction':rr['rate'],'mean_interval':rr['mean_interval'],
                'final_tail_loss':tail,'loss_ratio_vs_sync':tail/(ref_tail+1e-30),
                'stable':rr['stable']}))
    # Prediction 3: local first-order communication interval scales approximately sigma/(h ||KH||)
    # for small sigma, so event fraction should scale approximately linearly with h/sigma.
    print('SCALING_SWEEP')
    K=K_dense
    for h in [.04,.08,.16]:
        for sigma in [.02,.04,.08]:
            rr=run(h,K,sigma,steps=1000,seed=4,max_interval=1000)
            print(json.dumps({'kind':'scaling','h':h,'sigma':sigma,'events':rr['events'],
                'event_fraction':rr['rate'],'mean_interval':rr['mean_interval'],'stable':rr['stable']}))
    # Fixed-K LMI sanity check: report monotone margin as sigma increases.
    print('LMI_SWEEP')
    for name,K in [('diag',K_diag),('dense',K_dense)]:
        for sigma in [0.0,.02,.05,.1,.2]:
            margin,_=lmi_margin(Hs,K,alpha=.02,sigma=sigma,lam=1.0)
            print(json.dumps({'kind':'lmi','method':name,'sigma':sigma,'worst_eigenvalue':margin,'certified':margin<0}))

if __name__ == '__main__':
    main()