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

SEED = 1645
np.random.seed(SEED)
random.seed(SEED)

# Weakly normally hyperbolic cyclic optimizer on a quadratic.
# State x=(theta,m), phase advances by one fixed step per update.
def step(x, phi, eta0, amp, mu0, mu_amp, curvature=1.0, forcing=0.0):
    theta, mom = float(x[0]), float(x[1])
    eta = eta0 * (1.0 + amp * math.sin(phi))
    mu = mu0 + mu_amp * math.cos(phi)
    grad = curvature * theta + forcing
    mom2 = mu * mom + grad
    return np.array([theta - eta * mom2, mom2], dtype=float)

def jacobian(phi, eta0, amp, mu0, mu_amp, curvature=1.0):
    eta = eta0 * (1.0 + amp * math.sin(phi))
    mu = mu0 + mu_amp * math.cos(phi)
    # derivative of [theta-eta*(mu*m+curvature*theta), mu*m+curvature*theta]
    return np.array([[1.0 - eta * curvature, -eta * mu],
                     [curvature, mu]], dtype=float)

def cycle_map(x, pars, return_path=False):
    P = pars['P']; dphi = 2.0 * math.pi / P
    path = []
    y = np.asarray(x, dtype=float).copy()
    for k in range(P):
        if return_path: path.append(y.copy())
        y = step(y, k*dphi, **{q: pars[q] for q in ('eta0','amp','mu0','mu_amp','curvature','forcing')})
    return (y, path) if return_path else y

def monodromy(pars):
    P = pars['P']; dphi = 2.0*math.pi/P
    M = np.eye(2)
    # perturbations propagate J_k, hence M=J_{P-1}...J_0
    for k in range(P):
        M = jacobian(k*dphi, pars['eta0'], pars['amp'], pars['mu0'], pars['mu_amp'], pars['curvature']) @ M
    return M

def periodic_orbit(pars):
    # affine map is solved robustly by fixed-point iteration; fallback linear solve
    x = np.zeros(2)
    for _ in range(10000):
        y = cycle_map(x, pars)
        if np.linalg.norm(y-x) < 1e-13: return x
        x = y
    # Numerical finite difference around zero gives affine offset.
    c = cycle_map(np.zeros(2), pars)
    return np.linalg.solve(np.eye(2)-monodromy(pars), c)

def empirical_decay(pars, perturb=1e-6, cycles=12):
    xstar = periodic_orbit(pars)
    y0 = cycle_map(xstar, pars)
    # Start on the same phase with a transverse state perturbation.
    y = xstar + np.array([perturb, -0.7*perturb])
    norms = []
    for _ in range(cycles+1):
        norms.append(float(np.linalg.norm(y-xstar)))
        y = cycle_map(y, pars)
    # fit log norm, excluding any underflow tail
    vals = np.maximum(np.asarray(norms), 1e-300)
    slope = float(np.polyfit(np.arange(len(vals)), np.log(vals), 1)[0])
    return norms, slope

def noise_gain(pars, sigma=1e-5, trials=200, cycles=80):
    # IID gradient perturbation added at each update; compare stationary RMS.
    xstar = periodic_orbit(pars)
    rng = np.random.default_rng(SEED+7)
    vals=[]
    for t in range(trials):
        y=xstar.copy()
        for k in range(cycles*pars['P']):
            phi=2*math.pi*(k % pars['P'])/pars['P']
            # forcing is zero here; inject noise into gradient via equivalent forcing
            q=dict(pars); q['forcing']=float(rng.normal(0,sigma))
            y=step(y,phi, **{z:q[z] for z in ('eta0','amp','mu0','mu_amp','curvature','forcing')})
        vals.append(np.linalg.norm(y-xstar))
    return float(np.sqrt(np.mean(np.square(vals))))

