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

SEED = 2394
np.random.seed(SEED); random.seed(SEED)


def prox_l1(z, eta, lam, h):
    # Diagonal-metric prox for R(w)=lambda*||w||_1.
    return np.sign(z) * np.maximum(np.abs(z) - eta * lam * h, 0.0)


def adaptive_quadratic(A, h, w0, eta0=0.05, gamma_up=1.5, gamma_down=.5,
                       q=3, lam=0.0, steps=120):
    w = w0.copy(); eta = eta0; success = 0; accepts = 0; rejects = 0
    eta_hist=[]; obj_hist=[]; resid_hist=[]
    def f(x): return .5 * x @ A @ x
    for _ in range(steps):
        g = A @ w
        trial = prox_l1(w - eta * h * g, eta, lam, h)
        old, new = f(w), f(trial); s = trial-w
        # Objective Armijo safeguard. For R=0 this is the stated smooth merit.
        rhs = old + 1e-4 * (g @ s)
        if np.isfinite(new) and new <= rhs:
            w = trial; accepts += 1; success += 1
            if success >= q:
                eta *= gamma_up; success = 0
        else:
            rejects += 1; success = 0; eta *= gamma_down
        eta_hist.append(eta); obj_hist.append(f(w)); resid_hist.append(np.linalg.norm(A@w))
    return dict(w=w, eta=np.array(eta_hist), obj=np.array(obj_hist), residual=np.array(resid_hist), accepts=accepts, rejects=rejects)


def fixed_quadratic(A, h, w0, eta, steps=100):
    w=w0.copy(); vals=[]
    for _ in range(steps):
        vals.append(.5*w@A@w); w=w-eta*h*(A@w)
    vals.append(.5*w@A@w)
    return np.array(vals)


def run_stability_sweeps():
    # Diagonal H and A make the exact transformed eigenvalues transparent.
    dim=8; h=np.array([.5, .8, 1.0, 1.2, 1.5, .7, 1.1, .9])
    lambdas=np.array([1.,2.,4.,7.,10.,14.,18.,25.])
    A=np.diag(lambdas); w0=np.ones(dim)
    L=float(np.max(h*lambdas)); eta_c=2/L
    # Prediction 1: a fixed linear iteration is stable iff eta*L < 2.
    etas=np.linspace(.1*eta_c, 2.2*eta_c, 22)
    stable=[]
    for eta in etas:
        vals=fixed_quadratic(A,h,w0,eta,steps=80)
        stable.append(bool(np.all(np.isfinite(vals)) and vals[-1] < vals[0] and np.max(vals)<1e12))
    boundary=next((float(etas[i]) for i in range(len(etas)) if not stable[i]), float('nan'))
    # Refine empirical boundary by binary search on final norm.
    lo,hi=0.,3*eta_c
    for _ in range(45):
        mid=(lo+hi)/2
        vals=fixed_quadratic(A,h,w0,mid,steps=100)
        ok=np.all(np.isfinite(vals)) and vals[-1] < vals[0] and np.max(vals)<1e10
        if ok: lo=mid
        else: hi=mid
    empirical=float(lo)
    # Prediction 2: asymptotic contraction factor is max_i |1-eta*h_i*lambda_i|.
    eta_test=.55*eta_c
    pred_factor=float(np.max(np.abs(1-eta_test*h*lambdas)))
    vals=fixed_quadratic(A,h,w0,eta_test,steps=40)
    # Objective ratio tends to squared factor; estimate from late ratios.
    observed=float(np.median(np.sqrt(np.maximum(vals[-10:-1],1e-300)/np.maximum(vals[-11:-2],1e-300))))
    # Prediction 3: adaptive safeguard settles below the boundary and rejects unsafe trials.
    ad=adaptive_quadratic(A,h,w0,eta0=eta_c*.18,steps=180)
    late_eta=float(np.median(ad['eta'][-30:]))
    late_max=float(np.max(ad['eta'][-30:]))
    return {
      'L_exact':L, 'eta_critical_predicted':eta_c,
      'boundary_coarse_first_unstable':boundary, 'boundary_binary_observed':empirical,
      'boundary_relative_error':abs(empirical-eta_c)/eta_c,
      'contraction_eta':eta_test, 'contraction_factor_predicted':pred_factor,
      'contraction_factor_observed':observed, 'contraction_abs_error':abs(observed-pred_factor),
      'adaptive_late_median_eta':late_eta, 'adaptive_late_max_eta':late_max,
      'adaptive_eta_ratio_to_boundary':late_eta/eta_c,
      'adaptive_accepts':ad['accepts'], 'adaptive_rejects':ad['rejects'],
      'adaptive_final_objective':float(ad['obj'][-1]),
      'stable_sweep': [{'eta':float(e),'etaL':float(e*L),'stable':s} for e,s in zip(etas,stable)]
    }


