path: root/worldalign/view_anchor_battery.py

"""C5 second wave: fine-grained anchors at view level.

The node-level anchor wave failed because scene-level attributes are
coarse-redundant with the shared subspace. Region-level attributes are
not: a phrase names the color, relative size, lighting, or position of
one region, and the region's own pixels and box realize them. With
sixteen candidates per node, weak per-view anchors carry usable bits.

Per-view anchors, computed independently per modality:

- text, per phrase: color-lexicon histogram, size score, light score,
  horizontal/vertical position scores, numeral content;
- vision, per region: crop hue-band histogram, crop luminance, box area
  fraction, box center coordinates.

Scalar channels are rank-standardized within the node (sixteen values),
so they compare relative attributes inside one scene; lexicon-to-physics
maps are frozen world knowledge. Channels contribute only where the
phrase carries the attribute. Replayed view truth scores matching;
nothing here uses node pairs or the coarse frame.
"""

from __future__ import annotations

import argparse
import json
import re
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import numpy as np
import torch
from PIL import Image
from scipy.optimize import linear_sum_assignment

from .anchor_battery import (
    COLOR_BANDS,
    LIGHT_WORDS,
    NUMBER_WORDS,
    SIZE_WORDS,
    read_jsonl,
)
from .common import write_json
from .vg_view_probe import replay_view_permutations

HORIZONTAL_WORDS = {"left": -1.0, "right": 1.0}
VERTICAL_WORDS = {"top": -1.0, "upper": -1.0, "above": -1.0, "bottom": 1.0, "lower": 1.0, "below": 1.0}
# Twelve color classes: eight hue bands plus achromatic and brown classes
# with pixel rules on value and saturation. Frozen world knowledge.
COLOR_CLASSES = list(COLOR_BANDS) + ["white", "black", "gray", "brown"]
COLOR_SYNONYMS = {"grey": "gray", "tan": "brown", "beige": "brown", "golden": "yellow", "gold": "yellow", "silver": "gray"}


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("--image-cache", default="/tmp/yurenh2-worldalign-vg-images")
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument("--structure", default="artifacts/manifold_gate/attn_text_full.pt")
    parser.add_argument(
        "--vision-structure", default="artifacts/manifold_gate/attn_vision_full.pt"
    )
    parser.add_argument("--nodes", type=int, default=0)
    parser.add_argument(
        "--output", default="artifacts/manifold_gate/view_anchor_battery.json"
    )
    parser.add_argument(
        "--anchors-output", default="artifacts/manifold_gate/view_anchor_maps.pt"
    )
    return parser.parse_args()


def phrase_features(phrase: str) -> dict:
    tokens = [
        COLOR_SYNONYMS.get(token, token)
        for token in re.findall(r"[a-z]+|\d+", phrase.lower())
    ]
    color = np.zeros(len(COLOR_CLASSES))
    for index, name in enumerate(COLOR_CLASSES):
        color[index] = sum(token == name for token in tokens)
    numeral = 0.0
    for token in tokens:
        if token.isdigit() and int(token) <= 50:
            numeral += int(token)
        elif token in NUMBER_WORDS:
            numeral += NUMBER_WORDS[token]
    return {
        "color": color,
        "size": float(sum(SIZE_WORDS.get(token, 0.0) for token in tokens)),
        "light": float(sum(LIGHT_WORDS.get(token, 0.0) for token in tokens)),
        "horizontal": float(sum(HORIZONTAL_WORDS.get(token, 0.0) for token in tokens)),
        "vertical": float(sum(VERTICAL_WORDS.get(token, 0.0) for token in tokens)),
        "numeral": numeral,
    }


def region_features(record: dict, args: argparse.Namespace) -> dict:
    path = Path(args.image_cache, f"{record['source_image_id']}.jpg")
    with Image.open(path) as image:
        hsv = np.asarray(image.convert("HSV"), dtype=np.float64)
    height, width = hsv.shape[:2]
    hues, lums, areas, xs, ys = [], [], [], [], []
    for region in record["regions"]:
        x0 = max(0, min(int(region["x"]), width - 1))
        y0 = max(0, min(int(region["y"]), height - 1))
        x1 = max(x0 + 1, min(x0 + int(region["width"]), width))
        y1 = max(y0 + 1, min(y0 + int(region["height"]), height))
        crop = hsv[y0:y1, x0:x1]
        hue = crop[..., 0] * (360.0 / 255.0)
        saturation = crop[..., 1] / 255.0
        value = crop[..., 2] / 255.0
        white = (value > 0.75) & (saturation < 0.25)
        black = value < 0.25
        gray = (~white) & (~black) & (saturation < 0.25)
        brown_hue = np.minimum(np.abs(hue - 30.0), 360.0 - np.abs(hue - 30.0)) < 25.0
        brown = brown_hue & (saturation >= 0.25) & (value >= 0.2) & (value < 0.55)
        chromatic = (saturation >= 0.25) & (value >= 0.2) & (~brown)
        histogram = np.zeros(len(COLOR_CLASSES))
        for index, center in enumerate(COLOR_BANDS.values()):
            distance = np.minimum(np.abs(hue - center), 360.0 - np.abs(hue - center))
            histogram[index] = float(((distance < 25.0) & chromatic).mean())
        histogram[len(COLOR_BANDS) + 0] = float(white.mean())
        histogram[len(COLOR_BANDS) + 1] = float(black.mean())
        histogram[len(COLOR_BANDS) + 2] = float(gray.mean())
        histogram[len(COLOR_BANDS) + 3] = float(brown.mean())
        hues.append(histogram)
        lums.append(float(value.mean()))
        areas.append((x1 - x0) * (y1 - y0) / (width * height))
        xs.append((x0 + x1) / 2.0 / width)
        ys.append((y0 + y1) / 2.0 / height)
    return {
        "hue": np.stack(hues),
        "luminance": np.array(lums),
        "area": np.array(areas),
        "x": np.array(xs),
        "y": np.array(ys),
    }


