import math, time, itertools, json, random
import numpy as np


def random_regular(n, d, seed=0):
    if d < 0 or d >= n: raise ValueError("degree must satisfy 0 <= d < n")
    rng=np.random.default_rng(seed)
    if (n*d) % 2: raise ValueError("n*d must be even")
    # A randomly relabeled circulant is simple and exactly d-regular.
    # For odd d, add a perfect matching (requiring even n).
    if d % 2 and n % 2: raise ValueError("odd degree requires even n")
    adj=[set() for _ in range(n)]
    half=d//2
    for i in range(n):
        for z in range(1,half+1):
            adj[i].add((i+z)%n); adj[i].add((i-z)%n)
    if d % 2:
        perm=rng.permutation(n)
        for a,b in perm.reshape(-1,2):
            adj[int(a)].add(int(b)); adj[int(b)].add(int(a))
    perm=rng.permutation(n); inv=np.empty(n,dtype=int)
    for i,x in enumerate(perm): inv[x]=i
    out=[set() for _ in range(n)]
    for i in range(n):
        out[int(perm[i])]={int(perm[j]) for j in adj[i]}
    return [sorted(x) for x in out]

def ring(n, d):
    adj=[set() for _ in range(n)]
    for i in range(n):
        for z in range(1,d//2+1): adj[i].update(((i-z)%n,(i+z)%n))
    return [sorted(x) for x in adj]

def boundary(adj, U):
    us=set(U); out=set()
    for i in U:
        out.update(j for j in adj[i] if j not in us)
    return len(out)

def exact_stats(adj, k, eps):
    n=len(adj); rows=[]
    for s in range(k,n//2+1):
        vals=[]
        for U in itertools.combinations(range(n),s):
            b=boundary(adj,U); bound=eps*s/(math.log(3*s/k)**2)
            vals.append((b/s,b,bound))
        a=np.array(vals)
        rows.append((s,float(a[:,0].min()),float(a[:,0].mean()),float(a[:,0].max()),float(a[:,2].max())))
    return rows

def sampled_stats(adj,k,eps,seed=1,samples=1000):
    rng=np.random.default_rng(seed); n=len(adj); rows=[]
    for s in [k, min(2*k,n//2), min(4*k,n//2), n//2]:
        vals=[]
        for _ in range(samples):
            U=rng.choice(n,s,replace=False); b=boundary(adj,U)
            vals.append((b/s, eps/(math.log(3*s/k)**2)))
        rows.append((s,float(np.min([x[0] for x in vals])),float(np.mean([x[0] for x in vals])),float(np.mean([x[1] for x in vals]))))
    return rows

def bfs_depth(adj, seedset, target):
    seen=set(seedset); frontier=set(seedset); depth=0; history=[len(seen)]
    while len(seen)<target and frontier and depth<100:
        nxt=set()
        for i in frontier: nxt.update(adj[i])
        nxt-=seen; seen |= nxt; frontier=nxt; depth+=1; history.append(len(seen))
    return depth,history

def torch_benchmark(n=512,h=64,d=16,bs=8,iters=30):
    import torch
    device='cuda' if torch.cuda.is_available() else 'cpu'
    def run(kind):
        q=torch.randn(bs,n,h,device=device); k=torch.randn_like(q); v=torch.randn_like(q)
        if kind=='dense':
            fn=lambda: torch.softmax(torch.matmul(q,k.transpose(1,2))/math.sqrt(h),-1).matmul(v)
        else:
            adj=ring(n,d) if kind=='local' else random_regular(n,d,9)
            ii=np.repeat(np.arange(n),d); jj=np.array([j for row in adj for j in row])
            ti=torch.tensor(ii,device=device); tj=torch.tensor(jj,device=device)
            def fn():
                scores=(q[:,ti,:]*k[:,tj,:]).sum(-1)/math.sqrt(h)
                scores=scores.view(bs,n,d); att=torch.softmax(scores,-1)
                return (att.unsqueeze(-1)*v[:,tj,:].view(bs,n,d,h)).sum(2)
        for _ in range(5): fn()
        if device=='cuda': torch.cuda.synchronize()
        t=time.perf_counter()
        for _ in range(iters): fn()
        if device=='cuda': torch.cuda.synchronize()
        return (time.perf_counter()-t)/iters*1000
    try: return device,{x:run(x) for x in ['dense','local','expander']}
    except Exception as e:
        return 'cpu-fallback',{'error':str(e)}

def expansion_constant(adj, k):
    n=len(adj); vals=[]
    for size in range(k,n//2+1):
        for U in itertools.combinations(range(n),size):
            vals.append((boundary(adj,U)/size)*math.log(3*size/k)**2)
    return min(vals)

def max_exact_violation(adj, k, eps):
    n=len(adj); worst=0.0
    for size in range(k,n//2+1):
        bound=eps/(math.log(3*size/k)**2)
        for U in itertools.combinations(range(n),size):
            worst=max(worst, bound-boundary(adj,U)/size)
    return max(0.0,worst)

def predicted_steps(eps, k, target):
    x=float(k); steps=0
    while x < target and steps < 10000:
        x *= 1 + eps/(math.log(3*x/k)**2)
        steps += 1
    return steps, x

def bfs_to_target(adj, seedset, target):
    seen=set(seedset); frontier=set(seedset); depth=0
    while len(seen)<target and frontier and depth<10000:
        nxt=set()
        for i in frontier: nxt.update(adj[i])
        nxt-=seen; seen |= nxt; frontier=nxt; depth+=1
    return depth, len(seen)

def main():
    np.random.seed(0); random.seed(0)
    n=16; k=2
    expander=random_regular(n,4,3); ring_graph=ring(n,4)
    # The largest epsilon for which the definition holds is the minimum
    # of observed boundary_ratio * log^2(3|U|/k).
    c_exp=expansion_constant(expander,k); c_ring=expansion_constant(ring_graph,k)
    eps=0.8*c_exp
    exact=exact_stats(expander,k,eps)
    max_violation=max_exact_violation(expander,k,eps)
    # Prediction 1: every exact subset obeys the bound at eps <= epsilon_max.
    epsilon_scaling=[]
    for frac in [0.25,0.5,0.75,1.0,1.1]:
        e=frac*c_exp; rows=exact_stats(expander,k,e)
        viol=max_exact_violation(expander,k,e)
        epsilon_scaling.append({'fraction_of_max':frac,'epsilon':e,'max_violation':viol})
    # Prediction 2: better-mixed expander has stronger normalized expansion
    # than a same-degree local ring.
    family=[('ring',ring_graph),('random_regular',expander)]
    family_stats=[]
    for name,g in family:
        rows=sampled_stats(g,k,eps,samples=600)
        family_stats.append({'graph':name,'epsilon_max':expansion_constant(g,k),
                             'sampled_rows':rows})
    # Prediction 3: increasing epsilon in the theoretical recurrence decreases
    # the predicted number of layers to reach n/2; compare with BFS.
    growth=[]
    for d in [2,4,6,8]:
        gg=random_regular(64,d,100+d); cm=min(x[1] * math.log(3*x[0]/2)**2 for x in sampled_stats(gg,2,1.0,samples=1200))
        vals=[]
        for frac in [0.25,0.5,0.75,1.0]:
            e=frac*cm; pred,_=predicted_steps(e,2,32)
            obs,reached=bfs_to_target(gg,[0,1],32)
            vals.append({'fraction_of_max':frac,'predicted_layers_to_32':pred,
                         'observed_layers_to_32':obs,'reached':reached})
        growth.append({'d':d,'epsilon_max':cm,'sweep':vals})
    bench=torch_benchmark()
    result={'epsilon_max_expander':c_exp,'epsilon_max_ring':c_ring,
      'epsilon_used':eps,'exact_expansion_rows':exact,
      'max_bound_minus_observed':max_violation,
      'epsilon_scaling':epsilon_scaling,'graph_family':family_stats,
      'growth_sweep':growth,'attention_ms':bench[1],'device':bench[0]}
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()