"""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)


def solve_faq(visual, shuffled, **_):
    """Frank-Wolfe on the Birkhoff polytope: the canonical QAP relaxation.

    This is the convex-concave family the register listed and never reached.
    It optimises over doubly stochastic matrices and projects at the end,
    which is a different relaxation from both the spectral solvers and GW.
    """
    from scipy.optimize import quadratic_assignment

    result = quadratic_assignment(
        visual, shuffled, method="faq", options={"maximize": True, "n_init": 1}
    )
    return result.col_ind


def solve_faq_multi(visual, shuffled, **_):
    """FAQ from many random doubly stochastic starts, best by its own objective."""
    from scipy.optimize import quadratic_assignment

    result = quadratic_assignment(
        visual, shuffled, method="faq",
        options={"maximize": True, "n_init": 30, "rng": 0},
    )
    return result.col_ind


def solve_2opt(visual, shuffled, **_):
    from scipy.optimize import quadratic_assignment

    result = quadratic_assignment(
        visual, shuffled, method="2opt", options={"maximize": True, "rng": 0}
    )
    return result.col_ind


def solve_gw_multi(visual, shuffled, **_):
    """GW is non-convex; restart it and keep the best by the GW objective."""
    import ot

    size = len(visual)
    weights = ot.unif(size)
    best, best_value = None, np.inf
    generator = np.random.default_rng(0)
    for attempt in range(12):
        init = None
        if attempt:
            noise = generator.random((size, size)) + 1e-3
            init = noise / noise.sum()
        try:
            coupling, log = ot.gromov.gromov_wasserstein(
                visual, shuffled, weights, weights, "square_loss",
                G0=init, max_iter=200, tol_rel=1e-9, log=True,
            )
        except Exception:
            continue
        value = float(log["gw_dist"])
        if value < best_value:
            best, best_value = coupling, value
    if best is None:
        raise RuntimeError("all GW restarts failed")
    return coupling_to_permutation(best)


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

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


def solve_gw_deep_anneal(visual, shuffled, **_):
    """The winning solver, annealed further and started from the smooth end."""
    import ot

    size = len(visual)
    weights = ot.unif(size)
    coupling = None
    schedule = (2e-1, 1e-1, 5e-2, 3e-2, 2e-2, 1e-2, 7e-3, 5e-3, 3e-3, 2e-3, 1e-3)
    for epsilon in schedule:
        try:
            candidate = ot.gromov.entropic_gromov_wasserstein(
                visual, shuffled, weights, weights, "square_loss",
                epsilon=epsilon, G0=coupling, max_iter=300, tol=1e-9, solver="PGD",
            )
        except Exception:
            break
        if not np.isfinite(candidate).all():
            break
        coupling = candidate
    if coupling is None:
        raise RuntimeError("deep anneal diverged immediately")
    return coupling_to_permutation(coupling)


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

    size = len(visual)
    coupling = ot.gromov.semirelaxed_gromov_wasserstein(
        visual, shuffled, ot.unif(size), "square_loss", max_iter=200,
    )
    return coupling_to_permutation(coupling)


def _birkhoff_frank_wolfe(visual, shuffled, mu, coupling, steps):
    """Frank-Wolfe on the Birkhoff polytope for one concavity setting."""
    from scipy.optimize import linear_sum_assignment

    size = len(visual)
    for step in range(steps):
        gradient = 2.0 * (visual @ coupling @ shuffled) + 2.0 * mu * coupling
        _, columns = linear_sum_assignment(-gradient)
        vertex = np.zeros_like(coupling)
        vertex[np.arange(size), columns] = 1.0
        direction = vertex - coupling
        # exact line search on the quadratic along the segment
        quad = float((direction * (visual @ direction @ shuffled)).sum()) + mu * float(
            (direction * direction).sum()
        )
        linear = float((direction * gradient).sum())
        if abs(quad) < 1e-12:
            gamma = 1.0 if linear > 0 else 0.0
        else:
            gamma = float(np.clip(-linear / (2.0 * quad), 0.0, 1.0)) if quad < 0 else (
                1.0 if linear > 0 else 0.0
            )
        if gamma <= 1e-12:
            break
        coupling = coupling + gamma * direction
    return coupling


def solve_path(visual, shuffled, **_):
    """Convex-to-concave path following over the Birkhoff polytope.

    The relaxation the register named and never ran. Maximising a concave
    surrogate over the polytope has one smooth interior optimum; maximising a
    convex one attains its maximum at a vertex, which is a permutation. Adding
    mu*||P||^2 and sweeping mu from negative to positive walks continuously
    between the two, and the argument for it is exactly our situation: the
    smooth end has no local optima to be trapped by, and the solution is
    tracked into the sharp end rather than searched for there.
    """
    size = len(visual)
    scale = float(np.abs(np.linalg.eigvalsh(visual)).max()
                  * np.abs(np.linalg.eigvalsh(shuffled)).max())
    coupling = np.full((size, size), 1.0 / size)
    for fraction in np.linspace(-1.0, 1.0, 21):
        coupling = _birkhoff_frank_wolfe(
            visual, shuffled, fraction * scale, coupling, 30
        )
    return coupling_to_permutation(coupling)


def solve_moment_ladder(visual, shuffled, **_):
    """Match the eigenvalue moments of M, not just its sum.

    The energy is a moment functional in disguise: the pairwise term is the
    first moment of M = permuted-text * visual and the triangle term its
    third. The sequence {tr(M^k)} is what pins down M's spectrum, and we have
    only ever used two of its entries. This optimises a weighted combination
    over several k by Frank-Wolfe, which is the cheapest way to ask whether a
    longer moment sequence has a wider basin than a single term.
    """
    from scipy.optimize import linear_sum_assignment

    size = len(visual)
    coupling = np.full((size, size), 1.0 / size)
    for step in range(120):
        permuted = coupling @ shuffled @ coupling.T
        product = permuted * visual
        # d/dP of tr(M^k) terms, accumulated over the ladder
        gradient = 2.0 * (visual @ coupling @ shuffled)
        squared = product @ product
        gradient = gradient + 1.0 * (
            (visual * product) @ coupling @ shuffled
            + (visual * squared) @ coupling @ shuffled
        )
        _, columns = linear_sum_assignment(-gradient)
        vertex = np.zeros_like(coupling)
        vertex[np.arange(size), columns] = 1.0
        gamma = 2.0 / (step + 2.0)
        coupling = coupling + gamma * (vertex - coupling)
    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,
    "FAQ (Frank-Wolfe)": solve_faq,
    "FAQ x30 restarts": solve_faq_multi,
    "2-opt": solve_2opt,
    "GW x12 restarts": solve_gw_multi,
    "GW kl-loss": solve_gw_kl,
    "entropic GW deep anneal": solve_gw_deep_anneal,
    "semirelaxed GW": solve_semirelaxed,
    "PATH convex-concave": solve_path,
    "moment ladder tr(M^k)": solve_moment_ladder,
}


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()