def stability_sweep():
    base={'P':12,'eta0':0.8,'amp':0.0,'mu0':0.75,'mu_amp':0.0,'curvature':1.0,'forcing':0.0}
    # Prediction 1: analytic one-cycle Jacobian and finite differences agree.
    rows=[]
    for amp in [0.0,0.2,0.5,0.8]:
        p=dict(base,amp=amp,mu_amp=0.08)
        M=monodromy(p); rho=max(abs(np.linalg.eigvals(M)))
        eps=1e-7; x=periodic_orbit(p); c0=cycle_map(x,p)
        fd=np.column_stack([(cycle_map(x+eps*np.eye(2)[j],p)-c0)/eps for j in range(2)])
        rho_fd=max(abs(np.linalg.eigvals(fd)))
        rows.append({'amp':amp,'rho_analytic':float(rho),'rho_finite_difference':float(rho_fd),
                     'relative_error':float(abs(rho-rho_fd)/max(rho,1e-15))})

    # Prediction 2: rho<1 gives decay and rho>1 gives growth. Find predicted
    # boundary (rho=1) by bisection, then test both sides empirically.
    def rho_at(eta):
        p=dict(base,eta0=float(eta),amp=0.65,mu_amp=0.10)
        return float(max(abs(np.linalg.eigvals(monodromy(p)))))
    lo,hi=2.3,2.4
    for _ in range(50):
        mid=(lo+hi)/2
        if rho_at(mid)<1: lo=mid
        else: hi=mid
    predicted_boundary=(lo+hi)/2
    crossing=[]
    for eta in [predicted_boundary-0.05,predicted_boundary-0.01,
                predicted_boundary+0.01,predicted_boundary+0.05]:
        p=dict(base,eta0=float(eta),amp=0.65,mu_amp=0.10)
        rho=rho_at(eta); norms,slope=empirical_decay(p,cycles=12)
        crossing.append({'eta0':float(eta),'rho':rho,
                         'predicted_log_multiplier':float(math.log(rho)),
                         'observed_log_decay_per_cycle':slope,
                         'decays':bool(norms[-1] < norms[0])})

    # Prediction 3: near the stable cycle, noise amplification grows with the
    # resolvent scale 1/(1-rho). Vary momentum to span distinct rho values.
    ng=[]
    for mu in [0.10,0.25,0.40,0.55,0.70,0.82]:
        p=dict(base,eta0=0.35,amp=0.35,mu0=mu,mu_amp=0.0)
        rho=float(max(abs(np.linalg.eigvals(monodromy(p)))))
        if rho < .98:
            gain=noise_gain(p,sigma=2e-5,trials=150,cycles=80)
            ng.append({'mu0':mu,'rho':rho,'noise_rms':gain,
                       'predicted_resolvent':float(1/(1-rho))})
    # rank correlation is a scale-free check of monotonicity.
    order=np.argsort([x['rho'] for x in ng])
    noise_monotone=all(ng[order[i]]['noise_rms'] <= ng[order[i+1]]['noise_rms']
                       for i in range(len(order)-1))
    return rows,crossing,ng,{'predicted_eta_boundary_rho1':predicted_boundary,
                             'noise_monotone_with_rho':bool(noise_monotone)}

# Small nonlinear regression: same data, update count, and average LR.
def mlp_compare():
    rng=np.random.default_rng(SEED)
    X=rng.normal(size=(512,2)).astype(np.float64)
    y=(np.sin(X[:,0])+0.35*X[:,1]**2).astype(np.float64)
    split=384; Xtr,ytr=X[:split],y[:split]; Xte,yte=X[split:],y[split:]
    def run(cyclic):
        W1=rng.normal(0,.35,(2,16)); b1=np.zeros(16); W2=rng.normal(0,.2,16); b2=0.
        m=[np.zeros_like(W1),np.zeros_like(b1),np.zeros_like(W2),0.]
        eta0=.025; mu0=.85; P=20; amp=.55; muamp=.06
        for k in range(1800):
            ix=rng.choice(split,64,replace=False); a=Xtr[ix]; target=ytr[ix]
            h=np.tanh(a@W1+b1); pred=h@W2+b2; d=(2*(pred-target)/len(ix))
            dW2=h.T@d; db2=d.sum(); dh=d[:,None]*W2[None,:]; dz=dh*(1-h*h)
            dW1=a.T@dz; db1=dz.sum(0)
            if cyclic:
                ph=2*math.pi*(k%P)/P; eta=eta0*(1+amp*math.sin(ph)); mu=mu0+muamp*math.cos(ph)
            else: eta=eta0; mu=mu0
            for j,g in enumerate([dW1,db1,dW2,db2]):
                m[j]=mu*m[j]+g
            W1-=eta*m[0]; b1-=eta*m[1]; W2-=eta*m[2]; b2-=eta*m[3]
        pred=np.tanh(Xte@W1+b1)@W2+b2
        return float(np.mean((pred-yte)**2))
    # reset is intentionally deterministic via local seeds for fair pair.
    np.random.seed(SEED); random.seed(SEED); cyc=run(True)
    np.random.seed(SEED); random.seed(SEED); const=run(False)
    return {'constant_average_momentum_test_mse':const,'cyclic_test_mse':cyc}

def main():
    rows,crossing,ng,summary=stability_sweep()
    result={'seed':SEED,'predictions':[
      'Floquet multiplier from the analytic one-cycle Jacobian equals finite-difference cycle response.',
      'The decay/growth transition occurs at spectral radius rho=1.',
      'Stable noisy response increases with the predicted resolvent scale 1/(1-rho).'],
      'floquet_check':rows,'stability_boundary_sweep':crossing,'noise_scaling':ng,'summary':summary,
      'mlp_comparison':mlp_compare()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()