"""Decisive ORACLE test of the symNMF de-rotation / component-pruning candidate.

If the ORACLE version (perfect column match, perfect pruning) does not beat the
current blind state of the art (~0.83 on omit-size), the candidate is dead.
"""
import numpy as np, torch, time, warnings
warnings.filterwarnings('ignore')
from scipy.optimize import linear_sum_assignment

dev = torch.device('cuda:0')
DT = torch.float64

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(); N = len(V0)
rng = np.random.default_rng(777); sigma = rng.permutation(N)
T0s = T0[np.ix_(sigma, sigma)]; truth = np.argsort(sigma)
acc = lambda p: float((p == truth).mean())


def standardise(M):
    M = np.asarray(M, float); mask = ~np.eye(len(M), dtype=bool); v = M[mask]
    out = (M - v.mean()) / v.std(); np.fill_diagonal(out, 0.0); return out


def hung(A, B):
    C = ((A**2).sum(1)[:, None] + (B**2).sum(1)[None, :] - 2 * A @ B.T)
    r, c = linear_sum_assignment(C); return c


def sym_nmf_gpu(M, r, iters=3000, seed=0):
    Mg = torch.tensor(np.clip(M, 0, None), device=dev, dtype=DT)
    g = torch.Generator(device='cpu').manual_seed(seed)
    W = torch.abs(torch.randn(len(M), r, generator=g, dtype=DT)).to(dev) * float(np.sqrt(M.mean() / r))
    for _ in range(iters):
        W = W * (0.5 + 0.5 * (Mg @ W) / (W @ (W.T @ W) + 1e-12))
    return W.cpu().numpy()


# ---- energy + exact steepest descent (closed-form swap table) ----
Vs = standardise(V0); Ts = standardise(T0s)
CONST = float((Vs * Vs).sum() + (Ts * Ts).sum()); DEN = N * (N - 1)
Ag = torch.tensor(Vs, device=dev, dtype=DT); Tg = torch.tensor(Ts, device=dev, dtype=DT)


def energy(p):
    return (CONST - 2.0 * float((Ts[np.ix_(p, p)] * Vs).sum())) / DEN


def descend(p0, max_steps=4000, A=None, B=None):
    A = Ag if A is None else A
    Bfull = Tg if B is None else B
    p = torch.tensor(np.asarray(p0), device=dev, dtype=torch.long)
    iu = torch.triu_indices(N, N, offset=1, device=dev)
    for _ in range(max_steps):
        B = Bfull[p][:, p]
        C = A @ B
        dg = torch.diagonal(C)
        G = C + C.T - dg[:, None] - dg[None, :] + 2 * A * B
        vals = G[iu[0], iu[1]]
        k = int(vals.argmax())
        if float(vals[k]) <= 1e-12:
            break
        u, v = int(iu[0][k]), int(iu[1][k])
        p[u], p[v] = p[v].clone(), p[u].clone()
    return p.cpu().numpy()


def grampa(Vm, Tm, eta=0.2):
    lam, u = np.linalg.eigh(Vm); mu, v = np.linalg.eigh(Tm)
    ones = np.ones(len(Vm)); left = u.T @ ones; right = v.T @ ones
    w = np.outer(left, right) / ((lam[:, None] - mu[None, :]) ** 2 + eta**2)
    s = u @ w @ v.T
    r, c = linear_sum_assignment(-s); return c


print(f"E(truth) = {energy(truth):.6f}")
print("REF: GRAMPA on raw fields ->", end=" ")
p = grampa(Vs, Ts); pd = descend(p)
print(f"raw {acc(p):.4f} refined {acc(pd):.4f} E={energy(pd):.4f}")
rho_raw = np.corrcoef(Vs[~np.eye(N, dtype=bool)], Ts[np.ix_(truth, truth)][~np.eye(N, dtype=bool)])[0, 1]
print(f"REF: field correlation (aligned) = {rho_raw:.4f}")

