import json, math, random
import numpy as np
import torch
from torch import nn

SEED=1481
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def j_sign(j):
    return (-1)**(j*(j-1)//2)

def wedge2(x,y):
    return torch.outer(x,y)-torch.outer(y,x)

def fermionic_local(messages, require_even=True):
    # degree 0 and degree 2 truncation; odd incident degree is rejected
    n=len(messages)
    if require_even and n % 2: return torch.tensor(0., dtype=messages[0].dtype), torch.zeros((messages[0].numel(),messages[0].numel()))
    if n < 2: return torch.tensor(1., dtype=messages[0].dtype), torch.zeros((messages[0].numel(),messages[0].numel()))
    A=torch.zeros((messages[0].numel(),messages[0].numel()),dtype=messages[0].dtype)
    for i in range(n):
        for k in range(i+1,n): A += wedge2(messages[i],messages[k])
    return torch.tensor(1., dtype=messages[0].dtype), A

def math_checks():
    out={}
    # Prediction 1: every odd number of incident odd half-edges is exactly killed.
    odd_res=[]
    d=4
    for n in range(1,10,2):
        xs=[torch.randn(d) for _ in range(n)]
        s,A=fermionic_local(xs)
        odd_res.append(float(abs(s)+torch.linalg.norm(A)))
    out['odd_degree_predicted_zero']= {'predicted':0.0,'observed_max_residual':max(odd_res),'all_residuals':odd_res}
    # Prediction 2: parity sign is (-1)^|F|, swept over all masks and graph sizes.
    parity=[]
    for E in range(1,9):
        maxerr=0
        for mask in range(1<<E):
            m=mask.bit_count(); lhs=(-1)**sum(((2 if False else 0) for _ in []))
            # construct a cycle: each selected edge contributes to two endpoint degrees
            deg=[0]*E
            for e in range(E):
                if mask>>e&1: deg[e]+=1; deg[(e+1)%E]+=1
            # Odd local degrees vanish; compare signs only on surviving masks.
            if any(q % 2 for q in deg):
                continue
            rhs=(-1)**m
            local=1
            for q in deg: local*=(-1)**(q//2)
            maxerr=max(maxerr,abs(local-rhs))
        parity.append({'edges':E,'max_abs_error':maxerr})
    out['parity_predicted_exact']= {'predicted_error':0.0,'sweep':parity}
    # Prediction 3: random degree-2 wedge energy scales as d(d-1); normalized energy is stable.
    scaling=[]
    for d in [2,3,4,6,8,12]:
        vals=[]
        for _ in range(400):
            x=torch.randn(d); y=torch.randn(d); vals.append(float(torch.linalg.norm(wedge2(x,y))**2))
        mean=np.mean(vals); predicted=2*d*(d-1)
        scaling.append({'ell':d//2,'dimension':d,'predicted_mean':predicted,'observed_mean':float(mean),'normalized':float(mean/predicted)})
    out['wedge_energy_scaling']=scaling
    return out

class GraphSet:
    def __init__(self,ntrain=240, ntest=120, min_n=6,max_n=10):
        self.train=self.make(ntrain,min_n,max_n); self.test=self.make(ntest, max_n+1, max_n+5)
    def make(self,N,lo,hi):
        arr=[]
        for z in range(N):
            n=random.randint(lo,hi); y=z%2
            edges=[]
            if y==1: # cycle
                edges=[(i,(i+1)%n) for i in range(n)]
            else: # tree/path, same node count and degree statistics mostly
                edges=[(i,i+1) for i in range(n-1)]
                random.shuffle(edges)
            arr.append((n,edges,y))
        return arr

def batch_graph(g, max_n):
    n,edges,y=g
    # node scalar is constant; challenge is topology
    X=torch.ones(n,1)
    return X,edges,y

class GIN(nn.Module):
    def __init__(self,h=16, fermionic=False,d=4):
        super().__init__(); self.fermionic=fermionic; self.d=d
        self.inp=nn.Linear(1,h)
        self.e=nn.ModuleList([nn.Sequential(nn.Linear(h,h),nn.ReLU(),nn.Linear(h,h)) for _ in range(3)])
        if fermionic:
            self.odd=nn.ModuleList([nn.Linear(h,d) for _ in range(3)])
            self.mix=nn.ModuleList([nn.Sequential(nn.Linear(h+d*d,h),nn.ReLU(),nn.Linear(h,h)) for _ in range(3)])
        self.cls=nn.Sequential(nn.Linear(h, h),nn.ReLU(),nn.Linear(h,2))
    def forward(self,X,edges):
        h=self.inp(X); n=h.shape[0]
        adj=[[] for _ in range(n)]
        for u,v in edges: adj[u].append(v); adj[v].append(u)
        for t in range(3):
            agg=torch.zeros_like(h); odd=[[] for _ in range(n)]
            for v in range(n):
                for u in adj[v]:
                    agg[v]+=self.e[t](h[u])
                    if self.fermionic: odd[v].append(self.odd[t](h[u]))
            if self.fermionic:
                feats=[]
                for v in range(n):
                    A=fermionic_local(odd[v])[1].reshape(-1)
                    feats.append(torch.cat([agg[v],A]))
                h=h+self.mix[t](torch.stack(feats))
            else: h=h+agg
        return self.cls(h.mean(0))

def train_eval(fermionic, data):
    torch.manual_seed(SEED+int(fermionic)); model=GIN(fermionic=fermionic)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    for ep in range(12):
        random.shuffle(data.train); model.train()
        for g in data.train:
            X,e,y=batch_graph(g,16); loss=nn.functional.cross_entropy(model(X,e)[None,:],torch.tensor([y]))
            opt.zero_grad(); loss.backward(); opt.step()
    model.eval(); good=0; total=0
    with torch.no_grad():
        for g in data.test:
            X,e,y=batch_graph(g,16); good += int(model(X,e).argmax().item()==y); total+=1
    return good/total

def main():
    checks=math_checks(); data=GraphSet(); b=train_eval(False,data); f=train_eval(True,data)
    result={'seed':SEED,'checks':checks,'classification':{'baseline_gin_accuracy':b,'fermionic_accuracy':f,'test_graphs':len(data.test),'note':'test cycles have longer node counts than training'}}
    with open('results.json','w') as h: json.dump(result,h,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()