"""Selective recall / selective copying task generators."""
from __future__ import annotations

import numpy as np


def make_selective_copying_task_tokens(
    episode_len: int,
    num_episodes: int,
    d_in: int,
    *,
    num_info_tokens: int = 16,
    filler_token_id: int | None = None,
    write_token_id: int | None = None,
    seed: int | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Generate vector-token selective copying episodes.

    Args:
        episode_len: Sequence length per episode. Must be divisible by 3.
        num_episodes: Number of episodes to generate.
        d_in: Token embedding dimension.
        num_info_tokens: Number of informative token IDs.
        filler_token_id: Optional explicit filler ID.
        write_token_id: Optional explicit write ID.
        seed: RNG seed for deterministic generation.

    Returns:
        X: Input tensor with shape ``(num_episodes, episode_len, d_in)``.
        Y: Target tensor with shape ``(num_episodes, episode_len, d_in)``.
        token_table: Token embeddings, shape ``(vocab_size, d_in)``.
        ids: Integer token IDs, shape ``(num_episodes, episode_len)``.
    """
    if episode_len % 3 != 0:
        raise ValueError("episode_len must be divisible by 3")

    rng = np.random.default_rng(seed)

    filler_id = num_info_tokens if filler_token_id is None else int(filler_token_id)
    write_id = num_info_tokens + 1 if write_token_id is None else int(write_token_id)
    vocab_size = max(num_info_tokens, filler_id + 1, write_id + 1)

    first_len = (2 * episode_len) // 3
    second_len = episode_len // 3
    n_info = first_len // 2

    token_table = rng.standard_normal(size=(vocab_size, d_in)).astype(np.float32)
    token_table /= np.linalg.norm(token_table, axis=1, keepdims=True) + 1e-8

    X_all, Y_all, ids_all = [], [], []
    for _ in range(num_episodes):
        info_pos = rng.choice(first_len, size=n_info, replace=False)
        info_pos.sort()
        info_ids = rng.integers(low=0, high=num_info_tokens, size=n_info, dtype=np.int32)

        ids_first = np.full(first_len, filler_id, dtype=np.int32)
        ids_first[info_pos] = info_ids
        ids_second = np.full(second_len, write_id, dtype=np.int32)
        ids = np.concatenate([ids_first, ids_second], axis=0)

        X_ep = token_table[ids]
        Y_first = np.zeros((first_len, d_in), dtype=np.float32)
        Y_second = token_table[info_ids]
        Y_ep = np.concatenate([Y_first, Y_second], axis=0)

        X_all.append(X_ep)
        Y_all.append(Y_ep)
        ids_all.append(ids)

    return (
        np.stack(X_all, axis=0),
        np.stack(Y_all, axis=0),
        token_table,
        np.stack(ids_all, axis=0),
    )


__all__ = ["make_selective_copying_task_tokens"]