for r in (24, 42):
    t0 = time.time()
    WV = sym_nmf_gpu(V0, r, seed=1)
    WT = sym_nmf_gpu(T0s, r, seed=1)
    WV2 = sym_nmf_gpu(V0, r, seed=2)           # stability check: same matrix, new seed
    print(f"\n=== r={r}  [{time.time()-t0:.1f}s]  resid V {np.linalg.norm(V0-WV@WV.T)/np.linalg.norm(V0):.4f}"
          f"  T {np.linalg.norm(T0s-WT@WT.T)/np.linalg.norm(T0s):.4f}"
          f"  V(seed2) {np.linalg.norm(V0-WV2@WV2.T)/np.linalg.norm(V0):.4f}")

    nrm = lambda X: X / (np.linalg.norm(X, axis=0, keepdims=True) + 1e-12)
    An, Bn, A2n = nrm(WV), nrm(WT), nrm(WV2)

    # (i) UNIQUENESS: same matrix, two seeds
    S = An.T @ A2n; rr, cc = linear_sum_assignment(-S)
    cs = np.sort(S[rr, cc])[::-1]
    print(f"  [uniqueness] symNMF(V,seed1) vs symNMF(V,seed2) matched cos: "
          f"median {np.median(cs):.3f}  >0.9 {(cs>0.9).sum()}/{r}  >0.5 {(cs>0.5).sum()}/{r}  min {cs.min():.3f}")

    # (ii) ORACLE cross-modal column match -- full distribution
    Sx = An.T @ Bn[truth]; rr, cc = linear_sum_assignment(-Sx)
    ocs = Sx[rr, cc]; colmap = cc[np.argsort(rr)]
    srt = np.sort(ocs)[::-1]
    print(f"  [oracle cols] cos: >0.9 {(ocs>0.9).sum()}/{r}  >0.5 {(ocs>0.5).sum()}/{r}"
          f"  median {np.median(ocs):.3f}  min {ocs.min():.3f}")
    print(f"                sorted: {np.round(srt,2)}")
    print(f"                identity-map fraction (shared-seed artifact check): {(colmap==np.arange(r)).mean():.3f}")

    # (iii) row descriptors from oracle-matched columns
    Ar = WV / np.linalg.norm(WV, axis=1, keepdims=True).clip(1e-9)
    Br = WT[:, colmap] / np.linalg.norm(WT[:, colmap], axis=1, keepdims=True).clip(1e-9)
    p_desc = hung(Ar, Br)
    print(f"  [oracle-col row descriptors] Hungarian acc {acc(p_desc):.4f} -> descent {acc(descend(p_desc)):.4f}")

    # (iv) THE PROPOSAL, ORACLE FORM: prune unmatched V components, rebuild, re-solve
    for thr in (0.5, 0.7, 0.8):
        keep = np.where(ocs >= thr)[0]
        kv = rr[np.isin(np.arange(r), np.arange(r))]  # rr is identity-ordered by lsa
        keepV = rr[ocs >= thr]
        if len(keepV) < 4:
            print(f"  [prune thr={thr}] only {len(keepV)} kept, skip"); continue
        Vc = WV[:, keepV] @ WV[:, keepV].T
        Vcs = standardise(Vc)
        rho = np.corrcoef(Vcs[~np.eye(N, dtype=bool)], Ts[np.ix_(truth, truth)][~np.eye(N, dtype=bool)])[0, 1]
        # solve on the cleaned vision field against the original text field
        pg = grampa(Vcs, Ts)
        Acg = torch.tensor(Vcs, device=dev, dtype=DT)
        pgd = descend(pg, A=Acg)          # descent on the CLEANED objective
        pgd_true = descend(pgd)           # then polish on the TRUE objective
        print(f"  [prune thr={thr}] kept {len(keepV)}/{r}  rho(cleanV,T)={rho:.4f} (was {rho_raw:.4f})"
              f"  GRAMPA {acc(pg):.4f} -> clean-descent {acc(pgd):.4f} -> true-descent {acc(pgd_true):.4f}"
              f"  E={energy(pgd_true):.4f}")