#!/usr/bin/env python3
"""把 teacher 清洗输出转换成训练用 jsonl（teacher 标签为准）。

输入: clean_with_teacher.py 的 */latest_full.jsonl（含 teacher + v1）
输出: text / label(ad|normal) / ad_kind / scope / severity / severity_bucket / 教师置信度 / v1 对照
"""
import argparse
import json
from collections import Counter
from pathlib import Path


def bucket(score):
    if score is None:
        return -1
    if score < 0.5:
        return 0
    if score < 1.5:
        return 1
    if score < 2.5:
        return 2
    return 3


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--input", default="data/cleaned/archive_latest_full.jsonl")
    ap.add_argument("--out", default="data/cleaned/archive_latest_full_clean.jsonl")
    args = ap.parse_args()

    n = 0
    counts = Counter()
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    with open(args.input, encoding="utf-8") as f, open(out, "w", encoding="utf-8") as w:
        for line in f:
            if not line.strip():
                continue
            r = json.loads(line)
            t = r.get("teacher") or {}
            is_ad = (t.get("is_ad") or {}).get("choice") == "ad"
            kind = (t.get("ad_kind") or {}).get("choice") or "none"
            scope = (t.get("ad_scope") or {}).get("choice") or "message"
            sev = (t.get("severity") or {}).get("score")
            row = {
                "text": r.get("text"),
                "label": "ad" if is_ad else "normal",
                "ad_kind": kind if is_ad else "normal",
                "scope": scope if is_ad else "message",
                "severity": sev if is_ad else 0.0,
                "severity_bucket": bucket(sev) if is_ad else 0,
                "confidence": (t.get("is_ad") or {}).get("confidence"),
                "teacher": "jev-1.13",
                "v1_label": "ad" if r.get("v1_label") == "ad" else "normal",
                "v1_kind": r.get("v1_kind"),
                "v1_decider": r.get("v1_decider"),
                "v1_scope": r.get("v1_scope"),
                "v1_severity": r.get("v1_severity"),
                "chat_id": r.get("chat_id"),
                "message_id": r.get("message_id"),
            }
            w.write(json.dumps(row, ensure_ascii=False) + "\n")
            n += 1
            counts[(row["label"], row["ad_kind"])] += 1
    print("wrote", n, "->", out.resolve())
    print("label/kind:", dict(counts.most_common(20)))


if __name__ == "__main__":
    main()