"""World v2: free-form captions over a much richer procedural world.

Text is tokenized as raw UTF-8 bytes, so any string is accepted. What the model
*understands* is whatever this world teaches it: many colours and shapes, single-
object events (moving, falling, jumping, orbiting, fading, changing colour ...)
and two-object interactions (merging with colour mixing, splitting, colliding,
chasing, orbiting, swapping places), each described with many paraphrases.

Every clip comes with per-frame, per-object ground truth in `meta` for scoring.
"""
from __future__ import annotations

import math
import re
import zlib
from dataclasses import dataclass

import numpy as np
import torch

# ---------------------------------------------------------------- tokenizer

PAD, NULL, BOS = 256, 257, 258
BYTE_VOCAB = 259
MAX_BYTES = 256


class ByteTokenizer:
    """UTF-8 bytes -> ids. No vocabulary: every string is representable."""
    vocab_size = BYTE_VOCAB

    def __init__(self, max_len: int = MAX_BYTES):
        self.max_len = max_len

    def encode(self, caption: str) -> list[int]:
        return [BOS] + list(caption.encode("utf-8"))[: self.max_len - 1]

    def __call__(self, captions: list[str], pad_to_longest: bool = True) -> torch.Tensor:
        """[B, L] ids; L = longest caption in the batch (dynamic padding) or max_len."""
        encoded = [self.encode(c) for c in captions]
        L = max(map(len, encoded)) if pad_to_longest else self.max_len
        out = torch.full((len(captions), L), PAD, dtype=torch.long)
        for i, ids in enumerate(encoded):
            out[i, : len(ids)] = torch.tensor(ids)
        return out

    def null(self, n: int, length: int = 1) -> torch.Tensor:
        out = torch.full((n, length), PAD, dtype=torch.long)
        out[:, 0] = NULL
        return out


SEP = 259
HASH_OFFSET, HASH_BUCKETS = 260, 16384
HYBRID_VOCAB = HASH_OFFSET + HASH_BUCKETS
MAX_HYBRID = 320
_WORD = re.compile(r"[^\W_]+", re.UNICODE)


def word_bucket(word: str) -> int:
    """Stable hash of a case-folded word into one of HASH_BUCKETS ids (no vocabulary; any word maps somewhere)."""
    return HASH_OFFSET + zlib.crc32(word.casefold().encode("utf-8")) % HASH_BUCKETS


class HybridTokenizer(ByteTokenizer):
    """[BOS] utf-8 bytes [SEP] hashed words. Bytes keep spelling; hashed words give the encoder word identity
    from the first step. Still no vocabulary: every string encodes, nothing is UNK."""
    vocab_size = HYBRID_VOCAB

    def __init__(self, max_len: int = MAX_HYBRID, max_bytes: int = MAX_BYTES):
        super().__init__(max_len)
        self.max_bytes = max_bytes

    def encode(self, caption: str) -> list[int]:
        words = [word_bucket(w) for w in _WORD.findall(caption)]
        byte_ids = list(caption.encode("utf-8"))[: self.max_bytes - 1]
        room = self.max_len - 2 - len(byte_ids)
        return [BOS] + byte_ids + [SEP] + words[: max(0, room)]


# ---------------------------------------------------------------- the world's nouns

COLORS = {
    "red": (225, 35, 35), "orange": (250, 140, 20), "yellow": (245, 220, 40), "green": (40, 190, 60),
    "cyan": (40, 215, 230), "blue": (45, 85, 240), "purple": (145, 55, 215), "pink": (250, 125, 190),
    "brown": (135, 80, 35), "white": (245, 245, 245), "gray": (140, 140, 140), "black": (12, 12, 12),
}
COLOR_WORDS = {
    "red": ["red", "red", "red", "crimson", "scarlet"], "orange": ["orange"], "yellow": ["yellow", "yellow", "golden"],
    "green": ["green", "green", "lime"], "cyan": ["cyan", "cyan", "turquoise", "light blue"],
    "blue": ["blue", "blue", "dark blue"], "purple": ["purple", "purple", "violet"], "pink": ["pink"],
    "brown": ["brown"], "white": ["white"], "gray": ["gray", "grey", "silver"], "black": ["black"],
}
# Two colours that merge make this one (order-free). Pairs not listed blend to the nearest named colour.
MIX = {
    frozenset(["red", "blue"]): "purple", frozenset(["red", "yellow"]): "orange",
    frozenset(["blue", "yellow"]): "green", frozenset(["red", "white"]): "pink",
    frozenset(["black", "white"]): "gray", frozenset(["blue", "green"]): "cyan",
    frozenset(["red", "green"]): "brown", frozenset(["blue", "white"]): "cyan",
    frozenset(["purple", "white"]): "pink", frozenset(["orange", "blue"]): "brown",
}

SHAPES = ["circle", "square", "triangle", "star", "ring", "cross", "heart", "diamond", "hexagon"]
SHAPE_WORDS = {
    "circle": ["circle", "circle", "circle", "ball", "dot", "disc"], "square": ["square", "square", "box", "block"],
    "triangle": ["triangle"], "star": ["star"], "ring": ["ring", "ring", "donut"], "cross": ["cross", "plus sign"],
    "heart": ["heart"], "diamond": ["diamond"], "hexagon": ["hexagon"],
}
ROTATION_INVARIANT = {"circle", "ring"}

