#!/usr/bin/env python3
"""构建双头联合训练数据：每行同时带二分类标签和 7 类细类标签。

- label_bin: 0=ad, 1=normal
- label_kind: 0..6 细类；非广告或广告但无细类标签 = -100
- 训练池 = processed_v2/train；3% 分层切 dev（按 normal / 细类 / ad_unknown 分组）
- 测试 = processed_v2/val；source=archive 的子集即 827 干净 holdout
"""
import argparse
import collections
import json
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"]
MAP = {"contact": "porn"}
DROP = {"spam", "impersonate", "unspecified", "normal", "none", None, ""}


def load(path):
    return [json.loads(l) for l in open(path, encoding="utf-8") if l.strip()]


def clean(lab):
    if lab in MAP:
        lab = MAP[lab]
    return lab if lab in KINDS else None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--gold", default="data/raw/dataset/converted/gold.kind.jsonl")
    ap.add_argument("--archive", default="data/raw/archive/converted/archive.strong.binary.jsonl")
    ap.add_argument("--split_train", default="data/processed_v2/train.jsonl")
    ap.add_argument("--split_val", default="data/processed_v2/val.jsonl")
    ap.add_argument("--out", default="data/processed_dual")
    ap.add_argument("--dev_ratio", type=float, default=0.03)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    gold, arch = {}, {}
    gold_all, arch_all = set(), set()
    for r in load(args.gold):
        gold_all.add(r["text"].strip())
        lab = clean(r["label"])
        if lab:
            gold.setdefault(r["text"].strip(), lab)
    for r in load(args.archive):
        arch_all.add(r["text"].strip())
        if r.get("label") != "ad":
            continue
        lab = clean(r.get("ad_kind"))
        if lab:
            arch.setdefault(r["text"].strip(), lab)

    def kind_of(text):
        t = text.strip()
        return gold.get(t) or arch.get(t)

    def pack(r, source=None):
        text = r["text"]
        is_ad = r["label"] == "ad"
        k = kind_of(text) if is_ad else None
        row = {"text": text,
               "label_bin": 0 if is_ad else 1,
               "label_kind": KINDS.index(k) if k else -100}
        if source is not None:
            row["source"] = source
        return row

    train_rows = load(args.split_train)
    val_rows = load(args.split_val)
    train = [pack(r) for r in train_rows]
    val = []
    for r in val_rows:
        t = r["text"].strip()
        src = "archive" if t in arch_all else ("gold" if t in gold_all else "none")
        val.append(pack(r, src))

    # dev 分层键：normal / 细类 / ad_unknown
    keys = []
    for r in train:
        if r["label_bin"] == 1:
            keys.append("normal")
        elif r["label_kind"] >= 0:
            keys.append(KINDS[r["label_kind"]])
        else:
            keys.append("ad_unknown")
    idx = np.arange(len(train))
    tr_idx, dev_idx = train_test_split(idx, test_size=args.dev_ratio,
                                       stratify=np.array(keys), random_state=args.seed)
    fit = [train[i] for i in tr_idx]
    dev = [train[i] for i in dev_idx]
    holdout = [r for r in val if r.get("source") == "archive"]

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    for name, rows in [("train", fit), ("dev", dev), ("val", val), ("holdout827", holdout)]:
        with open(out / f"{name}.jsonl", "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        b = collections.Counter(r["label_bin"] for r in rows)
        k = collections.Counter(KINDS[r["label_kind"]] for r in rows if r["label_kind"] >= 0)
        print(f"{name:<12} n={len(rows):5d} bin={dict(b)} kind={dict(k)}")

    (out / "label_map.json").write_text(json.dumps({
        "binary": {"ad": 0, "normal": 1},
        "kind": {k: i for i, k in enumerate(KINDS)},
    }, ensure_ascii=False, indent=2), encoding="utf-8")
    print("输出:", out.resolve())


if __name__ == "__main__":
    main()