import json, math, random, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED = 2062
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
except Exception:
    device = 'cpu'


def cycle(n):
    g=[set() for _ in range(n)]
    for i in range(n):
        for j in ((i-1)%n,(i+1)%n): g[i].add(j)
    return g

def grid(side):
    n=side*side; g=[set() for _ in range(n)]
    for r in range(side):
      for c in range(side):
        i=r*side+c
        for rr,cc in ((r-1,c),(r+1,c),(r,c-1),(r,c+1)):
          if 0<=rr<side and 0<=cc<side: g[i].add(rr*side+cc)
    return g

def random_regular(n, d=3):
    # deterministic retry via networkx if available; fallback ring chords
    try:
      import networkx as nx
      G=nx.random_regular_graph(d,n,seed=SEED+n)
      return [set(G.neighbors(i)) for i in range(n)]
    except Exception:
      g=[set() for _ in range(n)]
      for i in range(n):
        for k in range(1,d//2+1): g[i].update(((i-k)%n,(i+k)%n))
      return g

def tree(d, depth):
    g=[]; levels=[[0]]; g.append(set())
    next_id=1
    for dep in range(depth):
      nxt=[]
      for p in levels[-1]:
        count=d if dep==0 else d-1
        for _ in range(count):
          x=next_id; next_id+=1
          while len(g)<=x: g.append(set())
          g[p].add(x); g[x].add(p); nxt.append(x)
      levels.append(nxt)
    return g, levels

def expand(g, Fset):
    out=set(Fset)
    for x in Fset: out.update(g[x])
    return out

def ratio(g,Fset): return len(expand(g,Fset))/max(1,len(Fset))

def graph_matrix(g):
    n=len(g); A=torch.zeros((n,n),dtype=torch.float32)
    for i,ns in enumerate(g):
      A[i,i]=1
      for j in ns: A[i,j]=1
    A=A/A.sum(1,keepdim=True).clamp_min(1)
    return A

class GNN(nn.Module):
    def __init__(self,din,hidden,layers, gated=False, delta=.35):
      super().__init__(); self.gated=gated; self.delta=delta
      self.inp=nn.Linear(din,hidden); self.ws=nn.ModuleList([nn.Linear(hidden,hidden) for _ in range(layers)])
      self.out=nn.Linear(hidden,2)
    def forward(self,x,A,rbar):
      h=F.relu(self.inp(x)); gs=[]
      for w in self.ws:
        msg=A@h
        if self.gated:
          g=torch.sigmoid(8*((1+self.delta)-rbar))
          # pooled skip is a stable global context, weighted by non-message fraction
          pooled=h.mean(0,keepdim=True).expand_as(h)
          h=h + g*w(msg) + (1-g)*0.25*w(pooled)
          gs.append(float(g.detach().cpu()))
        else:
          h=h+w(msg); gs.append(1.0)
        h=F.relu(h)
      return self.out(h),gs

def mechanism_checks():
    # Prediction 1: closed one-hop interval on a cycle has r=1+2/m.
    cyc=cycle(200); cycle_rows=[]
    for m in [2,4,8,16,32,64,100]:
      Fset=set(range(m)); obs=ratio(cyc,Fset); pred=1+2/m
      cycle_rows.append({'m':m,'predicted':pred,'observed':obs,'abs_error':abs(pred-obs)})
    # Prediction 2: d=3 tree balls have ratios tending to 2; exact ball sweep.
    tr,levels=tree(3,7); ball=set()
    tree_rows=[]
    for k in range(0,6):
      ball.update(levels[k]); obs=ratio(tr,ball)
      pred=(1+3*(2**(k+1)-1)/2 if False else None)
      # exact finite tree ratio from current ball and next level
      exact=(len(ball)+len(levels[k+1]) if k+1<len(levels) else len(expand(tr,ball)))/len(ball)
      tree_rows.append({'radius':k,'size':len(ball),'predicted_exact':exact,'observed':obs})
    # Prediction 3: sigmoid gate boundary is r=1+delta, independent of slope.
    delta=.35; boundary_rows=[]
    for slope in [2,4,8,16]:
      rs=np.linspace(1.0,2.0,1001); gs=1/(1+np.exp(-slope*((1+delta)-rs)))
      boundary_rows.append({'slope':slope,'predicted_r':1+delta,'observed_r_at_g=.5':float(rs[np.argmin(abs(gs-.5))])})
    # Persistence prediction: EMA crosses threshold iff sustained r exceeds it.
    persistence=[]
    for r in [1.1,1.3,1.6,2.0]:
      beta=.8; rb=1.; crossing=None
      for t in range(1,31):
        rb=beta*rb+(1-beta)*r
        if crossing is None and rb>1+delta: crossing=t
      predicted = (math.log((1+delta-1)/(r-1))/math.log(beta) if r>1+delta else None)
      persistence.append({'r':r,'cross_step_observed':crossing,'cross_step_continuous_prediction':None if predicted is None else predicted})
    return {'cycle_boundary_prediction':cycle_rows,'tree_expansion_prediction':tree_rows,'gate_boundary_prediction':boundary_rows,'ema_persistence_prediction':persistence}

def train_case(name,g,epochs=60):
    n=len(g); A=graph_matrix(g).to(device)
    # fixed random features; label is a local structural signal (degree parity plus index hash)
    deg=np.array([len(x) for x in g]); x=np.random.default_rng(SEED+n).normal(size=(n,8)).astype('float32')
    x[:,0]=deg/ max(1,deg.max())
    y=torch.tensor(((deg + np.arange(n))%2),dtype=torch.long,device=device)
    perm=np.random.default_rng(SEED).permutation(n); cut=max(2,int(.7*n)); tridx=torch.tensor(perm[:cut],device=device); vaidx=torch.tensor(perm[cut:],device=device)
    # sampled seed frontier and a stable EMA estimate, as the controller would see it online
    seeds=set(perm[:min(16,n)]); Fset=seeds; rb=1.; beta=.7; rs=[]; frontiers=[]
    for _ in range(6):
      Fset=expand(g,Fset); r=len(Fset)/max(1,len(seeds)); rb=beta*rb+(1-beta)*r; rs.append(rb); frontiers.append(len(Fset))
    rbar=torch.tensor(float(rb),device=device)
    result={}
    for gated in [False,True]:
      torch.manual_seed(SEED+int(gated)+n)
      model=GNN(8,24,6,gated=gated).to(device); opt=torch.optim.Adam(model.parameters(),lr=.025)
      t0=time.time()
      for _ in range(epochs):
        opt.zero_grad(); logits,_=model(torch.tensor(x,device=device),A,rbar)
        loss=F.cross_entropy(logits[tridx],y[tridx]); loss.backward(); opt.step()
      with torch.no_grad():
        logits,gs=model(torch.tensor(x,device=device),A,rbar); acc=(logits[vaidx].argmax(1)==y[vaidx]).float().mean().item()
        h=logits[vaidx]; cos=F.cosine_similarity(h[:,None,:],h[None,:,:],dim=-1).mean().item() if len(vaidx)>1 else 0
      result['gated' if gated else 'baseline']={'val_accuracy':acc,'mean_pairwise_cosine':cos,'seconds':time.time()-t0,'gate_mean':float(np.mean(gs))}
    result.update({'rbar':float(rb),'frontier_peak':max(frontiers),'raw_ratios':[float(x) for x in rs]})
    return result

def main():
    checks=mechanism_checks(); cases={}
    graphs={'cycle':cycle(100),'grid':grid(10),'random3':random_regular(100,3)}
    tr,_=tree(3,5); graphs['tree3']=tr
    for name,g in graphs.items():
      try: cases[name]=train_case(name,g)
      except Exception as e: cases[name]={'error':repr(e)}
    out={'seed':SEED,'device':device,'checks':checks,'experiments':cases}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()