SIZES = {"tiny": 0.085, "small": 0.115, "medium": 0.155, "big": 0.205, "huge": 0.26}
SIZE_WORDS = {"tiny": ["tiny"], "small": ["small", "little"], "medium": [""], "big": ["big", "large"],
              "huge": ["huge", "giant"]}

BACKGROUNDS = {  # name -> (top rgb, bottom rgb)
    "black": ((0, 0, 0), (0, 0, 0)), "white": ((250, 250, 250), (250, 250, 250)),
    "gray": ((120, 120, 120), (120, 120, 120)), "navy": ((10, 15, 70), (10, 15, 70)),
    "sky": ((110, 170, 240), (200, 230, 255)), "sunset": ((90, 40, 130), (250, 150, 60)),
    "grass": ((120, 190, 250), (40, 140, 40)),
}
BG_PHRASES = {
    "black": ["on a black background", "against a black background", "on black", "in the dark", "in darkness"],
    "white": ["on a white background", "against a white background", "on white", "on a blank white page"],
    "gray": ["on a gray background", "on a grey background", "against a gray wall"],
    "navy": ["at night", "in the night sky", "on a dark blue background", "on a navy background"],
    "sky": ["in the sky", "in a blue sky", "against the sky", "in the daytime sky"],
    "sunset": ["at sunset", "in a sunset sky", "against a sunset", "at dusk"],
    "grass": ["over a field of grass", "above green grass", "over a meadow", "in a grassy field"],
}

DIRS = {
    "left": (-1, 0), "right": (1, 0), "up": (0, -1), "down": (0, 1),
    "up-left": (-1, -1), "up-right": (1, -1), "down-left": (-1, 1), "down-right": (1, 1),
}
DIR_PHRASES = {
    "left": ["left", "to the left", "leftward", "toward the left", "from right to left"],
    "right": ["right", "to the right", "rightward", "toward the right", "from left to right"],
    "up": ["up", "upward", "upwards", "toward the top", "from bottom to top"],
    "down": ["down", "downward", "downwards", "toward the bottom", "from top to bottom"],
    "up-left": ["up and to the left", "diagonally up left", "toward the top left", "up-left"],
    "up-right": ["up and to the right", "diagonally up right", "toward the top right", "up-right"],
    "down-left": ["down and to the left", "diagonally down left", "toward the bottom left", "down-left"],
    "down-right": ["down and to the right", "diagonally down right", "toward the bottom right", "down-right"],
}
MOVE_VERBS = ["moves", "moves", "slides", "drifts", "travels", "goes", "glides", "floats", "heads"]
MOVE_VERBS_ING = ["moving", "sliding", "drifting", "traveling", "going", "gliding", "floating", "heading"]

SINGLE_EVENTS = {"move": 4, "fall": 1, "rise": 1, "jump": 1, "bounce": 1, "zigzag": 1, "circle": 1, "grow": 1,
                 "shrink": 1, "spin": 1, "still": 1, "fade_in": 1, "fade_out": 1, "blink": 1, "recolor": 1.5,
                 "teleport": 0.7}
PAIR_EVENTS = {"merge": 9, "split": 1.2, "collide": 1, "chase": 1, "orbit": 1, "swap": 1, "independent": 1.8}

# (event kind, shape) combinations never rendered in the training split: a generalisation test.
HELDOUT = {("merge", "star"), ("jump", "heart"), ("orbit", "hexagon"), ("fall", "diamond"), ("spin", "cross"),
           ("recolor", "triangle")}


@dataclass
class WorldConfig:
    frames: int = 16
    size: int = 32
    supersample: int = 2
    p_pair: float = 0.5           # with PAIR_EVENTS weights -> ~30% of clips are merges
    p_featured: float = 0.04      # the user's own example: a red and a blue circle merge into a purple circle
    merge_weight: float = 9.0     # relative weight of "merge" among PAIR_EVENTS (phase 3 lowers it: merges took over)
    p_mention_bg: float = 0.5
    p_default_black: float = 0.7  # background when the caption does not mention one


def _choice(rng, seq):
    return seq[rng.integers(len(seq))]


def _weighted(rng, weights: dict):
    keys = list(weights)
    p = np.array([weights[k] for k in keys], float)
    return keys[rng.choice(len(keys), p=p / p.sum())]


def _contrast_ok(rgb, bg: str) -> bool:
    top, bottom = (np.array(c, float) for c in BACKGROUNDS[bg])
    c = np.array(rgb, float)
    return min(np.linalg.norm(c - top), np.linalg.norm(c - bottom)) > 110


def nearest_color(rgb) -> str:
    c = np.array(rgb, float)
    return min(COLORS, key=lambda k: np.linalg.norm(c - np.array(COLORS[k], float)))


def mix_colors(a: str, b: str) -> str:
    if a == b:
        return a
    return MIX.get(frozenset([a, b])) or nearest_color((np.array(COLORS[a]) + np.array(COLORS[b])) / 2)


