import json, time
import numpy as np


def softmax(x):
    z = x - np.max(x, axis=-1, keepdims=True)
    e = np.exp(z)
    return e / e.sum(axis=-1, keepdims=True)


def init_probs(n, beta, seed=0, noise=0.15):
    rng = np.random.default_rng(seed)
    logits = 2*np.asarray(beta, float)[None, None, :] + noise*rng.normal(size=(n,n,3))
    logits = (logits + logits.transpose(1,0,2))/2
    p = softmax(logits)
    for i in range(n): p[i,i] = 1/3
    return p


def refine(p, beta, beta4, tau, iterations=8, damping=1.0):
    n = p.shape[0]
    beta = np.asarray(beta, float)
    q = p.copy()
    for _ in range(iterations):
        for i in range(n):
            for j in range(i+1, n):
                r = np.zeros(3)
                for k in range(n):
                    if k == i or k == j:
                        continue
                    for a in range(3):
                        o = [x for x in range(3) if x != a]
                        r[a] += q[i,k,o[0]]*q[j,k,o[1]] + q[i,k,o[1]]*q[j,k,o[0]]
                s = 2*beta + (beta4/n)*r
                new = softmax((s/tau)[None, :])[0]
                q[i,j] = q[j,i] = (1-damping)*q[i,j] + damping*new
    return q


def rainbow_density(p):
    n = p.shape[0]; total = 0.0; count = 0
    for i in range(n):
        for j in range(i+1,n):
            for k in range(j+1,n):
                x = 0.0
                for a in range(3):
                    o = [b for b in range(3) if b != a]
                    x += p[i,j,a]*(p[i,k,o[0]]*p[j,k,o[1]] + p[i,k,o[1]]*p[j,k,o[0]])
                total += x; count += 1
    return total/count


def entropy(p):
    return float((-p*np.log(np.maximum(p, 1e-12))).sum(axis=-1).mean())


def run():
    beta = np.array([0.10, -0.03, -0.07])
    out = {'math_checks': {}, 'mechanism': [], 'scaling': [], 'proxy': {}}
    # beta4=0 must make refinement identical to the unary router.
    p0 = init_probs(7, beta, seed=4)
    pzero = refine(p0, beta, 0.0, 0.7, iterations=5)
    unary = softmax((2*beta/0.7)[None, :])[0]
    out['math_checks']['beta4_zero_max_change_from_unary'] = float(np.max(np.abs(pzero[0,1]-unary)))
    out['math_checks']['simplex_max_error'] = float(np.max(np.abs(pzero.sum(-1)-1)))
    # Positive coupling response, using a symmetric unbiased initialization.
    for b4 in [0., 0.5, 1., 2., 4., 8.]:
        p = init_probs(8, np.zeros(3), seed=11, noise=0.0)
        q = refine(p, np.zeros(3), b4, 0.35, iterations=12, damping=0.7)
        out['mechanism'].append({'beta4': b4, 'rainbow_density': rainbow_density(q), 'entropy': entropy(q), 'max_p': float(q.max())})
    # Claimed 1/n scaling: compare motif-logit contribution at two graph sizes.
    for n in [6, 12, 20]:
        p = init_probs(n, np.zeros(3), seed=3, noise=0.0)
        # one explicit update's motif score magnitude
        rvals=[]
        for i in range(n):
            for j in range(i+1,n):
                r=np.zeros(3)
                for k in range(n):
                    if k in (i,j): continue
                    for a in range(3):
                        o=[x for x in range(3) if x!=a]
                        r[a]+=p[i,k,o[0]]*p[j,k,o[1]]+p[i,k,o[1]]*p[j,k,o[0]]
                rvals.append(np.max(np.abs(1.0*r/n)))
        out['scaling'].append({'n': n, 'median_normalized_score': float(np.median(rvals))})
    # Small fixed classification proxy: target is a noisy unary color; compare unary vs refined routing.
    rng=np.random.default_rng(21); n=10; d=5; samples=160
    X=rng.normal(size=(samples,n,d)); true_w=rng.normal(size=(d,3)); y=np.argmax(X@true_w,axis=2)
    # A router predicts the edge color from endpoint features; evaluate held-out edges.
    split=120; losses={'baseline':[],'idea':[]}; t0=time.perf_counter()
    for s in range(samples):
        logits=np.zeros((n,n,3))
        for i in range(n):
            for j in range(i+1,n):
                logits[i,j]=logits[j,i]=(X[s,i]+X[s,j])@true_w/2
        p=softmax(logits)
        if s>=split: losses['baseline'].append(-np.log(np.maximum(p[np.arange(n), (np.arange(n)+1)%n, y[s,np.arange(n)]],1e-12)).mean())
        q=refine(p, np.zeros(3), 2.0, 0.5, iterations=2, damping=0.5)
        if s>=split: losses['idea'].append(-np.log(np.maximum(q[np.arange(n), (np.arange(n)+1)%n, y[s,np.arange(n)]],1e-12)).mean())
    out['proxy']={'baseline_test_nll':float(np.mean(losses['baseline'])), 'idea_test_nll':float(np.mean(losses['idea'])), 'seconds':time.perf_counter()-t0}
    return out

if __name__ == '__main__':
    print(json.dumps(run(), indent=2))