import json, math, os
import numpy as np

SEED = 3057
np.random.seed(SEED)


def continuous_roots(lam, kappa, beta):
    return np.roots([1.0, lam + kappa + beta, beta * lam])


def fm_matrix(lam, kappa, beta, eta):
    rho = math.exp(-beta * eta)
    # theta' = (1-eta*(lambda+kappa))*theta + eta*kappa*m
    a = 1.0 - eta * (lam + kappa)
    b = eta * kappa
    # m' = rho*m + (1-rho)*theta'
    return np.array([[a, b], [(1-rho)*a, rho + (1-rho)*b]])


def spectral_radius(A):
    return float(np.max(np.abs(np.linalg.eigvals(A))))


def fm_rollout(lams, kappa, beta, eta, steps, x0):
    x = np.asarray(x0, dtype=float).copy()
    m = x.copy()
    xs, losses, forces = [], [], []
    rho = math.exp(-beta * eta)
    for _ in range(steps):
        g = lams * x
        force = kappa * (m - x)
        x = x - eta * g + eta * force
        m = rho * m + (1-rho) * x
        xs.append(x.copy()); losses.append(float(0.5*np.sum(lams*x*x)))
        forces.append(float(np.linalg.norm(force)))
    return np.asarray(xs), np.asarray(losses), np.asarray(forces)


def sgd_rollout(lams, eta, steps, x0):
    x = np.asarray(x0, dtype=float).copy(); xs=[]; losses=[]
    for _ in range(steps):
        x = x - eta*lams*x
        xs.append(x.copy()); losses.append(float(0.5*np.sum(lams*x*x)))
    return np.asarray(xs), np.asarray(losses)


def momentum_rollout(lams, eta, momentum, steps, x0):
    x=np.asarray(x0,dtype=float).copy(); v=np.zeros_like(x); xs=[]; losses=[]
    for _ in range(steps):
        v = momentum*v + lams*x
        x = x - eta*v
        xs.append(x.copy()); losses.append(float(0.5*np.sum(lams*x*x)))
    return np.asarray(xs), np.asarray(losses)


def main():
    # A known two-mode quadratic, plus a single-mode check for exact roots.
    lam, kappa, beta = 3.0, 2.0, 5.0
    roots = continuous_roots(lam, kappa, beta)
    root_residuals = [abs(r*r + (lam+kappa+beta)*r + beta*lam) for r in roots]

    # Predicted discrete boundary is spectral radius == 1; compare with rollout.
    grid = np.linspace(0.005, 1.2, 240)
    radii = np.array([spectral_radius(fm_matrix(lam,kappa,beta,e)) for e in grid])
    stable_pred = grid[radii < 1.0]
    predicted_eta_max = float(stable_pred[-1]) if len(stable_pred) else None
    # empirical stability criterion: bounded and decreasing from x=1 over 500 steps
    empirical = []
    for eta in grid:
        _, ls, _ = fm_rollout(np.array([lam]),kappa,beta,float(eta),500,[1.0])
        empirical.append(bool(np.all(np.isfinite(ls)) and np.max(ls) < 1e8 and ls[-1] < ls[0]))
    empirical = np.asarray(empirical)
    empirical_eta_max = float(grid[np.where(empirical)[0][-1]]) if np.any(empirical) else None

    # Matched tiny optimization benchmark: ill-conditioned quadratic, equal 250 steps.
    lams=np.array([1.0, 20.0]); x0=np.array([1.0,1.0]); steps=250
    configs = {
      'sgd': (lambda eta: sgd_rollout(lams,eta,steps,x0), [0.01,0.02,0.04,0.06,0.08]),
      'momentum': (lambda eta: momentum_rollout(lams,eta,0.9,steps,x0), [0.005,0.01,0.02,0.03,0.04]),
      'fading_memory': (lambda eta: fm_rollout(lams,2.0,5.0,eta,steps,x0), [0.005,0.01,0.02,0.03,0.04,0.05])
    }
    bench={}
    for name,(runner,etas) in configs.items():
        rows=[]
        for eta in etas:
            out=runner(eta)
            ls=out[1]
            finite=bool(np.all(np.isfinite(ls)) and np.max(ls)<1e12)
            rows.append({'eta':eta,'final_loss':float(ls[-1]) if finite else 1e12,
                         'min_loss':float(np.min(ls)) if finite else 1e12,
                         'stable':finite,'loss_at_50':float(ls[49]) if finite else 1e12})
        bench[name]=rows

    # Compare oscillation proxy on the fastest stable choices: sign changes of stiff mode.
    osc={}
    for name in ['sgd','momentum','fading_memory']:
        stable_rows=[r for r in bench[name] if r['stable']]
        best=min(stable_rows,key=lambda r:r['final_loss'])
        eta=best['eta']
        if name=='sgd': traj=sgd_rollout(lams,eta,steps,x0)[0]
        elif name=='momentum': traj=momentum_rollout(lams,eta,0.9,steps,x0)[0]
        else: traj=fm_rollout(lams,2.0,5.0,eta,steps,x0)[0]
        signs=np.sign(traj[:,1]); changes=int(np.sum(signs[1:]*signs[:-1]<0))
        osc[name]={'selected_eta':eta,'final_loss':best['final_loss'],'stiff_mode_sign_changes':changes}

    result={'seed':SEED,'continuous_roots':[[float(r.real),float(r.imag)] for r in roots],
      'root_residuals':root_residuals,'discrete_stability':{
        'lambda':lam,'kappa':kappa,'beta':beta,'predicted_eta_max':predicted_eta_max,
        'empirical_eta_max':empirical_eta_max,'relative_gap':abs(predicted_eta_max-empirical_eta_max)/predicted_eta_max},
      'benchmark':bench,'oscillation_comparison':osc}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()