# ---------------------------------------------------------------- phrasing

def noun_phrase(rng, obj: dict, article: str | None = None) -> str:
    size = _choice(rng, SIZE_WORDS[obj["size"]])
    color = _choice(rng, COLOR_WORDS[obj["color"]])
    shape = _choice(rng, SHAPE_WORDS[obj["shape"]])
    if article is None:
        article = _choice(rng, ["a", "a", "a", "the", "one", ""])
    if article.endswith("bigger"):
        size = ""
    if article == "a" and (size or color)[:1] in "aeiou":
        article = "an"
    return " ".join(w for w in (article, size, color, shape) if w)


def _move_phrase(rng, d: str) -> str:
    dp = _choice(rng, DIR_PHRASES[d])
    if d == "up" and rng.random() < 0.25:
        return _choice(rng, ["rises", "floats up", "goes up"])
    return f"{_choice(rng, MOVE_VERBS)} {dp}"


def describe_single(rng, obj: dict, ev: dict) -> str:
    np_ = noun_phrase(rng, obj)
    k = ev["kind"]
    if k == "move":
        vp = _move_phrase(rng, ev["dir"])
        if rng.random() < 0.2:
            vp = f"{_choice(rng, MOVE_VERBS_ING)} {_choice(rng, DIR_PHRASES[ev['dir']])}"
            return f"{np_} {vp}" if rng.random() < 0.5 else f"{np_}, {vp}"
        return f"{np_} {vp}"
    if k == "fall":
        return f"{np_} {_choice(rng, ['falls', 'falls down', 'drops', 'drops down', 'falls from the top'])}"
    if k == "rise":
        return f"{np_} {_choice(rng, ['rises from the bottom', 'rises up', 'floats up from the bottom', 'launches upward'])}"
    if k == "jump":
        d = ev.get("dir")
        tail = f" {_choice(rng, ['to the left', 'left'])}" if d == "left" else f" {_choice(rng, ['to the right', 'right'])}" if d == "right" else ""
        return f"{np_} {_choice(rng, ['jumps', 'hops', 'leaps'])}{tail}"
    if k == "bounce":
        return f"{np_} {_choice(rng, ['bounces around', 'bounces', 'bounces off the walls', 'ricochets around'])}"
    if k == "zigzag":
        return f"{np_} {_choice(rng, ['zigzags', 'moves in a zigzag'])} {_choice(rng, DIR_PHRASES[ev['dir']][:3])}"
    if k == "circle":
        return f"{np_} {_choice(rng, ['moves in a circle', 'goes around in a circle', 'circles around', 'travels in a loop'])}"
    if k == "grow":
        return f"{np_} {_choice(rng, ['grows', 'gets bigger', 'expands', 'grows larger', 'inflates'])}"
    if k == "shrink":
        return f"{np_} {_choice(rng, ['shrinks', 'gets smaller', 'shrinks away', 'deflates'])}"
    if k == "spin":
        return f"{np_} {_choice(rng, ['spins', 'rotates', 'spins in place', 'twirls'])}"
    if k == "still":
        return f"{np_} {_choice(rng, ['stays still', 'sits still', 'does not move', 'stays in place', 'is resting'])}"
    if k == "fade_in":
        return f"{np_} {_choice(rng, ['fades in', 'appears', 'slowly appears', 'materializes'])}"
    if k == "fade_out":
        return f"{np_} {_choice(rng, ['fades out', 'disappears', 'fades away', 'vanishes'])}"
    if k == "blink":
        return f"{np_} {_choice(rng, ['blinks', 'flashes', 'flickers', 'blinks on and off'])}"
    if k == "recolor":
        to = _choice(rng, COLOR_WORDS[ev["to"]])
        return f"{np_} {_choice(rng, ['turns', 'changes to', 'becomes', 'changes color to', 'turns into a'])} {to}".replace(
            "turns into a " + to, f"turns into a {to} {_choice(rng, SHAPE_WORDS[obj['shape']])}")
    if k == "teleport":
        return f"{np_} {_choice(rng, ['teleports', 'jumps to another spot', 'teleports across the screen'])}"
    raise ValueError(k)


