import sys, numpy as np, ot, torch
sys.path.insert(0,'/home/yurenh2/emm')
from scipy.optimize import linear_sum_assignment
from worldalign.synth_fast_gate import ClosedFormEnergy, fast_pair_descent
from worldalign.synth_triangle_gate import standardized
DEV=torch.device("cuda:1")
def standardise(m):
    mask=~np.eye(len(m),dtype=bool); v=m[mask]
    o=(m-v.mean())/v.std(); np.fill_diagonal(o,0.0); return o
def tlb(V,S):
    n=len(V); a=np.sort(V,1); b=np.sort(S,1)
    return (a*a).sum(1)[:,None]/n+(b*b).sum(1)[None,:]/n-2*(a@b.T)/n
st=torch.load("/home/yurenh2/emm/artifacts/synth_v1/omit_size.pt",map_location='cpu',weights_only=False)
V=standardise(st["visual_field"].double().numpy()); T=standardise(st["text_field"].double().numpy())
n=len(V); w=ot.unif(n)
rng=np.random.default_rng(0); hid=rng.permutation(n); S=T[np.ix_(hid,hid)]
Vg=standardized(torch.from_numpy(V).to(DEV)).double(); Sg=standardized(torch.from_numpy(S).to(DEV)).double()
en=ClosedFormEnergy(Sg,Vg,1.0,0.0,64)
M=tlb(V,S); M=M/M.max()
G=np.outer(w,w)
for a in (0.0,0.2):
    G=ot.gromov.fused_gromov_wasserstein(M,V,S,w,w,"square_loss",alpha=a,G0=G,max_iter=200,tol_rel=1e-9)
_,cols=linear_sum_assignment(-G)
fin=fast_pair_descent(Sg,Vg,torch.from_numpy(np.ascontiguousarray(cols)).to(DEV),2000).cpu().numpy()
truth=np.argsort(hid)
Ef=float(en.energy(torch.from_numpy(np.ascontiguousarray(fin)).to(DEV)[None])[0])
Et=float(en.energy(torch.from_numpy(np.ascontiguousarray(truth)).to(DEV)[None])[0])
acc=float((hid[fin]==np.arange(n)).mean())
print("accuracy=%.4f  E_found=%.10f  E_truth=%.10f  diff=%.3e"%(acc,Ef,Et,Ef-Et))
wrong=np.where(hid[fin]!=np.arange(n))[0]
print("n_wrong=%d"%len(wrong))
# for each wrong scene i: correlation between S-row of assigned node and S-row of true node
Sal=S  # in shuffled index space
corr=[]
vcorr=[]
for i in wrong:
    a_idx=fin[i]; t_idx=truth[i]
    x=np.delete(Sal[a_idx],[a_idx,t_idx]); y=np.delete(Sal[t_idx],[a_idx,t_idx])
    corr.append(np.corrcoef(x,y)[0,1])
    u=np.delete(V[i],[i]); 
    # visual side: is the visually-assigned scene similar to i? compare V rows of i and of the scene truly at a_idx
    j=hid[a_idx]
    v1=np.delete(V[i],[i,j]); v2=np.delete(V[j],[i,j])
    vcorr.append(np.corrcoef(v1,v2)[0,1])
rng2=np.random.default_rng(7); pairs=rng2.integers(0,n,(len(wrong),2))
base=[np.corrcoef(np.delete(Sal[p],[p,q]),np.delete(Sal[q],[p,q]))[0,1] for p,q in pairs if p!=q]
print("text-row corr of confused partners: median=%.3f mean=%.3f  | random pair baseline median=%.3f"%(np.median(corr),np.mean(corr),np.median(base)))
print("visual-row corr of confused partners: median=%.3f"%np.median(vcorr))
# cycle structure of the error
perm=hid[fin]
seen=set(); cyc=[]
for i in range(n):
    if i in seen or perm[i]==i: continue
    c=0; j=i
    while j not in seen:
        seen.add(j); j=perm[j]; c+=1
    cyc.append(c)
import collections
print("error cycle lengths:", dict(collections.Counter(cyc)))