import json, math, random
from pathlib import Path
import numpy as np

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


def matrix(k0, k1, k2, gamma=1.0):
    return np.array([[k0+k1, -k1], [-k2, k0+k2]], dtype=float) / gamma


def lyapunov_2x2(A, D):
    # Solve A S + S A.T = 2D for symmetric S=(a,b;b,c).
    M = np.array([[2*A[0,0], 2*A[0,1], 0],
                  [A[1,0], A[0,0]+A[1,1], A[0,1]],
                  [0, 2*A[1,0], 2*A[1,1]]])
    rhs = 2*np.array([D[0,0], D[0,1], D[1,1]])
    a,b,c = np.linalg.solve(M, rhs)
    return np.array([[a,b],[b,c]])


def ep_rate(A, D):
    S = lyapunov_2x2(A, D)
    invS = np.linalg.inv(S)
    B = D @ invS - A
    # E[(Bx)^T D^-1 (Bx)] = tr(B S B.T D^-1)
    ep = np.trace(B @ S @ B.T @ np.linalg.inv(D))
    return float(ep), S


def stability_checks():
    out = {}
    # Prediction 1: continuous stability boundary k0+k1+k2=0.
    k0, gamma = 1.0, 1.0
    vals = np.linspace(-1.4, 0.4, 19)
    rows = []
    for s in vals:
        A = matrix(k0, s/2, s/2, gamma)
        stable_pred = (k0+s) > 0
        stable_obs = np.min(np.real(np.linalg.eigvals(A))) > 0
        rows.append([float(s), bool(stable_pred), bool(stable_obs)])
    # discrete boundary for a representative asymmetric pair, spectral radius=1.
    eta_grid = np.linspace(.01, 1.5, 1500)
    k1,k2 = .75,.25
    A = matrix(k0,k1,k2,gamma)
    rhos = np.array([max(abs(np.linalg.eigvals(np.eye(2)-e*A))) for e in eta_grid])
    idx = np.where(rhos >= 1)[0][0]
    eta_obs = eta_grid[idx]
    eig = np.linalg.eigvals(A)
    eta_pred = min(2*np.real(z)/(abs(z)**2) for z in eig)
    # Direct deterministic rollout: classify as divergent if norm grows by 1e6.
    rollout = []
    for eta in np.linspace(.1, 1.4, 14):
        M = np.eye(2)-eta*A
        x=np.array([1., -1.]); initial=np.linalg.norm(x)
        for _ in range(300): x=M@x
        rollout.append({'eta':float(eta), 'rho':float(max(abs(np.linalg.eigvals(M)))), 'diverged':bool(np.linalg.norm(x)>1e6*initial)})
    out['stability'] = {'continuous_rows': rows, 'discrete_pred_eta': float(eta_pred), 'discrete_obs_eta': float(eta_obs), 'rho_at_pred': float(max(abs(np.linalg.eigvals(np.eye(2)-eta_pred*A)))), 'rollout':rollout}

    # Prediction 2: EP vanishes at delta=0 and is quadratic near reciprocity.
    D=np.eye(2) * .3
    s=.8
    deltas=np.linspace(-.7,.7,15)
    eps=[]
    for d in deltas:
        e,_=ep_rate(matrix(k0,(s+d)/2,(s-d)/2),D)
        eps.append(e)
    small=np.abs(deltas)<=.3
    coeff=np.polyfit(deltas[small]**2, np.array(eps)[small], 1)[0]
    zero=float(eps[len(eps)//2])
    # Ratio EP/delta^2 over small nonzero values.
    ratios=[eps[i]/deltas[i]**2 for i in range(len(deltas)) if small[i] and abs(deltas[i])>.05]
    out['ep']={'deltas':deltas.tolist(),'rates':eps,'zero_rate':zero,'quadratic_coeff':float(coeff),'small_ratio_mean':float(np.mean(ratios)),'small_ratio_cv':float(np.std(ratios)/np.mean(ratios))}

    # Prediction 3: covariance is finite only on stable side and grows toward boundary.
    cov_rows=[]
    for s in [-.8,-.5,0,.5,1.0]:
        A=matrix(k0,s/2,s/2)
        S=lyapunov_2x2(A,D)
        cov_rows.append({'s':s,'max_variance':float(np.max(np.linalg.eigvalsh(S))), 'pred_stable':bool(k0+s>0)})
    out['covariance']=cov_rows
    return out


def ml_experiment():
    # Tiny nonlinear problem with fixed full-batch gradients; explicit Langevin noise
    # makes the optimizer comparison deterministic and compute-matched.
    import torch
    torch.set_num_threads(4)
    torch.manual_seed(SEED)
    n=512
    x=torch.randn(n,2)
    y=((x[:,0]*x[:,1])>0).long()
    # fixed MLP parameter vector, functional forward
    shapes=[(2,24),(24,24),(24,2)]
    sizes=[a*b for a,b in shapes]+[24,24,2]
    total=sum(sizes)
    def unpack(v):
        p=[]; q=0
        for (a,b),sz in zip(shapes,sizes[:3]): p.append(v[q:q+sz].reshape(a,b)); q+=sz
        for sz in sizes[3:]: p.append(v[q:q+sz]); q+=sz
        return p
    def loss(v):
        w1,w2,w3,b1,b2,b3=unpack(v)
        z=torch.tanh(x@w1+b1); z=torch.tanh(z@w2+b2); logits=z@w3+b3
        return torch.nn.functional.cross_entropy(logits,y)
    def grad(v):
        v=v.detach().requires_grad_(True); l=loss(v); return l.detach(),torch.autograd.grad(l,v)[0].detach()
    init=torch.randn(total)*.15
    steps=250; eta=.08; T=.0008
    def run(kind):
        a=init.clone(); b=init.clone(); losses=[]
        for t in range(steps):
            if kind=='baseline':
                l,g=grad(a); a=a-eta*g
                losses.append(float(l))
            else:
                l1,g1=grad(a); l2,g2=grad(b)
                if kind=='reciprocal': k1=k2=.8
                else: k1,k2=.95,.25
                noise1=torch.randn_like(a)*math.sqrt(2*eta*T); noise2=torch.randn_like(a)*math.sqrt(2*eta*T)
                a=a-eta*(g1+k1*(a-b))-eta*.02*a+noise1
                b=b-eta*(g2+k2*(b-a))-eta*.02*b+noise2
                losses.append(float(loss((a+b)/2)))
        return float(loss((a+b)/2) if kind!='baseline' else loss(a)), float((losses[0]-losses[-1])), losses
    result={}
    for k in ['baseline','reciprocal','nonreciprocal']:
        final,drop,trace=run(k); result[k]={'final_loss':final,'loss_drop':drop,'trace_tail':trace[-10:]}
    return result

if __name__=='__main__':
    result={'seed':SEED,'checks':stability_checks()}
    try:
        result['ml']=ml_experiment()
    except Exception as e:
        result['ml_error']=repr(e)
    Path('results.json').write_text(json.dumps(result,indent=2))
    c=result['checks']
    print(json.dumps({'stability':c['stability'],'ep':c['ep'],'covariance':c['covariance'],'ml':result.get('ml',result.get('ml_error'))},indent=2))