"""View-level shared-structure probe for Visual Genome nodes.

The node-level manifold gate shows that pooled node states expose only a
low-dimensional coarse shared subspace, which cannot identify individual
assignments. This probe measures whether additional cross-modal signal
exists inside nodes, at the level of the sixteen region views, where both
modalities observe the same sixteen regions of the same scene.

The per-node view permutation is replayed from the preparation seed and the
private image IDs, verified exactly against the released text bundles, and
used only for evaluation. No cross-modal map is trained.

Three quantities are reported:

1. within-node relational alignment: correlation between the visual and the
   truth-aligned text view-relation fields, against a shuffled-view control;
2. coarse-frame view matching: population-context profiles make single
   views cross-modally comparable through the node-level frame; assignment
   accuracy is chance-normalized (16 candidates per node);
3. a within-node relational-only floor with no population context.
"""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path

import torch
import torch.nn.functional as F
from scipy.optimize import linear_sum_assignment

from .common import write_json
from .vg_prepare import load_config


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vg-dir", default="artifacts/vg_5k")
    parser.add_argument("--cache-dir", default="/tmp/yurenh2-worldalign-vg-hf")
    parser.add_argument("--seed", type=int, default=20260728)
    parser.add_argument("--views", type=int, default=16)
    parser.add_argument(
        "--context-nodes",
        type=int,
        default=1000,
        help="Node-level frame size for population-context profiles.",
    )
    parser.add_argument("--nodes", type=int, default=0, help="0 keeps all nodes.")
    parser.add_argument(
        "--skip-floor",
        action="store_true",
        help="Skip the slow within-node relational-only descent floor.",
    )
    parser.add_argument(
        "--output", default="artifacts/manifold_gate/vg_view_probe.json"
    )
    return parser.parse_args()


def replay_view_permutations(args: argparse.Namespace) -> dict:
    """Recover text-view -> vision-view mappings from the preparation RNG.

    ``random.shuffle`` consumes the generator identically for equal-length
    lists, so shuffling index lists replays the region subsample and the
    phrase shuffle exactly. Every node is verified against the released
    phrase bundle before use.
    """
    truth = [
        json.loads(line)
        for line in Path(args.vg_dir, "ground_truth.private.jsonl").read_text(
            encoding="utf-8"
        ).splitlines()
        if line.strip()
    ]
    text_nodes = {
        record["node_id"]: record
        for line in Path(args.vg_dir, "text_nodes.jsonl").read_text(
            encoding="utf-8"
        ).splitlines()
        if line.strip()
        for record in [json.loads(line)]
    }
    region_ds = load_config("region_descriptions_v1.2.0", args.cache_dir)
    image_rows = {
        int(image_id): row for row, image_id in enumerate(region_ds["image_id"])
    }
    mappings: dict[str, dict] = {}
    mismatched = 0
    for record in truth:
        image_id = int(record["source_image_id"])
        row = region_ds[image_rows[image_id]]
        regions = list(row["regions"])
        node_rng = random.Random(args.seed ^ image_id)
        order = list(range(len(regions)))
        node_rng.shuffle(order)
        selected = order[: args.views]
        phrase_perm = list(range(args.views))
        node_rng.shuffle(phrase_perm)
        phrases_view_order = [
            str(regions[index]["phrase"]).strip() for index in selected
        ]
        reconstructed = [phrases_view_order[j] for j in phrase_perm]
        released = text_nodes[record["text_node_id"]]["region_closed"]
        if reconstructed != released:
            mismatched += 1
            continue
        mappings[record["vision_node_id"]] = {
            "text_node_id": record["text_node_id"],
            # text view j describes vision view phrase_perm[j]
            "text_to_vision": phrase_perm,
        }
    if mismatched:
        print(json.dumps({"replay_mismatched_nodes": mismatched}))
    return mappings


def offdiag(matrix: torch.Tensor) -> torch.Tensor:
    mask = ~torch.eye(len(matrix), dtype=torch.bool)
    return matrix[mask]


def spearman(x: torch.Tensor, y: torch.Tensor) -> float:
    ranks = torch.stack(
        [x.argsort().argsort().double(), y.argsort().argsort().double()]
    )
    return float(torch.corrcoef(ranks)[0, 1])


