"""A matching benchmark, and the solvers we never tried on it.

Everything the project has used to search is one narrow family: spectral
initialisation on the pairwise field, exact steepest descent over
transpositions, tempering, and Sinkhorn. All of them act on the permutation
group with a pairwise objective in flat space. The measured cost of that
narrowness is large and specific -- on a field whose anchor bound is 0.997,
meaning the information is intact almost perfectly, blind matching returns
5.6%.

That makes a clean benchmark, because unlike natural data the answer is known
to be reachable. Three synthetic fields span the range: one that already
recovers, one whose captions omit a factor, one truncated to rank eight.
Natural data is included as the case where no method should be expected to
work, so a method that scores there is reporting a bug rather than a result.

Gromov-Wasserstein is the canonical formulation for exactly our problem --
align two metric spaces with no shared embedding, by matching their internal
distance structures -- and it has never been run on these fields. Its entropic
form with an annealed regulariser is the landscape-deformation method that the
rank ladder argued for and that we never reached. Each solver is scored blind
and then re-scored after exact refinement, since composition is what made the
spectral solver work.
"""

from __future__ import annotations

import argparse
import json
import time

import numpy as np
import torch

from .common import write_json
from .spectral_match import grampa, umeyama
from .synth_fast_gate import ClosedFormEnergy, fast_pair_descent
from .synth_triangle_gate import standardized


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", nargs="+", required=True)
    parser.add_argument("--labels", nargs="+", default=None)
    parser.add_argument("--trials", type=int, default=3)
    parser.add_argument("--iterations", type=int, default=600)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--methods", nargs="+", default=None)
    parser.add_argument("--output", default="artifacts/synth_v1/match_battery.json")
    return parser.parse_args()


def standardise(matrix: np.ndarray) -> np.ndarray:
    mask = ~np.eye(len(matrix), dtype=bool)
    values = matrix[mask]
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


def coupling_to_permutation(coupling: np.ndarray) -> np.ndarray:
    from scipy.optimize import linear_sum_assignment

    _, columns = linear_sum_assignment(-coupling)
    return columns


# --- solvers: each returns column assignment for rows of `visual` ---------

def solve_grampa(visual, shuffled, **_):
    return grampa(visual, shuffled, 1.0)


def solve_umeyama(visual, shuffled, **_):
    return umeyama(visual, shuffled)


def solve_gw(visual, shuffled, **_):
    import ot

    size = len(visual)
    weights = ot.unif(size)
    coupling = ot.gromov.gromov_wasserstein(
        visual, shuffled, weights, weights, "square_loss",
        max_iter=200, tol_rel=1e-9,
    )
    return coupling_to_permutation(coupling)


def solve_gw_annealed(visual, shuffled, **_):
    """Entropic GW with the regulariser annealed from smooth to sharp."""
    import ot

    size = len(visual)
    weights = ot.unif(size)
    coupling = None
    for epsilon in (5e-2, 2e-2, 1e-2, 5e-3, 2e-3):
        try:
            coupling = ot.gromov.entropic_gromov_wasserstein(
                visual, shuffled, weights, weights, "square_loss",
                epsilon=epsilon, G0=coupling, max_iter=200, tol=1e-9,
                solver="PGD",
            )
        except Exception:
            break
        if not np.isfinite(coupling).all():
            break
    if coupling is None or not np.isfinite(coupling).all():
        raise RuntimeError("entropic GW diverged at every regulariser")
    return coupling_to_permutation(coupling)


def solve_bapg(visual, shuffled, **_):
    """Bregman alternating projected gradient: a newer GW solver."""
    import ot

    size = len(visual)
    weights = ot.unif(size)
    coupling = ot.gromov.BAPG_gromov_wasserstein(
        visual, shuffled, weights, weights, "square_loss",
        epsilon=1e-1, max_iter=1000, tol=1e-9,
    )
    return coupling_to_permutation(coupling)


SOLVERS = {
    "grampa (current)": solve_grampa,
    "umeyama": solve_umeyama,
    "gromov-wasserstein": solve_gw,
    "entropic GW annealed": solve_gw_annealed,
    "BAPG GW": solve_bapg,
}


def main() -> None:
    args = parse_args()
    labels = args.labels or [p.split("/")[-1] for p in args.fields]
    chosen = args.methods or list(SOLVERS)
    device = torch.device(args.device)
    rows = []

    for path, label in zip(args.fields, labels):
        state = torch.load(path, map_location="cpu", weights_only=False)
        visual = standardise(state["visual_field"].double().numpy())
        text = standardise(state["text_field"].double().numpy())
        size = len(visual)
        text_gpu = standardized(torch.from_numpy(text).to(device)).double()
        visual_gpu = standardized(torch.from_numpy(visual).to(device)).double()
        print(f"\n=== {label}  (N={size}, chance={1/size:.4f})", flush=True)

        for name in chosen:
            solver = SOLVERS[name]
            raw, refined, seconds, failures = [], [], [], []
            for trial in range(args.trials):
                generator = np.random.default_rng(trial)
                hidden = generator.permutation(size)
                shuffled = text[np.ix_(hidden, hidden)]
                start = time.time()
                try:
                    columns = solver(visual=visual, shuffled=shuffled)
                except Exception as error:
                    failures.append(str(error)[:80])
                    continue
                seconds.append(time.time() - start)
                raw.append(float((hidden[columns] == np.arange(size)).mean()))

                shuffled_gpu = standardized(
                    torch.from_numpy(shuffled).to(device)
                ).double()
                final = fast_pair_descent(
                    shuffled_gpu, visual_gpu,
                    torch.from_numpy(np.ascontiguousarray(columns)).to(device),
                    args.iterations,
                )
                refined.append(
                    float((hidden[final.cpu().numpy()] == np.arange(size)).mean())
                )
            row = {
                "field": label, "method": name,
                "raw": float(np.mean(raw)) if raw else None,
                "refined": float(np.mean(refined)) if refined else None,
                "seconds": float(np.mean(seconds)) if seconds else None,
                "failures": failures,
            }
            rows.append(row)
            if raw:
                print("  %-22s raw=%.4f  refined=%.4f  (%.1fs)"
                      % (name, row["raw"], row["refined"], row["seconds"]), flush=True)
            else:
                print("  %-22s FAILED: %s" % (name, failures[0] if failures else "?"),
                      flush=True)

    write_json(args.output, {
        "protocol": (
            "Blind: the hidden permutation is generated per trial and read only "
            "for scoring. Each solver is scored as returned and again after "
            "exact steepest-descent refinement."
        ),
        "rows": rows,
    })
    print(json.dumps({"done": True}))


if __name__ == "__main__":
    main()