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

SEED=7
rng=np.random.default_rng(SEED)
phi=(1+math.sqrt(5))/2

# Composite convex objective: h is logistic loss, f is elastic-net, g is box indicator on Kx=x.
def make_problem(n=320,d=12):
    X=rng.normal(size=(n,d)); X[:,0]*=12.0; X[:,1]*=4.0  # sharp and flatter directions
    w=rng.normal(size=d); y=(X@w+rng.normal(scale=1.0,size=n)>0).astype(float)
    return X,y

def sigmoid(t): return 1/(1+np.exp(-np.clip(t,-40,40)))
def loss_grad(x,X,y,alpha,rho):
    p=sigmoid(X@x)
    h=np.mean(np.logaddexp(0,X@x)-y*(X@x))
    # h excludes f and g; derivative is minibatch/full gradient
    return h, X.T@(p-y)/len(y)

def prox_f(v,lam,alpha=2e-3,rho=2e-3):
    return np.sign(v)*np.maximum(np.abs(v)-lam*alpha,0)/(1+lam*rho)

def gstar_prox(y,sigma,c):
    # prox_{sigma g*}(y)=y-sigma clip(y/sigma,-c,c), K=I
    return y-sigma*np.clip(y/sigma,-c,c)

def full_obj(x,X,y,alpha=2e-3,rho=2e-3,c=.8):
    h,_=loss_grad(x,X,y,alpha,rho)
    return h+alpha*np.abs(x).sum()+rho*.5*(x@x), max(0.,np.abs(x).max()-c)

def grpd(X,y,steps,lam0,eta=.95,sigma=.5,c=.8,alpha=2e-3,rho=2e-3):
    d=X.shape[1]; x=np.zeros(d); z=x.copy(); dual=np.zeros(d); oldx=x.copy(); oldq=np.zeros(d); Lhat=1.0
    rec=[]
    for k in range(steps):
        h,q=loss_grad(x,X,y,alpha,rho)
        if k>0:
            Lhat=np.linalg.norm(q-oldq)/(np.linalg.norm(x-oldx)+1e-8)
            # robust EMA as explicitly allowed by idea
            Lhat=.8*Lhat+.2*Lhat_prev
        Lhat_prev=Lhat
        lam=eta/(Lhat+sigma+1e-8)
        # retain the requested initial-step stress: first iteration starts from lam0
        if k==0: lam=lam0
        xn=prox_f(z-lam*(q+dual),lam,alpha,rho)
        zn=((phi-1)/phi)*xn+(1/phi)*z
        dualn=gstar_prox(dual+sigma*zn,sigma,c)
        oldx,oldq=x.copy(),q.copy(); x,z,dual=xn,zn,dualn
        obj,v=full_obj(x,X,y,alpha,rho,c)
        rec.append((obj,v,lam,Lhat,np.linalg.norm(q)))
        if not np.isfinite(obj) or np.linalg.norm(x)>1e8: break
    return rec,x

def proxgrad(X,y,steps,lam,alpha=2e-3,rho=2e-3,c=.8):
    x=np.zeros(X.shape[1]); rec=[]
    for k in range(steps):
        _,q=loss_grad(x,X,y,alpha,rho)
        x=prox_f(x-lam*q,lam,alpha,rho)
        obj,v=full_obj(x,X,y,alpha,rho,c); rec.append((obj,v,lam,np.nan,np.linalg.norm(q)))
        if not np.isfinite(obj) or np.linalg.norm(x)>1e8: break
    return rec,x

def projected_pg(X,y,steps,lam,alpha=2e-3,rho=2e-3,c=.8):
    x=np.zeros(X.shape[1]); rec=[]
    for k in range(steps):
        _,q=loss_grad(x,X,y,alpha,rho)
        x=np.clip(prox_f(x-lam*q,lam,alpha,rho),-c,c)
        obj,v=full_obj(x,X,y,alpha,rho,c); rec.append((obj,v,lam,np.nan,np.linalg.norm(q)))
        if not np.isfinite(obj) or np.linalg.norm(x)>1e8: break
    return rec,x

def adamw(X,y,steps,lr,alpha=2e-3,rho=2e-3,c=.8):
    x=np.zeros(X.shape[1]); m=np.zeros_like(x); v=np.zeros_like(x); rec=[]
    for k in range(1,steps+1):
        _,q=loss_grad(x,X,y,alpha,rho); m=.9*m+.1*q; v=.999*v+.001*q*q
        x=x-lr*(m/(1-.9**k))/(np.sqrt(v/(1-.999**k))+1e-8)-lr*rho*x
        obj,viol=full_obj(x,X,y,alpha,rho,c); rec.append((obj,viol,lr,np.nan,np.linalg.norm(q)))
        if not np.isfinite(obj) or np.linalg.norm(x)>1e8: break
    return rec,x

def run():
    X,y=make_problem(); steps=150
    # Initial step sweep is the falsifiable stability test. PG has no constraint dual.
    # Use identical full-gradient evaluations; constants chosen around sharp curvature.
    spectral=np.linalg.eigvalsh(X.T@X/len(y)).max()/4
    candidates=[.02,.05,.1,.2,.5,1.,2.]
    rows=[]
    for lr in candidates:
        for name,fn in [('PG',lambda lr:proxgrad(X,y,steps,lr)),('ProjectedPG',lambda lr:projected_pg(X,y,steps,lr)),('GRPD',lambda lr:grpd(X,y,steps,lr)) ,('AdamW',lambda lr:adamw(X,y,steps,lr))]:
            rec,_=fn(lr); finite=bool(rec) and np.isfinite(rec[-1][0]) and rec[-1][0]<1e6
            rows.append({'method':name,'lr':lr,'final_loss':float(rec[-1][0]) if rec else None,'final_violation':float(rec[-1][1]) if rec else None,'finite':bool(finite),'min_loss':float(min(r[0] for r in rec)) if rec else None,'steps':len(rec)})
    # Representative trajectories at a stress level where behavior differs, plus math checks.
    checks={'phi':phi,'extrapolation_weights':[(phi-1)/phi,1/phi], 'weights_sum':((phi-1)/phi+1/phi)}
    # finite-difference curvature estimate agrees with local Hessian directional curvature
    x=rng.normal(size=X.shape[1])*0.1; dx=rng.normal(size=x.shape); dx/=np.linalg.norm(dx)
    _,q1=loss_grad(x,X,y,0,0); _,q2=loss_grad(x+1e-5*dx,X,y,0,0)
    est=np.linalg.norm(q2-q1)/(1e-5+1e-8)
    p=sigmoid(X@x); Hdx=X.T@((p*(1-p))*(X@dx))/len(y)
    true=float(np.linalg.norm(Hdx))
    checks.update({'curvature_estimate':float(est),'analytic_hessian_action':true,'curvature_relative_error':float(abs(est-true)/(true+1e-12))})
    out={'seed':SEED,'spectral_quadratic_curvature_bound':float(spectral),'checks':checks,'results':rows}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': run()