import sys, time, torch, numpy as np
sys.path.insert(0,'/home/yurenh2/emm')
from scipy.optimize import linear_sum_assignment
np.set_printoptions(precision=4, suppress=True, linewidth=200)
def standardise(M):
    M=np.asarray(M,dtype=np.float64); mask=~np.eye(len(M),dtype=bool); v=M[mask]
    out=(M-v.mean())/v.std(); np.fill_diagonal(out,0.0); return out
d=torch.load('/home/yurenh2/emm/artifacts/synth_v1/omit_size.pt',map_location='cpu')
V0=d['visual_field'].double().numpy(); T0=d['text_field'].double().numpy()
V=standardise(V0); T=standardise(T0); N=len(V)
truth=np.arange(N); acc=lambda p: float((p==truth).mean())
CONST=float((T*T).sum()+(V*V).sum())
def energy_from_align(S): return (CONST-2.0*S)/(N*(N-1))
def energy(p): return energy_from_align(float((T[np.ix_(p,p)]*V).sum()))
print("E(truth) =",energy(truth))

print("\n### 2. SPECTRAL CERTIFICATES (lower bounds on E over all permutations)")
lV=np.linalg.eigvalsh(V)[::-1]; lT=np.linalg.eigvalsh(T)[::-1]
print("  Hoffman-Wielandt / von Neumann bound  E >= ", energy_from_align(float((lV*lT).sum())))
# projected eigenvalue bound (Hadley-Rendl-Wolkowicz): project out the all-ones direction
Q,_=np.linalg.qr(np.hstack([np.ones((N,1))/np.sqrt(N), np.random.default_rng(0).standard_normal((N,N-1))]))
Vp=Q[:,1:].T@V@Q[:,1:]; Tp=Q[:,1:].T@T@Q[:,1:]
lVp=np.linalg.eigvalsh(Vp)[::-1]; lTp=np.linalg.eigvalsh(Tp)[::-1]
sV=V.sum(); sT=T.sum(); rV=V.sum(1); rT=T.sum(1)
# <V,PTP'> = (1/N)*?; exact decomposition: with x=ones/sqrt(N),
# tr(V P T P^T) = tr(Vp Pp Tp Pp^T) + 2 x^T V P T P^T x*? -- use the standard HRW split
cross = float(rV.sum()*rT.sum())/ (N*N)  # x^T V x * x^T T x term
# rank-1 cross terms bounded by sorted-inner-product of centred row sums
cV=np.sort(rV-rV.mean())[::-1]; cT=np.sort(rT-rT.mean())[::-1]
proj_bound = float((lVp*lTp).sum()) + cross + 2.0*float((cV*cT).sum())/N
print("  projected (HRW-style, upper bd on alignment) E >= ", energy_from_align(proj_bound))
print("  (both are LOWER bounds on E; E(truth)=%.4f, solvers report E(truth)+0.44=%.4f)"%(energy(truth),energy(truth)+0.44))

print("\n### 3. BLIND rotation-invariant node descriptors -> Hungarian seeds")
def report(name,DV,DT,topk=(12,25,50)):
    DV=DV/ (np.linalg.norm(DV,axis=1,keepdims=True)+1e-12); DT=DT/(np.linalg.norm(DT,axis=1,keepdims=True)+1e-12)
    C=((DV**2).sum(1)[:,None]+(DT**2).sum(1)[None,:]-2*DV@DT.T)
    r,c=linear_sum_assignment(C)
    a=acc(c)
    # confidence = margin between assigned cost and 2nd best in row
    Cm=C.copy(); Cm[np.arange(N),c]=np.inf
    margin=Cm.min(1)-C[np.arange(N),c]
    order=np.argsort(-margin)
    prec={k: float((c[order[:k]]==order[:k]).mean()) for k in topk}
    print(f"  {name:34s} full-acc {a:.3f} | precision@top-margin {prec}")
    return c
# (a) sorted row profile
report("sorted row profile (V vs T)", np.sort(V,1), np.sort(T,1))
# (b) row moments
def mom(M,K=8): return np.stack([ (M**k).mean(1) for k in range(1,K+1)],1)
report("row power moments k=1..8", mom(V), mom(T))
# (c) heat kernel signature on normalised Laplacian of the raw non-negative Gram
def lap_eigs(W):
    W=np.clip(W,0,None).copy(); np.fill_diagonal(W,0.0)
    dg=W.sum(1); Dm=1/np.sqrt(np.maximum(dg,1e-12))
    L=np.eye(len(W))-(Dm[:,None]*W*Dm[None,:])
    w,U=np.linalg.eigh(L); return w,U,dg
wV,UVl,dV=lap_eigs(V0); wT,UTl,dT=lap_eigs(T0)
ts=np.logspace(-2,1.5,24)
HV=np.stack([ (np.exp(-t*wV)[None,:]*UVl**2).sum(1) for t in ts],1)
HT=np.stack([ (np.exp(-t*wT)[None,:]*UTl**2).sum(1) for t in ts],1)
report("HKS (normalised Laplacian)", np.log(HV+1e-12), np.log(HT+1e-12))
# (d) wave kernel signature
def wks(w,U,M=24):
    lw=np.log(np.maximum(w,1e-6)); e=np.linspace(lw.min(),lw.max(),M); sig=(e[1]-e[0])*2
    return np.stack([ (np.exp(-((e_-lw)**2)/(2*sig**2))[None,:]*U**2).sum(1) for e_ in e],1)
report("WKS (normalised Laplacian)", wks(wV,UVl), wks(wT,UTl))
# (e) spectral graph wavelet (SGWT) coefficient energies
def sgwt(w,U,S=10):
    ss=np.logspace(-1.5,1.0,S)
    return np.stack([ ((s*w*np.exp(-s*w))[None,:]*U**2).sum(1) for s in ss],1)
report("SGWT band energies", sgwt(wV,UVl), sgwt(wT,UTl))
# (f) degree only
report("degree (row sum) only", dV[:,None], dT[:,None])