def describe_pair(rng, a: dict, b: dict, ev: dict) -> str:
    k = ev["kind"]
    A, B = noun_phrase(rng, a), noun_phrase(rng, b)
    if k == "merge":
        res = ev["result"]
        verb = _choice(rng, ["merge into", "merge into", "combine into", "fuse into", "blend into", "join to form",
                             "come together to form", "melt together into"])
        tail = f" {verb} {noun_phrase(rng, res, article=_choice(rng, ['a', 'one', 'a single', 'a bigger']))}"
        if rng.random() < 0.25:
            tail = f" {_choice(rng, ['merge', 'merge together', 'fuse together', 'combine'])}"
        return f"{A} and {B}{tail}"
    if k == "split":
        src = ev["source"]
        S = noun_phrase(rng, src)
        if rng.random() < 0.25:
            return f"{S} {_choice(rng, ['splits in two', 'splits into two', 'divides in half', 'breaks in two'])}"
        return f"{S} {_choice(rng, ['splits into', 'breaks into', 'divides into', 'separates into'])} {noun_phrase(rng, a, 'a')} and {noun_phrase(rng, b, 'a')}"
    if k == "collide":
        return f"{A} and {B} {_choice(rng, ['collide and bounce apart', 'crash into each other', 'bump into each other', 'collide', 'smash together and bounce off'])}"
    if k == "chase":
        if rng.random() < 0.3:
            return f"{B} is chased by {A}"
        return f"{A} {_choice(rng, ['chases', 'follows', 'runs after', 'is chasing'])} {B}"
    if k == "orbit":
        return f"{A} {_choice(rng, ['orbits', 'circles around', 'revolves around', 'goes around'])} {B}"
    if k == "swap":
        return f"{A} and {B} {_choice(rng, ['swap places', 'trade places', 'switch places', 'swap positions'])}"
    if k == "independent":
        return f"{describe_single(rng, a, ev['a'])} {_choice(rng, ['and', 'and', 'while', ', and'])} {describe_single(rng, b, ev['b'])}".replace(" , and", ", and")
    raise ValueError(k)


def _style(rng, caption: str, bg_phrase: str | None) -> str:
    if bg_phrase:
        caption = f"{bg_phrase}, {caption}" if rng.random() < 0.2 else f"{caption} {bg_phrase}"
    if rng.random() < 0.08:
        caption = _choice(rng, ["a video of ", "watch as ", "animation: ", "clip: "]) + caption
    if rng.random() < 0.3:
        caption = caption[0].upper() + caption[1:]
    if rng.random() < 0.25:
        caption += "."
    return caption.strip()


# ---------------------------------------------------------------- dynamics (normalised coords, t in [0,1])

def _uni(rng, lo, hi, size=None):
    """rng.uniform that tolerates an empty range (big objects in a small frame): collapses to the midpoint."""
    if hi < lo:
        lo = hi = (lo + hi) / 2
    return rng.uniform(lo, hi, size=size)


def _state(x, y, r, ang, rgb, alpha=1.0):
    return [float(x), float(y), float(r), float(ang), tuple(float(c) for c in rgb), float(alpha)]


