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

SEED = 421
np.random.seed(SEED)
torch.manual_seed(SEED)


def make_graph(n=600, k=3, p_in=.16, p_out=.035):
    labels = np.arange(n) % k
    rng = np.random.default_rng(SEED)
    A = np.zeros((n,n), dtype=np.uint8)  # A[u,v] means u -> v
    for u in range(n):
        same = labels == labels[u]
        probs = np.where(same, p_in, p_out)
        row = rng.random(n) < probs
        row[u] = False
        A[u] = row
    # ensure all vertices have a modest amount of incoming structure
    return A, labels


def signatures(A, landmarks):
    return A[np.asarray(landmarks), :].T.astype(np.float32)


def collision_count(Z):
    return int(Z.shape[0] - np.unique(Z, axis=0).shape[0])


def pairwise_c(A):
    # exact minimum symmetric difference of incoming neighborhoods
    B = A.T.astype(np.uint8)
    c = n_pairs = 10**9
    for i in range(len(B)):
        d = np.sum(np.bitwise_xor(B[i+1:], B[i]), axis=1) if i+1 < len(B) else np.array([])
        if len(d): c = min(c, int(d.min()))
    return c


def greedy_landmarks(A, s):
    n = len(A)
    incoming = A.T.astype(bool)
    # distinguish all pairs; represent only currently colliding pairs
    Z = np.zeros((n, 0), dtype=np.uint8)
    chosen = []
    remaining = set(range(n))
    for _ in range(min(s,n)):
        if collision_count(Z) == 0: break
        # Candidate score: number of currently equal-signature pairs split by q.
        groups = {}
        for i, row in enumerate(Z): groups.setdefault(row.tobytes(), []).append(i)
        best, bestscore = None, -1
        for q in remaining:
            score = 0
            bit = incoming[:,q]
            for inds in groups.values():
                if len(inds) > 1:
                    score += int(np.sum(bit[inds]) * (len(inds)-np.sum(bit[inds])))
            if score > bestscore: best, bestscore = q, score
        if best is None or bestscore <= 0: break
        chosen.append(best); remaining.remove(best)
        Z = np.column_stack([Z, incoming[:,best].astype(np.uint8)])
    return chosen


class MLP(nn.Module):
    def __init__(self, d, hidden=48, k=3):
        super().__init__(); self.net=nn.Sequential(nn.Linear(d,hidden),nn.ReLU(),nn.Linear(hidden,k))
    def forward(self,x): return self.net(x)


def train_eval(X, y, epochs=100):
    # Same fixed classifier for every representation, with fixed split and seed.
    torch.manual_seed(SEED)
    n=len(y); perm=np.random.default_rng(SEED).permutation(n)
    tr=perm[:int(.7*n)]; te=perm[int(.7*n):]
    model=MLP(X.shape[1], k=int(y.max()+1))
    opt=torch.optim.Adam(model.parameters(), lr=.015, weight_decay=1e-4)
    xt=torch.tensor(X); yt=torch.tensor(y, dtype=torch.long)
    t=time.perf_counter()
    for _ in range(epochs):
        opt.zero_grad(); loss=nn.functional.cross_entropy(model(xt[tr]),yt[tr]); loss.backward(); opt.step()
    with torch.no_grad():
        pred=model(xt[te]).argmax(1); acc=float((pred==yt[te]).float().mean())
    return acc, float(loss), time.perf_counter()-t


def math_check(A, s, draws=4000):
    n=len(A); incoming=A.T.astype(bool)
    rng=np.random.default_rng(SEED+1)
    # Select a representative pair and also evaluate all-pair aggregate prediction.
    pairs=[(0,1),(0,n//2),(1,n//2)]
    empirical=[]; exact=[]
    for u,v in pairs:
        diff=int(np.sum(np.logical_xor(incoming[u],incoming[v])))
        p=math.comb(n-diff,s)/math.comb(n,s) if n-diff>=s else 0.0
        hit=0
        for _ in range(draws):
            S=rng.choice(n,s,replace=False)
            if not np.any(incoming[u,S] != incoming[v,S]): hit+=1
        empirical.append(hit/draws); exact.append(p)
    # Aggregate exact bound and Monte Carlo mean collision count.
    c=pairwise_c(A)
    bound=math.comb(n,2)*(math.comb(n-c,s)/math.comb(n,s) if n-c>=s else 0.0)
    obs=[]
    for _ in range(200):
        S=rng.choice(n,s,replace=False); obs.append(collision_count(incoming[:,S].astype(np.uint8)))
    return {'n':n,'s':s,'c_exact':c,'pair_exact':exact,'pair_empirical':empirical,
            'expected_collision_upper_bound':bound,'sampled_mean_collisions':float(np.mean(obs)),
            'max_pair_abs_error':float(max(abs(a-b) for a,b in zip(empirical,exact)))}


def main():
    A,y=make_graph(); n=len(y)
    s=max(4, int(math.ceil(2*math.sqrt(n))))
    check=math_check(A,s)
    rng=np.random.default_rng(SEED+2)
    deg=A.sum(axis=1)+A.sum(axis=0)
    methods={
      'uniform':rng.choice(n,s,replace=False).tolist(),
      'degree':np.argsort(-deg)[:s].tolist(),
      'greedy':greedy_landmarks(A,s)
    }
    # noisy features are deliberately insufficient alone; signatures carry block structure.
    X0=rng.normal(0,1,(n,8)).astype(np.float32)
    results={}
    for name,S in methods.items():
        t=time.perf_counter(); Z=signatures(A,S); prep=time.perf_counter()-t
        # x + landmark binary positional signature, exactly the proposed interface
        acc,loss,train=train_eval(np.concatenate([X0,Z],axis=1),y)
        results[name]={'s':len(S),'collisions':collision_count(Z),'collision_rate':collision_count(Z)/n,
                       'accuracy':acc,'final_train_loss':loss,'preprocess_sec':prep,'train_sec':train}
    # Standard control: one-step full incoming-neighbor mean as graph-to-token interface.
    t=time.perf_counter(); Xfull=np.concatenate([X0, A.T.astype(np.float32)@X0/max(1,A.sum(axis=0).max())],axis=1); prep=time.perf_counter()-t
    acc,loss,train=train_eval(Xfull,y)
    results['full_1hop']={'features':Xfull.shape[1],'collisions':'n/a','accuracy':acc,'final_train_loss':loss,'preprocess_sec':prep,'train_sec':train}
    out={'seed':SEED,'graph':{'n':n,'edges':int(A.sum()),'classes':3},'math_check':check,'methods':results}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()