"""Data pipeline for the vi/ja -> en translation task.

The model is decoder-only, so translation is framed as plain language
modelling over a prompt/answer sequence:

    <bos> <vi> {vietnamese source} <en> {english target} <eos>

Only the tokens after ``<en>`` (the English target and the closing ``<eos>``)
contribute to the loss.  Conditioning tokens are still *read* by the model --
and still get routed through the MoE, which is what makes the per-language
expert heatmap meaningful -- they are simply not predicted.

Examples are packed back-to-back into fixed ``n_ctx``-token blocks rather than
padded individually.  A typical pair is ~60 tokens, so per-example padding to
256 would have thrown away three quarters of every batch, and "same number of
training tokens" would then have meant "same number of padding tokens".
"""

from __future__ import annotations

import numpy as np
import torch

from ..utils import RUNS, load_tokenizer, train_bpe

DATASET = "belumind/en-vi-ja-curated-500k-triplets"

PAD, BOS, EOS, UNK, VI, JA, EN = 0, 1, 2, 3, 4, 5, 6
LANG_VI, LANG_JA, LANG_EN, LANG_SPECIAL = 0, 1, 2, 3
LANG_NAMES = ["vi (source)", "ja (source)", "en (target)", "special"]


def load_raw(split: str, limit: int | None = None):
    from datasets import load_dataset
    ds = load_dataset(DATASET, split=split)
    if limit is not None and limit < len(ds):
        ds = ds.select(range(limit))
    return ds


# ---------------------------------------------------------------------------
# Tokenizer
# ---------------------------------------------------------------------------
def get_tokenizer(vocab_size: int = 16000, n_train_rows: int = 250_000,
                  force: bool = False):
    """Train (or load) the shared byte-level BPE over all three languages."""
    path = RUNS / "part1" / f"tokenizer_{vocab_size}.json"
    if path.exists() and not force:
        return load_tokenizer(path)

    ds = load_raw("train", limit=n_train_rows)

    def it():
        # Interleaving the three languages keeps any one script from
        # monopolising the merge budget.
        for row in ds:
            yield row["en"]
            yield row["vi"]
            yield row["ja"]

    print(f"[part1] training BPE (vocab={vocab_size}) on {3*len(ds):,} lines")
    return train_bpe(it(), vocab_size, path)


# ---------------------------------------------------------------------------
# Encoding into packed streams
# ---------------------------------------------------------------------------
def _encode_split(tok, split: str, limit: int | None):
    """Return (ids, is_target, lang) int arrays for one split, packed flat."""
    ds = load_raw(split, limit=limit)
    ids, tgt, lang = [], [], []

    en_all = tok.encode_batch([r["en"] for r in ds])
    vi_all = tok.encode_batch([r["vi"] for r in ds])
    ja_all = tok.encode_batch([r["ja"] for r in ds])

    for en, vi, ja in zip(en_all, vi_all, ja_all):
        e = en.ids
        for src_ids, marker, lang_id in ((vi.ids, VI, LANG_VI),
                                         (ja.ids, JA, LANG_JA)):
            seq = [BOS, marker] + src_ids + [EN] + e + [EOS]
            is_t = ([0, 0] + [0] * len(src_ids) + [0]
                    + [1] * len(e) + [1])
            lg = ([LANG_SPECIAL, LANG_SPECIAL] + [lang_id] * len(src_ids)
                  + [LANG_SPECIAL] + [LANG_EN] * len(e) + [LANG_SPECIAL])
            ids.extend(seq)
            tgt.extend(is_t)
            lang.extend(lg)

    return (np.asarray(ids, dtype=np.uint16),
            np.asarray(tgt, dtype=np.uint8),
            np.asarray(lang, dtype=np.uint8))


def get_streams(tok, split: str, limit: int | None = None, force: bool = False):
    """Tokenised streams for a split, cached on disk as .npy."""
    tag = f"{split}_{limit if limit is not None else 'all'}"
    cache = RUNS / "part1" / "cache"
    cache.mkdir(parents=True, exist_ok=True)
    files = {k: cache / f"{tag}_{k}.npy" for k in ("ids", "tgt", "lang")}

    if all(f.exists() for f in files.values()) and not force:
        return tuple(np.load(f) for f in files.values())

    print(f"[part1] encoding split={split} limit={limit}")
    arrs = _encode_split(tok, split, limit)
    for f, a in zip(files.values(), arrs):
        np.save(f, a)
    return arrs


class PackedBlocks(torch.utils.data.Dataset):
    """Chops a flat token stream into (x, y, loss_mask, lang) blocks."""

    def __init__(self, ids, tgt, lang, n_ctx: int):
        n = (len(ids) - 1) // n_ctx
        self.n_ctx, self.n = n_ctx, n
        self.ids, self.tgt, self.lang = ids, tgt, lang

    def __len__(self):
        return self.n

    def __getitem__(self, i):
        s = i * self.n_ctx
        e = s + self.n_ctx
        x = torch.from_numpy(self.ids[s:e].astype(np.int64))
        y = torch.from_numpy(self.ids[s + 1:e + 1].astype(np.int64))
        m = torch.from_numpy(self.tgt[s + 1:e + 1].astype(np.float32))
        lg = torch.from_numpy(self.lang[s:e].astype(np.int64))
        return x, y, m, lg


def make_loader(dataset, batch_size, shuffle, seed=1337, num_workers=2):
    g = torch.Generator().manual_seed(seed)
    return torch.utils.data.DataLoader(
        dataset, batch_size=batch_size, shuffle=shuffle, generator=g if shuffle else None,
        num_workers=num_workers, pin_memory=True, drop_last=True,
        persistent_workers=num_workers > 0)


# ---------------------------------------------------------------------------
# Raw pairs, used for BLEU (needs real sentence boundaries, not packed blocks)
# ---------------------------------------------------------------------------
def eval_pairs(split: str = "test", n: int = 500, seed: int = 0):
    """(lang_marker, source_text, reference_english) triples for generation."""
    ds = load_raw(split)
    rng = np.random.default_rng(seed)
    idx = rng.choice(len(ds), size=min(n, len(ds)), replace=False)
    out = []
    for i in idx:
        row = ds[int(i)]
        out.append(("vi", row["vi"], row["en"]))
        out.append(("ja", row["ja"], row["en"]))
    return out


def build_prompt(tok, lang: str, src_text: str, max_src: int = 128):
    marker = VI if lang == "vi" else JA
    src = tok.encode(src_text).ids[:max_src]
    return [BOS, marker] + src + [EN]