"""Procedural captioned video clips: coloured shapes moving on simple backgrounds.

Everything is generated on the fly from a seeded RNG, so there is no dataset on
disk and every clip comes with ground-truth metadata (per-frame trajectory) that
the evaluator can check generated videos against.
"""
from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import torch

# ---------------------------------------------------------------- vocabulary

SHAPES = ["circle", "square", "triangle", "star", "ring", "cross"]
COLORS = {
    "red": (230, 40, 40),
    "green": (40, 200, 60),
    "blue": (50, 90, 240),
    "yellow": (245, 220, 40),
    "cyan": (40, 220, 230),
    "magenta": (220, 50, 220),
    "white": (245, 245, 245),
    "orange": (250, 140, 20),
}
SIZES = {"small": 0.11, "medium": 0.16, "big": 0.23}  # radius as a fraction of frame size
BACKGROUNDS = ["black", "white", "gray", "navy", "gradient"]
# motion name -> caption words
MOTIONS = {
    "left": "moves left",
    "right": "moves right",
    "up": "moves up",
    "down": "moves down",
    "up-left": "moves up left",
    "up-right": "moves up right",
    "down-left": "moves down left",
    "down-right": "moves down right",
    "bounces": "bounces",
    "grows": "grows",
    "shrinks": "shrinks",
    "spins": "spins",
    "still": "stays still",
}
DIRECTIONS = {
    "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),
}
ROTATIONALLY_SYMMETRIC = {"circle", "ring"}  # spinning these is invisible, so never caption it

# Combos never rendered during training; used to test compositional generalisation.
HELDOUT_SHAPE_COLOR = {("star", "red"), ("circle", "yellow"), ("triangle", "cyan"), ("square", "magenta")}
HELDOUT_SHAPE_MOTION = {("ring", "down"), ("cross", "grows"), ("triangle", "bounces"), ("star", "left")}

SPECIALS = ["<pad>", "<unk>", "<null>"]
_WORDS = ["a", "and", "on", "background", "moves", "stays", "still", "left", "right", "up", "down",
          "bounces", "grows", "shrinks", "spins", *SIZES, *COLORS, *SHAPES, *BACKGROUNDS]
VOCAB = SPECIALS + sorted(set(_WORDS))
PAD, UNK, NULL = 0, 1, 2
MAX_TOKENS = 24


class Tokenizer:
    def __init__(self, max_len: int = MAX_TOKENS):
        self.max_len = max_len
        self.index = {w: i for i, w in enumerate(VOCAB)}

    def encode(self, caption: str) -> list[int]:
        ids = [self.index.get(w, UNK) for w in caption.lower().split()]
        return (ids or [NULL])[: self.max_len]

    def __call__(self, captions: list[str]) -> torch.Tensor:
        out = torch.full((len(captions), self.max_len), PAD, dtype=torch.long)
        for i, c in enumerate(captions):
            ids = self.encode(c)
            out[i, : len(ids)] = torch.tensor(ids)
        return out

    def null(self, n: int) -> torch.Tensor:
        """Token rows for the unconditional (classifier-free guidance) branch."""
        out = torch.full((n, self.max_len), PAD, dtype=torch.long)
        out[:, 0] = NULL
        return out


# ---------------------------------------------------------------- sampling captions

@dataclass
class ClipConfig:
    frames: int = 16
    size: int = 48
    p_two_objects: float = 0.3
    p_mention_bg: float = 0.85
    supersample: int = 2
    # Curriculum: restrict what gets sampled (None = everything).
    shapes: tuple[str, ...] | None = None
    colors: tuple[str, ...] | None = None
    sizes: tuple[str, ...] | None = None
    motions: tuple[str, ...] | None = None
    backgrounds: tuple[str, ...] | None = None

    @classmethod
    def simple(cls, **kw) -> "ClipConfig":
        """Stage-1 curriculum: one object, black background, 3 shapes, 4 colours, 5 motions."""
        base = dict(p_two_objects=0.0, p_mention_bg=0.0, shapes=("circle", "square", "triangle"),
                    colors=("red", "green", "blue", "yellow"), motions=("left", "right", "up", "down", "bounces"),
                    backgrounds=("black",))
        return cls(**{**base, **kw})


def _pick(rng: np.random.Generator, options, allowed):
    opts = [o for o in options if allowed is None or o in allowed]
    return opts[rng.integers(len(opts))]


def is_heldout(meta: dict) -> bool:
    return any((o["shape"], o["color"]) in HELDOUT_SHAPE_COLOR or (o["shape"], o["motion"]) in HELDOUT_SHAPE_MOTION
               for o in meta["objects"])


def _clashes(color: str, bg: str) -> bool:
    return (color, bg) in {("white", "white"), ("blue", "navy")}


