path: root/worldalign/blind_recovery.py

"""Blind unpaired recovery on gate-passing states.

The basin audit licenses this experiment: on content-projected VG states
the true assignment sits below every quench, so the remaining question is
search. Three arms attack the measured landscape shape (deep true basin,
glassy surroundings, smooth coarse spectrum):

A. entropic Sinkhorn annealing: soft coupling, temperature and entropy
   schedules, barycentric language states;
B. parallel tempering: replica-exchange Metropolis over permutations with
   the exact closed-form swap deltas, imported from spin-glass practice;
C. spectral-band homotopy: solve the coupling in a low-dimensional
   spectral band first, then warm-start progressively finer bands --
   a band-limited continuation in the sense of low-order Fourier terms on
   the symmetric group, instantiated through the content spectrum.

The text side enters through a hidden shuffle, so the optimizer's identity
carries no information. The hidden truth scores results and is used
nowhere else.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment

from .common import seed_everything, write_json
from .content_gate import flickr_states, vg_states
from .energy import log_sinkhorn
from .manifold_gate import standardize_relation


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", choices=["flickr", "vg"], default="vg")
    parser.add_argument("--manifest", default="artifacts/manifest.json")
    parser.add_argument("--vision", default="artifacts/vision.pt")
    parser.add_argument("--text", default="artifacts/text.pt")
    parser.add_argument("--text-orbits", default="artifacts/text_orbits_qwen0p5b.pt")
    parser.add_argument("--vg-vision", default="artifacts/vg_5k/vision_features.pt")
    parser.add_argument("--vg-text", default="artifacts/vg_5k/text_features.pt")
    parser.add_argument(
        "--vg-ground-truth", default="artifacts/vg_5k/ground_truth.private.jsonl"
    )
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=512)
    parser.add_argument("--subset-seed", type=int, default=0)
    parser.add_argument("--dims", type=int, default=256)
    parser.add_argument("--shrinkage", type=float, default=0.05)
    parser.add_argument("--arms", default="tempering,sinkhorn,homotopy")
    parser.add_argument("--replicas", type=int, default=8)
    parser.add_argument("--tempering-rounds", type=int, default=40000)
    parser.add_argument("--temp-high", type=float, default=3e-3)
    parser.add_argument("--temp-low", type=float, default=1e-5)
    parser.add_argument("--exchange-every", type=int, default=20)
    parser.add_argument("--sinkhorn-steps", type=int, default=2500)
    parser.add_argument("--sinkhorn-restarts", type=int, default=4)
    parser.add_argument("--homotopy-bands", default="4,8,16,32,64,full")
    parser.add_argument(
        "--unary-matrix",
        default="",
        help="Optional anchor similarity matrix (.pt) in ground-truth node "
        "order; injected as a unary field on the tempering energy.",
    )
    parser.add_argument("--unary-weight", type=float, default=2.0)
    parser.add_argument("--init", choices=["random", "unary"], default="random")
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--seed", type=int, default=20260730)
    parser.add_argument("--output", required=True)
    return parser.parse_args()


def relation_mse_soft(
    language_states: torch.Tensor, visual_standardized: torch.Tensor
) -> torch.Tensor:
    """Differentiable standardized-relation MSE for soft states."""
    normalized = F.normalize(language_states, dim=-1)
    relation = normalized @ normalized.T
    mask = ~torch.eye(len(relation), dtype=torch.bool, device=relation.device)
    values = relation[mask]
    standardized = (values - values.mean()) / values.std().clamp_min(1e-6)
    return (standardized - visual_standardized[mask]).square().mean()


def permutation_energy_batch(
    text_standardized: torch.Tensor,
    visual_standardized: torch.Tensor,
    permutations: torch.Tensor,
) -> torch.Tensor:
    fields = text_standardized[permutations[:, :, None], permutations[:, None, :]]
    mask = ~torch.eye(
        text_standardized.shape[-1], dtype=torch.bool, device=fields.device
    )
    return ((fields - visual_standardized) ** 2)[:, mask].mean(-1)


def proposal_swap_deltas(
    permuted_fields: torch.Tensor,
    visual_standardized: torch.Tensor,
    pairs_p: torch.Tensor,
    pairs_q: torch.Tensor,
) -> torch.Tensor:
    """Exact deltas for proposed swaps only, O(N) per proposal.

    permuted_fields: [R, N, N] text fields under each replica's current
    permutation; pairs_p/pairs_q: [R, P] proposal endpoints.
    """
    size = permuted_fields.shape[-1]
    count = size * (size - 1)
    replica_index = torch.arange(len(permuted_fields), device=permuted_fields.device)
    rows_p = permuted_fields[replica_index[:, None], pairs_p]  # [R, P, N]
    rows_q = permuted_fields[replica_index[:, None], pairs_q]
    visual_p = visual_standardized[pairs_p]
    visual_q = visual_standardized[pairs_q]
    self_p = (rows_p * visual_p).sum(-1)
    self_q = (rows_q * visual_q).sum(-1)
    cross_pq = (rows_p * visual_q).sum(-1)
    cross_qp = (rows_q * visual_p).sum(-1)
    direct = permuted_fields[replica_index[:, None], pairs_p, pairs_q]
    direct_visual = visual_standardized[pairs_p, pairs_q]
    total = self_p + self_q - cross_pq - cross_qp - 2.0 * direct * direct_visual
    delta = (4.0 / count) * total
    return delta.masked_fill(pairs_p == pairs_q, float("inf"))


def score(
    permutation: torch.Tensor, truth: torch.Tensor, energy: float, true_energy: float
) -> dict:
    return {
        "accuracy": float((permutation.cpu() == truth.cpu()).double().mean()),
        "energy": energy,
        "energy_over_true": energy / true_energy - 1.0,
    }


def coupling_metrics(coupling: torch.Tensor, truth: torch.Tensor) -> dict:
    n = len(coupling)
    truth = truth.to(coupling.device)
    mass_at_truth = float(coupling[torch.arange(n, device=coupling.device), truth].mean())
    rows, cols = linear_sum_assignment(-coupling.detach().cpu().numpy())
    permutation = torch.from_numpy(cols)
    forward = coupling.argmax(-1).cpu()
    backward = coupling.argmax(0).cpu()
    mutual = backward[forward] == torch.arange(n)
    sorted_mass = coupling.sort(-1, descending=True).values
    margin = (sorted_mass[:, 0] - sorted_mass[:, 1]).cpu()
    top = margin.argsort(descending=True)[: max(1, n // 10)]
    return {
        "mass_at_truth": mass_at_truth,
        "hungarian_accuracy": float((permutation == truth.cpu()).double().mean()),
        "mutual_nn_count": int(mutual.sum()),
        "mutual_nn_precision": float(
            (forward[mutual] == truth.cpu()[mutual]).double().mean()
        )
        if mutual.any()
        else None,
        "top_margin_decile_precision": float(
            (forward[top] == truth.cpu()[top]).double().mean()
        ),
    }


def arm_tempering(
    text_standardized: torch.Tensor,
    visual_standardized: torch.Tensor,
    truth: torch.Tensor,
    true_energy: float,
    args: argparse.Namespace,
    generator: torch.Generator,
    unary: torch.Tensor | None = None,
) -> dict:
    device = text_standardized.device
    size = text_standardized.shape[-1]
    replicas = args.replicas
    weight = args.unary_weight if unary is not None else 0.0

    def total_energy(permutations: torch.Tensor) -> torch.Tensor:
        energy = permutation_energy_batch(
            text_standardized, visual_standardized, permutations
        )
        if unary is not None:
            index = torch.arange(size, device=device)
            energy = energy - weight * unary[index, permutations].mean(-1)
        return energy

    temperatures = torch.logspace(
        torch.log10(torch.tensor(args.temp_low)),
        torch.log10(torch.tensor(args.temp_high)),
        replicas,
    ).to(device)
    if args.init == "unary" and unary is not None:
        rows, cols = linear_sum_assignment(-unary.cpu().numpy())
        start = torch.from_numpy(cols).to(device)
        permutations = torch.stack([start.clone() for _ in range(replicas)])
    else:
        permutations = torch.stack(
            [
                torch.randperm(size, generator=generator).to(device)
                for _ in range(replicas)
            ]
        )
    energies = total_energy(permutations)
    best = {"energy": float("inf"), "permutation": permutations[0].clone()}
    accepted = 0
    for round_index in range(args.tempering_rounds):
        fields = text_standardized[
            permutations[:, :, None], permutations[:, None, :]
        ]
        pairs_p = torch.randint(0, size, (replicas, 24), generator=generator).to(
            device
        )
        pairs_q = torch.randint(0, size, (replicas, 24), generator=generator).to(
            device
        )
        proposal_deltas = proposal_swap_deltas(
            fields, visual_standardized, pairs_p, pairs_q
        )
        if unary is not None:
            replica_index = torch.arange(replicas, device=device)[:, None]
            assigned_p = permutations[replica_index, pairs_p]
            assigned_q = permutations[replica_index, pairs_q]
            unary_delta = (
                unary[pairs_p, assigned_q]
                + unary[pairs_q, assigned_p]
                - unary[pairs_p, assigned_p]
                - unary[pairs_q, assigned_q]
            )
            proposal_deltas = proposal_deltas - (weight / size) * unary_delta
        noise = torch.rand(replicas, 24, generator=generator).to(device)
        acceptable = (
            proposal_deltas < -temperatures[:, None] * noise.clamp_min(1e-12).log()
        ) & torch.isfinite(proposal_deltas)
        for replica in range(replicas):
            hits = torch.nonzero(acceptable[replica])
            if not len(hits):
                continue
            first = int(hits[0, 0])
            p = int(pairs_p[replica, first])
            q = int(pairs_q[replica, first])
            permutations[replica][[p, q]] = permutations[replica][[q, p]]
            energies[replica] = energies[replica] + proposal_deltas[replica, first]
            accepted += 1
        if round_index % args.exchange_every == 0:
            for replica in range(replicas - 1):
                gap = (energies[replica] - energies[replica + 1]) * (
                    1.0 / temperatures[replica] - 1.0 / temperatures[replica + 1]
                )
                if gap > 0 or torch.rand(1, generator=generator).item() < float(
                    gap.exp()
                ):
                    permutations[[replica, replica + 1]] = permutations[
                        [replica + 1, replica]
                    ]
                    energies[[replica, replica + 1]] = energies[
                        [replica + 1, replica]
                    ]
        cold = int(energies.argmin())
        if float(energies[cold]) < best["energy"]:
            best = {
                "energy": float(energies[cold]),
                "permutation": permutations[cold].clone(),
            }
        if round_index % 5000 == 0:
            energies = total_energy(permutations)  # refresh against drift
            print(
                json.dumps(
                    {
                        "tempering_round": round_index,
                        "cold_energy": float(energies.min()),
                        "cold_accuracy": float(
                            (permutations[int(energies.argmin())].cpu() == truth)
                            .double()
                            .mean()
                        ),
                        "accepted": accepted,
                    }
                )
            )
    final = [
        score(
            permutations[r],
            truth,
            float(
                permutation_energy_batch(
                    text_standardized, visual_standardized, permutations[r : r + 1]
                )[0]
            ),
            true_energy,
        )
        for r in range(replicas)
    ]
    best_score = score(
        best["permutation"],
        truth,
        float(
            permutation_energy_batch(
                text_standardized, visual_standardized, best["permutation"][None]
            )[0]
        ),
        true_energy,
    )
    return {"replicas": final, "best": best_score, "accepted_moves": accepted}


def arm_sinkhorn(
    text_states: torch.Tensor,
    visual_standardized: torch.Tensor,
    truth: torch.Tensor,
    true_energy: float,
    args: argparse.Namespace,
    generator: torch.Generator,
    bands: list[int | None] | None = None,
) -> dict:
    device = text_states.device
    size = len(text_states)
    restarts = []
    for restart in range(args.sinkhorn_restarts if bands is None else 1):
        logits = torch.nn.Parameter(
            0.01
            * torch.randn(size, size, generator=generator).to(device)
        )
        optimizer = torch.optim.Adam([logits], lr=0.08)
        stage_states = text_states
        stages = bands if bands is not None else [None]
        for stage_index, band in enumerate(stages):
            if band is not None:
                keep = min(band, text_states.shape[-1])
                stage_states = text_states[:, :keep]
            steps = args.sinkhorn_steps // len(stages)
            for step in range(steps):
                progress = step / max(steps - 1, 1)
                temperature = 0.5 * (0.05 / 0.5) ** progress
                coupling = log_sinkhorn(logits, temperature, iterations=12)
                barycenter = coupling @ stage_states
                loss = relation_mse_soft(barycenter, visual_standardized)
                entropy = -(coupling * coupling.clamp_min(1e-12).log()).sum(-1).mean()
                loss = loss + (0.002 + 0.02 * progress) * entropy
                optimizer.zero_grad(set_to_none=True)
                loss.backward()
                optimizer.step()
        with torch.no_grad():
            coupling = log_sinkhorn(logits, 0.05, iterations=30)
        metrics = coupling_metrics(coupling, truth)
        rows, cols = linear_sum_assignment(-coupling.detach().cpu().numpy())
        permutation = torch.from_numpy(cols).to(device)
        rounded_energy = float(
            permutation_energy_batch(
                standardize_relation(
                    F.normalize(text_states, dim=-1)
                    @ F.normalize(text_states, dim=-1).T
                )[0][None].squeeze(0),
                visual_standardized,
                permutation[None],
            )[0]
        )
        metrics.update(
            score(permutation, truth, rounded_energy, true_energy)
        )
        restarts.append(metrics)
        print(json.dumps({"sinkhorn_restart": restart, **metrics}))
    return {"restarts": restarts}


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    if args.dataset == "flickr":
        text_states, visual_states = flickr_states(args)
    else:
        text_states, visual_states = vg_states(args)
    device = torch.device(args.device)
    generator = torch.Generator().manual_seed(args.seed)

    size = len(text_states)
    hidden = torch.randperm(size, generator=generator)
    truth = torch.argsort(hidden)  # input row j holds true counterpart hidden[j]
    text_input = text_states[hidden].float().to(device)
    visual_states = visual_states.float().to(device)

    unary = None
    if args.unary_matrix:
        anchor_state = torch.load(
            args.unary_matrix, map_location="cpu", weights_only=False
        )
        combined = anchor_state["combined"].double()
        subset_generator = torch.Generator().manual_seed(args.subset_seed)
        subset = torch.randperm(len(combined), generator=subset_generator)[
            : args.samples
        ]
        unary_subset = combined[subset][:, subset]
        unary_subset = (unary_subset - unary_subset.mean()) / unary_subset.std()
        unary = unary_subset[:, hidden].float().to(device)

    text_standardized = standardize_relation(
        F.normalize(text_input, dim=-1) @ F.normalize(text_input, dim=-1).T
    )[0]
    visual_standardized = standardize_relation(
        F.normalize(visual_states, dim=-1) @ F.normalize(visual_states, dim=-1).T
    )[0]
    true_energy = float(
        permutation_energy_batch(
            text_standardized, visual_standardized, truth[None].to(device)
        )[0]
    )
    report: dict = {
        "protocol": (
            "The text side enters through a hidden shuffle; the optimizer "
            "never sees pair, order, or truth information. Hidden truth "
            "scores outcomes only."
        ),
        "dataset": args.dataset,
        "dims": args.dims,
        "samples": size,
        "true_energy": true_energy,
        "chance_accuracy": 1.0 / size,
        "arms": {},
    }
    arms = [arm.strip() for arm in args.arms.split(",")]
    if "tempering" in arms:
        report["arms"]["tempering"] = arm_tempering(
            text_standardized,
            visual_standardized,
            truth,
            true_energy,
            args,
            generator,
            unary=unary,
        )
        if unary is not None:
            report["unary"] = {
                "weight": args.unary_weight,
                "matched_mean": float(
                    unary[torch.arange(size, device=unary.device), truth.to(unary.device)].mean()
                ),
                "grand_mean": float(unary.mean()),
                "init": args.init,
            }
    if "sinkhorn" in arms:
        report["arms"]["sinkhorn"] = arm_sinkhorn(
            text_input, visual_standardized, truth, true_energy, args, generator
        )
    if "homotopy" in arms:
        bands: list[int | None] = [
            None if band == "full" else int(band)
            for band in args.homotopy_bands.split(",")
        ]
        report["arms"]["homotopy"] = arm_sinkhorn(
            text_input,
            visual_standardized,
            truth,
            true_energy,
            args,
            generator,
            bands=bands,
        )
    summary = {}
    for arm, result in report["arms"].items():
        candidates = result.get("restarts") or result.get("replicas") or []
        if "best" in result:
            candidates = candidates + [result["best"]]
        best = max(candidates, key=lambda c: c["accuracy"], default=None)
        summary[arm] = best
    report["summary"] = summary
    print(json.dumps({"summary": summary}))
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()