"""World matching: spectral initialisation plus energy refinement.

The two solvers are complementary. A spectral solver reads the coarse
correspondence out of the eigenstructure of two relation fields in
polynomial time and without any search, but its rounding is noisy. Exact
steepest descent on the closed-form energy repairs the rounding but
cannot find the basin from a random start. Composed, they recover the
hidden assignment.

Neither stage sees a pair. The relation fields are built independently
per modality; the hidden permutation is applied to the text field and
read back only to score.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch

from .common import seed_everything, write_json
from .spectral_match import grampa, spectral_profile, umeyama
from .synth_fast_gate import ClosedFormEnergy, all_swaps, steepest_descent
from .synth_triangle_gate import standardized


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fields", required=True)
    parser.add_argument("--eta", type=float, default=1.0)
    parser.add_argument("--pair-weight", type=float, default=1.0)
    parser.add_argument("--triangle-weight", type=float, default=1.0)
    parser.add_argument("--refine-steps", type=int, default=4000)
    parser.add_argument("--chunk", type=int, default=256)
    parser.add_argument("--restarts", type=int, default=1)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--output", default="")
    return parser.parse_args()


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


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual = normalise(state["visual_field"].double().numpy().copy())
    text = normalise(state["text_field"].double().numpy().copy())
    size = len(visual)

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

    def accuracy(assignment: np.ndarray) -> float:
        return float((hidden[assignment] == truth).mean())

    device = torch.device(args.device)
    text_tensor = standardized(torch.from_numpy(shuffled).to(device)).float()
    visual_tensor = standardized(torch.from_numpy(visual).to(device)).float()
    energy = ClosedFormEnergy(
        text_tensor,
        visual_tensor,
        args.pair_weight,
        args.triangle_weight,
        args.chunk,
    )
    swaps = all_swaps(size, device)
    truth_permutation = torch.from_numpy(np.argsort(hidden)).to(device)
    true_energy = float(energy.energy(truth_permutation[None])[0])

    report = {
        "protocol": (
            "Fields are built per modality without pairs; the hidden "
            "permutation is applied to the text field and used only to "
            "score. Spectral initialisation is followed by exact "
            "steepest descent on the closed-form energy."
        ),
        "samples": size,
        "field_correlation_at_truth": float(
            np.corrcoef(
                visual[~np.eye(size, dtype=bool)],
                text[~np.eye(size, dtype=bool)],
            )[0, 1]
        ),
        "true_energy": true_energy,
        "chance": 1.0 / size,
        "spectra": {
            "visual": spectral_profile(visual),
            "text": spectral_profile(shuffled),
        },
        "stages": {},
    }

    spectral = grampa(visual, shuffled, args.eta)
    report["stages"]["spectral"] = {"accuracy": accuracy(spectral)}
    initial = torch.from_numpy(spectral.copy()).to(device)
    refined, value = steepest_descent(energy, initial, swaps, args.refine_steps)
    report["stages"]["spectral_plus_refinement"] = {
        "accuracy": accuracy(refined.cpu().numpy()),
        "energy": value,
        "energy_over_true": value / abs(true_energy) - true_energy / abs(true_energy),
    }

    quenches = []
    for restart in range(args.restarts):
        start = torch.from_numpy(generator.permutation(size)).to(device)
        final, final_value = steepest_descent(
            energy, start, swaps, args.refine_steps
        )
        quenches.append(
            {
                "accuracy": accuracy(final.cpu().numpy()),
                "energy": final_value,
            }
        )
    report["stages"]["random_start_refinement"] = quenches

    print(json.dumps({key: value for key, value in report["stages"].items()}))
    if args.output:
        write_json(args.output, report)
        print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()