path: root/worldalign/anchor_battery.py

"""C5 battery: gauge-free unary anchors from raw observations.

Two declared anchor classes, per `STRUCTURE_DESIGN.md`:

1. pure invariants -- quantities whose cross-modal identity needs no
   specification at all (numeral content of phrases);
2. minimally specified physical invariants -- fixed world-knowledge maps
   (color lexicon to hue bands, size words to box areas, luminance words
   to pixel luminance), frozen before evaluation and never tuned against
   pairs.

Anchors are unary and computed independently per modality: text from the
released phrase bundles, vision from cached image pixels and the region
boxes of the private preprocessing file (vision-side observables). Hidden
pairs score anchor quality; the anchor matrix itself is built without
them and feeds the tempering search as side information.
"""

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 .common import write_json

NUMBER_WORDS = {
    "one": 1, "two": 2, "three": 3, "four": 4, "five": 5, "six": 6,
    "seven": 7, "eight": 8, "nine": 9, "ten": 10, "eleven": 11,
    "twelve": 12, "dozen": 12,
}
# Hue band centers in degrees on the HSV wheel; frozen world knowledge.
COLOR_BANDS = {
    "red": 0.0, "orange": 30.0, "yellow": 60.0, "green": 120.0,
    "cyan": 180.0, "blue": 220.0, "purple": 275.0, "pink": 330.0,
}
ACHROMATIC = ("white", "black", "gray", "grey", "brown", "tan", "beige")
SIZE_WORDS = {
    "huge": 2.0, "large": 1.5, "big": 1.5, "tall": 1.0, "long": 1.0,
    "small": -1.0, "little": -1.0, "tiny": -2.0, "short": -0.5,
}
LIGHT_WORDS = {
    "white": 1.0, "bright": 1.0, "light": 0.5, "sunny": 1.0,
    "black": -1.0, "dark": -1.0, "shadow": -0.5, "night": -1.0,
}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--vg-dir", default="artifacts/vg_5k")
    parser.add_argument("--image-cache", default="/tmp/yurenh2-worldalign-vg-images")
    parser.add_argument("--image-size", type=int, default=224)
    parser.add_argument("--workers", type=int, default=16)
    parser.add_argument(
        "--output", default="artifacts/manifold_gate/anchor_battery.json"
    )
    parser.add_argument(
        "--matrix-output", default="artifacts/manifold_gate/anchor_unary.pt"
    )
    return parser.parse_args()


