"""Training mix for the coding fix (run rC): Beam's test-verified TACO solutions plus a replay sample of the effort
data that Dory-Beam-4B v2 was trained on, so the effort levels and the math are not forgotten.

  .venv/bin/python -I scripts/build_code_sft.py --beam 'data/code/beam/run-code-*.jsonl' \
      --verdicts data/code/verdicts.jsonl [--max-code 8000] [--replay-ratio 1.0]
  # -> data/code/sft_code_mix.jsonl, data/code/sft_code_mix_stats.json

Coding examples are Beam at medium effort, tagged `{reasoning effort: medium}` half the time and untagged otherwise,
as in scripts/build_effort_dataset.py. Replay examples are drawn uniformly from data/final_effort/sft_effort.jsonl.
"""

import argparse
import glob
import json
import random
from collections import Counter
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--beam", required=True)
    ap.add_argument("--verdicts", required=True)
    ap.add_argument("--max-code", type=int, default=0)
    ap.add_argument("--replay-ratio", type=float, default=1.0)
    ap.add_argument("--max-chars", type=int, default=40000, help="drop very long traces (~12k tokens)")
    ap.add_argument("--prefer-fn", action="store_true",
                    help="fill --max-code with function-style problems first (the HumanEval/MBPP format)")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    rng = random.Random(args.seed)

    passed = {v["key"] for v in map(json.loads, open(args.verdicts)) if v["passed"]}
    code = []
    for path in sorted(glob.glob(args.beam)):
        for line in open(path):
            r = json.loads(line)
            key = f"{r.get('id')}#{r.get('sample_idx', 0)}"
            if key not in passed or not r["reasoning"].strip():
                continue
            if len(r["reasoning"]) + len(r["content"]) > args.max_chars:
                continue
            code.append(r)
    rng.shuffle(code)
    if args.prefer_fn:
        fn = {t["id"] for t in map(json.loads, open(ROOT / "data" / "code" / "tests.jsonl")) if t["fn_name"]}
        code.sort(key=lambda r: r["id"] not in fn)  # stable: function-style first, each part still shuffled
    if args.max_code:
        code = code[:args.max_code]

    rows = []
    for r in code:
        tagged = rng.random() < 0.5
        user = r["prompt"] + ("\n\n{reasoning effort: medium}" if tagged else "")
        rows.append({"id": f"{r['id']}#code", "prompt_id": r["id"], "effort": "medium", "beam_effort": "medium",
                     "tagged": tagged, "domain": "code", "teacher": r["teacher"],
                     "messages": [{"role": "user", "content": user},
                                  {"role": "assistant", "reasoning_content": r["reasoning"], "content": r["content"]}],
                     "gold": "", "chars": len(r["reasoning"]) + len(r["content"])})
    effort = [json.loads(l) for l in open(ROOT / "data" / "final_effort" / "sft_effort.jsonl")]
    replay = rng.sample(effort, min(len(effort), round(len(rows) * args.replay_ratio)))
    mix = rows + replay
    rng.shuffle(mix)

    out = ROOT / "data" / "code" / "sft_code_mix.jsonl"
    with open(out, "w") as f:
        for r in mix:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    stats = {"code_examples": len(rows), "replay_examples": len(replay),
             "code_chars_mean": round(sum(r["chars"] for r in rows) / max(1, len(rows))),
             "replay_chars_mean": round(sum(r.get("chars", 0) for r in replay) / max(1, len(replay))),
             "replay_mix": Counter(f"{r['effort']}{'' if r['tagged'] else '-untagged'}" for r in replay),
             "code_sources": Counter(r["source"] for r in code),
             "code_function_style": sum("starter code" in r["prompt"] for r in code)}
    stats["approx_tokens_M"] = round((stats["code_chars_mean"] * len(rows) + stats["replay_chars_mean"] * len(replay))
                                     / 3.3 / 1e6, 1)
    (ROOT / "data" / "code" / "sft_code_mix_stats.json").write_text(json.dumps(stats, indent=1))
    print(json.dumps(stats, indent=1))


if __name__ == "__main__":
    main()