#!/usr/bin/env python3
"""Calibration corpus v3 (~15M tokens) for phase-2 full AQLM.

Writes /data/glm52-calib-v3/shard_NNNNN.npy (uint32 token ids, ~1M tokens
each) using the GLM tokenizer at /data/glm52. Deterministic (seed 42).

Mix (by tokens):
  ~40% code           (local vLLM sources + HF CodeFeedback + HF python code)
  ~25% tool-calling / agentic   (reused make_agentic_session, GLM chat template)
  ~15% instruction chat         (HF Alpaca + reused coding chat)
  ~10% medical        (MedQA textbook continuation + medical Q&A)
  ~10% prose          (vLLM docs markdown + HF Dolly general knowledge)

Reuses the generators in tools/collect_expert_stats_v2.py. Downloads three
ungated HF streaming datasets:
  - m-a-p/CodeFeedback-Filtered-Instruction   (code)
  - jtatman/python-code-dataset-500k          (code)
  - tatsu-lab/alpaca                          (instruction)
  - databricks/databricks-dolly-15k           (prose / general knowledge)

HELD-OUT EXCLUSIONS (these feed /data/glm52-heldout.txt; never include them):
  - do not read /data/glm52-heldout.txt
  - skip the last 25 vLLM docs/ markdown files (sorted)
  - skip the last 3 MedQA jsonl files (sorted)
  - skip vLLM python files at indices >= 400 of the seed-42 shuffle (the
    corpus builder's random.Random(42) order); we take only code_files[:400].
"""
import glob
import json
import os
import random
import sys

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(ROOT, "tools"))
import collect_expert_stats_v2 as cs   # noqa: E402  reuse the generators

TOK_DIR = "/data/glm52"
OUT = "/data/glm52-calib-v3"
HELDOUT = "/data/glm52-heldout.txt"
SHARD_TOKENS = 1_000_000
TARGET = 15_000_000
SEED = 42

BUDGET = {                     # target tokens per category
    "code": int(0.40 * TARGET),
    "agentic": int(0.25 * TARGET),
    "instruction": int(0.15 * TARGET),
    "medical": int(0.10 * TARGET),
    "prose": int(0.10 * TARGET),
}

MEDICAL_QA = [
    "Explain the differential diagnosis approach for acute chest pain, "
    "covering cardiac, pulmonary, gastrointestinal, and musculoskeletal "
    "causes, and which investigations discriminate between them.",
    "Describe how to systematically read a chest X-ray, common findings "
    "(consolidation, pneumothorax, effusion, cardiomegaly, nodules), and "
    "typical pitfalls in interpretation.",
    "Explain the pharmacology of beta-blockers: mechanism, receptor "
    "selectivity, indications, contraindications, and interactions.",
    "Walk through the pathophysiology of type 2 diabetes from insulin "
    "resistance to complications, and the mechanism of each major drug "
    "class used to treat it.",
    "Explain how CT and MRI imaging work physically, when each is "
    "preferred clinically, and their contraindications.",
    "Describe the interpretation of a full blood count and common "
    "patterns: microcytic vs macrocytic anaemia, neutrophilia, "
    "lymphopenia, thrombocytopenia, and their differential causes.",
    "Explain the staging and grading of solid tumours, TNM notation, "
    "and how imaging and histopathology contribute to each.",
    "Describe the physiology of the cardiac cycle and how it maps to "
    "ECG waveforms, heart sounds, and common arrhythmia mechanisms.",
    "Explain sepsis: definitions, pathophysiology, early recognition "
    "criteria, and the evidence behind initial management bundles.",
    "Explain acid-base disturbances and how to interpret an arterial "
    "blood gas step by step, with compensated and mixed examples.",
    "Describe the mechanisms and comparative effectiveness of the major "
    "vaccine platforms: live attenuated, inactivated, subunit, mRNA, "
    "and viral vector.",
]
GENERAL_QA = [
    "the history of the transistor", "the economics of shipping",
    "how compilers optimize loops", "how DNS resolution works",
    "what causes inflation", "the water cycle",
    "the development of the internet", "how vaccines are manufactured",
    "the physics of climate", "how GPS positioning works",
]


