import json, random, time
from itertools import combinations
import numpy as np

# Explicit orthogonality graph: vertices are representatives of nonzero {-1,0,1}^3 vectors.
V = np.array([(0,0,1),(0,1,-1),(0,1,0),(0,1,1),(1,-1,-1),(1,-1,0),
              (1,-1,1),(1,0,-1),(1,0,0),(1,0,1),(1,1,-1),(1,1,0),(1,1,1)], dtype=np.float32)
N = len(V)
E = [(i,j) for i in range(N) for j in range(i) if np.dot(V[i], V[j]) == 0]

# Small exact backtracking coloring check.
def colorable(k):
    adj = [set() for _ in range(N)]
    for i,j in E: adj[i].add(j); adj[j].add(i)
    c = [-1]*N
    def go(done):
        if done == N: return True
        u = max((i for i in range(N) if c[i] < 0), key=lambda i: len(adj[i]))
        for col in range(k):
            if all(c[w] != col for w in adj[u]):
                c[u] = col
                if go(done+1): return True
                c[u] = -1
        return False
    return go(0)

def graph_check():
    gram = V @ V.T
    edge_vals = np.array([gram[i,j] for i,j in E])
    # Unit-normalized vectors are the orthogonal-rank witness.
    U = V / np.linalg.norm(V, axis=1, keepdims=True)
    ug = U @ U.T
    violation = max(abs(ug[i,j]) for i,j in E)
    clique = max((len(C) for r in range(1,N+1) for C in combinations(range(N),r)
                  if all(((a,b) in E or (b,a) in E) for a,b in combinations(C,2))), default=0)
    return {'vertices':N, 'edges':len(E), 'explicit_dimension':3,
            'edge_dot_max_raw':float(max(abs(edge_vals))),
            'edge_violation_normalized':float(violation),
            'clique_lower_bound_xi':clique,
            'colorable_3':colorable(3), 'colorable_4':colorable(4)}

# Dataset: retain the first event through distractors and decide whether the final
# event is exclusive with it. This is a graph-defined relational language.
def make_data(num, T, seed):
    rng = np.random.default_rng(seed)
    # A learnable graph-walk-style promise language: the first event is the
    # context, fixed neutral distractors follow, and the final event is queried.
    # The target remains exactly edge compatibility with the initial event.
    x = np.zeros((num,T), dtype=np.int64)
    x[:,0] = rng.integers(0,N,size=num)
    x[:,-1] = rng.integers(0,N,size=num)
    edge = np.zeros((N,N), dtype=np.float32)
    for i,j in E: edge[i,j]=edge[j,i]=1
    y = 1.0 - edge[x[:,0], x[:,-1]]
    return x, y.astype(np.float32)


import torch
import torch.nn as nn

def train_one(name, width, orth_lambda, seed, device):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    class Recognizer(nn.Module):
        def __init__(self):
            super().__init__()
            self.emb = nn.Embedding(N,width)
            self.gru = nn.GRU(width,width,batch_first=True)
            self.out = nn.Linear(width,1)
        def forward(self,x):
            h,_ = self.gru(self.emb(x))
            return self.out(h[:,-1]).squeeze(-1)
        def orth_loss(self):
            u = self.emb.weight/(self.emb.weight.norm(dim=1,keepdim=True)+1e-8)
            dots = (u @ u.T)
            return sum(dots[i,j]**2 for i,j in E) / len(E)
    try:
        model=Recognizer().to(device)
        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        xb,yb=make_data(4000,16,seed+100)
        xt,yt=make_data(2000,64,seed+200)
        xb=torch.tensor(xb,device=device); yb=torch.tensor(yb,device=device)
        xt=torch.tensor(xt,device=device); yt=torch.tensor(yt,device=device)
        loss_fn=nn.BCEWithLogitsLoss()
        t0=time.time()
        for epoch in range(90):
            perm=torch.randperm(len(xb),device=device)
            for ix in perm.split(128):
                loss=loss_fn(model(xb[ix]),yb[ix])
                if orth_lambda: loss=loss+orth_lambda*model.orth_loss()
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            model.eval()
            def score(x,y):
                p=(torch.sigmoid(model(x))>.5).float()
                return float((p==y).float().mean()), float(loss_fn(model(x),y))
            a,l=score(xt,yt); tr,_=score(xb,yb)
            u=model.emb.weight/(model.emb.weight.norm(dim=1,keepdim=True)+1e-8)
            gv=max(abs(float((u[i]@u[j]).cpu())) for i,j in E)
        return {'name':name,'width':width,'train_acc':tr,'long_acc':a,'long_loss':l,
                'edge_violation':gv,'seconds':time.time()-t0,'params':sum(p.numel() for p in model.parameters())}
    except Exception as e:
        raise RuntimeError(f'{name}: {e}') from e

def main():
    check=graph_check()
    # CUDA is attempted but any runtime/device problem falls back to CPU.
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        if device=='cuda': torch.zeros(1,device='cuda')
    except Exception: device='cpu'
    results=[]
    # Standard width-4 GRU (classical color lower bound), width-3 controls,
    # and width-3 with the proposed edge orthogonality penalty. If CUDA has
    # any allocation/cuDNN failure, rerun the complete matched comparison CPU.
    configs=[('gru_width_chi',4,0.0),('gru_width_xi',3,0.0),('orthogonal_memory',3,3.0)]
    try:
        for args in configs: results.append(train_one(*args,seed=7,device=device))
    except Exception as exc:
        if device != 'cpu':
            device='cpu'; results=[]
            for args in configs: results.append(train_one(*args,seed=7,device=device))
        else: raise
    out={'device':device,'graph_check':check,'results':results,
         'note':'xi is only upper-bounded by the explicit 3D representation; clique size supplies a matching lower bound here.'}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()