import json
import numpy as np
from scipy.optimize import differential_evolution, minimize_scalar
SEED=2125


def M_interp(mu,L):
    return .5*np.array([[-mu*L,mu*L,mu,-L],[mu*L,-mu*L,-mu,L],
                        [mu,-mu,-1,1],[-L,L,1,-1]],float)


def interp_check(n=5000):
    rng=np.random.default_rng(SEED); vals=[]
    for _ in range(n):
        h=rng.uniform(.1,1.0); u,v=rng.normal(size=2)
        z=np.array([u,v,h*u,h*v]); vals.append(z@M_interp(.1,1)@z)
    vals=np.asarray(vals)
    return {'supplied_min':float(vals.min()),
            'supplied_negative_fraction':float(np.mean(vals < -1e-11)),
            'corrected_min':float((-vals).min()),
            'corrected_negative_fraction':float(np.mean(-vals < -1e-11)),
            'M_eigenvalues':np.linalg.eigvalsh(M_interp(.1,1)).tolist()}


def mat_recur(mu,A,c):
    # c=(a1,b0,b1), with a0=1-a1 (constant-gradient consistency).
    a1,b0,b1=c; a0=1-a1; T=np.zeros((4,4))
    T[0,0]=a0-b0*mu; T[0,1]=a1-b1*mu; T[0,2]=-b0*A; T[0,3]=-b1*A; T[1,0]=1
    T[2,0]=b0*A; T[2,1]=b1*A; T[2,2]=a0-b0*mu; T[2,3]=a1-b1*mu; T[3,2]=1
    return T


def rho(c,mu,A): return float(np.max(np.abs(np.linalg.eigvals(mat_recur(mu,A,c)))))

def optimize_coeff(mu,Amax):
    grid=np.linspace(0,Amax,16)
    def obj(c): return max(rho(c,mu,A) for A in grid)
    r=differential_evolution(obj,[(-1,1),(0,1),(-.5,.5)],seed=SEED,popsize=16,maxiter=100,polish=True)
    return r.x,float(r.fun)

def best_gda(mu,Amax):
    grid=np.linspace(0,Amax,16)
    def obj(g): return max(rho((0,g,0),mu,A) for A in grid)
    r=minimize_scalar(obj,bounds=(1e-5,1),method='bounded')
    return np.array([0,r.x,0]),float(r.fun)

def simulate(c,mu,A,steps=50):
    s=np.array([1.,-.7,-.7,1.]); T=mat_recur(mu,A,c); out=[]
    for _ in range(steps): out.append(float(np.linalg.norm(s))); s=T@s
    return out

def main():
    mu=.1; chk=interp_check(); synth,sobj=optimize_coeff(mu,1.5); gda,gobj=best_gda(mu,1.5)
    grid=np.linspace(0,2,41); rs=np.array([rho(synth,mu,A) for A in grid]); rg=np.array([rho(gda,mu,A) for A in grid])
    design=np.linspace(0,1.5,16); rdesign=np.array([rho(synth,mu,A) for A in design]); rgdesign=np.array([rho(gda,mu,A) for A in design])
    result={'seed':SEED,'mu':mu,'interpolation_check':chk,
      'coefficients':{'idea':[float(x) for x in synth],'baseline_gda':[float(x) for x in gda]},
      'design_worst_rho':{'idea':sobj,'baseline':gobj},
      'prediction_checks':{
       'P1_supplied_M_nonnegative':chk['supplied_negative_fraction']==0,
       'P1_sign_corrected_nonnegative':chk['corrected_negative_fraction']==0,
       'P2_stability_boundary_on_design_grid':{'idea':float(design[np.where(rdesign<1)[0][-1]]) if np.any(rdesign<1) else None,'baseline':float(design[np.where(rgdesign<1)[0][-1]]) if np.any(rgdesign<1) else None},
       'P3_rho_vs_coupling':{'idea_rho_A0':float(rs[0]),'idea_rho_A15':float(rs[30]),'baseline_rho_A0':float(rg[0]),'baseline_rho_A15':float(rg[30]),'idea_monotone':bool(np.all(np.diff(rs)>=-1e-7)),'baseline_monotone':bool(np.all(np.diff(rg)>=-1e-7))}},
      'sweep':{'coupling':grid.tolist(),'idea_rho':rs.tolist(),'baseline_rho':rg.tolist()},
      'trajectory_A1.2':{'idea':simulate(synth,mu,1.2),'baseline':simulate(gda,mu,1.2)}}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps({'check':chk,'idea':synth.tolist(),'baseline':gda.tolist(),'rho':(sobj,gobj)},indent=2))
if __name__=='__main__': main()