# Qubit 7B: data inspection, filtering, SFT conversion and tokenizer training.
# Runs on HF Jobs (CPU). Writes everything to OUT (bucket mount).
import os, sys, re, json, hashlib, random, subprocess, tempfile, time, glob, math
from collections import Counter

OUT = os.environ.get("OUT", "/ckpt/data")
QUBIT = os.environ.get("QUBIT_REPO", "/qubit")
N_TOK_TRAIN_WEB = int(os.environ.get("N_TOK_TRAIN_WEB", "600000"))
os.makedirs(OUT, exist_ok=True)
random.seed(1234)
stats = {"started": time.strftime("%Y-%m-%d %H:%M:%S")}

JA_RE = re.compile(r"[\u3040-\u309f\u30a0-\u30ff\u4e00-\u9fff\u3400-\u4dbf\uff66-\uff9f]")


def ja_ratio(s):
    s2 = re.sub(r"\s", "", s)
    if not s2:
        return 0.0
    return len(JA_RE.findall(s2)) / len(s2)


def h(s):
    return hashlib.sha1(s.encode("utf-8")).hexdigest()


def is_eval(key, pct=1.0):
    # deterministic held-out split by hash (no leakage)
    return int(h("eval:" + key)[:8], 16) % 10000 < pct * 100


def log(*a):
    print(*a, flush=True)


from datasets import load_dataset

# ------------------------------------------------------------------ existing tokenizers
import sentencepiece as spm
tok_report = {}
for p in sorted(glob.glob(os.path.join(QUBIT, "*.model"))):
    try:
        sp = spm.SentencePieceProcessor(model_file=p)
        sample = "日本の首都は東京です。Pythonでクイックソートを書いてください。"
        ids = sp.encode(sample)
        tok_report[os.path.basename(p)] = {
            "vocab": sp.get_piece_size(),
            "sample_tokens": len(ids),
            "unk_in_sample": sum(1 for i in ids if i == sp.unk_id()),
        }
    except Exception as e:
        tok_report[os.path.basename(p)] = {"error": str(e)[:200]}
stats["existing_tokenizers"] = tok_report
log("existing tokenizers:", json.dumps(tok_report, ensure_ascii=False))

# ------------------------------------------------------------------ helpers for SFT
def write_jsonl(path, rows):
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def split_and_dedupe(name, rows):
    seen, train, ev = set(), [], []
    dup = 0
    for r in rows:
        key = r["messages"][0]["content"] if r["messages"][0]["role"] != "system" else r["messages"][1]["content"]
        k = h(re.sub(r"\s+", "", key))
        if k in seen:
            dup += 1
            continue
        seen.add(k)
        (ev if is_eval(k) else train).append(r)
    random.shuffle(train)
    write_jsonl(os.path.join(OUT, f"sft_{name}_train.jsonl"), train)
    write_jsonl(os.path.join(OUT, f"sft_{name}_eval.jsonl"), ev)
    log(f"[{name}] train={len(train)} eval={len(ev)} dup_removed={dup}")
    return {"train": len(train), "eval": len(ev), "dup_removed": dup}


# ------------------------------------------------------------------ stage 2: conversation
conv_rows = []
c = Counter()
ds = load_dataset("shisa-ai/shisa-v2-sharegpt", split="train")
stats.setdefault("raw", {})["shisa-v2-sharegpt"] = len(ds)
role_map = {"human": "user", "user": "user", "gpt": "assistant", "assistant": "assistant", "system": "system"}
for ex in ds:
    msgs = []
    ok = True
    for t in ex["conversations"]:
        r = role_map.get(t["from"])
        if r is None or not (t["value"] or "").strip():
            ok = False
            break
        msgs.append({"role": r, "content": t["value"].strip()})
    if not ok or not any(m["role"] == "assistant" for m in msgs):
        c["shisa_bad_format"] += 1
        continue
    alltext = "".join(m["content"] for m in msgs if m["role"] != "system")
    if ja_ratio(alltext) < 0.25:
        c["shisa_non_ja"] += 1
        continue
    if len(alltext) > 16000:
        c["shisa_too_long"] += 1
        continue
    conv_rows.append({"source": "shisa", "messages": msgs})
