"""Data pipeline for Part 2: next-token pretraining on the human-AI corpus.

Every document in ``browndw/human-ai-parallel-corpus`` exists in several
parallel versions -- the human original (``@chunk_N``) and one continuation
per generator model (``@gpt-4o-*``).  Splitting rows at random would drop
near-paraphrases of the same underlying document into both train and test and
quietly inflate every validation number, so the split is made on the *base*
document id and all of a document's versions move together.
"""

from __future__ import annotations

import numpy as np
import torch

from ..utils import RUNS, load_tokenizer, train_bpe

DATASET = "browndw/human-ai-parallel-corpus"
PAD, BOS, EOS, UNK = 0, 1, 2, 3
SPECIALS = ["<pad>", "<bos>", "<eos>", "<unk>"]


def _base_id(doc_id: str) -> str:
    return doc_id.split("@")[0]


def load_splits(val_frac=0.02, test_frac=0.02, seed=1337):
    """Document-level train/val/test split of the corpus."""
    from datasets import load_dataset
    ds = load_dataset(DATASET, split="train")

    bases = sorted({_base_id(d) for d in ds["doc_id"]})
    rng = np.random.default_rng(seed)
    rng.shuffle(bases)
    n_val = int(len(bases) * val_frac)
    n_test = int(len(bases) * test_frac)
    val_b = set(bases[:n_val])
    test_b = set(bases[n_val:n_val + n_test])

    which = []
    for d in ds["doc_id"]:
        b = _base_id(d)
        which.append("val" if b in val_b else "test" if b in test_b else "train")
    which = np.array(which)
    idx = {s: np.flatnonzero(which == s) for s in ("train", "val", "test")}
    return ds, idx


def get_tokenizer(vocab_size: int = 16000, force: bool = False):
    path = RUNS / "part2" / f"tokenizer_{vocab_size}.json"
    if path.exists() and not force:
        return load_tokenizer(path)
    ds, idx = load_splits()
    train_texts = ds.select(idx["train"])["text"]
    print(f"[part2] training BPE (vocab={vocab_size}) on {len(train_texts):,} docs")
    return train_bpe(iter(train_texts), vocab_size, path, specials=SPECIALS)


def get_stream(tok, split: str, force: bool = False) -> np.ndarray:
    """Flat token stream for one split, documents joined by <eos>."""
    cache = RUNS / "part2" / "cache"
    cache.mkdir(parents=True, exist_ok=True)
    f = cache / f"{split}.npy"
    if f.exists() and not force:
        return np.load(f)

    ds, idx = load_splits()
    texts = ds.select(idx[split])["text"]
    print(f"[part2] encoding {split}: {len(texts):,} docs")
    ids = []
    B = 2000
    for s in range(0, len(texts), B):
        for e in tok.encode_batch(texts[s:s + B]):
            ids.append(BOS)
            ids.extend(e.ids)
            ids.append(EOS)
    arr = np.asarray(ids, dtype=np.uint16)
    np.save(f, arr)
    return arr


class Blocks(torch.utils.data.Dataset):
    """Contiguous (x, y) blocks over a flat token stream."""

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

    def __len__(self):
        return self.n

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


def make_loader(ds, batch_size, shuffle, seed=1337, num_workers=2):
    g = torch.Generator().manual_seed(seed)
    return torch.utils.data.DataLoader(
        ds, 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)


def bleu_prompts(tok, n=300, prompt_len=64, cont_len=48, seed=0):
    """(prompt_ids, reference_text) pairs from held-out test documents."""
    ds, idx = load_splits()
    rng = np.random.default_rng(seed)
    pick = rng.choice(idx["test"], size=min(n, len(idx["test"])), replace=False)
    out = []
    for i in pick:
        e = tok.encode(ds[int(i)]["text"]).ids
        if len(e) < prompt_len + cont_len:
            continue
        out.append(([BOS] + e[:prompt_len],
                    tok.decode(e[prompt_len:prompt_len + cont_len])))
    return out