import json, math, random
from pathlib import Path
import numpy as np
from scipy.optimize import minimize

# J_{a,b}(u)=(u^2-1)^2+a*u^2+b*u, u in [-2,2].
# At b=0: N=2 for a<2, N=1 for a>=2. Around a<2, r=sqrt(1-a/2),
# H(r)=4(2-a), so the positive-well displacement is |b|/[4(2-a)] + O(b^2).
def objective(u, a, b):
    x = float(np.asarray(u).reshape(-1)[0])
    return (x*x-1.0)**2 + a*x*x + b*x

def grad(x, a, b):
    x = float(x)
    return 4*x**3 + 2*(a-2)*x + b

def hess(x, a, b):
    x = float(x)
    return 12*x*x + 2*(a-2)

def audit(a, b, starts, eps=1e-3, grad_tol=2e-6, hess_tol=2e-6):
    ends=[]
    vals=[]
    for x0 in starts:
        r=minimize(lambda z: objective(z,a,b), [float(x0)], jac=lambda z: np.array([grad(z[0],a,b)]),
                   bounds=[(-2,2)], method='L-BFGS-B', options={'ftol':1e-14,'gtol':1e-10,'maxiter':300})
        x=float(r.x[0])
        if abs(grad(x,a,b)) <= grad_tol and hess(x,a,b) >= -hess_tol:
            ends.append(x); vals.append(objective(x,a,b))
    clusters=[]
    for x,v in sorted(zip(ends,vals)):
        if not clusters or abs(x-clusters[-1]['u']) > eps:
            clusters.append({'u':x,'J':v,'members':1})
        else:
            c=clusters[-1]; c['u']=(c['u']*c['members']+x)/(c['members']+1); c['J']=min(c['J'],v); c['members']+=1
    best=min(clusters,key=lambda c:c['J']) if clusters else {'u':float('nan'),'J':float('nan')}
    return {'N':len(clusters),'best_u':best['u'],'best_J':best['J'],'clusters':clusters}

def main():
    np.random.seed(7); random.seed(7)
    starts=np.linspace(-1.95,1.95,81)
    # Prediction 1: pitchfork count transition at a=2.
    count_rows=[]
    for a in np.linspace(0,3,13):
        out=audit(float(a),0.0,starts)
        predicted=2 if a < 2 else 1
        count_rows.append({'a':float(a),'N_observed':out['N'],'N_predicted':predicted})
    # Prediction 2: displacement is linear in |b| with slope 1/[4(2-a)].
    a=1.0; ref=audit(a,0.0,starts)['best_u']
    drift_rows=[]
    for b in np.linspace(0,0.16,9):
        out=audit(a,float(b),starts)
        # Track positive well, which is the best well for b>=0; reference is +1.
        observed=abs(out['best_u']-ref)
        predicted=abs(b)/(4*(2-a))
        drift_rows.append({'b':float(b),'e_observed':observed,'e_predicted_linear':predicted})
    # Prediction 3: larger curvature gap (2-a) suppresses drift; slope sweep.
    slope_rows=[]
    for aa in [0.5,1.0,1.5]:
        rr=audit(aa,0.0,starts); r=rr['best_u']; bs=np.array([0.02,0.04,0.06])
        es=[]
        for bb in bs: es.append(abs(audit(aa,float(bb),starts)['best_u']-r))
        fit=float(np.dot(bs,es)/np.dot(bs,bs))
        slope_rows.append({'a':aa,'slope_observed':fit,'slope_predicted':1/(4*(2-aa))})
    # Drift trajectory: nearly identical prediction loss, but surrogate decision drifts.
    traj=[]
    for t,b in enumerate(np.linspace(0,0.18,19)):
        out=audit(1.0,float(b),starts)
        e=abs(out['best_u']-ref)
        val_loss=1e-4*(1.0-b/0.18) # mildly improving but flat observational metric
        traj.append({'t':t,'b':float(b),'validation_loss':float(val_loss),'e':float(e),'N':out['N'],'J':float(out['best_J'])})
    tau_u=0.02; tau_J=0.03; nmax=2
    accepted=[r for r in traj if r['e']<=tau_u and abs(r['J']-traj[0]['J'])<=tau_J and r['N']<=nmax]
    # validation-only chooses last (slightly improving) checkpoint; audit chooses latest accepted.
    baseline=traj[-1]; audited=accepted[-1]
    result={'seed':7,'count_sweep':count_rows,'drift_sweep':drift_rows,'curvature_sweep':slope_rows,
            'trajectory':traj,'selection':{'baseline_validation_only':baseline,'audit_selected':audited,
             'tau_u':tau_u,'tau_J':tau_J,'N_max':nmax},
            'summary':{'count_all_correct':all(x['N_observed']==x['N_predicted'] for x in count_rows),
             'max_drift_linear_abs_error':max(abs(x['e_observed']-x['e_predicted_linear']) for x in drift_rows),
             'slope_relative_errors':[abs(x['slope_observed']-x['slope_predicted'])/x['slope_predicted'] for x in slope_rows]}}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result['summary'],indent=2))
    print('selection',json.dumps(result['selection'],indent=2))

if __name__=='__main__': main()