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

SEED = 17
np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    torch.tensor([1.0], device=DEVICE).sum().item()
except Exception:
    DEVICE = "cpu"


def core_sanity():
    # A residual with one nearly unidentifiable direction. Full inversion amplifies
    # its tiny residual, while the rank-safe inverse discards that direction.
    s = torch.tensor([1.0, 1e-9])
    J = torch.diag(s)
    r = torch.tensor([0.1, 0.1])
    full = -torch.linalg.pinv(J) @ r
    keep = s >= 1e-6 * s[0]
    safe = -(torch.diag(torch.where(keep, 1/s, torch.zeros_like(s))) @ r)
    # A first-order linear residual after each proposed step.
    full_res = r + J @ full
    safe_res = r + J @ safe
    return {"singular_values": s.tolist(), "full_step_norm": float(torch.linalg.norm(full)),
            "truncated_step_norm": float(torch.linalg.norm(safe)),
            "full_linearized_residual": float(torch.linalg.norm(full_res)),
            "truncated_linearized_residual": float(torch.linalg.norm(safe_res)),
            "kept_rank": int(keep.sum())}


def make_data(ntrain=96, ntest=256):
    xtr = torch.linspace(-1, 1, ntrain, device=DEVICE).reshape(-1,1)
    xte = torch.linspace(-1, 1, ntest, device=DEVICE).reshape(-1,1)
    def target(x):
        return torch.where(x < 0, torch.sin(4*x), 0.55*torch.sin(7*x)+0.25*x)
    ytr = target(xtr) + 0.025*torch.randn_like(xtr)
    return xtr, ytr, xte, target(xte)


def unpack(z, width):
    # tanh hidden layer, feature vector includes a constant output bias
    a = z[:width].reshape(width,1)
    b = z[width:2*width]
    return a, b

def features(z, x, width):
    a,b = unpack(z,width)
    h = torch.tanh(x @ a.T + b)
    return torch.cat([h, torch.ones((x.shape[0],1),device=x.device)], dim=1)

def projected(z, x, y, width, tau=1e-6):
    Phi = features(z,x,width)
    U,S,Vh = torch.linalg.svd(Phi, full_matrices=False)
    cutoff = tau*S[0] if S.numel() else 0.
    inv = torch.where(S >= cutoff, 1/S, torch.zeros_like(S))
    w = Vh.T @ (inv[:,None] * (U.T @ y))
    r = Phi @ w-y
    return w, r, S

def vp_gn(x,y,xt,yt,width=16,steps=35,rho=1e-4,tau=1e-6):
    z = (0.4*torch.randn(2*width,device=DEVICE)).requires_grad_()
    hist=[]; accepted=[]; conds=[]
    t0=time.perf_counter()
    for it in range(steps):
        w,r,S = projected(z,x,y,width,tau)
        loss=0.5*torch.sum(r*r)
        def residual(zz): return projected(zz,x,y,width,tau)[0] # not used; avoid differentiating solve
        # Jacobian of residual with output w held fixed (Gauss-Newton VP/envelope step)
        ww=w.detach()
        def fixed_res(zz): return (features(zz,x,width) @ ww-y).reshape(-1)
        J=torch.autograd.functional.jacobian(fixed_res,z,create_graph=False).reshape(x.shape[0], -1)
        uj,sj,vhj=torch.linalg.svd(J,full_matrices=False)
        conds.append(float(sj[0]/max(float(sj[-1]),1e-30)))
        keep=sj >= rho*sj[0]
        coeff = torch.where(keep, 1.0/sj, torch.zeros_like(sj))
        rhs = torch.mv(uj.T, r.reshape(-1))
        delta = -torch.mv(vhj.T, coeff * rhs)
        base=float(loss)
        alpha=1.0; ok=False
        for _ in range(9):
            with torch.no_grad():
                ztry=z+alpha*delta
                wtry,rtry,_=projected(ztry,x,y,width,tau)
                if float(0.5*torch.sum(rtry*rtry)) < base:
                    z=ztry.detach().requires_grad_(); ok=True; break
            alpha*=0.5
        if not ok: alpha=0.0
        accepted.append(alpha)
        with torch.no_grad():
            wp,rp,_=projected(z,x,y,width,tau)
            pred=features(z,xt,width)@wp
            hist.append((float(0.5*torch.sum(rp*rp)/len(y)),float(torch.mean((pred-yt)**2)),int(keep.sum())))
    return {"train_loss":hist[-1][0],"test_mse":hist[-1][1],"initial_train_loss":hist[0][0],
            "accepted_steps":sum(a>0 for a in accepted),"median_alpha":float(np.median(accepted)),
            "max_jacobian_condition":max(conds),"history":hist,"seconds":time.perf_counter()-t0}

def adam_run(x,y,xt,yt,width=16,steps=350):
    z=(0.4*torch.randn(2*width,device=DEVICE)).requires_grad_()
    w=(0.1*torch.randn(width+1,device=DEVICE)).requires_grad_()
    opt=torch.optim.Adam([z,w],lr=0.025)
    hist=[]; t0=time.perf_counter()
    for _ in range(steps):
        opt.zero_grad()
        pred=features(z,x,width)@w
        loss=torch.mean((pred-y)**2)/2
        loss.backward(); opt.step()
        with torch.no_grad(): hist.append((float(loss),float(torch.mean((features(z,xt,width)@w-yt)**2))))
    return {"train_loss":hist[-1][0],"test_mse":hist[-1][1],"initial_train_loss":hist[0][0],
            "history":hist,"seconds":time.perf_counter()-t0}

def main():
    sanity=core_sanity()
    x,y,xt,yt=make_data()
    # identical initial seed streams are controlled separately, but methods have
    # different parameterizations; report repeated fixed-seed single runs.
    torch.manual_seed(SEED+1); adam=adam_run(x,y,xt,yt)
    torch.manual_seed(SEED+2); vp=vp_gn(x,y,xt,yt)
    out={"device":DEVICE,"seed":SEED,"core_sanity":sanity,"adam":adam,"rank_safe_vp_gn":vp}
    with open("results.json","w") as f: json.dump(out,f,indent=2)
    print(json.dumps({k:v for k,v in out.items() if k not in ('adam','rank_safe_vp_gn')},indent=2))
    print(json.dumps({"adam":{k:adam[k] for k in adam if k!='history'},"rank_safe_vp_gn":{k:vp[k] for k in vp if k!='history'}},indent=2))
if __name__=='__main__': main()