log("shisa kept", len(conv_rows), dict(c))

ds = load_dataset("kunishou/databricks-dolly-15k-ja", split="train")
stats["raw"]["databricks-dolly-15k-ja"] = len(ds)
nd = 0
for ex in ds:
    ins, inp, out = (ex["instruction"] or "").strip(), (ex["input"] or "").strip(), (ex["output"] or "").strip()
    if not ins or not out or ja_ratio(ins + out) < 0.25:
        c["dolly_dropped"] += 1
        continue
    user = ins if not inp else f"{ins}\n\n{inp}"
    conv_rows.append({"source": "dolly", "messages": [{"role": "user", "content": user}, {"role": "assistant", "content": out}]})
    nd += 1
log("dolly kept", nd)
stats["filter_counts_conv"] = dict(c)
stats["sft_conv"] = split_and_dedupe("conv", conv_rows)

# ------------------------------------------------------------------ stage 3: code / logic
CODE_RE = re.compile(r"```|\bdef |\bimport |\bprint\(|\bclass |Python|python|コード|プログラム|関数を(書|作)|アルゴリズム")
MATH_RE = re.compile(r"\\\(|\\\[|方程式|計算|解いて|確率|整数|平方|微分|積分|\d+\s*[\+\-×÷\*/=]\s*\d+")
BAD_RE = re.compile(r"間違いがある可能性|再評価してみましょう|矛盾が(あり|生じ)|誤りがありました|申し訳ありませんが、")
code_rows, logic_rows = [], []
c = Counter()
ds = load_dataset("llm-jp/wizardlm8x22b-logical-math-coding-sft-ja", split="train")
stats["raw"]["wizardlm8x22b-logical-math-coding-sft-ja"] = len(ds)
for ex in ds:
    msgs = [{"role": m["role"], "content": (m["content"] or "").strip()} for m in ex["messages"]]
    if not msgs or msgs[0]["role"] not in ("user", "system") or not any(m["role"] == "assistant" for m in msgs):
        c["bad_format"] += 1
        continue
    if any(not m["content"] for m in msgs):
        c["empty_turn"] += 1
        continue
    asst = "\n".join(m["content"] for m in msgs if m["role"] == "assistant")
    alltext = "\n".join(m["content"] for m in msgs)
    is_code = bool(CODE_RE.search(alltext))
    if len(asst) < 30 or len(alltext) > 16000:
        c["length"] += 1
        continue
    if ja_ratio(alltext) < (0.10 if is_code else 0.25):
        c["non_ja"] += 1
        continue
    if BAD_RE.search(asst):
        c["self_contradiction"] += 1
        continue
    if is_code:
        code_rows.append({"source": "wizardlm_code", "messages": msgs})
    else:
        cat = "wizardlm_math" if MATH_RE.search(alltext) else "wizardlm_logic"
        c[cat] += 1
        logic_rows.append({"source": cat, "messages": msgs})
log("wizardlm code", len(code_rows), "logic/math", len(logic_rows), dict(c))
stats["filter_counts_wizardlm"] = dict(c)
stats["sft_code"] = split_and_dedupe("code", code_rows + logic_rows)
stats["sft_code"]["code_samples"] = len(code_rows)
stats["sft_code"]["logic_math_samples"] = len(logic_rows)

# ------------------------------------------------------------------ stage 4: math with verification
BOX_RE = re.compile(r"\\boxed\{([^{}]*)\}")
CODE_BLOCK_RE = re.compile(r"<llm-code>(.*?)</llm-code>", re.S)
NUM_RE = re.compile(r"-?\d+(?:,\d{3})*(?:\.\d+)?(?:[eE][-+]?\d+)?")


