"""Spectral matching of relation fields: Umeyama and GRAMPA.

Relation fields are N x N regardless of embedding dimension, so the two
modalities need no common representation dimension. What they do need is
comparable spectra: eigenvector matching degrades when effective ranks
differ or when eigenvalues cluster. GRAMPA is built for that regime -- it
weights every pair of eigenvectors by 1 / ((lambda_i - mu_j)^2 + eta^2)
instead of pairing them one to one -- so both solvers are provided along
with the spectral compatibility diagnostic that predicts whether either
can work.

No local search: these are polynomial-time solvers that sidestep the
glassy landscape entirely. Hidden pairs score the output only.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from scipy.optimize import linear_sum_assignment

from .common import read_json, seed_everything, write_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", help="Saved .pt with visual/text fields.")
    parser.add_argument("--data-dir", default="artifacts/synth_v0")
    parser.add_argument("--split", choices=["val", "test"], default="test")
    parser.add_argument("--samples", type=int, default=256)
    parser.add_argument("--merge-distance", type=float, default=30.0)
    parser.add_argument("--vision-views", type=int, default=4)
    parser.add_argument("--eta", default="0.05,0.1,0.2,0.5")
    parser.add_argument("--seed", type=int, default=20260731)
    parser.add_argument(
        "--output", default="artifacts/synth_v0/spectral_match.json"
    )
    parser.add_argument("--fields-output", default="")
    return parser.parse_args()


def spectral_profile(matrix: np.ndarray) -> dict:
    values = np.linalg.eigvalsh(matrix)[::-1]
    magnitude = np.abs(values)
    total = magnitude.sum()
    cumulative = np.cumsum(magnitude) / max(total, 1e-12)
    participation = (magnitude.sum() ** 2) / max((magnitude**2).sum(), 1e-12)
    gaps = np.abs(np.diff(values))
    return {
        "top_eigenvalues": values[:12].tolist(),
        "effective_rank_participation": float(participation),
        "rank_for_90_percent": int(np.searchsorted(cumulative, 0.90) + 1),
        "rank_for_99_percent": int(np.searchsorted(cumulative, 0.99) + 1),
        "median_relative_gap": float(
            np.median(gaps) / max(magnitude.max(), 1e-12)
        ),
        "min_relative_gap_top20": float(
            gaps[:20].min() / max(magnitude.max(), 1e-12)
        ),
    }


def umeyama(visual: np.ndarray, text: np.ndarray) -> np.ndarray:
    """Classic eigenvector-magnitude matching, sign ambiguity absorbed."""
    _, u = np.linalg.eigh(visual)
    _, v = np.linalg.eigh(text)
    score = np.abs(u) @ np.abs(v).T
    rows, cols = linear_sum_assignment(-score)
    return cols


def grampa(visual: np.ndarray, text: np.ndarray, eta: float) -> np.ndarray:
    """Pairwise eigen-alignment similarity, robust to clustered spectra."""
    lam, u = np.linalg.eigh(visual)
    mu, v = np.linalg.eigh(text)
    ones = np.ones(len(visual))
    left = u.T @ ones  # [N]
    right = v.T @ ones
    weight = np.outer(left, right) / ((lam[:, None] - mu[None, :]) ** 2 + eta**2)
    similarity = u @ weight @ v.T
    rows, cols = linear_sum_assignment(-similarity)
    return cols


def score_assignment(
    assignment: np.ndarray, hidden: np.ndarray, size: int
) -> dict:
    """assignment[i] is the shuffled-space index matched to visual row i.

    Shuffled index j denotes original node hidden[j], and visual row i
    denotes original node i, so the match is correct when
    hidden[assignment[i]] == i.
    """
    return {
        "accuracy": float((hidden[assignment] == np.arange(size)).mean()),
        "chance": 1.0 / size,
    }


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    if args.fields:
        state = torch.load(args.fields, map_location="cpu", weights_only=False)
        visual_field = state["visual_field"]
        text_field = state["text_field"]
    else:
        from .synth_triangle_gate import build_fields

        manifest = read_json(Path(args.data_dir, "manifest.json"))
        rows = manifest[args.split][: args.samples]
        visual_field, text_field = build_fields(args, rows, manifest)
        if args.fields_output:
            torch.save(
                {"visual_field": visual_field, "text_field": text_field},
                args.fields_output,
            )

    size = len(visual_field)
    visual = visual_field.double().numpy()
    text = text_field.double().numpy()
    # Center and scale: spectral methods compare shapes, not offsets.
    mask = ~np.eye(size, dtype=bool)
    for matrix in (visual, text):
        values = matrix[mask]
        matrix -= values.mean()
        matrix /= values.std()
        np.fill_diagonal(matrix, 0.0)

    generator = np.random.default_rng(args.seed)
    hidden = generator.permutation(size)
    text_shuffled = text[np.ix_(hidden, hidden)]

    report = {
        "protocol": (
            "Polynomial-time spectral solvers on N x N relation fields; "
            "embedding dimensions are irrelevant by construction. The "
            "hidden shuffle is applied to the text field and used only to "
            "score the returned assignment."
        ),
        "samples": size,
        "spectra": {
            "visual": spectral_profile(visual),
            "text": spectral_profile(text_shuffled),
        },
        "solvers": {},
    }
    # Both fields are built over the same row list, so they are aligned at
    # the truth without any permutation.
    correlation = float(np.corrcoef(visual[mask], text[mask])[0, 1])
    report["field_correlation_at_truth"] = correlation

    assignment = umeyama(visual, text_shuffled)
    report["solvers"]["umeyama"] = score_assignment(assignment, hidden, size)
    print(json.dumps({"umeyama": report["solvers"]["umeyama"]}))

    for eta in (float(value) for value in args.eta.split(",")):
        assignment = grampa(visual, text_shuffled, eta)
        result = score_assignment(assignment, hidden, size)
        report["solvers"][f"grampa_eta{eta}"] = result
        print(json.dumps({f"grampa_eta{eta}": result}))

    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()