"""Close the loop: recovered pairs to a map to generated descriptions.

Matching is transductive -- it aligns one fixed population. A usable
system needs a map, so the recovered assignment is treated as pseudo-pair
supervision for a small vision-to-text-state regressor, and the map is
then applied to scenes that took no part in the matching. Descriptions
are read out through the frozen text tower.

Two controls decide whether the output is image-conditioned: the same
pipeline with the recovered assignment replaced by a random one, and the
same map applied to a shuffled image. Hidden pairs score; nothing here
trains on them.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F

from .common import read_json, seed_everything, write_json


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--data-dir", default="artifacts/synth_v1")
    parser.add_argument("--states", required=True,
                        help=".pt with vision_states, text_states, rows for "
                             "the matched population and the held-out split.")
    parser.add_argument("--assignment", required=True,
                        help=".pt with the recovered permutation and truth.")
    parser.add_argument("--ridge", type=float, default=1.0)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--output", default="artifacts/synth_v1/decode.json")
    return parser.parse_args()


def fit_map(source: np.ndarray, target: np.ndarray, ridge: float) -> np.ndarray:
    source = np.concatenate([source, np.ones((len(source), 1))], axis=1)
    gram = source.T @ source + ridge * np.eye(source.shape[1])
    return np.linalg.solve(gram, source.T @ target)


def apply_map(weights: np.ndarray, source: np.ndarray) -> np.ndarray:
    source = np.concatenate([source, np.ones((len(source), 1))], axis=1)
    return source @ weights


def nearest_caption(
    predicted: np.ndarray, bank: np.ndarray
) -> np.ndarray:
    predicted = predicted / np.linalg.norm(predicted, axis=1, keepdims=True).clip(1e-9)
    bank = bank / np.linalg.norm(bank, axis=1, keepdims=True).clip(1e-9)
    return (predicted @ bank.T).argmax(1)


def factor_scores(
    predicted_rows: list[int], truth_rows: list[int], scenes: list[dict]
) -> dict:
    """Do the retrieved descriptions state the right world facts?"""
    colour_f1, count_ok, group_ok = [], [], []
    for predicted, truth in zip(predicted_rows, truth_rows):
        p, t = scenes[predicted], scenes[truth]
        pc = {g["color"] for g in p["groups"]}
        tc = {g["color"] for g in t["groups"]}
        overlap = len(pc & tc)
        precision = overlap / max(len(pc), 1)
        recall = overlap / max(len(tc), 1)
        colour_f1.append(
            0.0 if precision + recall == 0 else 2 * precision * recall / (precision + recall)
        )
        count_ok.append(
            sorted(g["count"] for g in p["groups"])
            == sorted(g["count"] for g in t["groups"])
        )
        group_ok.append(len(p["groups"]) == len(t["groups"]))
    return {
        "colour_set_f1": float(np.mean(colour_f1)),
        "count_multiset_exact": float(np.mean(count_ok)),
        "group_count_exact": float(np.mean(group_ok)),
    }


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
    states = torch.load(args.states, map_location="cpu", weights_only=False)
    recovered = torch.load(args.assignment, map_location="cpu", weights_only=False)

    match_vision = states["match_vision"].double().numpy()
    match_text = states["match_text"].double().numpy()
    held_vision = states["held_vision"].double().numpy()
    held_text = states["held_text"].double().numpy()
    held_rows = states["held_rows"]

    assignment = np.asarray(recovered["assignment"])
    generator = np.random.default_rng(args.seed)
    random_assignment = generator.permutation(len(assignment))

    report = {
        "protocol": (
            "The recovered assignment supplies pseudo-pairs for a ridge "
            "map from vision states to text states; the map is applied to "
            "held-out scenes that took no part in matching. Random "
            "assignment and shuffled-image controls bound the claim."
        ),
        "matched_population": len(assignment),
        "held_out": len(held_rows),
        "recovery_accuracy": float(
            (assignment == np.arange(len(assignment))).mean()
        ),
        "conditions": {},
    }

    for label, permutation in (
        ("recovered_pairs", assignment),
        ("random_pairs", random_assignment),
    ):
        weights = fit_map(match_vision, match_text[permutation], args.ridge)
        predicted = apply_map(weights, held_vision)
        retrieved = nearest_caption(predicted, held_text)
        exact = float((retrieved == np.arange(len(held_rows))).mean())
        entry = {
            "held_out_retrieval_exact": exact,
            "chance": 1.0 / len(held_rows),
            **factor_scores(
                [held_rows[i] for i in retrieved], held_rows, scenes
            ),
        }
        if label == "recovered_pairs":
            shuffled = generator.permutation(len(held_vision))
            predicted_shuffled = apply_map(weights, held_vision[shuffled])
            retrieved_shuffled = nearest_caption(predicted_shuffled, held_text)
            entry["shuffled_image_control"] = factor_scores(
                [held_rows[i] for i in retrieved_shuffled], held_rows, scenes
            )
        report["conditions"][label] = entry
        print(json.dumps({label: entry}))

    write_json(args.output, report)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()