def sparse_regression():
    rng=np.random.default_rng(SEED+7); n,d=256,24
    X=rng.normal(size=(n,d)); true=np.zeros(d); true[[1,5,9,15,20]]=rng.normal(size=5)
    y=X@true+.08*rng.normal(size=n); A=X.T@X/n; b=X.T@y/n; w0=np.zeros(d)
    lam=.025; L=np.linalg.eigvalsh(A).max(); h=np.ones(d); eta=1/L
    def obj(w): return .5*np.mean((X@w-y)**2)+lam*np.abs(w).sum()
    def run_sgd():
        w=w0.copy(); hist=[]
        for _ in range(250):
            ix=rng.choice(n,64,replace=False); g=X[ix].T@(X[ix]@w-y[ix])/64
            w-=.35*g; hist.append(obj(w))
        return w,hist
    def run_prox():
        w=w0.copy(); hist=[]
        for _ in range(250):
            w=prox_l1(w-eta*(A@w-b),eta,lam,h); hist.append(obj(w))
        return w,hist
    def run_ad():
        # Use deterministic full-batch gradients and the same safeguarded update.
        w=w0.copy(); H=np.ones(d); et=eta*.5; succ=0; hist=[]; acc=rej=0
        for _ in range(250):
            g=A@w-b; trial=prox_l1(w-et*H*g,et,lam,H)
            old=obj(w); new=obj(trial); s=trial-w
            if np.isfinite(new) and new <= old+1e-4*(g@s):
                oldw=w; oldg=g; w=trial; acc+=1; succ+=1
                # Diagonal inverse-BFGS-like secant update, clipped for robustness.
                yk=(A@w-b)-oldg; sk=w-oldw; den=yk*sk
                mask=den>1e-10
                H[mask]=np.clip(sk[mask]/den[mask],1e-4,100.)
                if succ>=3: et*=1.5; succ=0
            else: rej+=1; succ=0; et*=.5
            hist.append(obj(w))
        return w,hist,acc,rej
    ws,hs=run_sgd(); wp,hp=run_prox(); wa,ha,ac,re=run_ad()
    return {'sgd_final':float(hs[-1]),'prox_final':float(hp[-1]),'adaptive_final':float(ha[-1]),
            'sgd_nonzeros':int(np.sum(np.abs(ws)>1e-4)), 'prox_nonzeros':int(np.sum(np.abs(wp)>1e-4)),
            'adaptive_nonzeros':int(np.sum(np.abs(wa)>1e-4)), 'adaptive_accepts':ac,'adaptive_rejects':re,
            'iterations':250,'lambda':lam,'L':float(L)}


def main():
    out={'seed':SEED,'stability':run_stability_sweeps(),'sparse_regression':sparse_regression()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    s=out['stability']; r=out['sparse_regression']
    print(json.dumps({'stability':s,'sparse_regression':r},indent=2))

if __name__=='__main__': main()