def read_jsonl(path: Path) -> list[dict]:
    return [
        json.loads(line)
        for line in path.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


def text_anchors(phrases: list[str]) -> dict[str, np.ndarray | float]:
    tokens = [
        token
        for phrase in phrases
        for token in re.findall(r"[a-z]+|\d+", phrase.lower())
    ]
    numeral_total = 0.0
    for token in tokens:
        if token.isdigit() and int(token) <= 50:
            numeral_total += int(token)
        elif token in NUMBER_WORDS:
            numeral_total += NUMBER_WORDS[token]
    color_histogram = np.zeros(len(COLOR_BANDS))
    for index, color in enumerate(COLOR_BANDS):
        color_histogram[index] = sum(token == color for token in tokens)
    achromatic = sum(token in ACHROMATIC for token in tokens)
    chroma_score = color_histogram.sum() / max(color_histogram.sum() + achromatic, 1)
    size_score = float(sum(SIZE_WORDS.get(token, 0.0) for token in tokens))
    light_score = float(sum(LIGHT_WORDS.get(token, 0.0) for token in tokens))
    return {
        "numeral_total": float(numeral_total),
        "color_histogram": color_histogram,
        "chroma_score": float(chroma_score),
        "size_score": size_score,
        "light_score": light_score,
    }


def vision_anchors(record: dict, args: argparse.Namespace) -> dict:
    path = Path(args.image_cache, f"{record['source_image_id']}.jpg")
    with Image.open(path) as image:
        hsv = image.convert("HSV").resize(
            (args.image_size, args.image_size), Image.BILINEAR
        )
    values = np.asarray(hsv, dtype=np.float64)
    hue = values[..., 0] * (360.0 / 255.0)
    saturation = values[..., 1] / 255.0
    value = values[..., 2] / 255.0
    chromatic = (saturation > 0.25) & (value > 0.2)
    hue_histogram = np.zeros(len(COLOR_BANDS))
    for index, center in enumerate(COLOR_BANDS.values()):
        distance = np.minimum(np.abs(hue - center), 360.0 - np.abs(hue - center))
        hue_histogram[index] = float(((distance < 25.0) & chromatic).mean())
    width, height = record["width"], record["height"]
    areas = [
        (region["width"] * region["height"]) / max(width * height, 1)
        for region in record["regions"]
    ]
    return {
        "hue_histogram": hue_histogram,
        "chroma_fraction": float(chromatic.mean()),
        "luminance_mean": float(value.mean()),
        "mean_box_area": float(np.mean(areas)),
        "log_mean_box_area": float(np.log(np.mean(areas) + 1e-6)),
    }


def rank_z(x: np.ndarray) -> np.ndarray:
    order = x.argsort().argsort().astype(np.float64)
    order = (order - order.mean()) / max(order.std(), 1e-9)
    return order


def spearman(x: np.ndarray, y: np.ndarray) -> float:
    return float(np.corrcoef(rank_z(x), rank_z(y))[0, 1])


def main() -> None:
    args = parse_args()
    text_nodes = read_jsonl(Path(args.vg_dir, "text_nodes.jsonl"))
    vision_nodes = read_jsonl(Path(args.vg_dir, "vision_nodes.private.jsonl"))
    truth = read_jsonl(Path(args.vg_dir, "ground_truth.private.jsonl"))

    text_features = {
        record["node_id"]: text_anchors(record["region_closed"])
        for record in text_nodes
    }
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        vision_results = list(
            pool.map(lambda record: vision_anchors(record, args), vision_nodes)
        )
    vision_features = {
        record["node_id"]: result
        for record, result in zip(vision_nodes, vision_results)
    }

    # Aligned arrays in ground-truth order: index i pairs vision i, text i.
    vision_ids = [pair["vision_node_id"] for pair in truth]
    text_ids = [pair["text_node_id"] for pair in truth]
    n = len(truth)

    def text_array(key: str) -> np.ndarray:
        return np.array([text_features[node][key] for node in text_ids])

    def vision_array(key: str) -> np.ndarray:
        return np.array([vision_features[node][key] for node in vision_ids])

    declared_pairs = {
        "color": ("color_histogram", "hue_histogram"),
        "chroma": ("chroma_score", "chroma_fraction"),
        "light": ("light_score", "luminance_mean"),
        "size": ("size_score", "log_mean_box_area"),
        "numerals_vs_area": ("numeral_total", "log_mean_box_area"),
    }

    report: dict = {
        "protocol": (
            "Unary anchors, computed independently per modality; hidden "
            "pairs score anchor quality only. Lexicon-to-physics maps are "
            "frozen world knowledge, declared before evaluation."
        ),
        "nodes": n,
        "anchors": {},
    }
    channels: list[np.ndarray] = []
    channel_names: list[str] = []

    # Color: cosine between lexicon histogram and hue-band histogram under
    # the frozen band map, evaluated for every cross pair.
    text_color = np.stack([text_features[node]["color_histogram"] for node in text_ids])
    vision_color = np.stack(
        [vision_features[node]["hue_histogram"] for node in vision_ids]
    )
    text_color_norm = text_color / np.linalg.norm(text_color, axis=1, keepdims=True).clip(
        min=1e-9
    )
    vision_color_norm = vision_color / np.linalg.norm(
        vision_color, axis=1, keepdims=True
    ).clip(min=1e-9)
    color_similarity = vision_color_norm @ text_color_norm.T  # [vision, text]
    has_signal = (text_color.sum(1) > 0)[None, :] * np.ones((n, 1))
    matched_color = np.diag(color_similarity)
    informative = text_color.sum(1) > 0
    shuffle = np.random.default_rng(0).permutation(n)
    report["anchors"]["color_histogram"] = {
        "informative_text_fraction": float(informative.mean()),
        "matched_mean": float(matched_color[informative].mean()),
        "shuffled_mean": float(color_similarity[shuffle, np.arange(n)][informative].mean()),
        "matched_minus_shuffled_z": float(
            (matched_color[informative] - color_similarity[shuffle, np.arange(n)][informative]).mean()
            / (matched_color[informative] - color_similarity[shuffle, np.arange(n)][informative]).std()
            * np.sqrt(informative.sum())
        ),
    }
    channels.append(color_similarity * has_signal)
    channel_names.append("color")

    for name, (text_key, vision_key) in declared_pairs.items():
        if name == "color":
            continue
        t = text_array(text_key)
        v = vision_array(vision_key)
        rho = spearman(t, v)
        report["anchors"][name] = {"true_pairing_spearman": rho}
        similarity = -np.abs(rank_z(v)[:, None] - rank_z(t)[None, :])
        channels.append(similarity)
        channel_names.append(name)

    # Combined unary similarity: mean of per-channel rank-standardized maps.
    stacked = []
    for channel in channels:
        flat = channel.flatten()
        standardized = (channel - flat.mean()) / max(flat.std(), 1e-9)
        stacked.append(standardized)
    combined = np.mean(stacked, axis=0)
    matched = np.diag(combined)
    ranks = (combined >= matched[:, None]).sum(1)
    report["combined_unary"] = {
        "channels": channel_names,
        "matched_mean": float(matched.mean()),
        "grand_mean": float(combined.mean()),
        "true_z": float(
            (matched.mean() - combined.mean())
            / combined.std()
            * np.sqrt(n)
        ),
        "retrieval_r@10": float((ranks <= 10).mean()),
        "retrieval_r@100": float((ranks <= 100).mean()),
        "median_rank": float(np.median(ranks)),
        "chance_r@10": 10.0 / n,
    }
    torch.save(
        {
            "vision_node_ids": vision_ids,
            "text_node_ids": text_ids,
            "order_note": "row i = vision node truth[i]; column j = text node truth[j]",
            "channels": {name: torch.from_numpy(np.asarray(channel)) for name, channel in zip(channel_names, channels)},
            "combined": torch.from_numpy(combined),
        },
        args.matrix_output,
    )
    write_json(args.output, report)
    print(json.dumps(report, indent=2)[:2200])
    print(f"Wrote {args.output} and {args.matrix_output}")


if __name__ == "__main__":
    main()