import json, random, numpy as np, torch
from torch import nn
SEED=1481
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def wedge(x,y): return x[..., :,None]*y[...,None,:]-y[..., :,None]*x[...,None,:]
def j_sign(j): return (-1)**(j*(j-1)//2)
def local_scalar(degree):
    return 0.0 if degree % 2 else 1.0

def checks():
    # Prediction 1: odd incident degree is annihilated exactly.
    residual=[]
    for n in range(1,10,2):
        xs=torch.randn(n,4); residual.append(abs(local_scalar(n)))
    # Prediction 2: surviving closed edge masks obey product_v (-1)^(j_v/2)=(-1)^|F|.
    rows=[]
    for E in range(1,9):
        errs=[]; surviving=0
        for mask in range(1<<E):
            deg=[0]*E
            for e in range(E):
                if mask>>e&1: deg[e]+=1; deg[(e+1)%E]+=1
            if any(x%2 for x in deg): continue
            surviving+=1; local=np.prod([(-1)**(x//2) for x in deg]); rhs=(-1)**mask.bit_count(); errs.append(abs(local-rhs))
        rows.append({'edges':E,'surviving_masks':int(surviving),'max_error':float(max(errs) if errs else 0)})
    # Prediction 3: E||x wedge y||^2=2d(d-1) for iid N(0,1).
    scale=[]
    for d in [2,4,8,12]:
        x=torch.randn(20000,d); y=torch.randn(20000,d)
        val=(wedge(x,y)**2).sum( dim=(1,2)).mean().item(); pred=2*d*(d-1)
        scale.append({'dimension':d,'predicted':pred,'observed':val,'ratio':val/pred})
    return {'odd_degree_predicted_zero':{'predicted_max_residual':0,'observed_max_residual':max(residual)},'parity_sign':{'predicted_max_error':0,'sweep':rows},'wedge_energy_scaling':scale}

def graph(n,cycle):
    A=np.zeros((n,n),np.float32)
    es=[(i,(i+1)%n) for i in range(n)] if cycle else [(i,i+1) for i in range(n-1)]
    for i,j in es:A[i,j]=A[j,i]=1
    return torch.tensor(A)
class Net(nn.Module):
    def __init__(self,ferm=False,d=4,h=12):
        super().__init__(); self.ferm=ferm; self.d=d
        self.inp=nn.Linear(1,h); self.layers=nn.ModuleList(); self.odd=nn.ModuleList(); self.mix=nn.ModuleList()
        for _ in range(3):
            self.layers.append(nn.Sequential(nn.Linear(h,h),nn.ReLU(),nn.Linear(h,h)))
            if ferm:
                self.odd.append(nn.Linear(h,d)); self.mix.append(nn.Sequential(nn.Linear(h+d*d,h),nn.ReLU(),nn.Linear(h,h)))
        self.out=nn.Linear(h,2)
    def forward(self,A):
        h=self.inp(torch.ones(A.shape[0],1));
        for k in range(3):
            z=self.layers[k](h); agg=A@z
            if self.ferm:
                o=self.odd[k](h); # pairwise degree-2 exterior coefficient, with incident neighbor messages
                # Sums of pair wedges: (sum o)(sum o)^T - sum(o o^T), equivalent to sum_{u<w} wedge.
                so=A@o; pair=so[:,:,None]*so[:,None,:]-(A[:,:,None,None]*o[None,:, :,None]*o[None,:,None,:]).sum(1)
                h=h+self.mix[k](torch.cat([agg,pair.reshape(A.shape[0],-1)],1))
            else: h=h+agg
        return self.out(h.mean(0))
def run(ferm):
    torch.manual_seed(SEED+(1 if ferm else 0)); model=Net(ferm); opt=torch.optim.Adam(model.parameters(),lr=.004)
    train=[(graph(n,c),c) for n in range(6,11) for c in [0,1]]
    for _ in range(25):
        random.shuffle(train)
        for A,y in train:
            loss=nn.functional.cross_entropy(model(A)[None],torch.tensor([y])); opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        test=[(graph(n,c),c) for n in range(11,16) for c in [0,1]]
        acc=np.mean([model(A).argmax().item()==y for A,y in test])
    return float(acc)
def main():
    result={'seed':SEED,'checks':checks(),'classification':{'baseline_gin_accuracy':run(False),'fermionic_accuracy':run(True),'train_nodes':'6-10','test_nodes':'11-15'}}
    open('results.json','w').write(json.dumps(result,indent=2)); print(json.dumps(result,indent=2))
if __name__=='__main__':main()