def _single_states(rng, obj: dict, ev: dict, T: int, center=None) -> list:
    r = SIZES[obj["size"]] * _uni(rng, 0.92, 1.08)
    ang0 = _uni(rng, 0, 2 * math.pi)
    rgb = np.array(COLORS[obj["color"]], float)
    ts = np.linspace(0, 1, T)
    k = ev["kind"]
    m = r + 0.03

    def anywhere():
        return _uni(rng, m, 1 - m, size=2)

    if k == "move":
        dx, dy = DIRS[ev["dir"]]
        dist = _uni(rng, 0.4, 0.62) / (math.sqrt(2) if dx and dy else 1)
        dist = min(dist, (1 - 2 * m) * 0.98)

        def start(d):
            if d > 0:
                return _uni(rng, m, 1 - m - dist)
            if d < 0:
                return _uni(rng, m + dist, 1 - m)
            return _uni(rng, m, 1 - m)
        x0, y0 = start(dx), start(dy)
        return [_state(x0 + dx * dist * t, y0 + dy * dist * t, r, ang0, rgb) for t in ts]
    if k in ("fall", "rise"):
        x0 = _uni(rng, m, 1 - m)
        top, bot = m, 1 - m
        prog = ts ** 2 if k == "fall" else 1 - (1 - ts) ** 2
        return [_state(x0, top + (bot - top) * p if k == "fall" else bot - (bot - top) * p, r, ang0, rgb) for p in prog]
    if k == "jump":
        d = {"left": -1, "right": 1}.get(ev.get("dir"), 0)
        span = min(_uni(rng, 0.3, 0.5), max(0.0, 1 - 2 * m - 0.02)) * d
        x0 = _uni(rng, m + max(0, -span), 1 - m - max(0, span))
        ground = 1 - m
        height = _uni(rng, 0.35, min(0.6, 1 - 2 * m))
        return [_state(x0 + span * t, ground - height * 4 * t * (1 - t), r, ang0, rgb) for t in ts]
    if k == "bounce":
        lo, hi = m, 1 - m
        x, y = anywhere()
        th = _uni(rng, 0, 2 * math.pi)
        sp = _uni(rng, 1.6, 2.4)
        vx, vy = sp * math.cos(th), sp * math.sin(th)
        out, dt = [], 1 / (T - 1)
        for _ in range(T):
            out.append(_state(x, y, r, ang0, rgb))
            x, y = x + vx * dt, y + vy * dt
            if x < lo or x > hi:
                vx, x = -vx, float(np.clip(2 * np.clip(x, lo, hi) - x, lo, hi))
            if y < lo or y > hi:
                vy, y = -vy, float(np.clip(2 * np.clip(y, lo, hi) - y, lo, hi))
        return out
    if k == "zigzag":
        dx, dy = DIRS[ev["dir"]]
        dist = min(_uni(rng, 0.45, 0.6), 1 - 2 * m - 0.1)
        amp = _uni(rng, 0.08, 0.14)
        along0 = _uni(rng, m, 1 - m - dist) if (dx + dy) > 0 else _uni(rng, m + dist, 1 - m)
        across0 = _uni(rng, m + amp, 1 - m - amp)
        out = []
        for t in ts:
            along = along0 + (dx + dy) * dist * t
            zig = amp * (2 * abs((t * 3) % 1 - 0.5) * 2 - 1)
            x, y = (along, across0 + zig) if dx else (across0 + zig, along)
            out.append(_state(x, y, r, ang0, rgb))
        return out
    if k == "circle":
        R = _uni(rng, 0.18, max(0.19, 0.5 - m - 0.02))
        cx, cy = center if center is not None else (_uni(rng, m + R, 1 - m - R), _uni(rng, m + R, 1 - m - R))
        ph = _uni(rng, 0, 2 * math.pi)
        sgn = _choice(rng, [-1, 1])
        turns = _uni(rng, 0.8, 1.0)
        return [_state(cx + R * math.cos(ph + sgn * 2 * math.pi * turns * t), cy + R * math.sin(ph + sgn * 2 * math.pi * turns * t),
                       r, ang0, rgb) for t in ts]
    if k in ("grow", "shrink"):
        mm = r * 1.7 + 0.03
        x0, y0 = _uni(rng, mm, 1 - mm, size=2)
        f = (lambda t: 0.45 + 1.25 * t) if k == "grow" else (lambda t: 1.7 - 1.25 * t)
        return [_state(x0, y0, r * f(t), ang0, rgb) for t in ts]
    x0, y0 = anywhere()
    if k == "spin":
        sp = _choice(rng, [-1, 1]) * _uni(rng, 0.7, 1.1) * math.pi
        return [_state(x0, y0, r, ang0 + sp * t, rgb) for t in ts]
    if k == "still":
        return [_state(x0, y0, r, ang0, rgb) for t in ts]
    if k == "fade_in":
        return [_state(x0, y0, r, ang0, rgb, min(1, max(0, (t - 0.1) / 0.6))) for t in ts]
    if k == "fade_out":
        return [_state(x0, y0, r, ang0, rgb, min(1, max(0, (0.9 - t) / 0.6))) for t in ts]
    if k == "blink":
        period = _choice(rng, [4, 4, 6])
        return [_state(x0, y0, r, ang0, rgb, 1.0 if (i // (period // 2)) % 2 == 0 else 0.0) for i in range(T)]
    if k == "recolor":
        rgb2 = np.array(COLORS[ev["to"]], float)
        out = []
        for t in ts:
            w = min(1, max(0, (t - 0.3) / 0.4))
            out.append(_state(x0, y0, r, ang0, (1 - w) * rgb + w * rgb2))
        return out
    if k == "teleport":
        x1, y1 = max((anywhere() for _ in range(8)), key=lambda p: math.hypot(p[0] - x0, p[1] - y0))
        cut = rng.integers(T // 3, 2 * T // 3)
        return [_state(x0, y0, r, ang0, rgb) if i < cut else _state(x1, y1, r, ang0, rgb) for i in range(T)]
    raise ValueError(k)


def _pair_states(rng, a: dict, b: dict, ev: dict, T: int):
    ra = SIZES[a["size"]] * _uni(rng, 0.92, 1.08)
    rb = SIZES[b["size"]] * _uni(rng, 0.92, 1.08)
    ca, cb = np.array(COLORS[a["color"]], float), np.array(COLORS[b["color"]], float)
    angA, angB = _uni(rng, 0, 2 * math.pi, size=2)
    ts = np.linspace(0, 1, T)
    k = ev["kind"]
    if k in ("merge", "split"):
        res = ev["result"] if k == "merge" else ev["source"]
        rr = SIZES[res["size"]]
        cr = np.array(COLORS[res["color"]], float)
        mid = _uni(rng, 0.32, 0.68, size=2)
        th = _uni(rng, 0, math.pi)
        sep = _uni(rng, 0.62, 0.8)
        u = np.array([math.cos(th), math.sin(th)])
        pa = np.clip(mid - u * sep / 2, ra + 0.02, 1 - ra - 0.02)
        pb = np.clip(mid + u * sep / 2, rb + 0.02, 1 - rb - 0.02)
        SA, SB, SR = [], [], []
        for t in ts:
            tm = t if k == "merge" else 1 - t
            p = min(1, tm / 0.6)                       # approach
            f = min(1, max(0, (tm - 0.5) / 0.25))      # fuse
            xa = pa + (mid - pa) * p
            xb = pb + (mid - pb) * p
            SA.append(_state(*xa, ra * (1 - f) + rr * f, angA, ca * (1 - f) + cr * f, 1 - f))
            SB.append(_state(*xb, rb * (1 - f) + rr * f, angB, cb * (1 - f) + cr * f, 1 - f))
            SR.append(_state(*mid, rr * (0.6 + 0.4 * f), angA, cr, f))
        return [SA, SB, SR]
    if k == "collide":
        y = _uni(rng, 0.3, 0.7)
        vert = rng.random() < 0.3
        hit = _uni(rng, 0.4, 0.6)
        xa0, xb0 = _uni(rng, 0.05, 0.2) + ra, 1 - _uni(rng, 0.05, 0.2) - rb
        contact = (xa0 * rb + xb0 * ra) / (ra + rb)  # where they touch if both move at equal speed
        SA, SB = [], []
        for t in ts:
            if t <= hit:
                s = t / hit
                xa = xa0 + (contact - ra - xa0) * s
                xb = xb0 + (contact + rb - xb0) * s
            else:
                s = (t - hit) / (1 - hit)
                xa = contact - ra - (contact - ra - xa0) * s * 0.8
                xb = contact + rb + (xb0 - contact - rb) * s * 0.8
            pa, pb = (y, xa) if vert else (xa, y), (y, xb) if vert else (xb, y)
            SA.append(_state(*pa, ra, angA, ca))
            SB.append(_state(*pb, rb, angB, cb))
        return [SA, SB]
    if k == "chase":
        lead = _single_states(rng, b, {"kind": _choice(rng, ["move", "circle", "zigzag"]), "dir": _choice(rng, ["left", "right", "up", "down"])}, T)
        lag = 4
        SA = []
        for i in range(T):
            j = i - lag
            if j >= 0:
                x, y = lead[j][0], lead[j][1]
            else:  # extrapolate backwards along the leader's initial velocity
                vx, vy = lead[1][0] - lead[0][0], lead[1][1] - lead[0][1]
                x, y = lead[0][0] + vx * j, lead[0][1] + vy * j
            SA.append(_state(np.clip(x, ra, 1 - ra), np.clip(y, ra, 1 - ra), ra, angA, ca))
        SB = [[s[0], s[1], rb, angB, tuple(cb), 1.0] for s in lead]
        return [SA, SB]
    if k == "orbit":
        R = _uni(rng, 0.24, 0.32)
        c = (_uni(rng, 0.5 - (0.5 - R - ra - 0.02), 0.5 + (0.5 - R - ra - 0.02)),
             _uni(rng, 0.5 - (0.5 - R - ra - 0.02), 0.5 + (0.5 - R - ra - 0.02)))
        SA = _single_states(rng, a, {"kind": "circle"}, T, center=c)
        SA = [[s[0], s[1], ra, angA, tuple(ca), 1.0] for s in SA]
        SB = [_state(*c, rb, angB, cb) for _ in ts]
        return [SA, SB]
    if k == "swap":
        p0 = _uni(rng, 0.25, 0.75, size=2)
        th = _uni(rng, 0, 2 * math.pi)
        R = _uni(rng, 0.2, 0.28)
        c = np.clip(p0, R + 0.12, 1 - R - 0.12)
        SA, SB = [], []
        for t in ts:
            ph = th + math.pi * t
            SA.append(_state(c[0] + R * math.cos(ph), c[1] + R * math.sin(ph), ra, angA, ca))
            SB.append(_state(c[0] - R * math.cos(ph), c[1] - R * math.sin(ph), rb, angB, cb))
        return [SA, SB]
    if k == "independent":
        return [_single_states(rng, a, ev["a"], T), _single_states(rng, b, ev["b"], T)]
    raise ValueError(k)


# ---------------------------------------------------------------- sampling a scene

def _random_obj(rng, bg: str, shape=None, color=None, size=None) -> dict:
    while True:
        o = {"shape": shape or _choice(rng, SHAPES), "color": color or _choice(rng, list(COLORS)),
             "size": size or _choice(rng, list(SIZES))}
        if _contrast_ok(COLORS[o["color"]], bg):
            return o
        color = None


def _random_single_event(rng, obj: dict, bg: str) -> dict:
    while True:
        k = _weighted(rng, SINGLE_EVENTS)
        if k == "spin" and obj["shape"] in ROTATION_INVARIANT:
            continue
        ev = {"kind": k}
        if k in ("move", "zigzag"):
            ev["dir"] = _choice(rng, list(DIRS) if k == "move" else ["left", "right", "up", "down"])
        if k == "jump":
            ev["dir"] = _choice(rng, ["left", "right", None])
        if k == "recolor":
            to = _choice(rng, [c for c in COLORS if c != obj["color"] and _contrast_ok(COLORS[c], bg)])
            ev["to"] = to
        return ev


def is_heldout(scene: dict) -> bool:
    kinds = [scene["event"]["kind"]]
    if kinds[0] == "independent":
        kinds = [scene["event"]["a"]["kind"], scene["event"]["b"]["kind"]]
    shapes = [o["shape"] for o in scene["objects"]]
    res = scene["event"].get("result") or scene["event"].get("source")
    if res:
        shapes.append(res["shape"])
    return any((k, sh) in HELDOUT for k in kinds for sh in shapes)


def featured_scene(rng, bg: str = "black", mention_bg: bool = False) -> dict:
    """The user's example: a red circle and a blue circle merge into a purple circle."""
    small = ["tiny", "small", "small", "medium"]
    a = {"shape": "circle", "color": "red", "size": _choice(rng, small)}
    b = {"shape": "circle", "color": "blue", "size": _choice(rng, small)}
    if rng.random() < 0.5:
        a, b = b, a
    res = {"shape": "circle", "color": "purple", "size": _choice(rng, ["medium", "big"])}
    return {"objects": [a, b], "event": {"kind": "merge", "result": res}, "bg": bg, "mention_bg": mention_bg}


def random_scene(rng, cfg: WorldConfig, split: str = "train") -> dict:
    """split: 'train' (never a HELDOUT combo), 'heldout' (always one), 'any'."""
    while True:
        scene = _random_scene(rng, cfg)
        if split == "any" or (split == "heldout") == is_heldout(scene):
            return scene


def _random_scene(rng, cfg: WorldConfig) -> dict:
    if rng.random() < cfg.p_featured:
        mention_bg = rng.random() < cfg.p_mention_bg
        bg = _choice(rng, ["black", "white", "gray", "sky", "grass"]) if mention_bg else "black"
        return featured_scene(rng, bg, mention_bg)
    mention_bg = rng.random() < cfg.p_mention_bg
    bg = _choice(rng, list(BACKGROUNDS)) if (mention_bg or rng.random() > cfg.p_default_black) else "black"
    if rng.random() >= cfg.p_pair:
        o = _random_obj(rng, bg, size=_choice(rng, ["small", "medium", "medium", "big", "big", "huge", "tiny"]))
        return {"objects": [o], "event": _random_single_event(rng, o, bg), "bg": bg, "mention_bg": mention_bg}
    k = _weighted(rng, {**PAIR_EVENTS, "merge": cfg.merge_weight})
    small = ["tiny", "small", "small", "medium"]
    if k in ("merge", "split"):
        a = _random_obj(rng, bg, size=_choice(rng, small))
        partners = [next(iter(pair - {a["color"]})) for pair in MIX if a["color"] in pair]
        b_color = _choice(rng, partners) if partners and rng.random() < 0.8 else a["color"]
        if not _contrast_ok(COLORS[b_color], bg):
            return _random_scene(rng, cfg)
        b = _random_obj(rng, bg, color=b_color, size=_choice(rng, small))
        if rng.random() < 0.6:  # same shape most of the time, so the result's shape is unambiguous
            b["shape"] = a["shape"]
        mixed = mix_colors(a["color"], b["color"])
        if not _contrast_ok(COLORS[mixed], bg):
            return _random_scene(rng, cfg)
        res = {"shape": a["shape"] if a["shape"] == b["shape"] else "circle", "color": mixed,
               "size": _choice(rng, ["medium", "big"])}
        ev = {"kind": k, "result" if k == "merge" else "source": res}
        return {"objects": [a, b], "event": ev, "bg": bg, "mention_bg": mention_bg}
    a = _random_obj(rng, bg, size=_choice(rng, small))
    b = _random_obj(rng, bg, size=_choice(rng, small + (["big"] if k == "orbit" else [])))
    if a == b:
        return _random_scene(rng, cfg)
    ev = {"kind": k}
    if k == "independent":
        ev["a"], ev["b"] = _random_single_event(rng, a, bg), _random_single_event(rng, b, bg)
    return {"objects": [a, b], "event": ev, "bg": bg, "mention_bg": mention_bg}


def caption_scene(rng, scene: dict) -> str:
    objs, ev = scene["objects"], scene["event"]
    cap = describe_single(rng, objs[0], ev) if len(objs) == 1 else describe_pair(rng, objs[0], objs[1], ev)
    bgp = _choice(rng, BG_PHRASES[scene["bg"]]) if scene["mention_bg"] else None
    return _style(rng, cap, bgp)


# ---------------------------------------------------------------- rendering

def _poly_mask(px, py, verts):
    inside = np.zeros(px.shape, bool)
    n = len(verts)
    for i in range(n):
        x1, y1 = verts[i]
        x2, y2 = verts[(i + 1) % n]
        crosses = (y1 > py) != (y2 > py)
        xint = (x2 - x1) * (py - y1) / (y2 - y1 + 1e-12) + x1
        inside ^= crosses & (px < xint)
    return inside


def _regular(n, r, ang, cx, cy, inner=None):
    k = 2 * n if inner else n
    return [(cx + (inner if inner and i % 2 else r) * math.cos(ang - math.pi / 2 + 2 * math.pi * i / k),
             cy + (inner if inner and i % 2 else r) * math.sin(ang - math.pi / 2 + 2 * math.pi * i / k)) for i in range(k)]


def shape_mask(shape, px, py, cx, cy, r, ang):
    dx, dy = px - cx, py - cy
    if shape == "circle":
        return dx * dx + dy * dy <= r * r
    if shape == "ring":
        d2 = dx * dx + dy * dy
        return (d2 <= r * r) & (d2 >= (0.55 * r) ** 2)
    if shape == "square":
        return _poly_mask(px, py, _regular(4, r * 1.15, ang + math.pi / 4, cx, cy))
    if shape == "triangle":
        return _poly_mask(px, py, _regular(3, r * 1.2, ang, cx, cy))
    if shape == "star":
        return _poly_mask(px, py, _regular(5, r * 1.25, ang, cx, cy, inner=r * 0.5))
    if shape == "hexagon":
        return _poly_mask(px, py, _regular(6, r * 1.05, ang, cx, cy))
    c, s = math.cos(ang), math.sin(ang)
    u, v = (c * dx + s * dy) / r, (-s * dx + c * dy) / r
    if shape == "diamond":
        return np.abs(u) / 0.75 + np.abs(v) / 1.15 <= 1
    if shape == "cross":
        arm = 0.33
        return ((np.abs(u) <= 1) & (np.abs(v) <= arm)) | ((np.abs(v) <= 1) & (np.abs(u) <= arm))
    if shape == "heart":  # upright heart: small tilt only, so it stays recognisable
        c, s = math.cos(0.3 * math.sin(ang)), math.sin(0.3 * math.sin(ang))
        u, v = (c * dx + s * dy) / r, (-s * dx + c * dy) / r
        x, y = u * 1.15, -v * 1.15 + 0.15
        return (x * x + y * y - 1) ** 3 - x * x * y ** 3 <= 0
    raise ValueError(shape)


def _background(bg, S):
    top, bot = (np.array(c, np.float32) for c in BACKGROUNDS[bg])
    w = np.linspace(0, 1, S, dtype=np.float32)[:, None, None]
    return np.broadcast_to(top * (1 - w) + bot * w, (S, S, 3)).copy()


def render_scene(scene: dict, rng, cfg: WorldConfig) -> np.ndarray:
    T, S, ss = cfg.frames, cfg.size, cfg.supersample
    big = S * ss
    co = (np.arange(big, dtype=np.float32) + 0.5) / big
    px, py = np.meshgrid(co, co)
    objs, ev = scene["objects"], scene["event"]
    if len(objs) == 1:
        tracks = [_single_states(rng, objs[0], ev, T)]
        shapes = [objs[0]["shape"]]
    else:
        tracks = _pair_states(rng, objs[0], objs[1], ev, T)
        shapes = [objs[0]["shape"], objs[1]["shape"]]
        if ev["kind"] in ("merge", "split"):
            shapes.append((ev.get("result") or ev.get("source"))["shape"])
    scene["tracks"] = tracks
    scene["track_shapes"] = shapes
    first = [i for i, tr in enumerate(tracks) if tr[0][5] > 0.5 and tr[0][2] > 0.005]
    last = [i for i, tr in enumerate(tracks) if tr[T - 1][5] > 0.5 and tr[T - 1][2] > 0.005]
    scene["summary"] = {"initial_count": len(first), "final_count": len(last),
                        "initial_colors": [nearest_color(tracks[i][0][4]) for i in first],
                        "final_colors": [nearest_color(tracks[i][T - 1][4]) for i in last],
                        "final_shapes": [shapes[i] for i in last]}
    base = _background(scene["bg"], big)
    frames = np.empty((T, S, S, 3), np.uint8)
    for f in range(T):
        img = base.copy()
        for shape, tr in zip(shapes, tracks):
            x, y, r, ang, rgb, alpha = tr[f]
            if alpha <= 0.01 or r <= 0.005:
                continue
            m = shape_mask(shape, px, py, x, y, r, ang)
            img[m] = img[m] * (1 - alpha) + np.array(rgb, np.float32) * alpha
        img = img.reshape(S, ss, S, ss, 3).mean(axis=(1, 3))
        frames[f] = np.clip(img + 0.5, 0, 255).astype(np.uint8)
    return frames


def make_clip(rng, cfg: WorldConfig | None = None, split: str = "train"):
    cfg = cfg or WorldConfig()
    scene = random_scene(rng, cfg, split)
    caption = caption_scene(rng, scene)
    video = render_scene(scene, rng, cfg)
    scene["caption"] = caption
    return video, caption, scene


def sample_prompts(n: int, split: str = "train", seed: int = 0, cfg: WorldConfig | None = None) -> list[tuple[str, dict]]:
    """Captions + scenes. Call render_scene(scene, rng, cfg) to get ground-truth tracks/summary for a scene."""
    rng = np.random.default_rng(seed)
    cfg = cfg or WorldConfig()
    out = []
    for _ in range(n):
        scene = random_scene(rng, cfg, split)
        out.append((caption_scene(rng, scene), scene))
    return out


def to_tensor(video: np.ndarray) -> torch.Tensor:
    return torch.from_numpy(video).permute(3, 0, 1, 2).float().div(127.5).sub(1)


class WorldDataset(torch.utils.data.IterableDataset):
    """Infinite stream of (video float [3,T,H,W] in [-1,1], caption)."""

    def __init__(self, cfg: WorldConfig | None = None, seed: int = 0, split: str = "train"):
        self.cfg = cfg or WorldConfig()
        self.seed = seed
        self.split = split

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        rng = np.random.default_rng([self.seed, info.id if info else 0])
        while True:
            video, caption, _ = make_clip(rng, self.cfg, self.split)
            yield to_tensor(video), caption


def collate(batch):
    videos, captions = zip(*batch)
    return torch.stack(videos), list(captions)