"""
Build the v6 SFT corpus: the v5 three-way mix, with the synthetic corpus swapped for the larger
one that carries the relational pool.

This is `prepare_v5_sft.py` with `SYNTH_DIR` repointed and a fourth eval column added. Separate
file, same reason v5 was separate from v3: v5 is the shipped model and must stay reproducible.

The synth split comes from `split_synth_v6.py`, not `build_synth_dataset.py` — the v6 corpus is
20,139 rows larger, so a re-drawn shuffle would put 1,955 of v5's held-out rows into v6's train.
The eval rows are pinned to v5's instead, which keeps the synth cell comparable and, because
`eval.jsonl` is byte-identical to v5's, makes the mix-eval synth column the *same 400 rows* v5 was
scored on.

`eval_rel.jsonl` (400 relational rows, held out of train) is the column that actually measures
what v6 was built for — none of the legacy eval rows test a relational comparison.

    data/{train,eval}.jsonl            <- prepare_data.py       (Arabic_Reasoning_Dataset)
    data_gsm_sft/{train,eval}.jsonl    <- prepare_gsm_sft.py    (oddadmix/gsm8k-reasoning-ar)
    data_synth_v6_sft/{train,eval,eval_rel}.jsonl <- split_synth_v6.py (120,462-row corpus)

Why all three: v3 scores 77.3% on GSM8K-ar but 2.0% on the synth held-out set, and v4 (synth only)
inverts that — 35.6% synth, 19.5% GSM. Neither corpus alone covers the other's distribution, so v5
trains on all three at once.

Answer styles are NOT normalised — GSM8K ends in a bare numeral, Arabic_Reasoning and the synth
corpus in an "إذن، …" sentence. A `source` tag rides on every row so eval_reasoning.py scores each
on its own terms. The eval set keeps v1's 400 AR rows and v2's first 600 GSM rows *unchanged*, so
the AR and GSM cells stay directly comparable to v1/v2/v3.

Writes data_v6_sft/{train,eval}.jsonl.
"""
import json
import os
import random
from collections import Counter
from pathlib import Path

AR_DIR    = Path("data")             # v1 splits
GSM_DIR   = Path("data_gsm_sft")     # v2 splits
SYNTH_DIR = Path("data_synth_v6_sft")  # the 120,462-row corpus, relational pool included
OUT       = Path("data_v6_sft")

# Arabic_Reasoning is ~25x smaller than GSM8K, so it is repeated (same value v3 used).
REPEAT       = int(os.environ.get("REPEAT", 3))
SYNTH_REPEAT = int(os.environ.get("SYNTH_REPEAT", 1))
# Eval: v1's 400 AR rows are all kept and v2's first 600 GSM rows, exactly as v3 built them, plus
# synth rows so all three distributions are scored in one pass.
EVAL_GSM   = int(os.environ.get("EVAL_GSM", 600))
EVAL_SYNTH = int(os.environ.get("EVAL_SYNTH", 400))
EVAL_REL   = int(os.environ.get("EVAL_REL", 400))
SEED       = 42


def load(path):
    with open(path, encoding="utf-8") as fh:
        return [json.loads(line) for line in fh]


def tag(rows, source):
    return [{**r, "source": source} for r in rows]


def token_stats(rows, tok):
    per_source, lengths = Counter(), []
    for r in rows:
        text = (f"<|im_start|>user\n{r['instruction']}<|im_end|>\n<|im_start|>assistant\n"
                f"<think>\n{r['reasoning']}\n</think>\n{r['answer']}<|im_end|>")
        n = len(tok.encode(text, add_special_tokens=False))
        per_source[r["source"]] += n
        lengths.append(n)
    lengths.sort()
    pct = lambda p: lengths[min(int(len(lengths) * p / 100), len(lengths) - 1)]
    return per_source, {"p50": pct(50), "p95": pct(95), "p99": pct(99), "max": lengths[-1]}


def interleave(groups):
    """Round-robin proportional to each group's size, so a truncated eval run still covers all
    three sources instead of whichever landed first."""
    groups = [g for g in groups if g]
    if not groups:
        return []
    total = sum(len(g) for g in groups)
    out, idx = [], [0] * len(groups)
    for _ in range(total):
        # pick the group that is furthest behind its target share
        pick = min(range(len(groups)), key=lambda i: (idx[i] / len(groups[i])) if idx[i] < len(groups[i]) else 2.0)
        if idx[pick] >= len(groups[pick]):
            break
        out.append(groups[pick][idx[pick]])
        idx[pick] += 1
    for g, i in zip(groups, idx):          # anything the loop could not place
        out.extend(g[i:])
    return out


def main():
    for d in (AR_DIR, GSM_DIR, SYNTH_DIR):
        for split in ("train", "eval"):
            if not (d / f"{split}.jsonl").exists():
                raise SystemExit(f"missing {d/f'{split}.jsonl'}")

    ar_train    = tag(load(AR_DIR / "train.jsonl"), "arabic_reasoning")
    ar_eval     = tag(load(AR_DIR / "eval.jsonl"), "arabic_reasoning")
    gsm_train   = tag(load(GSM_DIR / "train.jsonl"), "gsm8k_ar")
    gsm_eval    = tag(load(GSM_DIR / "eval.jsonl"), "gsm8k_ar")
    synth_train = tag(load(SYNTH_DIR / "train.jsonl"), "synth_math_ar")
    synth_eval  = tag(load(SYNTH_DIR / "eval.jsonl"), "synth_math_ar")
    rel_eval    = tag(load(SYNTH_DIR / "eval_rel.jsonl"), "synth_relational_ar")

    train = gsm_train + ar_train * REPEAT + synth_train * SYNTH_REPEAT
    random.Random(SEED).shuffle(train)

    eval_rows = interleave([ar_eval, gsm_eval[:EVAL_GSM], synth_eval[:EVAL_SYNTH],
                            rel_eval[:EVAL_REL]])

    OUT.mkdir(exist_ok=True)
    for name, split in (("train", train), ("eval", eval_rows)):
        with open(OUT / f"{name}.jsonl", "w", encoding="utf-8") as fh:
            for r in split:
                fh.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"[+] {name}: {len(split):,} -> {OUT / f'{name}.jsonl'}  "
              f"{dict(Counter(r['source'] for r in split))}")

    try:
        from transformers import AutoTokenizer
        tok = AutoTokenizer.from_pretrained(os.environ.get("BASE_MODEL", "/notebooks/50M/50M-2048-Emhotob"))
        tok.add_special_tokens({"additional_special_tokens": ["<|im_start|>", "<|im_end|>", "<think>", "</think>"]})
        per_source, pct = token_stats(train, tok)
        total = sum(per_source.values())
        print(f"[*] train tokens: {total/1e6:.1f}M/epoch  (REPEAT={REPEAT}, SYNTH_REPEAT={SYNTH_REPEAT})")
        for src, n in per_source.most_common():
            print(f"      {src:<18} {n/1e6:6.2f}M  {100*n/total:5.1f}%")
        print(f"[*] sample length: p50 {pct['p50']}  p95 {pct['p95']}  p99 {pct['p99']}  max {pct['max']}")
    except Exception as exc:
        print(f"[!] token stats skipped: {exc}")


if __name__ == "__main__":
    main()