def rank_z(values: np.ndarray) -> np.ndarray:
    order = values.argsort().argsort().astype(np.float64)
    std = order.std()
    return (order - order.mean()) / (std if std > 1e-9 else 1.0)


def node_anchor_matrix(
    text: list[dict], vision: dict
) -> tuple[np.ndarray, np.ndarray]:
    """[V, V] anchor similarity (vision rows, text columns) and coverage."""
    views = len(text)
    channels = []
    coverage = np.zeros(5)

    text_color = np.stack([p["color"] for p in text])
    informative = text_color.sum(1) > 0
    if informative.any():
        tn = text_color / np.linalg.norm(text_color, axis=1, keepdims=True).clip(min=1e-9)
        vn = vision["hue"] / np.linalg.norm(vision["hue"], axis=1, keepdims=True).clip(min=1e-9)
        color = vn @ tn.T
        color[:, ~informative] = 0.0
        flat = color[:, informative]
        scale = flat.std()
        channels.append(color / (scale if scale > 1e-9 else 1.0))
        coverage[0] = informative.mean()

    for slot, (text_key, vision_values, sign) in enumerate(
        (
            ("size", rank_z(vision["area"]), 1.0),
            ("light", rank_z(vision["luminance"]), 1.0),
            ("horizontal", rank_z(vision["x"]), 1.0),
            ("vertical", rank_z(vision["y"]), 1.0),
        ),
        start=1,
    ):
        scores = np.array([p[text_key] for p in text])
        informative = scores != 0
        if informative.sum() < 2:
            continue
        text_rank = rank_z(scores)
        similarity = -np.abs(vision_values[:, None] - sign * text_rank[None, :])
        similarity[:, ~informative] = 0.0
        flat = similarity[:, informative]
        scale = flat.std()
        channels.append(similarity / (scale if scale > 1e-9 else 1.0))
        coverage[slot] = informative.mean()
    if not channels:
        return np.zeros((views, views)), coverage
    return np.mean(channels, axis=0), coverage


def structure_signature_matrix(
    text_views: torch.Tensor, visual_views: torch.Tensor
) -> np.ndarray:
    """Frame-free structural channel: sorted within-node relation rows."""
    def signatures(views: torch.Tensor) -> torch.Tensor:
        views = torch.nn.functional.normalize(views.double(), dim=-1)
        relation = views @ views.T
        size = len(relation)
        mask = ~torch.eye(size, dtype=torch.bool)
        rows = relation.masked_select(mask).reshape(size, size - 1)
        rows = (rows - rows.mean()) / rows.std().clamp_min(1e-9)
        return rows.sort(-1).values

    a = signatures(visual_views)
    b = signatures(text_views)
    cost = ((a[:, None, :] - b[None, :, :]) ** 2).sum(-1)
    similarity = -cost
    return ((similarity - similarity.mean()) / similarity.std().clamp_min(1e-9)).numpy()


