"""Does choosing mutually dissimilar scenes buy anything, and at what price?

If the cross-modal agreement lives in a narrow subspace, then which scenes
are matched matters as much as how many. Scenes that sit close together in
that subspace are the decoys: nothing in either field separates them, so a
solver has no way to tell them apart. Spreading the population out should
help more than enlarging it.

Selection has to be one-sided to stay legal -- picking diverse images and
diverse captions independently yields two sets that are not in bijection,
which the permutation formulation cannot express. So this runs as a
**diagnostic**: scenes are chosen by visual dissimilarity alone and their
partners are taken from the hidden pairing. That is a protocol violation and
the result is an upper bound, not a method. It is run first because it is
cheap and it prices the idea: if diversity does not help even when the
counterpart set is handed over, the legal unbalanced version is not worth
building.

Reported against a random-subset control of the same size and the same
pipeline.
"""

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, 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/natural_pipeline_w64.pt")
    parser.add_argument("--size", type=int, default=128)
    parser.add_argument("--trials", type=int, default=5)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument("--output", default="artifacts/vg_5k/diversity_subset.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 farthest_point(field: np.ndarray, size: int, start: int) -> np.ndarray:
    """Greedy farthest-point selection under the visual field's similarity.

    High field entry means similar, so dissimilarity is the negated entry and
    each new scene is the one whose closest already-chosen neighbour is as far
    as possible.
    """
    chosen = [start]
    closest = field[start].copy()
    for _ in range(size - 1):
        candidate = int(np.argmin(np.where(np.isin(np.arange(len(field)), chosen),
                                           np.inf, closest)))
        chosen.append(candidate)
        closest = np.maximum(closest, field[candidate])
    return np.array(chosen)


def evaluate(visual: np.ndarray, text: np.ndarray, device, trials: int) -> dict:
    size = len(visual)
    visual = normalise(visual)
    text = normalise(text)
    mask = ~np.eye(size, dtype=bool)
    rho = float(np.corrcoef(visual[mask], text[mask])[0, 1])
    swaps = all_swaps(size, device)
    accuracies = []
    for trial in range(trials):
        generator = np.random.default_rng(trial)
        hidden = generator.permutation(size)
        shuffled = text[np.ix_(hidden, hidden)]
        columns = grampa(visual, shuffled, 1.0)
        energy = ClosedFormEnergy(
            standardized(torch.from_numpy(shuffled).to(device)).float(),
            standardized(torch.from_numpy(visual).to(device)).float(),
            1.0, 1.0, 256,
        )
        final, _ = steepest_descent(
            energy, torch.from_numpy(columns.copy()).to(device), swaps, 3000
        )
        accuracies.append(float((hidden[final.cpu().numpy()] == np.arange(size)).mean()))
    # how much of the correlation survives stripping the leading directions
    symmetric = (visual + visual.T) / 2.0
    values, vectors = np.linalg.eigh(symmetric)
    order = np.argsort(np.abs(values))[::-1][10:]
    stripped_visual = (vectors[:, order] * values[order]) @ vectors[:, order].T
    symmetric_text = (text + text.T) / 2.0
    values_t, vectors_t = np.linalg.eigh(symmetric_text)
    order_t = np.argsort(np.abs(values_t))[::-1][10:]
    stripped_text = (vectors_t[:, order_t] * values_t[order_t]) @ vectors_t[:, order_t].T
    retained = float(
        np.corrcoef(offdiagonal(stripped_visual), offdiagonal(stripped_text))[0, 1]
    )
    return {
        "correlation": rho,
        "correlation_after_strip10": retained,
        "recovery": float(np.mean(accuracies)),
        "chance": 1.0 / size,
    }


def main() -> None:
    args = parse_args()
    state = torch.load(args.fields, map_location="cpu", weights_only=False)
    visual_full = state["visual_field"].double().numpy()
    text_full = state["text_field"].double().numpy()
    population = len(visual_full)
    device = torch.device(args.device)
    size = min(args.size, population)

    rows = []
    for start in range(3):
        index = farthest_point(visual_full, size, start * 37 % population)
        result = evaluate(
            visual_full[np.ix_(index, index)], text_full[np.ix_(index, index)],
            device, args.trials,
        )
        result["selection"] = f"diverse (start {start})"
        rows.append(result)
        print(json.dumps(result), flush=True)

    for trial in range(3):
        generator = np.random.default_rng(100 + trial)
        index = generator.choice(population, size, replace=False)
        result = evaluate(
            visual_full[np.ix_(index, index)], text_full[np.ix_(index, index)],
            device, args.trials,
        )
        result["selection"] = f"random (seed {trial})"
        rows.append(result)
        print(json.dumps(result), flush=True)

    diverse = [r for r in rows if r["selection"].startswith("diverse")]
    random_rows = [r for r in rows if r["selection"].startswith("random")]
    summary = {
        "protocol": (
            "DIAGNOSTIC, NOT A METHOD: scenes selected by visual dissimilarity "
            "and their partners taken from the hidden pairing. One-sided "
            "selection cannot be realised without the pairing, so this is an "
            "upper bound used to price the idea before building the legal "
            "unbalanced version."
        ),
        "population": population,
        "subset_size": size,
        "rows": rows,
        "diverse_mean_correlation": float(np.mean([r["correlation"] for r in diverse])),
        "random_mean_correlation": float(np.mean([r["correlation"] for r in random_rows])),
        "diverse_mean_recovery": float(np.mean([r["recovery"] for r in diverse])),
        "random_mean_recovery": float(np.mean([r["recovery"] for r in random_rows])),
    }
    print(json.dumps({k: v for k, v in summary.items() if k != "rows"}))
    write_json(args.output, summary)


if __name__ == "__main__":
    main()