code/bundle/binding_task.py

"""Synthetic attribute-binding and same/different probes (Appendix H)."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Literal

import numpy as np
from PIL import Image, ImageDraw

SHAPES = ("circle", "square", "triangle")
COLORS = {
    "red": (220, 60, 60),
    "green": (60, 180, 60),
    "blue": (60, 60, 220),
    "yellow": (220, 220, 60),
    "purple": (160, 60, 200),
    "cyan": (60, 200, 200),
}
COLOR_NAMES = tuple(COLORS.keys())
BG = (200, 200, 200)
IMAGE_SIZE = 224
SHAPE_SIZE = IMAGE_SIZE // 6  # 37


def _center(side: Literal["left", "right"]) -> tuple[int, int]:
    x = IMAGE_SIZE // 4 if side == "left" else (3 * IMAGE_SIZE) // 4
    return x, IMAGE_SIZE // 2


def draw_shape(
    draw: ImageDraw.ImageDraw,
    shape: str,
    color: tuple[int, int, int],
    center: tuple[int, int],
    size: int = SHAPE_SIZE,
) -> None:
    cx, cy = center
    if shape == "circle":
        draw.ellipse((cx - size, cy - size, cx + size, cy + size), fill=color)
    elif shape == "square":
        draw.rectangle((cx - size, cy - size, cx + size, cy + size), fill=color)
    elif shape == "triangle":
        h = int(size * math.sqrt(3))
        pts = [(cx, cy - (2 * h) // 3), (cx - size, cy + h // 3), (cx + size, cy + h // 3)]
        draw.polygon(pts, fill=color)
    else:
        raise ValueError(shape)


def render_scene(
    left_shape: str,
    left_color: str,
    right_shape: str,
    right_color: str,
) -> Image.Image:
    img = Image.new("RGB", (IMAGE_SIZE, IMAGE_SIZE), BG)
    draw = ImageDraw.Draw(img)
    draw_shape(draw, left_shape, COLORS[left_color], _center("left"))
    draw_shape(draw, right_shape, COLORS[right_color], _center("right"))
    return img


@dataclass
class BindingTrial:
    query: Image.Image
    candidates: list[Image.Image]
    target_index: int
    meta: dict


@dataclass
class SameDiffTrial:
    image_a: Image.Image
    image_b: Image.Image
    same: bool
    meta: dict


def generate_binding_trials(n: int = 500, seed: int = 42) -> list[BindingTrial]:
    rng = np.random.default_rng(seed)
    trials: list[BindingTrial] = []
    for _ in range(n):
        s1, s2 = rng.choice(SHAPES, size=2, replace=False)
        cq = list(rng.choice(COLOR_NAMES, size=2, replace=False))
        remaining = [c for c in COLOR_NAMES if c not in cq]
        ct = list(rng.choice(remaining, size=2, replace=False))
        # Extra colors for distractor D2 if needed
        rem2 = [c for c in remaining if c not in ct]
        if len(rem2) >= 2:
            ct2 = list(rng.choice(rem2, size=2, replace=False))
        else:
            # Fall back to a permutation of target colors for D2
            ct2 = [ct[1], ct[0]]

        query = render_scene(s1, cq[0], s2, cq[1])
        target = render_scene(s1, ct[0], s2, ct[1])
        d1 = render_scene(s2, ct[0], s1, ct[1])  # shape swap
        d2 = render_scene(s2, ct2[0], s1, ct2[1])  # shape swap, different colors
        d3 = render_scene(s1, ct[0], s1, ct[1])  # partial: correct left, wrong right shape

        candidates = [target, d1, d2, d3]
        order = rng.permutation(4)
        shuffled = [candidates[i] for i in order]
        target_index = int(np.where(order == 0)[0][0])
        trials.append(
            BindingTrial(
                query=query,
                candidates=shuffled,
                target_index=target_index,
                meta={
                    "shapes": (s1, s2),
                    "query_colors": tuple(cq),
                    "target_colors": tuple(ct),
                    "disjoint": len(set(cq) & set(ct)) == 0,
                },
            )
        )
    return trials


def generate_samediff_trials(n: int = 500, seed: int = 42) -> list[SameDiffTrial]:
    rng = np.random.default_rng(seed + 1)
    trials: list[SameDiffTrial] = []
    n_same = n // 2
    for i in range(n):
        same = i < n_same
        s1, s2 = rng.choice(SHAPES, size=2, replace=False)
        c1, c2 = rng.choice(COLOR_NAMES, size=2, replace=False)
        image_a = render_scene(s1, c1, s2, c2)
        if same:
            remaining = [c for c in COLOR_NAMES if c not in (c1, c2)]
            c1p, c2p = rng.choice(remaining, size=2, replace=False)
            image_b = render_scene(s1, c1p, s2, c2p)
        else:
            if rng.random() < 0.5:
                # Shape swap
                image_b = render_scene(s2, c1, s1, c2)
            else:
                # Shape change on left
                others = [s for s in SHAPES if s != s1]
                s1p = rng.choice(others)
                image_b = render_scene(s1p, c1, s2, c2)
        trials.append(SameDiffTrial(image_a=image_a, image_b=image_b, same=same, meta={"same": same}))
    rng.shuffle(trials)
    return trials


def verify_disjoint_colors(trials: list[BindingTrial]) -> float:
    return float(np.mean([t.meta["disjoint"] for t in trials]))