import json, math, time
import numpy as np
from spectral_edge import SpectralEdgeController

SEED=2744

def toy_sweeps():
    rng=np.random.default_rng(SEED); n=160
    A=rng.normal(size=(n,n)); W=(A+A.T)/2/np.sqrt(n)
    lam=float(np.linalg.eigvalsh(W)[-1])
    ratios=np.array([.70,.85,.95,.99,1.01,1.10]); T=260
    v=rng.normal(size=n); v/=np.linalg.norm(v)
    coeff=abs(np.linalg.eigh(W)[1][:,-1] @ v)
    boundary=[]
    for r in ratios:
        x=coeff*r**np.arange(T)
        slope=float(np.polyfit(np.arange(40,T),np.log(x[40:]),1)[0])
        boundary.append({'ratio':float(r),'slope':slope,'predicted_slope':math.log(r),
          'xi_observed':(-1/slope if slope<0 else None),
          'xi_predicted':(-1/math.log(r) if r<1 else None),
          'norm_ratio':float(x[-1]/x[0])})
    scaling=[]
    for sigma in [.35,.55,.8,1.1]:
        B=rng.normal(size=(n,n)); B=(B+B.T)/2*np.sqrt(2)*sigma/np.sqrt(n)
        L=float(np.linalg.eigvalsh(B)[-1])
        scaling.append({'sigma':sigma,'lambda_edge':L,'predicted_edge':2*sigma,
          'critical_gain':1/L,'predicted_gain':1/(2*sigma)})
    target=.90; alpha=.7; g=2.0/lam; hist=[]
    for _ in range(35):
        hist.append(g*lam); g*=math.exp(alpha*(target-g*lam))
    return {'lambda_edge':lam,'boundary':boundary,'scaling':scaling,
      'controller':{'target':target,'initial_edge':hist[0],'final_edge':hist[-1],'history':hist}}

def nonlinear_comparison():
    rng=np.random.default_rng(SEED+9); n=48; T=80
    A=rng.normal(size=(n,n)); W=(A+A.T)/2
    raw_edge=float(np.linalg.eigvalsh(W)[-1]); W=W/raw_edge
    xs=rng.normal(0,.08,size=(T,n))
    def rollout(g, online=False):
        h=np.zeros(n); states=[]; gains=[]; edges=[]
        ctl=SpectralEdgeController(target=.90,alpha=.30,g_min=.01,g_max=2.)
        ctl.g=g
        for x in xs:
            z=g*(W@h+x); h=np.tanh(z); d=1-h*h
            edge=ctl.edge_estimate(W,d,iterations=5,rng=rng)
            if online: g=ctl.update(edge)
            gains.append(g); edges.append(g*edge); states.append(h.copy())
        # exact product for the realized trajectory with final fixed gain is diagnostic;
        # use per-step gains for a more faithful sensitivity product.
        P=np.eye(n)
        for h,gg in zip(states,gains): P=np.diag(1-h*h)@(gg*W)@P
        grad=float(np.linalg.svd(P,compute_uv=False)[0])
        return {'initial_gain':gains[0],'final_gain':gains[-1],
          'mean_effective_edge':float(np.mean(edges)),'last_effective_edge':float(edges[-1]),
          'gradient_norm':grad,'final_state_norm':float(np.linalg.norm(states[-1]))}
    # Same W and inputs; baseline is standard unit recurrent gain, idea starts at same gain.
    baseline=rollout(1.0,False); idea=rollout(1.0,True)
    return {'raw_lambda_edge':raw_edge,'normalized_lambda_edge':1.,
      'baseline':baseline,'idea':idea,'target_edge':.90}

def main():
    t=time.time(); out={'seed':SEED,'toy':toy_sweeps(),
      'nonlinear':nonlinear_comparison()}; out['runtime_sec']=time.time()-t
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()