class ShardWriter:
    def __init__(self, out, shard_tokens):
        self.out = out
        self.shard_tokens = shard_tokens
        os.makedirs(out, exist_ok=True)
        self.buf = []          # list of np.uint32 arrays
        self.buf_n = 0
        self.shard = 0
        self.total = 0
        self.by_cat = {}

    def add(self, ids, cat):
        if len(ids) == 0:
            return
        a = np.asarray(ids, dtype=np.uint32)
        self.buf.append(a)
        self.buf_n += len(a)
        self.total += len(a)
        self.by_cat[cat] = self.by_cat.get(cat, 0) + len(a)
        while self.buf_n >= self.shard_tokens:
            self._flush(self.shard_tokens)

    def _flush(self, n):
        cat = np.concatenate(self.buf)
        head, tail = cat[:n], cat[n:]
        p = os.path.join(self.out, f"shard_{self.shard:05d}.npy")
        np.save(p, head)
        assert head.size > 0, f"empty shard {p}"
        print(f"  wrote {p} ({head.size} tokens)", flush=True)
        self.shard += 1
        self.buf = [tail] if tail.size else []
        self.buf_n = tail.size

    def finalize(self):
        if self.buf_n > 0:
            cat = np.concatenate(self.buf)
            p = os.path.join(self.out, f"shard_{self.shard:05d}.npy")
            np.save(p, cat)
            assert cat.size > 0, f"empty shard {p}"
            print(f"  wrote {p} ({cat.size} tokens)", flush=True)
            self.shard += 1
            self.buf, self.buf_n = [], 0


# ---------------------------------------------------------------- generators
def gen_code_local(art):
    """Raw local code, restricted to the seed-42 shuffle's first 400 files
    (held-out uses vLLM python at indices >= 400)."""
    for p in art["code_files"][:400]:
        t = cs._read(p, 60000)
        if len(t) > 300:
            yield f"# ==== {os.path.relpath(p, ROOT)} ====\n{t}\n"


def gen_code_hf():
    """HF code: CodeFeedback (query+answer) and python-code-dataset."""
    from datasets import load_dataset
    cf = iter(load_dataset("m-a-p/CodeFeedback-Filtered-Instruction",
                           split="train", streaming=True))
    py = iter(load_dataset("jtatman/python-code-dataset-500k",
                           split="train", streaming=True))
    while True:
        got = False
        try:
            r = next(cf)
            yield f"### Task:\n{r['query']}\n\n### Solution:\n{r['answer']}\n"
            got = True
        except StopIteration:
            pass
        try:
            r = next(py)
            yield (f"# {r.get('instruction', '')}\n{r.get('output', '')}\n")
            got = True
        except StopIteration:
            pass
        if not got:
            return


def gen_agentic(rng, art, tok):
    while True:
        yield cs.make_agentic_session(rng, art, tok)


def gen_instruction(tok):
    from datasets import load_dataset
    ds = load_dataset("tatsu-lab/alpaca", split="train", streaming=True)
    for r in ds:
        instr = r["instruction"]
        if r.get("input"):
            instr = f"{instr}\n\n{r['input']}"
        msgs = [{"role": "user", "content": instr},
                {"role": "assistant", "content": r["output"]}]
        yield tok.apply_chat_template(msgs, tokenize=False)


def gen_coding_chat(rng, art, tok):
    while True:
        yield cs.make_coding_chat(rng, art, tok)


def gen_medical(rng, tok):
    # MedQA textbook continuation, excluding the last 3 (sorted) jsonl files
    files = sorted(glob.glob("/tmp/medqa/**/*.jsonl", recursive=True))
    kept = files[:-3] if len(files) > 3 else files
    buf = ""
    for f in kept:
        for line in open(f, errors="ignore"):
            try:
                buf += json.loads(line)["text"] + "\n\n"
            except (json.JSONDecodeError, KeyError):
                continue
            while len(buf) >= cs.MAX_TOKENS * 4:
                yield buf[: cs.MAX_TOKENS * 4]
                buf = buf[cs.MAX_TOKENS * 4:]
    if buf.strip():
        yield buf
    # medical Q&A prompts (chat-templated)
    for q in MEDICAL_QA:
        msgs = [{"role": "user", "content": q + " Be thorough and detailed."}]
        yield tok.apply_chat_template(msgs, tokenize=False,
                                      add_generation_prompt=True)


