"""Blind recovery on natural data, run because the old gate was wrong.

The standing gate said the correlation had to reach 0.9 and natural data sat
at 0.656, so no recovery run was warranted and none was made. The rank
experiment retires that rule: a synthetic field truncated to rank eight has a
correlation of 0.906 and recovers 13%, while the same field at rank sixteen
recovers 96%. Correlation alone neither predicts nor forbids recovery, and the
width of the shared spectrum is the variable that moves it.

By the width measure the best natural configuration now sits at sixteen shared
directions, inside the interval where the synthetic ladder crosses from 13% to
96%. Its correlation is far lower, so the joint condition is probably not met
-- but the honest way to find out is to run the search rather than to infer
the answer from a statistic that has just been shown not to govern it.

Hidden pairs are generated per trial and read only for scoring.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json
from .spectral_match import grampa, 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", default="artifacts/vg_5k/np_n256.pt")
    parser.add_argument("--label", default="natural-best")
    parser.add_argument("--trials", type=int, default=5)
    parser.add_argument("--etas", type=float, nargs="+", default=[0.2, 0.5, 1.0, 2.0])
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--output", default="artifacts/vg_5k/natural_recovery.json")
    return parser.parse_args()


def offdiagonal(matrix: np.ndarray) -> np.ndarray:
    return matrix[~np.eye(len(matrix), dtype=bool)]


def normalise(matrix: np.ndarray) -> np.ndarray:
    values = offdiagonal(matrix)
    out = (matrix - values.mean()) / values.std()
    np.fill_diagonal(out, 0.0)
    return out


def main() -> None:
    args = parse_args()
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual = normalise(state["visual_field"].double().numpy())
    text = normalise(state["text_field"].double().numpy())
    size = len(visual)
    device = torch.device(args.device)
    swaps = all_swaps(size, device)
    mask = ~np.eye(size, dtype=bool)
    rho = float(np.corrcoef(visual[mask], text[mask])[0, 1])

    rows = []
    for trial in range(args.trials):
        generator = np.random.default_rng(trial)
        hidden = generator.permutation(size)
        shuffled = text[np.ix_(hidden, hidden)]
        energy = ClosedFormEnergy(
            standardized(torch.from_numpy(shuffled).to(device)).float(),
            standardized(torch.from_numpy(visual).to(device)).float(),
            1.0, 1.0, 256,
        )

        def score(columns: np.ndarray) -> float:
            return float((hidden[columns] == np.arange(size)).mean())

        # several spectral starts, keep the one the energy prefers -- model
        # selection by energy only, never by accuracy
        best_energy, best_columns, best_tag = np.inf, None, ""
        candidates = [("umeyama", umeyama(visual, shuffled))]
        candidates += [(f"grampa eta={eta}", grampa(visual, shuffled, eta))
                       for eta in args.etas]
        for tag, columns in candidates:
            start = torch.from_numpy(columns.copy()).to(device)
            value = float(energy.energy(start[None])[0])
            if value < best_energy:
                best_energy, best_columns, best_tag = value, columns, tag

        initial = score(best_columns)
        final, _ = steepest_descent(
            energy, torch.from_numpy(best_columns.copy()).to(device), swaps, 5000
        )
        refined = score(final.cpu().numpy())
        rows.append(
            {
                "trial": trial,
                "chosen_start": best_tag,
                "spectral_accuracy": initial,
                "refined_accuracy": refined,
            }
        )
        print(
            f"trial {trial}: start={best_tag:<14} spectral={initial:.4f} "
            f"refined={refined:.4f}  (chance {1/size:.4f})",
            flush=True,
        )

    summary = {
        "protocol": (
            "Blind: the hidden permutation is generated per trial and used "
            "only for scoring. The spectral start is chosen by energy, never "
            "by accuracy."
        ),
        "label": args.label,
        "size": size,
        "field_correlation_at_truth": rho,
        "chance": 1.0 / size,
        "rows": rows,
        "mean_spectral": float(np.mean([r["spectral_accuracy"] for r in rows])),
        "mean_refined": float(np.mean([r["refined_accuracy"] for r in rows])),
    }
    print(json.dumps({k: v for k, v in summary.items() if k != "rows"}))
    write_json(args.output, summary)


if __name__ == "__main__":
    main()