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

SEED = 563
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)


def torus(n, m):
    # vertices are (x,y), directed edges use displacement (+x,-x,+y,-y)
    N=n*m; edges=[]; rel=[]
    def ix(x,y): return (x%n)*m+(y%m)
    for x in range(n):
        for y in range(m):
            u=ix(x,y)
            for r,(dx,dy) in enumerate(((1,0),(-1,0),(0,1),(0,-1))):
                edges.append((u,ix(x+dx,y+dy))); rel.append(r)
    return N, np.asarray(edges,np.int64), np.asarray(rel,np.int64)


def quotient_check(n=5,m=7):
    # Signed cycle incidence vectors in Z^2. Every elementary torus cycle sums to 0.
    # The wrap cycles demonstrate finite quotient relations n*a_x=0 and m*a_y=0.
    cycles=[]
    for y in range(m):
        cycles.append(np.array([n,0],dtype=np.int64))
    for x in range(n):
        cycles.append(np.array([0,m],dtype=np.int64))
    for x in range(n):
        for y in range(m):
            cycles.append(np.array([0,0],dtype=np.int64)) # square commutator
    C=np.stack(cycles)
    local_max=int(np.max(np.abs(C[:n+m])))
    assert np.all(C[n+m:] == 0)
    # Labels in A=Z_n x Z_m; every edge displacement agrees modulo (n,m).
    N,E,R=torus(n,m); labels=np.array([(x,y) for x in range(n) for y in range(m)],dtype=np.int64)
    gens=np.array([[1,0],[-1,0],[0,1],[0,-1]],dtype=np.int64)
    residual=(labels[E[:,1]]-labels[E[:,0]]-gens[R])
    ok=np.all((residual[:,0]%n==0)&(residual[:,1]%m==0))
    # path independence: two paths' displacement difference is a cycle relation
    assert ok
    return {"group":f"Z_{n} x Z_{m}","cycle_rows":len(C),"max_wrap_incidence":local_max,
            "all_edge_label_differences_zero_in_quotient":bool(ok),
            "relation_classes":4,"directed_edges":int(len(E))}


class RelationMP(nn.Module):
    def __init__(self, d=16):
        super().__init__(); self.W=nn.Parameter(torch.randn(4,d,d)*.08)
        self.out=nn.Sequential(nn.Linear(d,d),nn.Tanh(),nn.Linear(d,1))
    def forward(self,x,edge,rel):
        z=torch.zeros_like(x)
        for r in range(4):
            mask=(rel==r); src=edge[mask,0]; dst=edge[mask,1]
            z.index_add_(0,dst,x[src]@self.W[r].T)
        return self.out(z)

class GCN(nn.Module):
    def __init__(self,d=16):
        super().__init__(); self.W=nn.Parameter(torch.randn(d,d)*.08)
        self.out=nn.Sequential(nn.Linear(d,d),nn.Tanh(),nn.Linear(d,1))
    def forward(self,x,edge,rel):
        z=torch.zeros_like(x); z.index_add_(0,edge[:,1],x[edge[:,0]]@self.W.T)
        deg=torch.bincount(edge[:,1],minlength=x.shape[0]).float().unsqueeze(1)
        return self.out(z/deg)

def data(n,m,seed):
    rng=np.random.default_rng(seed); N,E,R=torus(n,m)
    x=torch.tensor(rng.normal(size=(N,16)),dtype=torch.float32)
    # A relation-dependent one-hop convolution, with independent inputs per graph.
    coeff=torch.tensor([1.0,-0.7,0.45,-1.15])
    y=torch.zeros(N)
    for r in range(4): y.index_add_(0,torch.tensor(E[R==r,1]),coeff[r]*x[torch.tensor(E[R==r,0]),0])
    return x,torch.tensor(E),torch.tensor(R),y[:,None]

def fit(model, n=6,m=6,steps=700):
    opt=torch.optim.Adam(model.parameters(),lr=.025)
    model.train()
    for t in range(steps):
        x,e,r,y=data(n,m,1000+t)
        pred=model(x,e,r); loss=((pred-y)**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    return model

def evaluate(model,n,m,count=20):
    model.eval(); vals=[]
    with torch.no_grad():
        for j in range(count):
            x,e,r,y=data(n,m,9000+j); vals.append(float(((model(x,e,r)-y)**2).mean()))
    return float(np.mean(vals))

def main():
    math_result=quotient_check()
    # Same architecture scale, source training and larger-geometry transfer.
    rel=fit(RelationMP()); gcn=fit(GCN())
    results={
      "source_6x6": {"GCN":evaluate(gcn,6,6),"quotient_tied":evaluate(rel,6,6)},
      "transfer_8x8": {"GCN":evaluate(gcn,8,8),"quotient_tied":evaluate(rel,8,8)},
      "parameters":{"GCN":sum(p.numel() for p in gcn.parameters()),"quotient_tied":sum(p.numel() for p in rel.parameters()),"untied_relation_bank_equivalent":4*16*16},
      "compression_edges_over_generators":(2*6*6*4)/4,
      "math":math_result}
    Path('results.json').write_text(json.dumps(results,indent=2))
    print(json.dumps(results,indent=2))

if __name__=='__main__': main()