#!/usr/bin/env python3
"""Odzysk tokenow edu z blendu fork-B (RB4 runda 1) + naprawa separatora.

Blend fork-B = przetasowane bloki: arcmix[0:~1.15B] (separator dokumentu 12285, zero 12287)
oraz dokumenty edu zakonczone 12287 (<|im_end|>, blad buildu: eot=vocab-1).

Dzielimy strumien na kawalki konczace sie 12287. Kawalek bez 12285 i bez 12286 = czysty
dokument edu -> zapis z separatorem podmienionym na 12285. Kawalek zawierajacy 12285/12286
niesie blok arcmixu -> odrzucony w calosci (traci sie co najwyzej jeden dokument edu na blok).

Wyjscie: edu.bin (uint16) + edu.json (liczby do weryfikacji wobec logu buildu).
"""
import argparse
import hashlib
import json

import numpy as np

EOS, IM_START, IM_END = 12285, 12286, 12287
CHUNK = 200_000_000


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--blend", required=True)
    ap.add_argument("--out", required=True)
    a = ap.parse_args()

    b = np.memmap(a.blend, dtype=np.uint16, mode="r")
    n = len(b)
    out = open(a.out, "wb")
    sha = hashlib.sha256()
    kept_docs = kept_tok = drop_pieces = drop_tok = 0
    carry = np.empty(0, dtype=np.uint16)
    for lo in range(0, n, CHUNK):
        x = np.concatenate([carry, np.asarray(b[lo:lo + CHUNK])])
        ends = np.flatnonzero(x == IM_END)
        if lo + CHUNK >= n and (len(ends) == 0 or ends[-1] != len(x) - 1):
            tail = len(x) - (ends[-1] + 1 if len(ends) else 0)
            drop_pieces += 1
            drop_tok += tail
        if len(ends) == 0:
            carry = x
            continue
        starts = np.concatenate([[0], ends[:-1] + 1])
        bad = np.zeros(len(ends), dtype=bool)
        for code in (EOS, IM_START):
            pos = np.flatnonzero(x[:ends[-1] + 1] == code)
            bad[np.searchsorted(ends, pos)] = True
        seg = x[:ends[-1] + 1].copy()
        seg[ends] = EOS
        keep_mask = np.repeat(~bad, ends - starts + 1)
        kept = seg[keep_mask]
        kept.tofile(out)
        sha.update(kept.tobytes())
        kept_docs += int((~bad).sum())
        kept_tok += int(len(kept))
        drop_pieces += int(bad.sum())
        drop_tok += int((ends - starts + 1)[bad].sum())
        carry = x[ends[-1] + 1:]
    out.close()
    stats = {"blend_tokens": int(n), "edu_docs": kept_docs, "edu_tokens": kept_tok,
             "dropped_pieces": drop_pieces, "dropped_tokens": drop_tok,
             "separator": EOS, "sha256": sha.hexdigest()}
    json.dump(stats, open(a.out.rsplit(".", 1)[0] + ".json", "w"), indent=1)
    print(json.dumps(stats), flush=True)


if __name__ == "__main__":
    main()