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

SEED = 2043

def seed_all(seed=SEED):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)


def ego_intersections(n, edges):
    """Exact n=1 closed ego sets S[u]={u and its neighbors}; undirected edges."""
    nbr = [set([i]) for i in range(n)]
    for u,v in edges:
        nbr[u].add(v); nbr[v].add(u)
    stats = {}
    for u,v in edges:
        I = nbr[u] & nbr[v]
        # Formula in prompt, with epsilon only in denominator.
        C = 2.0*len(I)/(len(nbr[u])+len(nbr[v])+1e-12)
        r = len(I)/(len(nbr[u])*len(nbr[v])+1e-12)
        stats[(u,v)] = (sorted(I), C, r, len(nbr[u]), len(nbr[v]))
        stats[(v,u)] = stats[(u,v)]
    return nbr, stats


def math_verification():
    # Prediction 1: for a regular edge, adding t common neighbors gives
    # C=(2(t+2))/(2(k+1))=(t+2)/(k+1), hence slope 1/(k+1).
    k=6
    rows=[]
    for t in range(0, k-1):
        # endpoints have k-1 other neighbors, t shared among them, plus edge.
        # Build only sets directly to avoid irrelevant graph edges.
        su=set(range(k+1)); sv={k+1}; sv.update(range(0,t+1)); sv.add(k+2)
        # normalize desired sizes: endpoint sets each k+1 and intersection is edge endpoints + t
        su={0}|set(range(2,k+1)); sv={1}|set(range(2, t+2))|set(range(k+2, 2*k-t+1))
        # u=0, v=1; force both sizes k+1 and common {0,1,2..t+1}
        # easier use abstract sets with exact sizes
        su={0,1}|set(range(2,t+2))|set(range(100,100+(k-1-t)))
        sv={0,1}|set(range(2,t+2))|set(range(200,200+(k-1-t)))
        C=2*len(su&sv)/(len(su)+len(sv)+1e-12)
        pred=(t+2)/(k+1)
        rows.append((t,C,pred))
    max_err=max(abs(x-y) for _,x,y in rows)
    slope=np.polyfit([t for t,_,_ in rows],[c for _,c,_ in rows],1)[0]
    slope_pred=1/(k+1)

    # Prediction 2: alpha=0.5 at C*=-(a0+a2*r)/a1 and is increasing in C.
    a0,a1,a2=-1.2,4.0,0.7
    r=0.08
    cstar=-(a0+a2*r)/a1
    Cs=np.linspace(0,1,101)
    alphas=1/(1+np.exp(-(a0+a1*Cs+a2*r)))
    crossing=float(Cs[np.argmin(abs(alphas-.5))])
    monotone=bool(np.all(np.diff(alphas)>0))

    # Prediction 3: for fixed ordinary o and overlap q, update is affine in alpha,
    # with endpoint values o and q and slope q-o.
    o=np.array([1.3,-0.4,0.7]); q=np.array([-0.2,0.8,0.1])
    aa=np.linspace(0,1,11)
    updates=np.array([(1-x)*o+x*q for x in aa])
    fit=np.polyfit(aa, updates[:,0], 1)[0]
    predicted_slope=float(q[0]-o[0])
    return {
      'closure_sweep': {'k':k,'rows':rows,'max_abs_error':float(max_err),
                        'observed_slope':float(slope),'predicted_slope':float(slope_pred)},
      'gate_sweep': {'a0':a0,'a1':a1,'a2':a2,'r':r,'predicted_C_at_alpha_half':float(cstar),
                     'observed_C_at_alpha_half':crossing,'monotone_in_C':monotone,
                     'alpha_at_0':float(alphas[0]),'alpha_at_1':float(alphas[-1])},
      'update_sweep': {'predicted_slope_coordinate0':predicted_slope,'observed_slope_coordinate0':float(fit),
                       'endpoint_error':float(np.max(np.abs(updates[0]-o))+np.max(np.abs(updates[-1]-q)))}
    }