def _random_object(rng: np.random.Generator, bg: str, cfg: ClipConfig) -> dict:
    while True:
        shape = _pick(rng, SHAPES, cfg.shapes)
        color = _pick(rng, COLORS, cfg.colors)
        motion = _pick(rng, MOTIONS, cfg.motions)
        size = _pick(rng, SIZES, cfg.sizes)
        if not _clashes(color, bg) and not (motion == "spins" and shape in ROTATIONALLY_SYMMETRIC):
            return {"shape": shape, "color": color, "size": size, "motion": motion}


def _random_meta(rng: np.random.Generator, cfg: ClipConfig, split: str) -> dict:
    """split: 'train' (no held-out combos), 'heldout' (one object, held-out combo), 'any'."""
    while True:
        bg = _pick(rng, BACKGROUNDS, cfg.backgrounds)
        if split == "heldout":
            obj = _random_object(rng, bg, cfg)
            if rng.random() < 0.5:
                obj["shape"], obj["color"] = sorted(HELDOUT_SHAPE_COLOR)[rng.integers(len(HELDOUT_SHAPE_COLOR))]
            else:
                obj["shape"], obj["motion"] = sorted(HELDOUT_SHAPE_MOTION)[rng.integers(len(HELDOUT_SHAPE_MOTION))]
            objects = [obj]
        else:
            n = 2 if rng.random() < cfg.p_two_objects else 1
            objects = [_random_object(rng, bg, cfg) for _ in range(n)]
        meta = {"objects": objects, "bg": bg, "mention_bg": bool(rng.random() < cfg.p_mention_bg)}
        if any(_clashes(o["color"], bg) for o in objects):
            continue
        if split == "train" and is_heldout(meta):
            continue
        return meta


def caption_of(meta: dict) -> str:
    parts = []
    for o in meta["objects"]:
        size = "" if o["size"] == "medium" else o["size"] + " "
        parts.append(f"a {size}{o['color']} {o['shape']} {MOTIONS[o['motion']]}")
    cap = " and ".join(parts)
    if meta.get("mention_bg", True):
        cap += f" on a {meta['bg']} background"
    return cap


def parse_caption(caption: str) -> dict:
    """Inverse of caption_of (without trajectories). Unmentioned size -> 'medium', bg -> None."""
    words = caption.lower().split()
    bg = None
    if "on" in words:
        i = words.index("on")
        bg = words[i + 2] if i + 2 < len(words) else None
        words = words[:i]
    phrases, cur = [], []
    for w in words:
        if w == "and":
            phrases.append(cur)
            cur = []
        else:
            cur.append(w)
    phrases.append(cur)
    motion_by_words = {v: k for k, v in MOTIONS.items()}
    objects = []
    for p in phrases:
        p = [w for w in p if w != "a"]
        size = next((w for w in p if w in SIZES), "medium")
        color = next((w for w in p if w in COLORS), None)
        shape = next((w for w in p if w in SHAPES), None)
        rest = " ".join(p[p.index(shape) + 1:]) if shape in p else ""
        objects.append({"shape": shape, "color": color, "size": size, "motion": motion_by_words.get(rest)})
    return {"objects": objects, "bg": bg, "mention_bg": bg is not None}


def sample_prompts(n: int, split: str = "train", seed: int = 0, cfg: ClipConfig | None = None) -> list[tuple[str, dict]]:
    rng = np.random.default_rng(seed)
    cfg = cfg or ClipConfig()
    out = []
    for _ in range(n):
        meta = _random_meta(rng, cfg, split)
        out.append((caption_of(meta), meta))
    return out


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

def _background(bg: str, S: int) -> np.ndarray:
    img = np.empty((S, S, 3), np.float32)
    if bg == "gradient":
        ramp = np.linspace(20, 120, S, dtype=np.float32)[:, None, None]
        img[:] = ramp * np.array([0.6, 0.8, 1.0], np.float32)
    else:
        img[:] = {"black": (0, 0, 0), "white": (250, 250, 250), "gray": (128, 128, 128), "navy": (15, 20, 80)}[bg]
    return img


def _polygon_mask(px, py, verts) -> np.ndarray:
    """Even-odd point-in-polygon test, vectorised over pixels."""
    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: int, r: float, angle: float, cx: float, cy: float, inner: float | None = None):
    pts = []
    k = 2 * n if inner else n
    for i in range(k):
        a = angle - math.pi / 2 + 2 * math.pi * i / k
        rr = (inner if i % 2 else r) if inner else r
        pts.append((cx + rr * math.cos(a), cy + rr * math.sin(a)))
    return pts


