#!/usr/bin/env python3
"""RB4 blend builder (runda 3+), zastepuje blend_edu.py z rundy 1.

Poprawki wobec rundy 1 (confoundy z 2026-09-25):
- separator dokumentu brany z tokenizera (<|endoftext|>), nigdy vocab-1;
- zrodlo probkowane z CALEGO pliku na granicach dokumentow (seedowany wybor), nie z glowy;
- manifest z tokenami per zrodlo i per region (dokumenty czatu/QA = zawieraja <|im_start|>).

Tryby:
  index   --src arcmix.bin --out arcmix_index.npz
          granice dokumentow + flaga QA (dokument zawiera <|im_start|>)
  sample  --src arcmix.bin --index arcmix_index.npz --tokens N [--qa-weight W] --out X.bin
          seedowany wybor dokumentow (dokument QA wystepuje W razy w puli) az do N tokenow,
          zapis w przetasowanej kolejnosci
  mix     --a A.bin --b B.bin --a-tokens N --b-tokens M --out X.bin
          dwa pliki z tym samym separatorem -> przetasowanie na poziomie dokumentow

Kazdy tryb zapisuje <out>.json (manifest) obok wyjscia.
"""
import argparse
import hashlib
import json
from pathlib import Path

import numpy as np

CHUNK = 250_000_000


def special_ids(tokenizer):
    added = {t["content"]: t["id"] for t in json.load(open(tokenizer, encoding="utf-8"))["added_tokens"]}
    return added["<|endoftext|>"], added["<|im_start|>"]


def doc_bounds(arr, eos):
    """(starts, lens) dokumentow konczacych sie eos; ogon bez eos pomijany."""
    ends = []
    for lo in range(0, len(arr), CHUNK):
        ends.append(np.flatnonzero(np.asarray(arr[lo:lo + CHUNK]) == eos) + lo)
    ends = np.concatenate(ends).astype(np.int64)
    starts = np.concatenate([[0], ends[:-1] + 1])
    return starts, ends - starts + 1


def write_docs(out_path, parts):
    """parts: iterowalne (array, start, len); zapis + sha256."""
    sha, total, buf, buflen = hashlib.sha256(), 0, [], 0
    with open(out_path, "wb") as out:
        for arr, s, n in parts:
            buf.append(np.asarray(arr[s:s + n]))
            buflen += n
            if buflen >= 50_000_000:
                x = np.concatenate(buf)
                x.tofile(out)
                sha.update(x.tobytes())
                total += len(x)
                buf, buflen = [], 0
        if buf:
            x = np.concatenate(buf)
            x.tofile(out)
            sha.update(x.tobytes())
            total += len(x)
    return total, sha.hexdigest()


def cmd_index(a):
    eos, im_start = special_ids(a.tokenizer)
    arr = np.memmap(a.src, dtype=np.uint16, mode="r")
    starts, lens = doc_bounds(arr, eos)
    qa_pos = np.concatenate([np.flatnonzero(np.asarray(arr[lo:lo + CHUNK]) == im_start) + lo
                             for lo in range(0, len(arr), CHUNK)])
    qa = np.zeros(len(starts), dtype=bool)
    qa[np.searchsorted(starts, qa_pos, side="right") - 1] = True
    np.savez(a.out, starts=starts, lens=lens, qa=qa)
    meta = {"src": a.src, "src_tokens": int(len(arr)), "eos": eos, "docs": int(len(starts)),
            "doc_tokens": int(lens.sum()), "qa_docs": int(qa.sum()), "qa_tokens": int(lens[qa].sum())}
    Path(a.out + ".json").write_text(json.dumps(meta, indent=1))
    print(json.dumps(meta), flush=True)


def cmd_sample(a):
    eos, _ = special_ids(a.tokenizer)
    arr = np.memmap(a.src, dtype=np.uint16, mode="r")
    idx = np.load(a.index)
    starts, lens, qa = idx["starts"], idx["lens"], idx["qa"]
    pool = np.concatenate([np.arange(len(starts))] + [np.flatnonzero(qa)] * (a.qa_weight - 1))
    rng = np.random.default_rng(a.seed)
    rng.shuffle(pool)
    take = pool[:np.searchsorted(np.cumsum(lens[pool]), a.tokens) + 1]
    total, sha = write_docs(a.out, ((arr, int(starts[i]), int(lens[i])) for i in take))
    qa_tok = int(lens[take][qa[take]].sum())
    meta = {"mode": "sample", "src": a.src, "seed": a.seed, "eos": eos, "qa_weight": a.qa_weight,
            "docs": int(len(take)), "unique_docs": int(len(np.unique(take))), "tokens": total,
            "qa_tokens": qa_tok, "qa_share": qa_tok / total,
            "src_qa_share": float(lens[qa].sum() / lens.sum()), "sha256": sha}
    Path(a.out + ".json").write_text(json.dumps(meta, indent=1))
    print(json.dumps(meta), flush=True)


def cmd_mix(a):
    eos, _ = special_ids(a.tokenizer)
    A = np.memmap(a.a, dtype=np.uint16, mode="r")
    B = np.memmap(a.b, dtype=np.uint16, mode="r")
    sa, la = doc_bounds(A, eos)
    sb, lb = doc_bounds(B, eos)
    rng = np.random.default_rng(a.seed)
    oa, ob = rng.permutation(len(sa)), rng.permutation(len(sb))
    oa = oa[:np.searchsorted(np.cumsum(la[oa]), a.a_tokens) + 1]
    ob = ob[:np.searchsorted(np.cumsum(lb[ob]), a.b_tokens) + 1]
    tagged = np.concatenate([np.stack([np.zeros_like(oa), oa], 1), np.stack([np.ones_like(ob), ob], 1)])
    rng.shuffle(tagged)
    src = ((A, int(sa[i]), int(la[i])) if t == 0 else (B, int(sb[i]), int(lb[i])) for t, i in tagged)
    total, sha = write_docs(a.out, src)
    at, bt = int(la[oa].sum()), int(lb[ob].sum())
    meta = {"mode": "mix", "a": a.a, "b": a.b, "seed": a.seed, "eos": eos,
            "a_docs": int(len(oa)), "a_tokens": at, "b_docs": int(len(ob)), "b_tokens": bt,
            "tokens": total, "a_share": at / total, "sha256": sha}
    Path(a.out + ".json").write_text(json.dumps(meta, indent=1))
    print(json.dumps(meta), flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tokenizer", required=True)
    ap.add_argument("--seed", type=int, default=1337)
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("index"); p.add_argument("--src", required=True); p.add_argument("--out", required=True)
    p = sub.add_parser("sample"); p.add_argument("--src", required=True); p.add_argument("--index", required=True)
    p.add_argument("--tokens", type=int, required=True); p.add_argument("--qa-weight", type=int, default=1)
    p.add_argument("--out", required=True)
    p = sub.add_parser("mix"); p.add_argument("--a", required=True); p.add_argument("--b", required=True)
    p.add_argument("--a-tokens", type=int, required=True); p.add_argument("--b-tokens", type=int, required=True)
    p.add_argument("--out", required=True)
    a = ap.parse_args()
    {"index": cmd_index, "sample": cmd_sample, "mix": cmd_mix}[a.cmd](a)


if __name__ == "__main__":
    main()