#!/usr/bin/env python3
"""挑蒸馏用的双向 SFT 检查点：val_log.csv 里 cond 的 mse（各 t 取平均）最低的已存 step；与最优差 < TOL 的里取最晚的
（val 噪声内多训的更好）。只看 step ≥ --min_step（v8 的 SFT 教师是 step-6000）。
usage: pick_sft.py <sft_out> [--min_step 6000] [--tol 0.005]   → stdout 打印 ckpt 绝对路径，<sft_out>/sft_pick.json 记理由"""
import argparse
import csv
import collections
import json
import os
import sys


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("out")
    ap.add_argument("--min_step", type=int, default=6000)
    ap.add_argument("--tol", type=float, default=0.005)
    a = ap.parse_args()
    rows = list(csv.DictReader(open(os.path.join(a.out, "val_log.csv"))))
    by = collections.defaultdict(list)
    for r in rows:
        by[int(r["step"])].append(float(r["mse"]))
    cand = {s: sum(v) / len(v) for s, v in by.items()
            if s >= a.min_step and os.path.exists(os.path.join(a.out, f"step-{s}.safetensors"))}
    if not cand:
        sys.exit(f"pick_sft: {a.out} 里没有 step ≥ {a.min_step} 且有 val 记录的检查点")
    best = min(cand.values())
    pick = max(s for s, v in cand.items() if v <= best * (1 + a.tol))
    path = os.path.abspath(os.path.join(a.out, f"step-{pick}.safetensors"))
    json.dump({"pick": pick, "path": path, "val_mse_mean": cand[pick], "best": best, "tol": a.tol,
               "candidates": dict(sorted(cand.items()))}, open(os.path.join(a.out, "sft_pick.json"), "w"), indent=1)
    print(path)


if __name__ == "__main__":
    main()