import json, random
from pathlib import Path
import numpy as np
import torch
from torch import nn

SEED = 425
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
try: torch.set_num_threads(4)
except Exception: pass

def ring_ops(n):
    D = np.zeros((n,n), dtype=np.float64)
    for i in range(n):
        D[i,i] = -1.; D[i,(i+1)%n] = 1.
    return D, D.T @ D

def riesz_matrix(n, lam, eps=1e-6):
    D,L = ring_ops(n)
    A = np.eye(n) + lam*lam*(L + eps*np.eye(n))
    return lam * D @ np.linalg.inv(A), D, L

def sanity():
    out=[]
    for n in (32,64,128):
        lam=1.0
        R,D,L=riesz_matrix(n,lam)
        direct=lam*D
        rng=np.random.default_rng(SEED+n)
        # noisy piecewise-smooth signal, measuring derivative noise amplification
        x=np.sin(2*np.pi*np.arange(n)/n); x[n//2:]+=0.8
        noise=rng.normal(size=n)
        clean=np.linalg.norm(R@x); noisy=np.linalg.norm(R@(x+.5*noise))
        dclean=np.linalg.norm(direct@x); dnoisy=np.linalg.norm(direct@(x+.5*noise))
        op=np.linalg.svd(R,compute_uv=False)[0]
        dop=np.linalg.svd(direct,compute_uv=False)[0]
        # Isolate perturbation response, rather than conflating it with signal energy.
        r_noise=np.linalg.norm(R@(0.5*noise))/np.linalg.norm(0.5*noise)
        d_noise=np.linalg.norm(direct@(0.5*noise))/np.linalg.norm(0.5*noise)
        out.append({'n':n,'resolvent_operator_norm':float(op),'direct_operator_norm':float(dop),
                    'noise_ratio_resolvent':float(noisy/clean),'noise_ratio_direct':float(dnoisy/dclean),
                    'noise_amplification_resolvent':float(r_noise),'noise_amplification_direct':float(d_noise)})
    return out

class Branch(nn.Module):
    def __init__(self,n,d,kind):
        super().__init__(); self.kind=kind; self.n=n
        R,D,L=riesz_matrix(n,1.0)
        M=R if kind=='riesz' else 1.0*D if kind=='direct' else np.zeros((n,n))
        self.register_buffer('M',torch.tensor(M,dtype=torch.float32))
        self.base=nn.Linear(d,16); self.out=nn.Linear(32 if kind!='plain' else 16,2)
        self.g=nn.Parameter(torch.tensor(0.1 if kind=='riesz' else 1.0))
    def forward(self,x):
        # x: B,N,d; branch is applied to node dimension
        h=torch.relu(self.base(x))
        if self.kind!='plain':
            r=torch.einsum('nm,bmd->bnd',self.M,h)
            r=r/(r.square().mean(dim=(1,2),keepdim=True).sqrt()+1e-5)
            h=h+self.g*r
        
        pooled=h.mean(1)
        if self.kind!='plain': pooled=torch.cat([pooled, r.abs().mean(1)],dim=1)
        return self.out(pooled)

def dataset(n=64, samples=1200, noise=.10):
    rng=np.random.default_rng(SEED)
    X=np.zeros((samples,n,1),np.float32); y=np.zeros(samples,np.int64)
    t=np.arange(n)
    for k in range(samples):
        cls=k%2; y[k]=cls
        phase=rng.uniform(0,2*np.pi)
        if cls==0:
            # Smooth oscillation and a random phase: same pointwise marginal as class 1.
            sig=np.sin(2*np.pi*t/n*2+phase)
        else:
            # High-frequency alternating signal, with matched amplitude and random phase.
            sig=np.sign(np.sin(2*np.pi*t/n*16+phase))
        X[k,:,0]=sig+rng.normal(0,noise,n)
    p=rng.permutation(samples); return X[p],y[p]

def train(kind, Xtr,ytr,Xv,yv,n,epochs=35):
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model=Branch(n,1,kind).to(dev)
        opt=torch.optim.Adam(model.parameters(),lr=3e-3); lossfn=nn.CrossEntropyLoss()
        xt=torch.tensor(Xtr,device=dev); yt=torch.tensor(ytr,device=dev)
        xv=torch.tensor(Xv,device=dev); yvt=torch.tensor(yv,device=dev)
        for _ in range(epochs):
            for s in range(0,len(xt),128):
                opt.zero_grad(); loss=lossfn(model(xt[s:s+128]),yt[s:s+128]); loss.backward(); opt.step()
        with torch.no_grad():
            acc=(model(xv).argmax(1)==yvt).float().mean().item()
            pert=(model(xv+0.10*torch.randn_like(xv)).argmax(1)==yvt).float().mean().item()
            val=lossfn(model(xv),yvt).item()
        return {'val_loss':val,'accuracy':acc,'perturbed_accuracy':pert}
    except Exception as e:
        if dev=='cuda':
            torch.cuda.empty_cache(); torch.set_default_device('cpu')
            return train(kind,Xtr,ytr,Xv,yv,n,epochs)
        raise

def main():
    check=sanity(); X,y=dataset(); Xtr,ytr,Xv,yv=X[:900],y[:900],X[900:],y[900:]
    results={k:train(k,Xtr,ytr,Xv,yv,64) for k in ('plain','direct','riesz')}
    report={'seed':SEED,'sanity':check,'results':results}
    Path('results.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report,indent=2))
if __name__=='__main__': main()