'''
data.py defines the synthetic benchmark data pipeline for the spectral world model project.

Specifies vocabulary and action labels for text-conditioned dynamics
Generates moving shape sequences with image frames, text tokens, actions, and rewards
Saves train/val/test splits to a compressed .npz benchmark file
Provides TransitionDataset for frame-to-frame transition learning
Provides SequenceDataset for full-sequence rollout evaluation
'''
# Andrew Kiruluta, UC Berkeley, June 2026

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import json
import math
import numpy as np
import torch
from torch.utils.data import Dataset


VOCAB = [
    "<pad>",
    "square", "circle",
    "stay", "up", "down", "left", "right",
    "slow", "fast",
    "x0", "x1", "x2", "x3",
    "y0", "y1", "y2", "y3",
]
TOKEN_TO_ID = {tok: i for i, tok in enumerate(VOCAB)}
ACTION_NAMES = ["stay", "up", "down", "left", "right"]


@dataclass
class DatasetConfig:
    image_size: int = 32
    seq_len: int = 8
    max_text_len: int = 5
    train_sequences: int = 256
    val_sequences: int = 64
    test_sequences: int = 64
    seed: int = 7


def _draw_square(canvas: np.ndarray, x: int, y: int, size: int) -> None:
    h, w = canvas.shape
    x0 = max(0, x - size // 2); x1 = min(w, x + size // 2 + 1)
    y0 = max(0, y - size // 2); y1 = min(h, y + size // 2 + 1)
    canvas[y0:y1, x0:x1] = 1.0


def _draw_circle(canvas: np.ndarray, x: int, y: int, radius: int) -> None:
    h, w = canvas.shape
    yy, xx = np.ogrid[:h, :w]
    mask = (xx - x) ** 2 + (yy - y) ** 2 <= radius**2
    canvas[mask] = 1.0


def _text_tokens(shape: str, action: int, speed: int, x: int, y: int, image_size: int, max_len: int) -> np.ndarray:
    xb = min(3, max(0, int(4 * x / image_size)))
    yb = min(3, max(0, int(4 * y / image_size)))
    words = [shape, ACTION_NAMES[action], "fast" if speed == 2 else "slow", f"x{xb}", f"y{yb}"]
    ids = [TOKEN_TO_ID[w] for w in words][:max_len]
    ids += [0] * (max_len - len(ids))
    return np.asarray(ids, dtype=np.int64)


def _apply_action(x: int, y: int, action: int, speed: int, image_size: int, radius: int) -> tuple[int, int]:
    dx, dy = 0, 0
    if action == 1: dy = -speed
    elif action == 2: dy = speed
    elif action == 3: dx = -speed
    elif action == 4: dx = speed
    x = int(np.clip(x + dx, radius + 1, image_size - radius - 2))
    y = int(np.clip(y + dy, radius + 1, image_size - radius - 2))
    return x, y


def generate_sequences(n_sequences: int, cfg: DatasetConfig, rng: np.random.Generator) -> dict[str, np.ndarray]:
    images = np.zeros((n_sequences, cfg.seq_len, 1, cfg.image_size, cfg.image_size), dtype=np.float32)
    tokens = np.zeros((n_sequences, cfg.seq_len, cfg.max_text_len), dtype=np.int64)
    actions = np.zeros((n_sequences, cfg.seq_len - 1), dtype=np.int64)
    rewards = np.zeros((n_sequences, cfg.seq_len - 1), dtype=np.float32)

    for n in range(n_sequences):
        shape = "square" if rng.random() < 0.5 else "circle"
        speed = int(rng.choice([1, 2]))
        radius = 3 if shape == "circle" else 4
        x = int(rng.integers(radius + 2, cfg.image_size - radius - 2))
        y = int(rng.integers(radius + 2, cfg.image_size - radius - 2))
        last_action = 0
        for t in range(cfg.seq_len):
            canvas = np.zeros((cfg.image_size, cfg.image_size), dtype=np.float32)
            if shape == "square":
                _draw_square(canvas, x, y, 6)
            else:
                _draw_circle(canvas, x, y, radius)
            # Add tiny structured sensor noise to make denoising useful.
            canvas += rng.normal(0.0, 0.01, size=canvas.shape).astype(np.float32)
            canvas = np.clip(canvas, 0.0, 1.0)
            images[n, t, 0] = canvas
            tokens[n, t] = _text_tokens(shape, last_action, speed, x, y, cfg.image_size, cfg.max_text_len)
            if t < cfg.seq_len - 1:
                action = int(rng.integers(0, len(ACTION_NAMES)))
                actions[n, t] = action
                x, y = _apply_action(x, y, action, speed, cfg.image_size, radius)
                last_action = action
                # Reward favors reaching the upper-right quadrant.
                rewards[n, t] = float((x > cfg.image_size * 0.65) and (y < cfg.image_size * 0.35))
    return {"images": images, "tokens": tokens, "actions": actions, "rewards": rewards}


def generate_benchmark_npz(path: str | Path, cfg: DatasetConfig = DatasetConfig()) -> None:
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(cfg.seed)
    train = generate_sequences(cfg.train_sequences, cfg, rng)
    val = generate_sequences(cfg.val_sequences, cfg, rng)
    test = generate_sequences(cfg.test_sequences, cfg, rng)
    np.savez_compressed(
        path,
        train_images=train["images"], train_tokens=train["tokens"], train_actions=train["actions"], train_rewards=train["rewards"],
        val_images=val["images"], val_tokens=val["tokens"], val_actions=val["actions"], val_rewards=val["rewards"],
        test_images=test["images"], test_tokens=test["tokens"], test_actions=test["actions"], test_rewards=test["rewards"],
        vocab=np.asarray(VOCAB), action_names=np.asarray(ACTION_NAMES),
        config=json.dumps(cfg.__dict__),
    )


class TransitionDataset(Dataset):
    def __init__(self, npz_path: str | Path, split: str):
        data = np.load(npz_path, allow_pickle=True)
        self.images = data[f"{split}_images"]
        self.tokens = data[f"{split}_tokens"]
        self.actions = data[f"{split}_actions"]
        self.seq_len = self.images.shape[1]
        self.index = [(i, t) for i in range(self.images.shape[0]) for t in range(self.seq_len - 1)]

    def __len__(self) -> int:
        return len(self.index)

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        i, t = self.index[idx]
        return {
            "image": torch.from_numpy(self.images[i, t]),
            "tokens": torch.from_numpy(self.tokens[i, t]).long(),
            "action": torch.tensor(self.actions[i, t]).long(),
            "next_image": torch.from_numpy(self.images[i, t + 1]),
            "next_tokens": torch.from_numpy(self.tokens[i, t + 1]).long(),
        }


class SequenceDataset(Dataset):
    def __init__(self, npz_path: str | Path, split: str):
        data = np.load(npz_path, allow_pickle=True)
        self.images = data[f"{split}_images"]
        self.tokens = data[f"{split}_tokens"]
        self.actions = data[f"{split}_actions"]
        self.rewards = data[f"{split}_rewards"]

    def __len__(self) -> int:
        return self.images.shape[0]

    def __getitem__(self, idx: int) -> dict[str, torch.Tensor]:
        return {
            "images": torch.from_numpy(self.images[idx]),
            "tokens": torch.from_numpy(self.tokens[idx]).long(),
            "actions": torch.from_numpy(self.actions[idx]).long(),
            "rewards": torch.from_numpy(self.rewards[idx]),
        }