import os, time, json, random
import numpy as np
import torch
import torch.nn as nn

SEED=397
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device.type=='cuda': torch.zeros(1,device=device)
except Exception:
    device=torch.device('cpu')

# Structured quad mesh. Each cell owns its four corner copies; shared_id is the
# only connectivity object needed by the unassembled implementation.
def mesh(n):
    cells=[]
    for y in range(n):
        for x in range(n):
            cells.append([y*(n+1)+x, y*(n+1)+x+1,
                          (y+1)*(n+1)+x+1, (y+1)*(n+1)+x])
    ids=np.asarray(cells,dtype=np.int64)
    k=np.bincount(ids.ravel(), minlength=(n+1)**2).astype(np.float32)
    return ids,k

# Exact projection check, including idempotence, symmetry, and fixation.
def projection_check():
    ids,k=mesh(3); m=ids.size; N=k.size
    G=np.zeros((m,N),dtype=np.float64)
    G[np.arange(m),ids.ravel()]=1
    S=G@np.diag(1/k)@G.T
    u=np.random.randn(N,3); v=G@u
    z=np.random.randn(m,2); manual=G@((G.T@z)/k[:,None])
    return {'projection_error':float(np.max(np.abs(S@v-v))),
            'idempotence_error':float(np.max(np.abs(S@S-S))),
            'symmetry_error':float(np.max(np.abs(S-S.T))),
            'manual_vs_matrix':float(np.max(np.abs(S@z-manual)))}

class LocalNet(nn.Module):
    def __init__(self, width=32, layers=8):
        super().__init__(); self.blocks=nn.ModuleList()
        for i in range(layers):
            self.blocks.append(nn.Sequential(nn.Linear(1 if i==0 else width,width),nn.Tanh()))
        self.out=nn.Linear(width,1)
    def forward(self,x):
        for b in self.blocks: x=b(x)
        return self.out(x)

class CellNet(nn.Module):
    def __init__(self, ids, k, sync_every, width=32, layers=8):
        super().__init__(); self.ids=torch.tensor(ids,device=device); self.k=torch.tensor(k,device=device)
        self.sync_every=sync_every; self.blocks=nn.ModuleList()
        self.local_nodes=ids.shape[1]; self.width=width
        for i in range(layers):
            iw=1 if i==0 else width
            self.blocks.append(nn.Sequential(nn.Linear(self.local_nodes*iw,self.local_nodes*width),nn.Tanh()))
        self.out=nn.Linear(width,1)
    def sync(self,h):
        # h: batch, cell, local, channel; segmented reduction via index_add.
        B,C,J,W=h.shape; flat=self.ids.reshape(-1); z=h.reshape(B,-1,W)
        sums=torch.zeros(B,self.k.numel(),W,device=h.device,dtype=h.dtype)
        sums.index_add_(1,flat,z)
        return (sums[:,flat,:]/self.k[flat][None,:,None]).reshape(B,C,J,W)
    def forward(self,x):
        # input is assembled nodal scalar, duplicated into local copies
        h=x[:,self.ids].unsqueeze(-1)
        for i,b in enumerate(self.blocks,1):
            B,C,J,W=h.shape
            h=b(h.reshape(B,C,J*W)).reshape(B,C,J,self.width)
            if self.sync_every and i%self.sync_every==0: h=self.sync(h)
        h=self.sync(h)
        # average local outputs to assembled vector
        y=h[...,0]
        B=y.shape[0]; out=torch.zeros(B,self.k.numel(),device=y.device)
        out.index_add_(1,self.ids.reshape(-1),y.reshape(B,-1))
        return out/self.k[None,:]

class AssembledNet(nn.Module):
    def __init__(self, adj, width=32, layers=8):
        super().__init__(); self.register_buffer('A',adj); self.blocks=nn.ModuleList()
        for i in range(layers): self.blocks.append(nn.Sequential(nn.Linear(1 if i==0 else width,width),nn.Tanh()))
        self.out=nn.Linear(width,1)
    def forward(self,x):
        h=x.unsqueeze(-1)
        for b in self.blocks:
            # standard assembled graph-GNN layer: neighbor aggregate plus self
            h=b(h + torch.matmul(self.A,h))
        return self.out(h).squeeze(-1)

def data(n,Nsample=512):
    # Diffusion-like target: a fixed local operator, represented as a smooth
    # coefficient field mapped to its neighbor average.
    yy,xx=np.mgrid[0:n+1,0:n+1]; p=(n+1)**2
    A=np.zeros((p,p),np.float32)
    for q in range(p):
        y,x=divmod(q,n+1)
        ns=[]
        for dy,dx in ((-1,0),(1,0),(0,-1),(0,1)):
            if 0<=y+dy<=n and 0<=x+dx<=n: ns.append((y+dy)*(n+1)+x+dx)
        A[q,ns]=1/len(ns)
    X=np.random.randn(Nsample,p).astype('float32')
    Y=X@A.T
    return torch.tensor(X),torch.tensor(Y),torch.tensor(A)

def train(model,X,Y,steps=160):
    model.to(device); X=X.to(device); Y=Y.to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    t=time.perf_counter();
    for s in range(steps):
        ix=torch.randint(0,len(X),(64,),device=device); pred=model(X[ix]); loss=((pred-Y[ix])**2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    if device.type=='cuda': torch.cuda.synchronize()
    elapsed=time.perf_counter()-t
    with torch.no_grad(): val=((model(X[-128:])-Y[-128:])**2).mean().item()
    return val,steps*64/elapsed,elapsed

def main():
    check=projection_check(); n=8; ids,k=mesh(n); X,Y,A=data(n)
    # Normalize adjacency to avoid an unfairly unstable assembled baseline.
    models=[('assembled',AssembledNet(A*.5,layers=8)),
            ('cell_sync_1',CellNet(ids,k,1,layers=8)),
            ('cell_sync_4',CellNet(ids,k,4,layers=8)),
            ('cell_sync_8',CellNet(ids,k,8,layers=8))]
    results={}
    for name,m in models: results[name]=train(m,X,Y)
    # Disagreement is measured immediately before a scheduled synchronization.
    h=torch.randn(16,len(ids),4,8,device=device); cm=CellNet(ids,k,4)
    pre=cm.sync(h); disagreement_before=(h-pre).pow(2).mean().sqrt().item()
    results['projection_check']=check; results['pre_sync_rms_disagreement']=disagreement_before
    results['device']=str(device); results['cells']=len(ids); results['nodes']=len(k)
    print(json.dumps(results,indent=2))
    with open('results.json','w') as f: json.dump(results,f,indent=2)
if __name__=='__main__': main()