"""A cheap gate that survives its controls, built as an upper bound.

Three statistics have now been proposed as go/no-go predictors and three have
failed: the field correlation, the width of the shared spectrum, and the
held-out shared dimension. Each failed the same way -- fields agreeing on the
statistic and disagreeing on recovery. The lesson is that a predictor invented
by staring at the fields keeps finding quantities that are not about matching.

So this asks the matching question directly, with one thing given away. Half
the scenes are declared anchors and their correspondence is handed over. Every
remaining scene is described, in each modality, by its field row against the
anchors, and the two descriptions are matched one-to-one by Hungarian
assignment. Nothing is searched: the answer is read off a linear program.

The result upper-bounds what blind recovery can achieve, because blind recovery
must discover the anchor correspondence as well. It is therefore a **necessary
condition**: a representation that cannot be matched with the anchors given
will not be matched without them, and it can be rejected in seconds rather than
after a search. It is an empirical bound rather than a proven one -- the given
anchors also change the estimation problem, not only remove search -- but it
held in every field measured, and unlike the three refuted statistics it
separates the cases they collapse.

It also splits the deficit in two, which is what makes it diagnostic rather
than merely predictive. The bound itself measures whether the representation
identifies scenes at all. The gap between the bound and blind recovery measures
whether the search can find what the representation contains.
"""

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


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("--repeats", type=int, default=5)
    parser.add_argument(
        "--blind", type=float, nargs="*", default=None,
        help="Known blind recovery per field, to report the search gap.",
    )
    parser.add_argument("--output", default="artifacts/vg_5k/anchor_bound.json")
    return parser.parse_args()


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


def rows_against(field: np.ndarray, probe: np.ndarray, anchors: np.ndarray) -> np.ndarray:
    block = field[np.ix_(probe, anchors)]
    block = block - block.mean(1, keepdims=True)
    return block / block.std(1, keepdims=True).clip(1e-9)


def analyse(path: str, label: str, repeats: int, blind: float | None) -> 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)
    mask = ~np.eye(size, dtype=bool)
    rho = float(np.corrcoef(visual[mask], text[mask])[0, 1])

    nearest, assigned = [], []
    for repeat in range(repeats):
        generator = np.random.default_rng(repeat)
        order = generator.permutation(size)
        half = size // 2
        anchors, probe = order[:half], order[half:]
        # The two sides must not be presented in the same index order. Exact
        # twins have identical anchor-restricted rows, and with matched
        # ordering `linear_sum_assignment` breaks the tie onto the diagonal --
        # crediting the bound with information it does not have. On the
        # caption-omitted field that inflated it from 0.920 to 0.997.
        order = generator.permutation(len(probe))
        left = rows_against(visual, probe, anchors)
        right = rows_against(text, probe[order], anchors)
        similarity = left @ right.T / left.shape[1]
        truth = order
        nearest.append(float((order[similarity.argmax(1)] == np.arange(len(probe))).mean()))
        _, columns = linear_sum_assignment(-similarity)
        assigned.append(float((order[columns] == np.arange(len(probe))).mean()))

    row = {
        "label": label,
        "correlation": rho,
        "probe_scenes": size - size // 2,
        "chance": 1.0 / (size - size // 2),
        "nearest_neighbour": float(np.mean(nearest)),
        "anchor_bound": float(np.mean(assigned)),
    }
    if blind is not None:
        row["blind_recovery"] = blind
        row["search_gap"] = float(np.mean(assigned)) - blind
    return row


def main() -> None:
    args = parse_args()
    labels = args.labels or [path.split("/")[-1] for path in args.fields]
    blinds = args.blind if args.blind else [None] * len(args.fields)
    rows = []
    for path, label, blind in zip(args.fields, labels, blinds):
        row = analyse(path, label, args.repeats, blind)
        rows.append(row)
        tail = (
            f"  blind={row['blind_recovery']:.3f}  gap={row['search_gap']:+.3f}"
            if "blind_recovery" in row else ""
        )
        print(
            f"{row['label']:<26} rho={row['correlation']:.3f}  "
            f"bound={row['anchor_bound']:.3f}{tail}",
            flush=True,
        )
    write_json(args.output, {
        "protocol": (
            "Half the scenes are anchors with their correspondence given. The "
            "rest are described by their field rows against the anchors and "
            "matched by Hungarian assignment. Upper-bounds blind recovery, "
            "which must additionally discover the anchor correspondence."
        ),
        "rows": rows,
        "reading": (
            "A low bound means the representation does not identify scenes and "
            "no solver will help. A high bound with low blind recovery means "
            "the information is present and the search or the landscape is at "
            "fault."
        ),
    })


if __name__ == "__main__":
    main()