def to_num(s):
    if s is None:
        return None
    m = NUM_RE.findall(str(s).replace("，", ","))
    if not m:
        return None
    try:
        return float(m[-1].replace(",", ""))
    except Exception:
        return None


def close(a, b):
    return a is not None and b is not None and abs(a - b) <= max(1e-2, 1e-4 * abs(b))


def run_code(code):
    with tempfile.NamedTemporaryFile("w", suffix=".py", delete=False) as f:
        f.write(code)
        p = f.name
    try:
        r = subprocess.run([sys.executable, "-I", p], capture_output=True, text=True, timeout=5,
                           cwd=tempfile.gettempdir(), env={"PATH": "/usr/bin:/bin"})
        if r.returncode != 0:
            return None
        lines = [l for l in r.stdout.strip().splitlines() if l.strip()]
        return to_num(lines[-1]) if lines else None
    except Exception:
        return None
    finally:
        os.unlink(p)


math_rows = []
c = Counter()
ds = load_dataset("HachiML/alpaca_jp_math", split="v1.0_cleaned")
stats["raw"]["alpaca_jp_math(v1.0_cleaned)"] = len(ds)
for ex in ds:
    ins, inp, out = (ex["instruction"] or "").strip(), (ex["input"] or "").strip(), (ex["output"] or "").strip()
    if not ins or not out:
        c["empty"] += 1
        continue
    if str(ex.get("clean")).lower() != "true":
        c["not_clean"] += 1
        continue
    boxed = BOX_RE.findall(out)
    ans = to_num(boxed[-1]) if boxed else None
    blocks = CODE_BLOCK_RE.findall(out)
    verified = False
    if blocks:
        res = run_code("\n".join(blocks))
        if res is not None and ans is not None:
            verified = close(res, ans)
            c["exec_match" if verified else "exec_mismatch"] += 1
        else:
            c["exec_failed_or_no_answer"] += 1
    if not verified and close(ex.get("code_result"), ex.get("text_result")) and ans is not None and close(ans, ex.get("text_result")):
        verified = True
        c["meta_match"] += 1
    if not verified:
        c["dropped_unverified"] += 1
        continue
    # make code blocks readable markdown
    out2 = out.replace("<llm-code>", "```python\n").replace("</llm-code>", "\n```").replace("<llm-code-output>", "実行結果:\n```\n").replace("</llm-code-output>", "\n```")
    user = ins if not inp else f"{ins}\n\n{inp}"
    math_rows.append({"source": "alpaca_jp_math", "messages": [{"role": "user", "content": user}, {"role": "assistant", "content": out2}]})
log("math kept", len(math_rows), dict(c))
stats["filter_counts_math"] = dict(c)
stats["sft_math"] = split_and_dedupe("math", math_rows)

# ------------------------------------------------------------------ fineweb sample for tokenizer + quality stats
def web_ok(ex, seen):
    t = ex["text"]
    if len(t) < 200:
        return "short"
    if ja_ratio(t) < 0.5:
        return "low_ja"
    lines = [l for l in t.splitlines() if l.strip()]
    if lines and len(set(lines)) / len(lines) < 0.6:
        return "repetitive"
    k = h(t[:2000])
    if k in seen:
        return "dup"
    seen.add(k)
    return None

c = Counter()
seen = set()
tok_txt = os.path.join(tempfile.gettempdir(), "spm_train.txt")
n_web = 0
web_chars = 0
ds = load_dataset("hotchpotch/fineweb-2-edu-japanese", "sample_10BT", split="train", streaming=True)
with open(tok_txt, "w", encoding="utf-8") as f:
    for ex in ds:
        why = web_ok(ex, seen)
        if why:
            c[why] += 1
            continue
        c["kept"] += 1
        t = ex["text"][:4000].replace("\n", " ")
        f.write(t + "\n")
        web_chars += len(t)
        n_web += 1
        if n_web >= N_TOK_TRAIN_WEB:
            break
    for name in ("conv", "code", "math"):
        for line in open(os.path.join(OUT, f"sft_{name}_train.jsonl"), encoding="utf-8"):
            r = json.loads(line)
            for m in r["messages"]:
                f.write(m["content"][:4000].replace("\n", " ") + "\n")
