import json, math, random
import numpy as np
import torch


def ridge(S, b, lam):
    return np.linalg.solve(S + lam*np.eye(S.shape[0]), b)


def estimate_lambda(w, S, b, lo=1e-8, hi=10.0):
    grid = np.r_[0.0, np.logspace(math.log10(lo), math.log10(hi), 240)]
    scores = np.array([np.sum((w-ridge(S,b,lam))**2) for lam in grid])
    j = int(np.argmin(scores))
    if 0 < j < len(grid)-1:
        a, z = max(grid[j-1], grid[1]), grid[j+1]
        for lam in np.logspace(np.log10(a), np.log10(z), 80):
            score=np.sum((w-ridge(S,b,lam))**2)
            if score < scores[j]: scores[j], grid[j] = score, lam
    return float(grid[j])


def math_check(seed=7):
    rng=np.random.default_rng(seed); d=8; n=30000
    A=rng.normal(size=(d,d)); S=A@A.T/d+.2*np.eye(d)
    theta=rng.normal(size=d); true_lam=.37
    X=rng.multivariate_normal(np.zeros(d),S,n); y=X@theta
    b=S@theta; w=ridge(S,b,true_lam)
    Sh=np.zeros((d,d)); bh=np.zeros(d); alpha=.97
    for i in range(0,n,100):
        H=X[i:i+100]; yy=y[i:i+100]
        Sh=alpha*Sh+(1-alpha)*H.T@H/len(H)
        bh=alpha*bh+(1-alpha)*H.T@yy/len(H)
    est=estimate_lambda(w,Sh,bh,hi=4)
    return {'true_lambda':true_lam,'estimated_lambda':est,
            'abs_error':abs(est-true_lam),
            'ridge_residual':float(np.linalg.norm(w-ridge(S,b,true_lam)))}


def run(seed, mode, steps=500):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        d,N,batch=12,1800,64; eps=.18; lr=.035
        gen=np.random.default_rng(seed); scales=np.linspace(.45,2.0,d)
        X=gen.normal(size=(N,d))*scales; theta=gen.normal(size=d); theta/=np.linalg.norm(theta)
        y=X@theta+.15*gen.normal(size=N)
        Xt=torch.tensor(X,dtype=torch.float32,device=device); yt=torch.tensor(y,dtype=torch.float32,device=device)
        w=torch.zeros(d,device=device,requires_grad=True)
        wd=.035 if mode=='fixed_weight_decay' else 0.
        opt=torch.optim.SGD([w],lr=lr,weight_decay=wd)
        Sema=torch.zeros((d,d),device=device); bema=torch.zeros(d,device=device)
        alpha=.96; lambda0=.035; rho=.28; smoothed=0.; updates=0; hist=[]
        for t in range(steps):
            ix=torch.randint(0,N,(batch,),device=device); H=Xt[ix]; yy=yt[ix]
            opt.zero_grad(set_to_none=True); residual=H@w-yy
            loss=.5*(torch.abs(residual)+eps*torch.linalg.vector_norm(w)).square().mean()
            loss.backward()
            with torch.no_grad():
                Sema=alpha*Sema+(1-alpha)*H.T@H/batch
                bema=alpha*bema+(1-alpha)*H.T@yy/batch
                if mode=='controller' and t>=60 and t%10==0:
                    Sn=Sema.cpu().numpy(); bn=bema.cpu().numpy(); wn=w.cpu().numpy()
                    current=estimate_lambda(wn,Sn,bn,hi=3.)
                    smoothed=current if updates==0 else .85*smoothed+.15*current
                    # Do not react to a >10x jump in the noisy early estimate.
                    if updates==0 or (.1*max(smoothed,1e-8)<=current<=10*max(smoothed,1e-8)):
                        updates+=1; hist.append(smoothed)
                if mode=='controller' and updates>=2:
                    w.grad.add_(rho*(smoothed-lambda0)*w)
            opt.step()
        with torch.no_grad():
            r=Xt@w-yt; clean=.5*r.square().mean(); robust=.5*(torch.abs(r)+eps*torch.linalg.vector_norm(w)).square().mean()
            return {'clean_loss':float(clean.cpu()),'robust_loss':float(robust.cpu()),
                    'weight_norm':float(torch.linalg.vector_norm(w).cpu()),
                    'lambda_last':float(hist[-1]) if hist else None,'device':device}
    except Exception:
        if device=='cuda':
            torch.cuda.empty_cache()
            # Explicit CPU retry, preserving deterministic seed and settings.
            return run_cpu(seed,mode,steps)
        raise


def run_cpu(seed,mode,steps):
    # CUDA fallback uses the same numerical path on CPU.
    available=torch.cuda.is_available
    torch.cuda.is_available=lambda:False
    try: return run(seed,mode,steps)
    finally: torch.cuda.is_available=available


def main():
    out={'math_check':math_check(),'runs':{}}
    for mode in ['adv_sgd','fixed_weight_decay','controller']:
        vals=[run(s,mode) for s in [11,23,41,59,71]]; out['runs'][mode]=vals
        out['summary_'+mode]={
            'clean_loss':float(np.mean([v['clean_loss'] for v in vals])),
            'robust_loss':float(np.mean([v['robust_loss'] for v in vals])),
            'weight_norm':float(np.mean([v['weight_norm'] for v in vals])),
            'robust_std':float(np.std([v['robust_loss'] for v in vals]))}
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()