def _shape_mask(shape: str, px, py, cx, cy, r, angle) -> np.ndarray:
    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 _polygon_mask(px, py, _regular(4, r * 1.15, angle + math.pi / 4, cx, cy))
    if shape == "triangle":
        return _polygon_mask(px, py, _regular(3, r * 1.2, angle, cx, cy))
    if shape == "star":
        return _polygon_mask(px, py, _regular(5, r * 1.25, angle, cx, cy, inner=r * 0.5))
    if shape == "cross":
        c, s = math.cos(angle), math.sin(angle)
        u, v = c * dx + s * dy, -s * dx + c * dy
        arm = 0.33 * r
        return ((np.abs(u) <= r) & (np.abs(v) <= arm)) | ((np.abs(v) <= r) & (np.abs(u) <= arm))
    raise ValueError(shape)


def _trajectory(rng: np.random.Generator, obj: dict, T: int) -> list[tuple[float, float, float, float]]:
    """Per-frame (cx, cy, r, angle) in normalised [0,1] frame coordinates."""
    r0 = SIZES[obj["size"]] * rng.uniform(0.9, 1.1)
    angle0 = rng.uniform(0, 2 * math.pi)
    m = obj["motion"]
    ts = np.linspace(0, 1, T)
    if m in DIRECTIONS:
        dx, dy = DIRECTIONS[m]
        dist = rng.uniform(0.35, 0.6) / (math.sqrt(2) if dx and dy else 1)
        margin = r0 + 0.02

        def start(d):
            if d > 0:
                return rng.uniform(margin, max(margin, 1 - margin - dist))
            if d < 0:
                return rng.uniform(min(1 - margin, margin + dist), 1 - margin)
            return rng.uniform(margin, 1 - margin)

        x0, y0 = start(dx), start(dy)
        return [(x0 + dx * dist * t, y0 + dy * dist * t, r0, angle0) for t in ts]
    if m == "bounces":
        theta = rng.uniform(0, 2 * math.pi)
        speed = rng.uniform(1.4, 2.2)
        lo, hi = r0, 1 - r0
        x, y = rng.uniform(lo, hi, size=2)
        vx, vy = speed * math.cos(theta), speed * math.sin(theta)
        out, dt = [], 1 / (T - 1)
        for _ in range(T):
            out.append((x, y, r0, angle0))
            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
    margin = r0 * (1.6 if m in ("grows", "shrinks") else 1.0) + 0.02
    cx, cy = rng.uniform(margin, 1 - margin, size=2)
    if m == "grows":
        return [(cx, cy, r0 * (0.5 + 1.1 * t), angle0) for t in ts]
    if m == "shrinks":
        return [(cx, cy, r0 * (1.6 - 1.1 * t), angle0) for t in ts]
    if m == "spins":
        spin = rng.choice([-1, 1]) * rng.uniform(0.6, 1.0) * math.pi
        return [(cx, cy, r0, angle0 + spin * t) for t in ts]
    return [(cx, cy, r0, angle0) for _ in ts]  # still


def render(meta: dict, rng: np.random.Generator, cfg: ClipConfig) -> np.ndarray:
    """Render meta (adds meta['objects'][i]['traj']) to uint8 [T, H, W, 3]."""
    T, S, ss = cfg.frames, cfg.size, cfg.supersample
    big = S * ss
    coords = (np.arange(big, dtype=np.float32) + 0.5) / big
    px, py = np.meshgrid(coords, coords)
    bg = _background(meta["bg"], big)
    for o in meta["objects"]:
        o["traj"] = _trajectory(rng, o, T)
    frames = np.empty((T, S, S, 3), np.uint8)
    for f in range(T):
        img = bg.copy()
        for o in meta["objects"]:
            cx, cy, r, ang = o["traj"][f]
            mask = _shape_mask(o["shape"], px, py, cx, cy, r, ang)
            img[mask] = COLORS[o["color"]]
        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: np.random.Generator, cfg: ClipConfig | None = None, split: str = "train"):
    cfg = cfg or ClipConfig()
    meta = _random_meta(rng, cfg, split)
    video = render(meta, rng, cfg)
    return video, caption_of(meta), meta


# ---------------------------------------------------------------- torch dataset

def to_tensor(video: np.ndarray) -> torch.Tensor:
    """uint8 [T,H,W,3] -> float [3,T,H,W] in [-1,1]."""
    return torch.from_numpy(video).permute(3, 0, 1, 2).float().div(127.5).sub(1)


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

    def __init__(self, cfg: ClipConfig | None = None, seed: int = 0, include_heldout: bool = False):
        self.cfg = cfg or ClipConfig()
        self.seed = seed
        self.split = "any" if include_heldout else "train"

    def __iter__(self):
        info = torch.utils.data.get_worker_info()
        wid = info.id if info else 0
        rng = np.random.default_rng([self.seed, wid])
        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)