def main() -> None:
    args = parse_args()
    mappings = replay_view_permutations(args)

    vision = torch.load(
        Path(args.vg_dir, "vision_features.pt"), map_location="cpu", weights_only=False
    )
    text = torch.load(
        Path(args.vg_dir, "text_features.pt"), map_location="cpu", weights_only=False
    )
    vision_index = {node: i for i, node in enumerate(vision["node_ids"])}
    text_index = {node: i for i, node in enumerate(text["node_ids"])}
    visual_views = F.normalize(vision["region_features"].double(), dim=-1)
    text_views = F.normalize(text["region_features"].double(), dim=-1)

    node_ids = sorted(mappings)
    if args.nodes:
        node_ids = node_ids[: args.nodes]

    # Node-level coarse frame: view-mean states over a fixed context set,
    # truth-aligned so both profile axes index the same underlying scenes.
    # The frame is evaluation-side scaffolding for an upper bound; a blind
    # system would substitute its recovered coarse alignment here.
    context_ids = node_ids[: args.context_nodes]
    visual_frame = F.normalize(
        torch.stack(
            [visual_views[vision_index[node]].mean(0) for node in context_ids]
        ),
        dim=-1,
    )
    text_frame = F.normalize(
        torch.stack(
            [
                text_views[text_index[mappings[node]["text_node_id"]]].mean(0)
                for node in context_ids
            ]
        ),
        dim=-1,
    )

    generator = torch.Generator().manual_seed(args.seed)
    relation_matched: list[float] = []
    relation_shuffled: list[float] = []
    profile_hits = 0
    profile_top3 = 0
    profile_total = 0
    floor_hits = 0
    floor_total = 0
    matched_profile_corr: list[float] = []
    mismatched_profile_corr: list[float] = []

    for node in node_ids:
        mapping = mappings[node]
        v = visual_views[vision_index[node]]
        t = text_views[text_index[mapping["text_node_id"]]]
        text_to_vision = torch.tensor(mapping["text_to_vision"])
        vision_to_text = torch.empty_like(text_to_vision)
        vision_to_text[text_to_vision] = torch.arange(len(text_to_vision))

        relation_visual = v @ v.T
        relation_text = t @ t.T
        aligned = relation_text[vision_to_text][:, vision_to_text]
        relation_matched.append(spearman(offdiag(relation_visual), offdiag(aligned)))
        shuffle = torch.randperm(len(v), generator=generator)
        shuffled = relation_text[shuffle][:, shuffle]
        relation_shuffled.append(spearman(offdiag(relation_visual), offdiag(shuffled)))

        # Population-context profiles through the coarse frame.
        profile_visual = v @ visual_frame.T
        profile_text = t @ text_frame.T
        rank_visual = profile_visual.argsort(-1).argsort(-1).double()
        rank_text = profile_text.argsort(-1).argsort(-1).double()
        rank_visual = rank_visual - rank_visual.mean(-1, keepdim=True)
        rank_text = rank_text - rank_text.mean(-1, keepdim=True)
        rank_visual = F.normalize(rank_visual, dim=-1)
        rank_text = F.normalize(rank_text, dim=-1)
        similarity = rank_visual @ rank_text.T  # [vision view, text view]
        rows, cols = linear_sum_assignment(-similarity.numpy())
        truth_cols = vision_to_text.numpy()
        profile_hits += int((cols == truth_cols).sum())
        ranks = (
            similarity
            >= similarity.gather(1, vision_to_text[:, None])
        ).sum(-1)
        profile_top3 += int((ranks <= 3).sum())
        profile_total += len(rows)
        matched_profile_corr.extend(
            similarity.gather(1, vision_to_text[:, None]).squeeze(1).tolist()
        )
        mismatched_profile_corr.extend(
            similarity[~torch.eye(len(v), dtype=torch.bool)][:64].tolist()
        )

        # Relational-only floor: 2-swap descent on within-node fields from a
        # random start, no population context.
        if args.skip_floor:
            continue
        state = torch.randperm(len(v), generator=generator)
        for _ in range(200):
            improved = False
            current = relation_text[state][:, state]
            base = float(((relation_visual - current) ** 2).mean())
            for p in range(len(v)):
                for q in range(p + 1, len(v)):
                    trial = state.clone()
                    trial[[p, q]] = trial[[q, p]]
                    candidate = relation_text[trial][:, trial]
                    if float(((relation_visual - candidate) ** 2).mean()) < base - 1e-12:
                        state = trial
                        improved = True
                        break
                if improved:
                    break
            if not improved:
                break
        floor_hits += int((state == vision_to_text).sum())
        floor_total += len(v)

    matched = torch.tensor(relation_matched)
    shuffled = torch.tensor(relation_shuffled)
    matched_corr = torch.tensor(matched_profile_corr)
    mismatched_corr = torch.tensor(mismatched_profile_corr)
    report = {
        "protocol": (
            "View permutations are replayed from the preparation seed and "
            "verified against released phrase bundles; they are used only "
            "for evaluation. The coarse frame uses truth-aligned node "
            "states, so profile matching is an upper bound on what a "
            "recovered node alignment would enable."
        ),
        "nodes": len(node_ids),
        "verified_fraction": len(mappings) / 5000.0,
        "within_node_relation": {
            "matched_spearman_mean": float(matched.mean()),
            "matched_spearman_std": float(matched.std()),
            "shuffled_spearman_mean": float(shuffled.mean()),
            "shuffled_spearman_std": float(shuffled.std()),
            "gap_z_across_nodes": float(
                (matched.mean() - shuffled.mean())
                / (matched - shuffled).std().clamp_min(1e-12)
                * len(matched) ** 0.5
            ),
        },
        "coarse_frame_view_matching": {
            "context_nodes": len(context_ids),
            "accuracy": profile_hits / max(profile_total, 1),
            "top3_accuracy": profile_top3 / max(profile_total, 1),
            "chance": 1.0 / args.views,
            "matched_profile_corr_mean": float(matched_corr.mean()),
            "mismatched_profile_corr_mean": float(mismatched_corr.mean()),
        },
        "within_node_relational_floor": {
            "accuracy": floor_hits / max(floor_total, 1),
            "chance": 1.0 / args.views,
        },
    }
    write_json(args.output, report)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()