import json
from pathlib import Path
import numpy as np


def gamma_formula(k):
    k = float(k)
    return (4-k*k + k*np.sqrt(k*k+8*k))/(2*(k*k+2*k+2))


def paper_gains(lam2, lamN):
    # Article theorem: kappa = lambda2/lambdaN.
    k = lam2 / lamN
    g = gamma_formula(k)
    return g, 1-g, 2*(1-g)/lam2


def idea_gains(lam2, lamN):
    # Literal task statement: kappa = lambdaN/lambda2 with the same formula.
    k = lamN / lam2
    g = gamma_formula(k)
    return g, 1-g, 2*(1-g)/lam2


def modal_matrix(lam, alpha, eps):
    q = 1.0 - eps*lam
    return np.array([[q, -alpha], [-eps*lam, q-alpha]], dtype=float)


def pole_radius(lam, alpha, eps):
    return float(np.max(np.abs(np.linalg.eigvals(modal_matrix(lam, alpha, eps)))))


def worst_radius(lam2, lamN, alpha, eps, n=10001):
    ls = np.linspace(lam2, lamN, n)
    rs = np.array([pole_radius(x, alpha, eps) for x in ls])
    i = int(np.argmax(rs))
    return float(rs[i]), float(ls[i])


def run():
    pole=[]
    for ratio in [1, 2, 5, 10, 30, 100]:
        l2, lN = 1., float(ratio)
        pg, pa, pe = paper_gains(l2,lN)
        ig, ia, ie = idea_gains(l2,lN)
        pr, pl = worst_radius(l2,lN,pa,pe)
        ir, il = worst_radius(l2,lN,ia,ie)
        pole.append({'lambdaN_over_lambda2':ratio,
            'paper_predicted_gamma':pg, 'paper_measured_radius':pr,
            'paper_abs_error':abs(pr-pg), 'paper_worst_lambda':pl,
            'idea_predicted_gamma':ig, 'idea_measured_radius':ir,
            'idea_abs_error':abs(ir-ig), 'idea_worst_lambda':il})

    endpoint=[]
    for ratio in [2,5,10,30]:
        g,a,e=paper_gains(1.,ratio)
        endpoint.append({'ratio':ratio,'predicted_gamma':g,
            'radius_lambda2':pole_radius(1.,a,e),
            'radius_lambdaN':pole_radius(float(ratio),a,e)})

    scaling=[]
    for l2,lN in [(0.5,5),(1,10),(2,20),(4,40)]:
        g,a,e=paper_gains(l2,lN)
        scaling.append({'lambda2':l2,'lambdaN':lN,'epsilon':e,
                        'epsilon_times_lambda2':e*l2,'alpha':a,'gamma':g})

    boundary=[]
    l2,lN=1.,10.; g,a,e=paper_gains(l2,lN)
    for mult in [.5,.8,1.,1.2]:
        ep=mult*e
        r,w=worst_radius(l2,lN,a,ep)
        boundary.append({'epsilon_over_lower_bound':mult,'radius':r,
                         'predicted_at_or_below_gamma':bool(r<=g+1e-8),'worst_lambda':w})

    rng=np.random.default_rng(7); n=8
    W=np.zeros((n,n))
    for i in range(n-1): W[i,i+1]=W[i+1,i]=1.
    L=np.diag(W.sum(1))-W
    ev=np.linalg.eigvalsh(L); l2,lN=ev[1],ev[-1]
    def simulate(alpha,eps,steps=100):
        c=rng.normal(size=n); x=rng.normal(size=n); gg=x-c; y=gg.copy(); vals=[]
        for _ in range(steps):
            xn=x-eps*(L@x)-alpha*y; gn=xn-c
            yn=y-eps*(L@y)+gn-gg; x,y,gg=xn,yn,gn
            val=float(np.mean(.5*(x-c)**2)) if np.all(np.isfinite(x)) else float('inf')
            vals.append(val)
            if not np.isfinite(val) or val>1e100: break
        return float(vals[-1]), bool(np.isfinite(vals[-1]))
    pg,pa,pe=paper_gains(l2,lN); ig,ia,ie=idea_gains(l2,lN)
    rng=np.random.default_rng(7); bl, bok=simulate(.1,.1)
    rng=np.random.default_rng(7); pl, pok=simulate(pa,pe)
    rng=np.random.default_rng(7); il, iok=simulate(ia,ie)
    comparison={'graph_lambda2':float(l2),'graph_lambdaN':float(lN),
      'baseline_final_loss':bl,'baseline_finite':bok,
      'paper_tuned_final_loss':pl,'paper_tuned_finite':pok,
      'idea_literal_final_loss':il,'idea_literal_finite':iok}
    out={'pole_convention_sweep':pole,'paper_endpoint_prediction':endpoint,
         'epsilon_scaling_prediction':scaling,'lower_bound_prediction':boundary,
         'quadratic_comparison':comparison,
         'conclusion':'The authoritative idea reverses the paper theorem kappa definition; literal gains are unstable for condition number > 1.'}
    text=json.dumps(out,indent=2)
    Path('results.json').write_text(text)
    print(text)

if __name__=='__main__': run()