stats["fineweb_tokenizer_sample"] = {"kept": n_web, "chars": web_chars, "filter": dict(c)}
log("fineweb sample", stats["fineweb_tokenizer_sample"])

# fineweb validation set from the dedicated *test* split (no overlap with train split)
# streaming=True: a non-streaming load would prepare the whole 326 GB train split too
dv = load_dataset("hotchpotch/fineweb-2-edu-japanese", "default", split="test", streaming=True)
val = []
for ex in dv:
    if ja_ratio(ex["text"]) >= 0.5 and len(ex["text"]) >= 200:
        val.append(ex["text"])
        if len(val) >= 2000:
            break
log("pretrain val docs", len(val))
write_jsonl(os.path.join(OUT, "pretrain_val.jsonl"), [{"text": t} for t in val])
stats["pretrain_val_docs"] = len(val)

# ------------------------------------------------------------------ train new 32k tokenizer
SPECIAL = ["<|system|>", "<|user|>", "<|assistant|>", "<|end|>"]
prefix = os.path.join(OUT, "qubit7b_ja_spm32k")
spm.SentencePieceTrainer.train(
    input=tok_txt, model_prefix=prefix, vocab_size=32000, model_type="unigram",
    character_coverage=0.9995, byte_fallback=True, split_digits=True,
    pad_id=0, unk_id=1, bos_id=2, eos_id=3, user_defined_symbols=SPECIAL,
    input_sentence_size=3_000_000, shuffle_input_sentence=True, num_threads=os.cpu_count(),
    max_sentence_length=16384, normalization_rule_name="identity", remove_extra_whitespaces=False,
    allow_whitespace_only_pieces=True, train_extremely_large_corpus=False,
)
sp = spm.SentencePieceProcessor(model_file=prefix + ".model")
assert sp.get_piece_size() == 32000
samples = ["日本の首都はどこですか？", "こんにちは。あなたは何ができますか？", "125 × 48 を計算してください。",
           "Pythonでクイックソートを書いてください。", "なぜ飛行機は空を飛べるのですか？",
           "def quicksort(arr):\n    return arr"]
stats["new_tokenizer"] = {
    "vocab": sp.get_piece_size(),
    "special_ids": {s: sp.piece_to_id(s) for s in SPECIAL + ["<pad>", "<unk>", "<s>", "</s>"]},
    "samples": {s: sp.encode(s, out_type=str) for s in samples},
    "roundtrip_ok": all(sp.decode(sp.encode(s)) == s for s in samples),
}
# chars per token on val
tot_c = tot_t = 0
for t in val[:500]:
    tot_c += len(t)
    tot_t += len(sp.encode(t))
stats["new_tokenizer"]["chars_per_token_val"] = round(tot_c / tot_t, 3)
for name, rep in tok_report.items():
    if "error" in rep:
        continue
    o = spm.SentencePieceProcessor(model_file=os.path.join(QUBIT, name))
    tc = sum(len(o.encode(t)) for t in val[:500])
    rep["chars_per_token_val"] = round(tot_c / tc, 3)
# SFT token counts
for name in ("conv", "code", "math"):
    n = 0
    for line in open(os.path.join(OUT, f"sft_{name}_train.jsonl"), encoding="utf-8"):
        r = json.loads(line)
        n += sum(len(sp.encode(m["content"])) + 3 for m in r["messages"])
    stats[f"sft_{name}"]["train_tokens"] = n
stats["finished"] = time.strftime("%Y-%m-%d %H:%M:%S")
json.dump(stats, open(os.path.join(OUT, "prep_stats.json"), "w"), ensure_ascii=False, indent=2)
log(json.dumps(stats, ensure_ascii=False, indent=1))
log("PREP_DONE")