def gen_prose(tok):
    # vLLM docs markdown, excluding the last 25 (sorted) files
    md = sorted(glob.glob(f"{ROOT}/vllm/docs/**/*.md", recursive=True))
    kept = md[:-25] if len(md) > 25 else md
    for p in kept:
        t = cs._read(p, 60000)
        if len(t) > 200:
            yield t + "\n"
    for topic in GENERAL_QA:
        msgs = [{"role": "user", "content":
                 f"Give me a long, detailed explanation of {topic}, with "
                 "history, fundamentals, examples, and common misconceptions."}]
        yield tok.apply_chat_template(msgs, tokenize=False,
                                      add_generation_prompt=True)
    # HF Dolly general knowledge (prose contexts + responses)
    from datasets import load_dataset
    ds = load_dataset("databricks/databricks-dolly-15k", split="train",
                      streaming=True)
    for r in ds:
        parts = [r.get("instruction", ""), r.get("context", ""),
                 r.get("response", "")]
        yield "\n\n".join(p for p in parts if p) + "\n"


def drive(writer, tok, gen, cat, budget):
    """Pull texts from `gen` until the category token budget is met."""
    start = writer.by_cat.get(cat, 0)
    for text in gen:
        if not text:
            continue
        ids = tok.encode(text[:200000], add_special_tokens=False)
        writer.add(ids, cat)
        if writer.by_cat.get(cat, 0) - start >= budget:
            break
    got = writer.by_cat.get(cat, 0) - start
    print(f"[{cat}] {got} tokens (budget {budget})", flush=True)


def main():
    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(TOK_DIR, trust_remote_code=True)

    rng = random.Random(SEED)          # FIRST consumer -> matches seed-42 shuffle
    art = cs.gather_artifacts(rng)     # defines the code_files[:400] cutoff
    assert len(art["code_files"]) >= 400, "need >=400 code files for the cutoff"

    writer = ShardWriter(OUT, SHARD_TOKENS)

    # code: interleave local + HF sources round-robin to hit the budget
    def code_stream():
        loc = gen_code_local(art)
        hf = gen_code_hf()
        while True:
            emitted = False
            for g in (loc, hf, hf):    # weight HF ~2x (local is finite)
                try:
                    yield next(g)
                    emitted = True
                except StopIteration:
                    pass
            if not emitted:
                return

    drive(writer, tok, code_stream(), "code", BUDGET["code"])
    drive(writer, tok, gen_agentic(rng, art, tok), "agentic",
          BUDGET["agentic"])

    def instr_stream():
        al = gen_instruction(tok)
        cc = gen_coding_chat(rng, art, tok)
        while True:
            emitted = False
            for g in (al, al, cc):     # mostly alpaca, some coding chat
                try:
                    yield next(g)
                    emitted = True
                except StopIteration:
                    pass
            if not emitted:
                return

    drive(writer, tok, instr_stream(), "instruction", BUDGET["instruction"])
    drive(writer, tok, gen_medical(rng, tok), "medical", BUDGET["medical"])
    drive(writer, tok, gen_prose(tok), "prose", BUDGET["prose"])

    writer.finalize()

    # ---------------------------------------------------------- verification
    shards = sorted(glob.glob(os.path.join(OUT, "shard_*.npy")))
    assert shards, "no shards written"
    sizes = [np.load(s, mmap_mode="r").shape[0] for s in shards]
    assert all(sz > 0 for sz in sizes), "an empty shard was written"
    total = sum(sizes)
    print("\n==== calib-v3 summary ====")
    print(f"total tokens: {total}")
    print(f"shards: {len(shards)} (sizes {min(sizes)}..{max(sizes)})")
    print("composition (tokens / %):")
    for c, n in writer.by_cat.items():
        print(f"  {c:12s} {n:>10d}  {100*n/total:5.1f}%")

    # tokenizer round-trips a sample
    samp = np.load(shards[0])[1000:1300].tolist()
    txt = tok.decode(samp)
    re_ids = tok.encode(txt, add_special_tokens=False)
    print(f"\nround-trip sample decode ({len(txt)} chars): {txt[:160]!r}")
    print(f"round-trip re-encode matches: {re_ids == samp} "
          f"(len {len(re_ids)} vs {len(samp)})")
    # held-out guard
    assert not any(os.path.samefile(s, HELDOUT) for s in shards
                   if os.path.exists(HELDOUT) and os.path.exists(s))
    print("done.")


if __name__ == "__main__":
    main()