#!/usr/bin/env python3
import argparse, gzip, json, random, time
from pathlib import Path
import numpy as np
from transformers import AutoTokenizer

BUDGETS_1B = {
    "english_web": 250_000_000,
    "english_edu": 200_000_000,
    "chinese_clean": 250_000_000,
    "code": 150_000_000,
    "math": 100_000_000,
    "science": 30_000_000,
    "qa_as_text": 20_000_000,
}

def open_text(path):
    return gzip.open(path, "rt", encoding="utf-8") if path.suffix == ".gz" else path.open("r", encoding="utf-8")

def flush(out_dir, split, shard_idx, arrays):
    if not arrays: return None
    arr = np.stack(arrays, axis=0)
    path = out_dir / f"{split}_{shard_idx:05d}.npy"
    np.save(path, arr)
    return {"path": path.name, "blocks": int(arr.shape[0]), "tokens": int(arr.size)}

def main():
    ap=argparse.ArgumentParser()
    ap.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
    ap.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
    ap.add_argument("--docmix-dir", default="data/cpt_docmix_5b_sources_8192_20260614")
    ap.add_argument("--out-dir", default="data/cpt_packed_5b_seq8192_seed42_stratified_20260614")
    ap.add_argument("--scale", type=float, default=5.0)
    ap.add_argument("--seq-len", type=int, default=8192)
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--eval-blocks", type=int, default=2048)
    ap.add_argument("--eval-rate", type=float, default=0.01)
    ap.add_argument("--shard-blocks", type=int, default=2048)
    ap.add_argument("--log-every-docs", type=int, default=20000)
    args=ap.parse_args()
    base=Path(args.base_dir); docmix=base/args.docmix_dir; out_dir=base/args.out_dir; out_dir.mkdir(parents=True, exist_ok=True)
    sources={cat: docmix/"documents"/f"{cat}.jsonl.gz" for cat in BUDGETS_1B}
    missing=[str(p) for p in sources.values() if not p.exists()]
    if missing: raise FileNotFoundError(missing)
    budgets={k:int(v*args.scale) for k,v in BUDGETS_1B.items()}
    tok=AutoTokenizer.from_pretrained(base/args.tokenizer, trust_remote_code=True); eos=tok.eos_token_id
    rng=random.Random(args.seed)
    total_budget=sum(budgets.values())
    eval_quota={k:round(args.eval_blocks*v/total_budget) for k,v in budgets.items()}
    diff=args.eval_blocks-sum(eval_quota.values())
    if diff: eval_quota["english_web"] += diff
    stats={"docs_seen":0,"train_blocks":0,"eval_blocks":0,"train_tokens":0,"eval_tokens":0,"source_docs":{},"source_tokens":{},"train_blocks_by_category":{},"eval_blocks_by_category":{},"leftover_tokens_by_category":{},"start_time":time.time()}
    shards={"train":[],"eval":[]}; shard_idx={"train":0,"eval":0}; train_arrays=[]; eval_arrays=[]
    def add_block(block, category):
        if stats["eval_blocks_by_category"].get(category,0) < eval_quota.get(category,0) and rng.random() < args.eval_rate:
            eval_arrays.append(block); stats["eval_blocks"] += 1; stats["eval_tokens"] += args.seq_len
            stats["eval_blocks_by_category"][category]=stats["eval_blocks_by_category"].get(category,0)+1
            if len(eval_arrays) >= args.shard_blocks:
                rec=flush(out_dir,"eval",shard_idx["eval"],eval_arrays); shards["eval"].append(rec); shard_idx["eval"] += 1; eval_arrays.clear()
        else:
            train_arrays.append(block); stats["train_blocks"] += 1; stats["train_tokens"] += args.seq_len
            stats["train_blocks_by_category"][category]=stats["train_blocks_by_category"].get(category,0)+1
            if len(train_arrays) >= args.shard_blocks:
                rec=flush(out_dir,"train",shard_idx["train"],train_arrays); shards["train"].append(rec); shard_idx["train"] += 1; train_arrays.clear()
    for category,path in sources.items():
        target=budgets[category]; cat_tokens=0; cat_docs=0; buffer=[]
        with open_text(path) as f:
            for line in f:
                if cat_tokens >= target: break
                if not line.strip(): continue
                row=json.loads(line); text=row.get("text") or ""
                if not text: continue
                ids=tok.encode(text, add_special_tokens=False)
                if not ids: continue
                ids.append(eos)
                if cat_tokens + len(ids) > target and cat_tokens > 0:
                    break
                buffer.extend(ids); cat_tokens += len(ids); cat_docs += 1; stats["docs_seen"] += 1
                stats["source_docs"][category]=cat_docs; stats["source_tokens"][category]=cat_tokens
                while len(buffer) >= args.seq_len:
                    block=np.asarray(buffer[:args.seq_len], dtype=np.uint32); del buffer[:args.seq_len]; add_block(block, category)
                if stats["docs_seen"] % args.log_every_docs == 0:
                    rec={k:stats[k] for k in ["docs_seen","train_blocks","eval_blocks","train_tokens","eval_tokens"]}; rec.update({"category":category,"category_tokens":cat_tokens,"elapsed_sec":time.time()-stats["start_time"]})
                    print(json.dumps(rec, ensure_ascii=False), flush=True)
        while len(buffer) >= args.seq_len:
            block=np.asarray(buffer[:args.seq_len], dtype=np.uint32); del buffer[:args.seq_len]; add_block(block, category)
        stats["leftover_tokens_by_category"][category]=len(buffer)
    rec=flush(out_dir,"train",shard_idx["train"],train_arrays)
    if rec: shards["train"].append(rec)
    rec=flush(out_dir,"eval",shard_idx["eval"],eval_arrays)
    if rec: shards["eval"].append(rec)
    manifest={**stats,"tokenizer":str(base/args.tokenizer),"seq_len":args.seq_len,"seed":args.seed,"scale":args.scale,"budgets":budgets,"sources":{k:str(v) for k,v in sources.items()},"eval_quota_blocks":eval_quota,"train_shards":shards["train"],"eval_shards":shards["eval"],"elapsed_sec":time.time()-stats["start_time"]}
    (out_dir/"manifest.json").write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(out_dir/"manifest.json")
if __name__=="__main__": main()