#!/usr/bin/env python3
"""把 teacher 清洗后的干净数据转成四头训练格式，并分层切分。

输入: data/cleaned/archive_latest_full_clean.jsonl（build_cleaned_dataset.py 输出）
输出: data/processed_teacher/{train,dev,val,holdout827}.jsonl
字段: text / label_bin / label_kind / label_scope / label_sev
- normal: kind=-100, scope=1(message), sev=0
- ad: kind 0..6（none -> -100）; scope account=0/message=1; sev bucket 0..3
"""
import argparse
import json
from collections import Counter
from pathlib import Path

import numpy as np
from sklearn.model_selection import train_test_split

KINDS = ["crypto", "gambling", "porn", "porn_bait", "promo", "scam", "spam_flood"]


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

    rows = []
    for line in open(args.input, encoding="utf-8"):
        r = json.loads(line)
        is_ad = r.get("label") == "ad"
        kind = r.get("ad_kind")
        rows.append({
            "text": r.get("text"),
            "label_bin": 0 if is_ad else 1,
            "label_kind": KINDS.index(kind) if (is_ad and kind in KINDS) else -100,
            "label_scope": (0 if (is_ad and r.get("scope") == "account") else 1),
            "label_sev": (int(r.get("severity_bucket", 0)) if is_ad else 0),
        })
    labels = np.array([r["label_bin"] for r in rows])
    idx = np.arange(len(rows))
    tr, rest = train_test_split(idx, test_size=0.10, stratify=labels, random_state=args.seed)
    dev, rest2 = train_test_split(rest, test_size=0.50, stratify=labels[rest], random_state=args.seed)
    val, hold = train_test_split(rest2, test_size=0.50, stratify=labels[rest2], random_state=args.seed)

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    for name, ids in [("train", tr), ("dev", dev), ("val", val), ("holdout827", hold)]:
        data = [rows[i] for i in ids]
        with open(out / f"{name}.jsonl", "w", encoding="utf-8") as f:
            for x in data:
                f.write(json.dumps(x, ensure_ascii=False) + "\n")
        print(name, len(data),
              "bin", dict(Counter(x["label_bin"] for x in data)),
              "kind", dict(Counter(x["label_kind"] for x in data)),
              "scope", dict(Counter(x["label_scope"] for x in data)),
              "sev", dict(Counter(x["label_sev"] for x in data)))
    print("out", out.resolve())


if __name__ == "__main__":
    main()