"""第 22 章：自编短文本的清洗、去重分组和溯源；不抓取外部数据。"""

import argparse
from hashlib import sha256
from html.parser import HTMLParser
from itertools import combinations
import json
from pathlib import Path
from random import Random
import re
import unicodedata

PROTECTED = "模型评测集的问题不应进入训练数据。"


def records():
    texts = [
        "打印报告前，请检查纸张余量，并确认打印机已经连接。",
        "<p>打印报告前，请检查纸张余量，并确认打印机已经连接。</p>",
        "打印报告前，请检查纸张余量，并确认打印机已经连接！",
        "训练模型之前，先把数据分为训练集、验证集和测试集。",
        "这份材料的授权情况尚未确认，请先保留来源记录。",
        "如需反馈打印问题，请联系 demo@example.org，并写明错误信息。",
        "你好！",
        "！？？！？？！？？！？？！？？！？？",
        "示例试题：" + PROTECTED,
        "图像分类器接收像素数组，输出每个候选类别的分数。",
        "计算平均损失时，应统计实际参与预测的词元数量。",
        "保存模型时，同时记录词表、网络配置和参数文件。",
        "训练模型之前，先把数据划分为训练集、验证集和测试集。",
        "中文文本通常没有明确的空格词边界，切分方式需要统一。",
    ]
    return [dict(id=f"r{i:02}", source="本书自编示例", domain="说明文本",
                 rights="待确认（模拟）" if i == 5 else "作者提供的教学文本", text=text)
            for i, text in enumerate(texts, 1)]


class ExtractText(HTMLParser):
    def __init__(self):
        super().__init__(convert_charrefs=True)
        self.parts, self.hidden = [], 0

    def handle_starttag(self, tag, attrs):
        if tag in {"script", "style"}:
            self.hidden += 1
        if tag in {"p", "div", "br"} and not self.hidden:
            self.parts.append(" ")

    def handle_endtag(self, tag):
        if tag in {"script", "style"}:
            self.hidden = max(0, self.hidden-1)
        if tag in {"p", "div"} and not self.hidden:
            self.parts.append(" ")

    def handle_data(self, data):
        if not self.hidden:
            self.parts.append(data)


def clean(text):
    parser = ExtractText()
    parser.feed(text)
    text = unicodedata.normalize("NFC", "".join(parser.parts))
    text = " ".join(text.split())
    # 只演示一种明显格式；不能替代真实隐私审查。
    return re.subn(r"[\w.+-]+@[\w.-]+\.[A-Za-z]{2,}", "[EMAIL]", text)


def shingles(text, size=3):
    return {text[i:i+size] for i in range(max(0, len(text)-size+1))}


def jaccard(left, right):
    a, b = shingles(left), shingles(right)
    return len(a & b)/len(a | b) if a | b else 0.0


def pipeline():
    raw = records()
    stages, rejected, kept, by_hash = [("原始记录", len(raw))], [], [], {}
    permitted = [row for row in raw if row["rights"] == "作者提供的教学文本"]
    rejected.extend((row["id"], "授权待确认") for row in raw if row not in permitted)
    stages.append(("来源筛选后", len(permitted)))
    quality = []
    for row in permitted:
        text, redactions = clean(row["text"])
        ratio = sum(char.isalnum() for char in text)/max(1, len(text))
        if len(text) < 12 or ratio < 0.6:
            rejected.append((row["id"], "长度/字符比例筛选"))
        else:
            quality.append({**row, "text": text, "redactions": redactions})
    stages.append(("格式筛选后", len(quality)))
    screened = []
    for row in quality:
        if PROTECTED in row["text"]:
            rejected.append((row["id"], "命中受保护评测文本"))
        else:
            screened.append(row)
    stages.append(("污染筛查后", len(screened)))
    for row in screened:
        digest = sha256(row["text"].encode()).hexdigest()
        if digest in by_hash:
            by_hash[digest]["original_ids"].append(row["id"])
            rejected.append((row["id"], "规范化后完全重复"))
        else:
            item = {**row, "sha256": digest, "original_ids": [row["id"]]}
            by_hash[digest] = item
            kept.append(item)
    stages.append(("完全去重后", len(kept)))
    # 相近文本保留，但整个连通组进入同一个数据划分。
    parent = {row["id"]: row["id"] for row in kept}

    def root(key):
        while key != parent[key]:
            key = parent[key]
        return key

    pairs = []
    for left, right in combinations(kept, 2):
        score = jaccard(left["text"], right["text"])
        if score >= 0.72:
            parent[root(right["id"])] = root(left["id"])
            pairs.append((left["id"], right["id"], score))
    groups = sorted({root(row["id"]) for row in kept})
    Random(2026).shuffle(groups)
    train_end = int(0.7*len(groups))
    assignments = {group: "训练" if i < train_end else "验证" if i == train_end else "测试"
                   for i, group in enumerate(groups)}
    for row in kept:
        row["group"] = root(row["id"])
        row["split"] = assignments[row["group"]]
    stages.append(("近重复分组数", len(groups)))
    return kept, stages, sorted(rejected), pairs


def mixture(counts, alpha):
    weights = [count**alpha for count in counts]
    return [weight/sum(weights) for weight in weights]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", type=Path, help="可选：保存处理后 JSONL 与同名 .audit.json")
    args = parser.parse_args()
    kept, stages, rejected, pairs = pipeline()
    for name, count in stages:
        print(f"{name}：{count}")
    print("未保留：" + "；".join(f"{key} {reason}" for key, reason in rejected))
    print("近重复：" + "；".join(f"{a}/{b}={score:.3f}" for a, b, score in pairs))
    print("邮箱替换次数：" + str(sum(row["redactions"] for row in kept)))
    for split in ["训练", "验证", "测试"]:
        print(f"{split}：" + ", ".join(row["id"] for row in kept if row["split"] == split))
    print("配比示例的词元量：900、90、10（与上述小语料无关）")
    for alpha in [1.0, 0.5]:
        print(f"alpha={alpha:.1f}：" + "、".join(f"{p:.1%}" for p in mixture([900, 90, 10], alpha)))
    if args.out:
        args.out.parent.mkdir(parents=True, exist_ok=True)
        args.out.write_text("".join(json.dumps(row, ensure_ascii=False)+"\n" for row in kept))
        audit = dict(version=1, stages=stages, rejected=rejected, near_pairs=pairs,
                     raw=records(), near_threshold=0.72, split_seed=2026)
        args.out.with_suffix(".audit.json").write_text(json.dumps(audit, ensure_ascii=False, indent=2))
        print(f"已保存处理结果与审计记录：{args.out}")


if __name__ == "__main__":
    main()