import json, math, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED = 17
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

class CrossChannelScattering(nn.Module):
    """Fixed complex Morlet analysis, cross-channel amplitude products, averaging."""
    def __init__(self, channels, edges, bands=(0.10, 0.20), kernel=25, pool=32, eps=1e-8):
        super().__init__()
        self.channels, self.edges, self.bands = channels, edges, tuple(bands)
        self.pool, self.eps = pool, eps
        t = torch.arange(kernel, dtype=torch.float32) - (kernel-1)/2
        sigma = kernel / 6.0
        ws = []
        for f in bands:
            g = torch.exp(-0.5*(t/sigma)**2)
            # zero-mean approximately analytic Morlet pair
            ws.append(torch.stack([g*torch.cos(2*math.pi*f*t), g*torch.sin(2*math.pi*f*t)]))
        w = torch.stack(ws) # bands, 2, kernel
        self.register_buffer('wr', w[:,0].view(len(bands),1,kernel))
        self.register_buffer('wi', w[:,1].view(len(bands),1,kernel))

    def forward(self, x):
        B,C,T = x.shape
        xr = x.reshape(B*C,1,T)
        pad = self.wr.shape[-1]//2
        zr = F.conv1d(xr, self.wr, padding=pad).reshape(B,C,len(self.bands),T)
        zi = F.conv1d(xr, self.wi, padding=pad).reshape(B,C,len(self.bands),T)
        out=[]
        for m,n in self.edges:
            # q = (zr_m+i zi_m) conj(zr_n+i zi_n)
            qr = zr[:,m]*zr[:,n] + zi[:,m]*zi[:,n]
            qi = zi[:,m]*zr[:,n] - zr[:,m]*zi[:,n]
            u = torch.sqrt(qr.square()+qi.square()+self.eps)
            out.append(F.avg_pool1d(u.reshape(B,len(self.bands),T), self.pool, stride=self.pool))
        return torch.cat(out, dim=1)

def verify():
    torch.manual_seed(SEED)
    z1 = torch.randn(3, 2, 20) + 1j*torch.randn(3,2,20)
    z2 = torch.randn(3, 2, 20) + 1j*torch.randn(3,2,20)
    q=z1*z2.conj()
    identity=(q.abs()-(z1.abs()*z2.abs())).abs().max().item()
    phase=torch.rand(3,2,20)*2*math.pi
    phase_err=((z1*torch.exp(1j*phase)).abs()-z1.abs()).abs().max().item()
    # end-to-end layer agrees with direct magnitude product (apart from eps)
    layer=CrossChannelScattering(2,[(0,1)],bands=(.12,),kernel=17,pool=8)
    x=torch.randn(2,2,64)
    got=layer(x)
    # independently reproduce its convolution and pooling
    pad=8; zr=F.conv1d(x.reshape(4,1,64),layer.wr,padding=pad).reshape(2,2,1,64)
    zi=F.conv1d(x.reshape(4,1,64),layer.wi,padding=pad).reshape(2,2,1,64)
    direct=F.avg_pool1d(torch.sqrt((zr[:,0]*zr[:,1]+zi[:,0]*zi[:,1]).square()+(zi[:,0]*zr[:,1]-zr[:,0]*zi[:,1]).square()+1e-8).reshape(2,1,64),8,stride=8)
    implementation_err=(got-direct).abs().max().item()
    return {'complex_magnitude_identity_max_error':identity,'phase_invariance_max_error':phase_err,'layer_direct_formula_max_error':implementation_err}

def make_data(n, T=256):
    # Class 1: neighboring sensors share a slowly varying amplitude envelope;
    # phases are independent, so coupling is amplitude- rather than phase-based.
    t=np.arange(T)/T; X=np.zeros((n,4,T),np.float32); y=np.arange(n)%2
    for i,c in enumerate(y):
        freq=.12 + .008*np.random.randn()
        envs=[]
        if c==1:
            env=0.55+0.45*(np.sin(2*np.pi*(1.5*t+np.random.rand()))**2)
            envs=[env,env, .6+.4*np.random.rand()*np.ones(T), .6+.4*np.random.rand()*np.ones(T)]
        else:
            envs=[0.55+0.45*np.sin(2*np.pi*(1.5*t+np.random.rand()))**2 for _ in range(4)]
        for ch in range(4):
            ph=np.random.rand()*2*np.pi
            X[i,ch]=envs[ch]*np.sin(2*np.pi*freq*np.arange(T)+ph)+.65*np.random.randn(T)
    return torch.tensor(X),torch.tensor(y,dtype=torch.long)

class RawCNN(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Conv1d(4,12,15,padding=7),nn.ReLU(),nn.AvgPool1d(8),nn.Conv1d(12,12,9,padding=4),nn.ReLU(),nn.AdaptiveAvgPool1d(1))
        self.fc=nn.Linear(12,2)
    def forward(self,x): return self.fc(self.net(x).squeeze(-1))
class SNSTModel(nn.Module):
    def __init__(self):
        super().__init__(); self.s=CrossChannelScattering(4,[(0,1),(1,2),(2,3)],bands=(.10,.20),kernel=25,pool=32); self.fc=nn.Sequential(nn.Flatten(),nn.Linear(3*2*8,16),nn.ReLU(),nn.Linear(16,2))
    def forward(self,x): return self.fc(self.s(x))

def train(model, Xtr,ytr,Xte,yte, epochs=35):
    model=model.to(DEVICE); Xtr=Xtr.to(DEVICE); ytr=ytr.to(DEVICE); Xte=Xte.to(DEVICE); yte=yte.to(DEVICE)
    opt=torch.optim.Adam(model.parameters(),lr=3e-3,weight_decay=1e-3)
    for _ in range(epochs):
        model.train(); p=torch.randperm(len(Xtr),device=DEVICE)
        for j in range(0,len(Xtr),64):
            ix=p[j:j+64]; loss=F.cross_entropy(model(Xtr[ix]),ytr[ix]); opt.zero_grad(); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        pred=model(Xte).argmax(1); acc=(pred==yte).float().mean().item()
    return acc

def experiment():
    # 20% labels: deliberately small training set, fixed held-out test.
    X,y=make_data(800); Xtr,ytr=X[:160],y[:160]; Xte,yte=X[160:],y[160:]
    # normalize using training statistics, as required for a fair front end
    mu=Xtr.mean((0,2),keepdim=True); sd=Xtr.std((0,2),keepdim=True)+1e-5
    Xtr=(Xtr-mu)/sd; Xte=(Xte-mu)/sd
    results=[]
    for kind in ('raw','snst'):
        vals=[]
        for seed in (3,7,11):
            torch.manual_seed(seed); np.random.seed(seed)
            model=RawCNN() if kind=='raw' else SNSTModel()
            vals.append(train(model,Xtr,ytr,Xte,yte))
        results.append((kind,vals,float(np.mean(vals)),float(np.std(vals))))
    return results

if __name__=='__main__':
    try:
        if DEVICE=='cuda': torch.cuda.init()
    except Exception as e:
        DEVICE='cpu'; print('CUDA fallback:',repr(e))
    check=verify(); print('DEVICE',DEVICE); print('VERIFY',json.dumps(check))
    try:
        results=experiment()
    except Exception as e:
        if DEVICE != 'cpu':
            print('Runtime fallback to CPU:',repr(e))
            DEVICE='cpu'
            results=experiment()
        else:
            raise
    print('RESULTS',json.dumps(results))