"""Amplification: the instance that beat fifteen solvers, solved by not solving it.

Fifteen methods returned under 5% on the caption-omitted field -- spectral,
Birkhoff, five Gromov-Wasserstein variants, convex-concave path following. The
measurements said why, and the answer was not a sixteenth solver.

Three facts fit together. The information is intact (anchor bound 0.997). The
truth is a strict local minimum whose basin is wide. And exact descent
**amplifies**: a start that is 10% correct comes out 42% correct, one that is
20% correct comes out 78%. What no initialiser could do was produce a start
correct enough to enter that regime, because every one of them lands in the
same systematically wrong region.

So the move is to stop asking any single initialiser to be right. Run a diverse
pool of cheap descents, let them vote on pairings, round the vote matrix to a
permutation with Hungarian assignment -- rounding, not taking the argmax, which
is worth several points on its own -- descend from that, and rebuild the pool by
perturbing the result. Each round feeds the amplifier a better start than the
last. Nothing here is a better optimiser; it is a ladder built out of the one
the field already had.

Ladders vary: on the hard field one run plateaus near 0.36 while another climbs
to 0.72. Several are run and the best is chosen **by energy**, never by
accuracy, so the selection stays blind.
"""

from __future__ import annotations

import argparse
import json

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

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("--pool", type=int, default=60)
    parser.add_argument("--rounds", type=int, default=8)
    parser.add_argument("--ladders", type=int, default=3)
    parser.add_argument("--trials", type=int, default=3)
    parser.add_argument("--iterations", type=int, default=2000)
    parser.add_argument("--fresh", type=float, default=0.25,
                        help="Fraction of each rebuilt pool that stays fully random.")
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--output", default="artifacts/synth_v1/amplify.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 ladder(text_gpu, visual_gpu, visual, shuffled, energy, args, generator, device):
    """One amplification ladder. Returns (permutation, energy, per-round trace)."""
    size = len(visual)
    pool = [grampa(visual, shuffled, eta) for eta in (0.2, 0.5, 1.0, 2.0, 5.0)]
    pool += [generator.permutation(size) for _ in range(args.pool - 5)]

    best, best_energy, trace = None, np.inf, []
    for round_index in range(args.rounds):
        votes = np.zeros((size, size))
        for start in pool:
            final = fast_pair_descent(
                text_gpu, visual_gpu,
                torch.from_numpy(np.ascontiguousarray(start)).to(device),
                args.iterations,
            ).cpu().numpy()
            votes[np.arange(size), final] += 1.0

        _, columns = linear_sum_assignment(-votes)
        descended = fast_pair_descent(
            text_gpu, visual_gpu,
            torch.from_numpy(np.ascontiguousarray(columns)).to(device),
            args.iterations,
        )
        value = float(energy.energy(descended[None])[0])
        current = descended.cpu().numpy()
        if value < best_energy:
            best, best_energy = current, value
        trace.append(value)

        # rebuild the pool around the current answer, keeping some of it random
        # so a plateaued ladder still has a way out
        fresh = int(args.fresh * args.pool)
        pool = [current.copy()]
        for _ in range(args.pool - fresh - 1):
            perturbed = current.copy()
            for _ in range(int(generator.integers(10, 120))):
                i, j = generator.integers(0, size, 2)
                perturbed[[i, j]] = perturbed[[j, i]]
            pool.append(perturbed)
        pool += [generator.permutation(size) for _ in range(fresh)]
    return best, best_energy, trace


def main() -> None:
    args = parse_args()
    labels = args.labels or [p.split("/")[-1] for p in args.fields]
    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)
        visual_gpu = standardized(torch.from_numpy(visual).to(device)).double()
        print(f"\n=== {label} (N={size}, chance={1/size:.4f})", flush=True)

        chosen, oracle_best = [], []
        for trial in range(args.trials):
            hidden = np.random.default_rng(trial).permutation(size)
            shuffled = text[np.ix_(hidden, hidden)]
            text_gpu = standardized(torch.from_numpy(shuffled).to(device)).double()
            energy = ClosedFormEnergy(text_gpu, visual_gpu, 1.0, 0.0, 256)

            results = []
            for index in range(args.ladders):
                generator = np.random.default_rng(1000 * trial + index)
                permutation, value, _ = ladder(
                    text_gpu, visual_gpu, visual, shuffled, energy,
                    args, generator, device,
                )
                accuracy = float((hidden[permutation] == np.arange(size)).mean())
                results.append((value, accuracy))
            # selection by energy only -- the hidden pairing is never consulted
            picked = min(results, key=lambda item: item[0])[1]
            chosen.append(picked)
            oracle_best.append(max(r[1] for r in results))
            print(f"  trial {trial}: ladders={[round(a, 3) for _, a in results]}  "
                  f"picked-by-energy={picked:.4f}", flush=True)

        row = {
            "field": label, "size": size, "chance": 1.0 / size,
            "accuracy_selected_by_energy": float(np.mean(chosen)),
            "accuracy_if_oracle_picked_best_ladder": float(np.mean(oracle_best)),
        }
        rows.append(row)
        print(f"  MEAN selected-by-energy={row['accuracy_selected_by_energy']:.4f}  "
              f"(oracle ladder choice would give "
              f"{row['accuracy_if_oracle_picked_best_ladder']:.4f})", flush=True)

    write_json(args.output, {
        "protocol": (
            "Blind. Hidden permutation generated per trial and read only for "
            "scoring; the ladder is chosen by energy, never by accuracy."
        ),
        "pool": args.pool, "rounds": args.rounds, "ladders": args.ladders,
        "rows": rows,
    })
    print(json.dumps({"done": True}))


if __name__ == "__main__":
    main()