import math, json, time
import numpy as np
from scipy.optimize import linear_sum_assignment
import torch
import torch.nn as nn
import torch.nn.functional as F

SEED=2140
np.random.seed(SEED); torch.manual_seed(SEED)

def sinkhorn(C, a=None, b=None, eps=0.1, iters=60):
    C=np.asarray(C,float); n,m=C.shape
    if a is None: a=np.ones(n)/n
    if b is None: b=np.ones(m)/m
    a=np.asarray(a); b=np.asarray(b)
    eps=max(float(eps),1e-10)
    logK=-C/eps
    la=np.log(a); lb=np.log(b)
    u=np.zeros(n); v=np.zeros(m)
    for _ in range(iters):
        u=la-np.logaddexp.reduce(logK+v[None,:],axis=1)
        v=lb-np.logaddexp.reduce(logK+u[:,None],axis=0)
    P=np.exp(logK+u[:,None]+v[None,:])
    return P

def lifted_cost(W,V,P):
    C=(W[:,:,None,None]-V[None,None,:,:])**2
    return float(np.einsum('ik,jl,ijkl->',P,P,C,optimize=True))

def quotient_sinkhorn(W,V,eps_scale=.03):
    n=W.shape[0]; # node-pair costs
    C=np.zeros((n,n))
    # quadratic assignment cost used to obtain a node coupling; this is the
    # standard lifted objective evaluated approximately by alternating updates.
    P=np.ones((n,n))/n
    for _ in range(5):
        C=np.einsum('jl,ikjl->ik',P,(W[:,:,None,None]-V[None,None,:,:])**2)*n
        med=np.median(C[C>1e-12]) if np.any(C>1e-12) else 1.
        P=sinkhorn(C,eps=max(eps_scale*med,1e-5),iters=80)
    return P, lifted_cost(W,V,P)

def make_graph(n=6):
    x=np.random.rand(n,n); return (x+x.T)/2

def toy_checks():
    W=make_graph(); perm=np.random.permutation(len(W)); V=W[np.ix_(perm,perm)]
    # identity and optimized lifted costs
    I=np.eye(len(W))/len(W)
    fixed=lifted_cost(W,V,I)
    inv=np.argsort(perm); Prelabel=np.zeros_like(I); Prelabel[np.arange(len(W)),inv]=1/len(W)
    exact_relabel=lifted_cost(W,V,Prelabel)
    P,q=quotient_sinkhorn(W,V)
    marg=max(np.max(abs(P.sum(1)-1/len(W))),np.max(abs(P.sum(0)-1/len(W))))
    # noise scaling prediction: d_Q is proportional to sigma for p=2
    sigmas=np.array([0,.01,.03,.06,.12,.24])
    vals=[]
    for s in sigmas:
        reps=[]
        for r in range(4):
            N=np.random.randn(*W.shape); N=(N+N.T)/2
            Vn=W+float(s)*N
            # Compare the unpermuted graph to its noisy version under identity.
            reps.append(math.sqrt(lifted_cost(W,Vn,I)))
        vals.append(np.mean(reps))
    slope=float(np.dot(sigmas,vals)/max(np.dot(sigmas,sigmas),1e-12))
    pred_intercept=vals[0]
    # regularization prediction: entropy rises with epsilon
    base=(W[:,:,None,None]-V[None,None,:,:])**2
    med=np.median(base)
    epses=med*np.array([.003,.01,.03,.1,.3,1.0])
    ents=[]; costs=[]
    for e in epses:
        # pairwise node cost from uniform lifted current iterate
        P=np.ones((len(W),len(W)))/len(W)
        for _ in range(5):
            C=np.einsum('jl,ikjl->ik',P,base)*len(W)
            P=sinkhorn(C,eps=e,iters=80)
        ent=float(-(P[P>0]*np.log(P[P>0])).sum())
        ents.append(ent); costs.append(lifted_cost(W,V,P))
    assert lifted_cost(W,W,I) < 1e-12
    return dict(identity_zero=math.sqrt(lifted_cost(W,W,I)), exact_relabel_root=math.sqrt(exact_relabel), fixed_root=math.sqrt(fixed), quotient_root=math.sqrt(q), marginal_error=marg,
                sigmas=sigmas.tolist(), noise_roots=np.array(vals).tolist(), noise_slope=slope,
                entropy_eps=epses.tolist(), entropies=ents, reg_costs=costs)

class OrderEquivariantNet(nn.Module):
    def __init__(self,n):
        super().__init__(); self.node=nn.Sequential(nn.Linear(n,24),nn.ReLU(),nn.Linear(24,12),nn.ReLU()); self.cls=nn.Linear(12,2)
    def forward(self,A):
        one = A.ndim == 2
        if one: A = A.unsqueeze(0)
        H=self.node(A) # each row sees its relational neighborhood
        out=self.cls(H.mean(1))
        return (out[0] if one else out), H

def data(n=8, count=100):
    As=[]; ys=[]
    for y in range(2):
        for _ in range(count//2):
            p=.25 if y==0 else .65
            X=(np.random.rand(n,n)<p).astype('float32'); X=np.triu(X,1); X=X+X.T
            np.fill_diagonal(X,0); As.append(X); ys.append(y)
    ix=np.random.permutation(len(ys)); return np.array(As)[ix],np.array(ys)[ix]

def train_compare():
    np.random.seed(SEED); torch.manual_seed(SEED)
    A,y=data(); n=A.shape[1]; cut=int(.7*len(A))
    At=torch.tensor(A[:cut]); yt=torch.tensor(y[:cut]); Av=torch.tensor(A[cut:]); yv=torch.tensor(y[cut:])
    def run(align):
        torch.manual_seed(SEED+int(align)); model=OrderEquivariantNet(n); opt=torch.optim.Adam(model.parameters(),lr=.01)
        t0=time.time()
        for step in range(35):
            ix=torch.randperm(cut)[:16]; X=At[ix]; target=yt[ix]
            if align:
                X2=[]; Ps=[]
                for x in X.numpy():
                    p=np.random.permutation(n); X2.append(x[np.ix_(p,p)])
                    # coupling is computed between relational rows, with detached Sinkhorn
                    C=((x[:,None,:]-X2[-1][None,:,:])**2).mean(2)
                    Ps.append(sinkhorn(C,eps=.1*max(np.median(C),1e-3),iters=35))
                X2=torch.tensor(np.array(X2),dtype=torch.float32)
                logits,H=model(X); logits2,H2=model(X2)
                align_loss=0.0
                for h1,h2,P in zip(H,H2,Ps):
                    Pt=torch.tensor(P,dtype=h1.dtype)
                    align_loss=align_loss+((h1[:,None,:]-h2[None,:,:])**2*Pt[:,:,None]).sum()
                align_loss=align_loss/len(H)
                loss=F.cross_entropy(logits,target)+.15*F.mse_loss(logits,logits2)+.05*align_loss
            else:
                logits,H=model(X); loss=F.cross_entropy(logits,target)
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=model(Av)[0].argmax(1); acc=float((pred==yv).float().mean())
            # prediction variance over 12 arbitrary node permutations
            vars=[]
            for x in Av[:20]:
                ls=[]
                for _ in range(5):
                    p=torch.randperm(n); ls.append(torch.softmax(model(x[p][:,p])[0],0)[1].item())
                vars.append(np.var(ls))
        return dict(acc=acc,perm_prob_variance=float(np.mean(vars)),seconds=time.time()-t0)
    return run(False),run(True)

def main():
    toy=toy_checks(); base,idea=train_compare()
    out={'toy':toy,'baseline':base,'idea':idea}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()