"""Builds the SFT set from teach_gen.py output: what the module system knows how to do, as plain model outputs.

Per problem (one example each, so easy problems do not dominate):
  * "short": the shortest correct natural solution among the N samples (efficient reasoning that still ends right);
  * "stop":  if no natural sample is shorter, a checked stop example (thinking cut where the answer was already
             known and stable + </think> + the model's own closing answer) — teaches ending the thinking on time.
usage: build_sft.py --inp results/teach/pilot.jsonl --items data/teach/*.jsonl --out data/teach/sft_pilot.jsonl
"""
import argparse
import collections
import glob
import json
import zlib

from transformers import AutoTokenizer

EOS = 124900
ap = argparse.ArgumentParser()
ap.add_argument("--inp", nargs="+", required=True)
ap.add_argument("--items", default="data/teach/*.jsonl")
ap.add_argument("--out", required=True)
ap.add_argument("--tokenizer", default="models/LFM3-1.7B")
ap.add_argument("--max-len", type=int, default=7000)
ap.add_argument("--mode", default="short", choices=["short", "correctness"],
                help="correctness: only problems the model solves sometimes (one random correct sample each) and "
                     "unsolved problems rescued by a checked stop; no length-only examples")
args = ap.parse_args()

tok = AutoTokenizer.from_pretrained(args.tokenizer)
items = {}
for p in glob.glob(args.items):
    if "sft_" in p:
        continue
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        items[r["id"]] = r
by = collections.defaultdict(list)
for p in args.inp:
    for l in open(p, encoding="utf-8"):
        r = json.loads(l)
        by[r["id"]].append(r)

out, stats = [], collections.Counter()
for pid, recs in by.items():
    it = items[pid]
    good = [r for r in recs if r["correct"]]
    stops = [r["stop"] for r in recs if r.get("stop", {}).get("ok")]
    best_short = min(good, key=lambda r: r["n_tokens"]) if good else None
    best_stop = min(stops, key=lambda s: len(s["tokens"])) if stops else None
    p_ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
                                    add_generation_prompt=True, return_dict=False)
    if args.mode == "correctness":
        rescue = [r["stop"] for r in recs if r.get("stop", {}).get("ok") and not r["stop"].get("from_correct")]
        if good and len(good) < len(recs):
            pick = good[zlib.crc32(pid.encode()) % len(good)]
            comp, kind = [t for t in pick["tokens"] if t != EOS] + [EOS], "partial"
        elif not good and rescue:
            comp, kind = [t for t in rescue[0]["tokens"] if t != EOS] + [EOS], "rescue"
        else:
            stats["skipped"] += 1
            continue
        if len(p_ids) + len(comp) > args.max_len:
            stats["too_long"] += 1
            continue
        out.append({"id": pid, "task": it["task"], "kind": kind, "prompt_ids": p_ids, "completion_ids": comp,
                    "solved": f"{len(good)}/{len(recs)}"})
        stats[f"{it['task']}_{kind}"] += 1
        continue
    if best_stop and (not best_short or len(best_stop["tokens"]) < best_short["n_tokens"]):
        comp, kind = [t for t in best_stop["tokens"] if t != EOS] + [EOS], "stop"
    elif best_short:
        comp, kind = [t for t in best_short["tokens"] if t != EOS] + [EOS], "short"
    else:
        stats["no_example"] += 1
        continue
    if len(p_ids) + len(comp) > args.max_len:
        stats["too_long"] += 1
        continue
    out.append({"id": pid, "task": it["task"], "kind": kind, "prompt_ids": p_ids, "completion_ids": comp,
                "solved": f"{len(good)}/{len(recs)}"})
    stats[f"{it['task']}_{kind}"] += 1
    stats[f"len_{kind}"] += len(comp)
    stats[f"natural_mean_{kind}"] += sum(r["n_tokens"] for r in recs) / len(recs)

with open(args.out, "w", encoding="utf-8") as f:
    for r in out:
        f.write(json.dumps(r) + "\n")
for k in ("short", "stop"):
    n = sum(v for kk, v in stats.items() if kk.endswith("_" + k) and not kk.startswith(("len", "natural")))
    if n:
        print(f"{k}: {n} examples, mean length {stats['len_' + k] / n:.0f} vs mean sample length "
              f"{stats['natural_mean_' + k] / n:.0f}")
print(dict((k, v) for k, v in stats.items() if not k.startswith(("len", "natural"))))