"""Replication with BLIND energy selection over the (m,K) grid, two fresh hidden perms."""
import sys, time, numpy as np, torch
sys.path.insert(0, '/home/yurenh2/emm')
from scipy.optimize import linear_sum_assignment
from scipy.stats import ortho_group
from worldalign.synth_fast_gate import fast_pair_descent

dev = 'cuda:0'
def standardise(M):
    M = np.asarray(M, dtype=np.float64); m = ~np.eye(len(M), dtype=bool); v = M[m]
    o = (M - v.mean())/v.std(); np.fill_diagonal(o, 0.0); return o
d = torch.load('/home/yurenh2/emm/artifacts/synth_v1/omit_size.pt', map_location='cpu', weights_only=False)
V = standardise(d['visual_field']); T0 = standardise(d['text_field']); N = len(V)
Vt = torch.tensor(V, dtype=torch.float32, device=dev)
wV, UV = np.linalg.eigh(V); wV = wV[::-1]; UV = UV[:, ::-1]

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 icp(XV, XT, O, iters=25):
    for _ in range(iters):
        p = hung(XV, XT@O); u, s, vt = np.linalg.svd(XT[p].T@XV); On = u@vt
        if np.allclose(On, O, atol=1e-10): O = On; break
        O = On
    return hung(XV, XT@O)

for SEED in (11, 23):
    rng = np.random.default_rng(SEED); sigma = rng.permutation(N)
    T = T0[np.ix_(sigma, sigma)]; truth = np.argsort(sigma)
    acc = lambda p: float((np.asarray(p) == truth).mean())
    Tt = torch.tensor(T, dtype=torch.float32, device=dev)
    CONST = float((Tt*Tt).sum() + (Vt*Vt).sum())
    def energy(p):
        P = torch.as_tensor(np.asarray(p), dtype=torch.long, device=dev)
        return (CONST - 2.0*float((Tt[P[:, None], P[None, :]]*Vt).sum()))/(N*(N-1))
    def descend(p, steps=4000):
        P = torch.as_tensor(np.ascontiguousarray(p), dtype=torch.long, device=dev)
        return fast_pair_descent(Tt, Vt, P, steps).cpu().numpy()
    wT, UT = np.linalg.eigh(T); wT = wT[::-1]; UT = UT[:, ::-1]
    t0 = time.time(); sols = []
    for r in (8, 10, 12, 14, 16):
        XV = UV[:, :r]*np.sqrt(np.abs(wV[:r])); XT = UT[:, :r]*np.sqrt(np.abs(wT[:r]))
        for _ in range(30):
            p = icp(XV, XT, ortho_group.rvs(r, random_state=int(rng.integers(1 << 30))))
            sols.append((energy(p), p, r))
    sols.sort(key=lambda z: z[0])
    print(f"\n##### hidden seed {SEED}: {len(sols)} ICP sols in {time.time()-t0:.0f}s, "
          f"E(truth)={energy(truth):.4f}, best member E={sols[0][0]:.4f} acc={acc(sols[0][1]):.3f}", flush=True)
    A = [descend(p) for _, p, _ in sols[:5]]
    bA = min(A, key=energy)
    print(f"  ROUTE A (descend 5 lowest-E, blind pick): E={energy(bA):.4f} acc={acc(bA):.3f}", flush=True)

    def expand(aL, aR):
        pr = np.setdiff1d(np.arange(N), aL); prR = np.setdiff1d(np.arange(N), aR)
        L = V[np.ix_(pr, aL)]; R = T[np.ix_(prR, aR)]
        L = (L-L.mean(1, keepdims=True))/L.std(1, keepdims=True).clip(1e-9)
        R = (R-R.mean(1, keepdims=True))/R.std(1, keepdims=True).clip(1e-9)
        S = L@R.T/L.shape[1]; _, c = linear_sum_assignment(-S)
        full = np.zeros(N, dtype=int); full[aL] = aR; full[pr] = prR[c]; return full
    grid = []
    for m in (10, 30, 60):
        votes = np.zeros((N, N))
        for _, p, _ in sols[:m]: votes[np.arange(N), p] += 1
        conf = votes.max(1) - np.partition(votes, -2, axis=1)[:, -2]
        order = np.argsort(-conf); pm = votes.argmax(1)
        for K in (12, 25, 50, 100):
            sel = order[:K]; part = pm[sel]; keep = np.zeros(len(sel), bool); seen = set()
            for i, c in enumerate(part):
                if c not in seen: seen.add(c); keep[i] = True
            sel, part = sel[keep], part[keep]
            prec = float((part == truth[sel]).mean())
            e = expand(sel, part); pd = descend(e)
            grid.append((energy(pd), acc(pd), m, K, prec, acc(e)))
            print(f"   m={m:2d} K={K:3d} prec={prec:.3f} expand={acc(e):.3f} -> descent acc={acc(pd):.3f} E={energy(pd):.4f}", flush=True)
    grid.sort()
    print(f"  ROUTE B (blind pick by energy over the 12 cells): m={grid[0][2]} K={grid[0][3]} "
          f"E={grid[0][0]:.4f} acc={grid[0][1]:.3f}   [best cell by accuracy was {max(g[1] for g in grid):.3f}]", flush=True)
    print(f"  wall {time.time()-t0:.0f}s", flush=True)