"""Control: is the FGW warm-homotopy win driven by TLB *information*, or by any
linear term at all (generic symmetry-breaking)?  Three linear costs:

  tlb      -- the proposed distance-profile cost
  tlbperm  -- the same matrix with its columns randomly permuted: identical
              value distribution, correspondence information destroyed
  rand     -- iid uniform noise cost

Plus 3 trials, cold and warm sweeps, blind selection by E(P).
"""
from __future__ import annotations
import json, sys, time
sys.path.insert(0, "/home/yurenh2/emm")
import numpy as np, ot, torch
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")
ALPHAS = [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 0.95, 1.0]


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


def main():
    path = "/home/yurenh2/emm/artifacts/synth_v1/omit_size.pt"
    st = torch.load(path, 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)
    rows = []
    for trial in range(3):
        rng = np.random.default_rng(trial); 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)
        et = float(en.energy(torch.from_numpy(
            np.ascontiguousarray(np.argsort(hid))).to(DEV)[None])[0])

        def acc(c): return float((hid[c] == np.arange(n)).mean())

        def refine(c):
            f = fast_pair_descent(Sg, Vg,
                                  torch.from_numpy(np.ascontiguousarray(c)).to(DEV),
                                  600).cpu().numpy()
            e = float(en.energy(torch.from_numpy(np.ascontiguousarray(f)).to(DEV)[None])[0])
            return e, acc(f)

        base = tlb(V, S)
        shuf = np.random.default_rng(100 + trial).permutation(n)
        variants = {
            "tlb": base / base.max(),
            "tlbperm": base[:, shuf] / base.max(),
            "rand": np.random.default_rng(200 + trial).random((n, n)),
        }
        for vname, M in variants.items():
            for mode in ("cold", "warm"):
                G = np.outer(w, w); recs = []
                for a in ALPHAS:
                    g0 = None if mode == "cold" else G
                    G2 = ot.gromov.fused_gromov_wasserstein(
                        M, V, S, w, w, "square_loss", alpha=a, G0=g0,
                        max_iter=200, tol_rel=1e-9)
                    if mode == "warm":
                        G = G2
                    _, cols = linear_sum_assignment(-G2)
                    e, ar = refine(cols)
                    recs.append({"alpha": a, "raw": acc(cols), "refined": ar, "E": e})
                best = min(recs, key=lambda x: x["E"])
                row = {"trial": trial, "cost": vname, "mode": mode,
                       "E_truth": et, "blind_refined": best["refined"],
                       "blind_alpha": best["alpha"], "blind_E": best["E"],
                       "oracle_refined": max(x["refined"] for x in recs),
                       "curve": [round(x["refined"], 4) for x in recs]}
                rows.append(row)
                print(json.dumps(row), flush=True)
    scratch = ("/tmp/claude-1273071/-home-yurenh2-emm/"
               "bb97201d-d9b5-46e6-82fc-981d37e7ab98/scratchpad/")
    json.dump(rows, open(scratch + "control.json", "w"))
    print(json.dumps({"done": True}))


if __name__ == "__main__":
    main()