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

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

# ---------- Stage 1: transfer-function sanity check ----------
def lockin_demo():
    # Stable diagonal A, driven by sinusoidal input. Exact H is known.
    lambdas=np.array([0.7, 2.0, 5.0])
    B=np.array([1.0, .4, .2]); C=np.array([.8, -.3, .5])
    w=1.3; eps=.03; dt=.01; periods=25
    T=2*np.pi/w*periods; n=int(T/dt); t=np.arange(n)*dt
    x=np.zeros(3); ys=[]
    for ti in t:
        u=eps*np.sin(w*ti)
        x += dt*(-lambdas*x+B*u)
        ys.append(C@x)
    # discard first periods as transient and correlate on integer-period tail
    cut=n//5; tt=t[cut:]; yy=np.asarray(ys[cut:])
    # For y=Re(H eps exp(i wt)), 2 mean(y exp(-iwt))/eps estimates H.
    Hhat=1j*(2*np.mean(yy*np.exp(-1j*w*tt))/eps)
    Hexact=np.sum(C*B/(lambdas+1j*w))
    rel=abs(Hhat-Hexact)/abs(Hexact)
    # discrete-gradient stability transition is independently checked
    eta_grid=np.linspace(.02, .55, 107)
    stable=[abs(1-eta*max(lambdas))<1 for eta in eta_grid]
    empirical=max(eta for eta,s in zip(eta_grid,stable) if s)
    boundary=2/max(lambdas)
    # Direct quadratic mode test: z_{t+1}=(1-eta*lambda)z_t, with lambda_max=5.
    # Below 2/lambda_max it decays; above it grows (alternating when eta*lambda>1).
    quad=[]
    for eta in [0.30, 0.39, 0.41, 0.50]:
        z=np.array([1.0, 0.3, -0.2]); norms=[]
        for _ in range(80):
            z=(1-eta*lambdas)*z; norms.append(float(np.linalg.norm(z)))
        quad.append({'eta':eta,'final_norm':norms[-1], 'decays':bool(norms[-1]<norms[0])})
    return {'H_est': [float(Hhat.real),float(Hhat.imag)], 'H_exact':[float(Hexact.real),float(Hexact.imag)],
            'relative_transfer_error':float(rel), 'theory_eta_critical':boundary,
            'empirical_eta_critical_grid':float(empirical), 'quadratic_mode_sweep':quad, 'stability_claim_observed':bool(empirical<boundary and abs(empirical-boundary)<.02),
            'phase_convention': 'corrected sine-input lock-in by multiplying raw estimate by i'}

# ---------- Stage 2: tiny real training comparison ----------
def run_training():
    import torch
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split
    from sklearn.preprocessing import StandardScaler
    torch.manual_seed(SEED); np.random.seed(SEED)
    X,y=load_digits(return_X_y=True)
    X=StandardScaler().fit_transform(X).astype('float32')
    Xtr,Xte,ytr,yte=train_test_split(X,y,test_size=.25,random_state=SEED,stratify=y)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        device=torch.device(dev)
        if dev=='cuda': torch.zeros(1,device=device)
    except Exception:
        device=torch.device('cpu')
    xt=torch.tensor(Xtr,device=device); yt=torch.tensor(ytr,device=device)
    xev=torch.tensor(Xte,device=device); yev=torch.tensor(yte,device=device)
    def make():
        torch.manual_seed(SEED)
        return torch.nn.Sequential(torch.nn.Linear(64,48),torch.nn.Tanh(),torch.nn.Linear(48,10)).to(device)
    def loss_acc(model, x, y):
        with torch.no_grad():
            z=model(x); return float(torch.nn.functional.cross_entropy(z,y)), float((z.argmax(1)==y).float().mean())
    def hessian_top(model, xb, yb, iters=4):
        # Hessian-vector power iteration: an inexpensive empirical local curvature estimate.
        params=[p for p in model.parameters() if p.requires_grad]
        out=model(xb); l=torch.nn.functional.cross_entropy(out,yb)
        g=torch.autograd.grad(l,params,create_graph=True)
        v=[torch.randn_like(p) for p in params]; norm=torch.sqrt(sum((q*q).sum() for q in v)); v=[q/norm for q in v]
        val=0.
        for _ in range(iters):
            gv=sum((a*b).sum() for a,b in zip(g,v))
            hv=torch.autograd.grad(gv,params,retain_graph=True)
            norm=torch.sqrt(sum((q*q).sum() for q in hv)).detach().clamp_min(1e-8)
            v=[q/norm for q in hv]
            val=float(sum((a*b).sum() for a,b in zip(v,hv)).detach())
        return max(val,1e-3)
    results={}
    steps=700; bs=96; K=35
    for mode in ['fixed','cosine','impedance']:
        model=make(); opt=torch.optim.SGD(model.parameters(),lr=.08)
        eta=.08; max_loss=0.; t0=time.time(); hist=[]
        for step in range(steps):
            idx=torch.randint(0,len(xt),(bs,),device=device); xb,yb=xt[idx],yt[idx]
            opt.zero_grad(set_to_none=True); l=torch.nn.functional.cross_entropy(model(xb),yb); l.backward(); opt.step()
            if mode=='cosine':
                eta=.08*.5*(1+math.cos(math.pi*step/steps));
                for g in opt.param_groups:g['lr']=eta
            elif mode=='impedance' and step%K==0:
                # Probe/estimate the local dominant relaxation pole (curvature lambda_max),
                # then apply eta=gamma*(2/lambda_max), clipped to avoid abrupt jumps.
                lam=hessian_top(model,xt[:bs],yt[:bs]); target=.7*2/lam
                eta=float(np.clip(min(eta*1.1,target),.005,.35))
                for g in opt.param_groups:g['lr']=eta
            if step%K==0:
                q,_=loss_acc(model,xt[:min(256,len(xt))],yt[:min(256,len(yt))]); max_loss=max(max_loss,q); hist.append((step,eta,q))
        tr,ta=loss_acc(model,xt,yt); te,ea=loss_acc(model,xev,yev)
        results[mode]={'train_loss':tr,'test_loss':te,'test_accuracy':ea,'max_sampled_loss':max_loss,'final_lr':eta,'history':hist,'seconds':time.time()-t0}
    # fixed-rate sweep exposes the alternating/divergent side of the claimed boundary.
    sweep=[]
    for eta0 in [.03,.08,.16,.30,.45,.60]:
        m=make(); o=torch.optim.SGD(m.parameters(),lr=eta0); losses=[]
        for step in range(250):
            idx=torch.randint(0,len(xt),(bs,),device=device); o.zero_grad(set_to_none=True)
            ll=torch.nn.functional.cross_entropy(m(xt[idx]),yt[idx]); ll.backward(); o.step()
            if step%25==0: losses.append(loss_acc(m,xt[:256],yt[:256])[0])
        sweep.append({'eta':eta0,'losses':losses,'finite':bool(np.isfinite(losses).all()),'final_loss':losses[-1]})
    return {'device':str(device),'comparisons':results,'fixed_rate_sweep':sweep}

if __name__=='__main__':
    out={'seed':SEED,'lockin':lockin_demo()}
    try: out['training']=run_training()
    except Exception as e: out['training_error']=repr(e)
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))