"""R3 battery evaluation: do model-internal relations align across modalities?

Compares, at the replayed true view correspondence, the cross-modal
alignment of relation fields read from inside the frozen models
(cross-phrase attention per layer group, in-context state cosines) against
the isolated-encoding cosine baseline (matched rho 0.128, z 46.6).
Spearman is rank-based, so monotone attention transforms are immaterial.
View truth is evaluation-only.
"""

from __future__ import annotations

import argparse
import json

import torch
import torch.nn.functional as F

from .common import write_json
from .vg_view_probe import offdiag, replay_view_permutations, spearman


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("--text-attn", default="artifacts/manifold_gate/attn_text_full.pt")
    parser.add_argument(
        "--vision-attn", default="artifacts/manifold_gate/attn_vision_full.pt"
    )
    parser.add_argument("--nodes", type=int, default=0)
    parser.add_argument(
        "--output", default="artifacts/manifold_gate/attention_probe.json"
    )
    return parser.parse_args()


def channel_fields(state: dict, side: str) -> dict[str, torch.Tensor]:
    """Named [N, V, V] relation fields for one side."""
    attention = state["attention_channels"].double()
    fields = {
        f"{side}_attn_g{group}": attention[:, group]
        for group in range(attention.shape[1])
    }
    context = F.normalize(state["context_states"].double(), dim=-1)
    fields[f"{side}_ctx_cos"] = context @ context.transpose(-2, -1)
    return fields


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

    text_state = torch.load(args.text_attn, map_location="cpu", weights_only=False)
    vision_state = torch.load(args.vision_attn, map_location="cpu", weights_only=False)
    text_index = {node: i for i, node in enumerate(text_state["node_ids"])}
    vision_index = {node: i for i, node in enumerate(vision_state["node_ids"])}

    baseline_vision = torch.load(
        f"{args.vg_dir}/vision_features.pt", map_location="cpu", weights_only=False
    )
    baseline_text = torch.load(
        f"{args.vg_dir}/text_features.pt", map_location="cpu", weights_only=False
    )
    baseline_vision_index = {
        node: i for i, node in enumerate(baseline_vision["node_ids"])
    }
    baseline_text_index = {node: i for i, node in enumerate(baseline_text["node_ids"])}

    text_fields = channel_fields(text_state, "text")
    vision_fields = channel_fields(vision_state, "vision")

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

    pairs = [
        (t_name, v_name) for t_name in text_fields for v_name in vision_fields
    ]
    matched: dict[tuple[str, str], list[float]] = {pair: [] for pair in pairs}
    shuffled: dict[tuple[str, str], list[float]] = {pair: [] for pair in pairs}
    baseline_matched: list[float] = []
    baseline_shuffled: list[float] = []
    generator = torch.Generator().manual_seed(args.seed)

    for node in node_ids:
        mapping = mappings[node]
        text_row = text_index[mapping["text_node_id"]]
        vision_row = vision_index[node]
        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))
        shuffle = torch.randperm(args.views, generator=generator)

        for t_name, v_name in pairs:
            text_field = text_fields[t_name][text_row]
            aligned = text_field[vision_to_text][:, vision_to_text]
            visual_field = vision_fields[v_name][vision_row]
            matched[(t_name, v_name)].append(
                spearman(offdiag(visual_field), offdiag(aligned))
            )
            scrambled = text_field[shuffle][:, shuffle]
            shuffled[(t_name, v_name)].append(
                spearman(offdiag(visual_field), offdiag(scrambled))
            )

        crops = F.normalize(
            baseline_vision["region_features"][
                baseline_vision_index[node]
            ].double(),
            dim=-1,
        )
        phrases = F.normalize(
            baseline_text["region_features"][
                baseline_text_index[mapping["text_node_id"]]
            ].double(),
            dim=-1,
        )
        crop_field = crops @ crops.T
        phrase_field = (phrases @ phrases.T)[vision_to_text][:, vision_to_text]
        baseline_matched.append(spearman(offdiag(crop_field), offdiag(phrase_field)))
        scrambled = (phrases @ phrases.T)[shuffle][:, shuffle]
        baseline_shuffled.append(spearman(offdiag(crop_field), offdiag(scrambled)))

    def summarize(matched_values: list[float], shuffled_values: list[float]) -> dict:
        m = torch.tensor(matched_values)
        s = torch.tensor(shuffled_values)
        return {
            "matched_mean": float(m.mean()),
            "shuffled_mean": float(s.mean()),
            "gap_z": float(
                (m.mean() - s.mean())
                / (m - s).std().clamp_min(1e-12)
                * len(m) ** 0.5
            ),
        }

    report = {
        "protocol": (
            "Replayed view truth, evaluation-only. Each cell is the "
            "within-node cross-modal Spearman of relation fields at the "
            "true view correspondence, against a shuffled-view control."
        ),
        "nodes": len(node_ids),
        "baseline_isolated_cosine": summarize(baseline_matched, baseline_shuffled),
        "channels": {
            f"{t} x {v}": summarize(matched[(t, v)], shuffled[(t, v)])
            for t, v in pairs
        },
    }
    write_json(args.output, report)
    best = max(
        report["channels"].items(), key=lambda item: item[1]["matched_mean"]
    )
    print(
        json.dumps(
            {
                "baseline": report["baseline_isolated_cosine"],
                "best_channel": {best[0]: best[1]},
                "nodes": len(node_ids),
            }
        )
    )
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()