def main() -> None:
    args = parse_args()
    mappings = replay_view_permutations(args)
    text_nodes = {
        record["node_id"]: record
        for record in read_jsonl(Path(args.vg_dir, "text_nodes.jsonl"))
    }
    vision_nodes = {
        record["node_id"]: record
        for record in read_jsonl(Path(args.vg_dir, "vision_nodes.private.jsonl"))
    }
    text_structure = torch.load(args.structure, map_location="cpu", weights_only=False)
    vision_structure = torch.load(
        args.vision_structure, map_location="cpu", weights_only=False
    )
    text_index = {node: i for i, node in enumerate(text_structure["node_ids"])}
    vision_index = {node: i for i, node in enumerate(vision_structure["node_ids"])}

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

    vision_records = [vision_nodes[node] for node in node_ids]
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        vision_results = list(
            pool.map(lambda record: region_features(record, args), vision_records)
        )

    generator = np.random.default_rng(args.seed)
    accuracies = {"anchors": [], "structure": [], "anchors_plus_structure": []}
    matched_scores, shuffled_scores = [], []
    coverage_totals = np.zeros(5)
    anchor_maps = {}
    informative_view_hits: list[bool] = []
    seed_pool: list[tuple[float, bool]] = []
    text_view_features: dict[str, dict] = {}
    vision_view_features: dict[str, dict] = {}

    for node, vision_features_one in zip(node_ids, vision_results):
        mapping = mappings[node]
        phrases = text_nodes[mapping["text_node_id"]]["region_closed"]
        text_features_list = [phrase_features(p) for p in phrases]
        anchors, coverage = node_anchor_matrix(text_features_list, vision_features_one)
        coverage_totals += coverage
        text_to_vision = np.array(mapping["text_to_vision"])
        vision_to_text = np.empty_like(text_to_vision)
        vision_to_text[text_to_vision] = np.arange(len(text_to_vision))

        structure = structure_signature_matrix(
            torch.as_tensor(
                text_structure["context_states"][text_index[mapping["text_node_id"]]]
            ),
            torch.as_tensor(vision_structure["context_states"][vision_index[node]]),
        )
        combined = anchors + structure

        informative_columns = np.abs(anchors).sum(0) > 1e-9
        for name, matrix in (
            ("anchors", anchors),
            ("structure", structure),
            ("anchors_plus_structure", combined),
        ):
            rows, cols = linear_sum_assignment(-matrix)
            accuracies[name].append(float((cols == vision_to_text).mean()))
            if name == "anchors" and informative_columns.any():
                assigned_text = cols  # vision row i -> text column
                correct = assigned_text == vision_to_text
                text_informative_hit = informative_columns[assigned_text]
                informative_view_hits.extend(
                    correct[text_informative_hit].tolist()
                )
                sorted_scores = np.sort(matrix, axis=1)[:, ::-1]
                margins = sorted_scores[:, 0] - sorted_scores[:, 1]
                for row in range(len(matrix)):
                    if informative_columns[assigned_text[row]]:
                        seed_pool.append(
                            (float(margins[row]), bool(correct[row]))
                        )
        matched = anchors[np.arange(len(anchors)), vision_to_text]
        shuffle = generator.permutation(len(anchors))
        matched_scores.append(float(matched.mean()))
        shuffled_scores.append(
            float(anchors[np.arange(len(anchors)), shuffle].mean())
        )
        anchor_maps[node] = torch.from_numpy(anchors).float()
        text_view_features[mapping["text_node_id"]] = {
            "color": np.stack([p["color"] for p in text_features_list]),
            "size": np.array([p["size"] for p in text_features_list]),
            "light": np.array([p["light"] for p in text_features_list]),
            "horizontal": np.array([p["horizontal"] for p in text_features_list]),
            "vertical": np.array([p["vertical"] for p in text_features_list]),
            "numeral": np.array([p["numeral"] for p in text_features_list]),
        }
        vision_view_features[node] = vision_features_one

    matched_arr = np.array(matched_scores)
    shuffled_arr = np.array(shuffled_scores)
    report = {
        "protocol": (
            "Per-view anchors and frame-free structural signatures; "
            "replayed view truth scores matching only. No node pairs, no "
            "coarse frame."
        ),
        "nodes": len(node_ids),
        "chance": 1.0 / args.views,
        "true_frame_profile_baseline": 0.1448,
        "channel_coverage_mean": {
            "color": float(coverage_totals[0] / len(node_ids)),
            "size": float(coverage_totals[1] / len(node_ids)),
            "light": float(coverage_totals[2] / len(node_ids)),
            "horizontal": float(coverage_totals[3] / len(node_ids)),
            "vertical": float(coverage_totals[4] / len(node_ids)),
        },
        "anchor_matched_vs_shuffled_z": float(
            (matched_arr - shuffled_arr).mean()
            / (matched_arr - shuffled_arr).std().clip(min=1e-12)
            * np.sqrt(len(matched_arr))
        ),
        "view_matching_accuracy": {
            name: float(np.mean(values)) for name, values in accuracies.items()
        },
        "informative_view_accuracy": float(np.mean(informative_view_hits))
        if informative_view_hits
        else None,
        "informative_view_count": len(informative_view_hits),
    }
    if seed_pool:
        seed_pool.sort(key=lambda item: -item[0])
        calibration = {}
        for fraction in (0.01, 0.05, 0.10, 0.25):
            k = max(1, int(len(seed_pool) * fraction))
            calibration[f"top_{fraction:.0%}_margin"] = {
                "count": k,
                "precision": float(np.mean([hit for _, hit in seed_pool[:k]])),
            }
        report["seed_calibration"] = calibration
    torch.save(
        {
            "anchor_maps": anchor_maps,
            "text_view_features": text_view_features,
            "vision_view_features": vision_view_features,
            "color_classes": COLOR_CLASSES,
            "order_note": "vision rows, text columns, released view order",
        },
        args.anchors_output,
    )
    write_json(args.output, report)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()