"""Is the truth the optimum, tied with it, or beaten? Three descents, not seven hundred.

The question that decides the next move is narrow, and the first attempt at it
was not. Suppressing a caption factor leaves the information intact -- the
anchor bound is 0.997 -- while blind matching returns 5.6%, and there are three
possible reasons with three different responses:

    E(truth) <  E(best found)   the searcher simply missed it   -> better optimiser
    E(truth) == E(best found)   the truth is one of many minima  -> quotient the symmetry
    E(truth) >  E(best found)   the truth is not the optimum     -> information is gone

Answering it does not need a strong searcher. It needs the energy of the truth
and the energy of whatever the standard pipeline actually converges to, which
is one descent per trial rather than the sixty-one the first version ran -- at
N=256 a single descent over all 32,640 transpositions costs minutes, so the
difference is hours against seconds.
"""

from __future__ import annotations

import argparse
import json

import numpy as np
import torch

from .common import write_json
from .spectral_match import grampa
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("--tolerance", type=float, default=1e-3)
    parser.add_argument("--output", default="artifacts/synth_v1/degeneracy_gate.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 analyse(path: str, label: str, args: argparse.Namespace) -> dict:
    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)
    device = torch.device(args.device)

    gaps, accuracies, truths = [], [], []
    for trial in range(args.trials):
        generator = np.random.default_rng(trial)
        hidden = generator.permutation(size)
        shuffled = text[np.ix_(hidden, hidden)]
        shuffled_gpu = standardized(torch.from_numpy(shuffled).to(device)).double()
        visual_gpu = standardized(torch.from_numpy(visual).to(device)).double()
        energy = ClosedFormEnergy(shuffled_gpu, visual_gpu, 1.0, 0.0, 256)
        truth = torch.from_numpy(np.argsort(hidden).copy()).to(device)
        truth_energy = float(energy.energy(truth[None])[0])

        start = torch.from_numpy(grampa(visual, shuffled, 1.0).copy()).to(device)
        final = fast_pair_descent(shuffled_gpu, visual_gpu, start, args.iterations)
        found_energy = float(energy.energy(final[None])[0])
        # descent from the truth itself: if it moves, the truth is not a local
        # minimum either, which separates "not optimal" from "merely not found"
        from_truth = fast_pair_descent(
            shuffled_gpu, visual_gpu, truth.clone(), args.iterations
        )
        truth_local = float(energy.energy(from_truth[None])[0])

        gaps.append(found_energy - truth_energy)
        truths.append(truth_energy - truth_local)
        accuracies.append(float((hidden[final.cpu().numpy()] == np.arange(size)).mean()))

    gap = float(np.mean(gaps))
    scale = abs(float(np.mean([g for g in gaps]))) + 1.0
    if gap < -args.tolerance:
        verdict = "truth is NOT the optimum -- information deficit"
    elif abs(gap) <= args.tolerance:
        verdict = "truth is TIED with what is found -- degenerate, quotient the symmetry"
    else:
        verdict = "truth is DEEPER than what is found -- optimiser failure"
    return {
        "label": label,
        "size": size,
        "energy_gap_found_minus_truth": gap,
        "truth_descends_further_by": float(np.mean(truths)),
        "blind_accuracy": float(np.mean(accuracies)),
        "verdict": verdict,
    }


def main() -> None:
    args = parse_args()
    labels = args.labels or [p.split("/")[-1] for p in args.fields]
    rows = []
    for path, label in zip(args.fields, labels):
        row = analyse(path, label, args)
        rows.append(row)
        print(
            f"{row['label']:<20} gap(found-truth)={row['energy_gap_found_minus_truth']:+.4f}  "
            f"truth-descends={row['truth_descends_further_by']:+.4f}  "
            f"acc={row['blind_accuracy']:.3f}\n"
            f"{'':<20} -> {row['verdict']}",
            flush=True,
        )
    write_json(args.output, {
        "protocol": (
            "One spectral-start descent per trial, plus one descent started at "
            "the truth. A positive gap means the searcher stopped above the "
            "truth; zero means it found something the energy cannot "
            "distinguish from the truth; negative means the truth is beaten."
        ),
        "rows": rows,
    })
    print(json.dumps({"done": True}))


if __name__ == "__main__":
    main()