"""
Re-split the merged v6 synthetic corpus so v5 -> v6 stays a fair comparison.

build_synth_dataset.py shuffles with SEED=42 over whatever rows it is given. The v6 corpus has
20,139 more rows than the v5 one, so that shuffle lands differently and **1,955 of v5's 2,000
held-out synth rows fall into v6's train split**. Training on them and then reporting the synth
cell would be scoring memorisation.

So the eval split is not re-drawn, it is *pinned*: `data_synth_sft/eval.jsonl` (v5's rows, in v5's
order) is copied through verbatim, and every one of those instructions is removed from train. The
first 1,000 of them are the same rows v4 and v5 were scored on, so the cell stays comparable
across all three models.

A second held-out set, `eval_rel.jsonl`, is carved from the relational pool (task_id >= 1,000,000)
— the whole point of v6 is a capability v5 lacks, and none of the legacy eval rows test it.

Writes data_synth_v6_sft/{train,eval,eval_rel}.jsonl.
"""
import json
import random
from pathlib import Path

import pyarrow.parquet as pq

CORPUS   = Path("out_merged_v6/arabic_math_reasoning_synth.parquet")
LEGACY   = Path("data_synth_sft/eval.jsonl")      # v5's held-out synth rows — pinned, not redrawn
OUT      = Path("data_synth_v6_sft")
EVAL_REL = 400
REL_MIN_TASK_ID = 1_000_000
SEED = 42

FIELDS = ("instruction", "reasoning", "answer")


def sft(row, source="synth_math_ar"):
    return {**{k: row[k] for k in FIELDS}, "source": source}


def main():
    rows = pq.read_table(CORPUS).to_pylist()
    by_instruction = {r["instruction"]: r for r in rows}
    print(f"[*] corpus {len(rows):,} rows")

    legacy = [json.loads(l) for l in open(LEGACY, encoding="utf-8")]
    missing = [r for r in legacy if r["instruction"] not in by_instruction]
    print(f"[*] pinned eval {len(legacy):,} rows, {len(missing)} no longer in the corpus")

    held = {r["instruction"] for r in legacy}

    # relational held-out: deterministic sample of the new pool, also excluded from train
    rel = [r for r in rows if r["task_id"] >= REL_MIN_TASK_ID and r["instruction"] not in held]
    rel.sort(key=lambda r: (r["task_id"], r["instruction"]))       # parquet order is shuffled
    eval_rel = random.Random(SEED).sample(rel, min(EVAL_REL, len(rel)))
    held |= {r["instruction"] for r in eval_rel}
    print(f"[*] relational rows {len(rel):,}, holding out {len(eval_rel):,}")

    train = [r for r in rows if r["instruction"] not in held]
    n_rel_train = sum(1 for r in train if r["task_id"] >= REL_MIN_TASK_ID)

    OUT.mkdir(exist_ok=True)
    for name, split in (("train", [sft(r) for r in train]),
                        ("eval", legacy),                 # verbatim, v5's order
                        ("eval_rel", [sft(r) for r in eval_rel])):
        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'}")
    print(f"[*] train carries {n_rel_train:,} relational rows ({n_rel_train/len(train):.1%})")

    leak = sum(1 for r in train if r["instruction"] in held)
    print(f"[{'+' if leak == 0 else '!'}] contamination check: {leak} held-out rows in train")


if __name__ == "__main__":
    main()