def make_graph(n=48, p_in=.30, p_out=.045, triangle_bias=0.55):
    y=np.array([i < n//2 for i in range(n)],dtype=np.int64)
    edges=[]
    rng=np.random.RandomState(SEED+int(1000*p_in))
    for i in range(n):
        for j in range(i+1,n):
            p=p_in if y[i]==y[j] else p_out
            if rng.rand()<p: edges.append((i,j))
    # Add triangles within each class, making high closure but modest graph size.
    for i in range(0,n,3):
        group=list(range(i,min(i+3,n)))
        for a in range(len(group)):
            for b in range(a+1,len(group)):
                e=(group[a],group[b])
                if e not in edges: edges.append(e)
    return edges,y


def features(n,y,edges):
    rng=np.random.RandomState(99)
    x=rng.randn(n,8).astype(np.float32)*0.9
    x[:,0] += (2*y-1)*0.35
    return torch.tensor(x), torch.tensor(y,dtype=torch.long)

class GeccLayer(nn.Module):
    def __init__(self, d):
        super().__init__(); self.w0=nn.Linear(d,d,bias=False); self.w1=nn.Linear(d,d,bias=False); self.w2=nn.Linear(d,d,bias=False)
        self.a=nn.Parameter(torch.tensor([-1.,2.,.2]))
    def forward(self,h, edges, stats, force_gate=None):
        out=self.w0(h); ordinary=torch.zeros_like(h); corr=torch.zeros_like(h)
        for u,v in edges:
            for dst,src in ((u,v),(v,u)):
                I,C,r,_,_=stats[(dst,src)]
                alpha=torch.sigmoid(self.a[0]+self.a[1]*C+self.a[2]*r)
                if force_gate is not None: alpha=torch.tensor(force_gate,device=h.device)
                ordinary[dst]+=self.w1(h[src])
                q=h[I].mean(0) if I else torch.zeros(h.shape[1],device=h.device)
                corr[dst]+=self.w2(q)
                out[dst] += (1-alpha)*self.w1(h[src]) + alpha*self.w2(q) - self.w0(h[dst])*0
        return torch.relu(out)

class GCN(nn.Module):
    def __init__(self,d=8,hid=16):
        super().__init__(); self.l1=nn.Linear(d,hid); self.l2=nn.Linear(hid,2)
    def forward(self,x,edges):
        n=x.shape[0]; z=torch.zeros_like(x); deg=torch.zeros(n,device=x.device)
        for u,v in edges: z[u]+=x[v]; z[v]+=x[u]; deg[u]+=1; deg[v]+=1
        z=z/(deg[:,None]+1); return self.l2(torch.relu(self.l1(x+z)))

class GECC(nn.Module):
    def __init__(self,d=8,hid=16):
        super().__init__(); self.inp=nn.Linear(d,hid); self.msg=GeccLayer(hid); self.out=nn.Linear(hid,2)
    def forward(self,x,edges,stats): return self.out(self.msg(torch.relu(self.inp(x)),edges,stats))


def train_one(model,x,y,edges,stats,train_idx,steps=60):
    opt=torch.optim.Adam(model.parameters(),lr=.025,weight_decay=2e-3)
    for _ in range(steps):
        opt.zero_grad(); logits=model(x,edges,stats) if isinstance(model,GECC) else model(x,edges)
        loss=nn.functional.cross_entropy(logits[train_idx],y[train_idx]); loss.backward(); opt.step()
    with torch.no_grad():
        logits=model(x,edges,stats) if isinstance(model,GECC) else model(x,edges)
        pred=logits.argmax(1); return float((pred!=y).float().mean()), float((pred[~train_idx]!=y[~train_idx]).float().mean())


def mini_experiment():
    seed_all(); edges,y_np=make_graph(); n=len(y_np); _,stats=ego_intersections(n,edges); x,y=features(n,y_np,edges)
    train_idx=torch.zeros(n,dtype=torch.bool); train_idx[:n//2:2]=True; train_idx[n//2::2]=True
    base=train_one(GCN(),x,y,edges,stats,train_idx); idea=train_one(GECC(),x,y,edges,stats,train_idx)
    Cs=[v[1] for v in stats.values()]; return {'n':n,'edges':len(edges),'mean_C_directed':float(np.mean(Cs)), 'train_nodes':int(train_idx.sum()),'baseline_error_all_and_test':base,'idea_error_all_and_test':idea}

if __name__=='__main__':
    seed_all(); result={'math_verification':math_verification(),'mini_experiment':mini_experiment()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))