"""
cleaning.py —— 第 3 讲：一条最小的预训练数据清洗流水线
=====================================================================

【这个文件要回答的问题】
    slides 把网页数据的处理分成三步：内容被发到网上 → 爬虫下载一部分 → **过滤和清洗**。
    这里实现第三步里最常见的几个零件（都只用 Python 标准库）：

        extract_all / extract_main   从 HTML 里取出文字：全部可见文字（≈ Common Crawl 的 WET） vs 只取正文
        LangID                       字符 n-gram 朴素贝叶斯语言识别（≈ fastText 语言识别的最简版）
        c4_rules                     C4 的启发式规则（Raffel et al. 2020，slides 第 32 页）
        gopher_rules                 Gopher / MassiveText 的质量规则（Rae et al. 2021）的一个子集
        dedup_span3                  C4 的"三句话片段"级去重
        exact_key / MinHash / lsh_buckets   精确去重与 MinHash + LSH 近似去重
        ngram_set / contaminated     基于 13-gram 重叠的测试集污染检测（GPT-3 的做法）
        to_tokens                    把清洗后的文字切回花园世界的词，供第 05 号脚本训练语言模型

    每个函数都返回"为什么删"的理由，好让脚本统计"每条规则删了什么"。

【LLM 联系】
    这些规则看起来很粗糙，但正是 C4、RefinedWeb、FineWeb、Dolma 等真实数据集的主干。
    它们的共同问题（本讲的实验会量化）：规则是为"标准英语网页"设计的，删掉的东西里
    总有一部分是好文字——而且删错的比例在不同人群、方言、语言之间并不均匀。
"""

import hashlib
import math
import random
import re
from collections import Counter, defaultdict

# =====================================================================
# 1. 从 HTML 取出文字
# =====================================================================

_BLOCK_TAGS = r"</?(p|div|nav|main|aside|footer|header|h1|h2|h3|li|ul|br|title|body|html|head)[^>]*>"


def extract_all(html):
    """
    取出页面上**全部**可见文字（去掉 <script>、<style> 的内容和所有标签），块级元素之间换行。

    【返回】str，一行一个"块"（导航栏、cookie 提示、段落、侧栏、页脚各占一行）。
    【思路】这大致就是 Common Crawl 的 WET 文件给出的东西：没有正文抽取，样板文字全在里面。
    【边界】只处理本讲网页生成器产出的这种规整 HTML；真实网页要用专门的解析器（如 resiliparse、trafilatura）。
    """
    html = re.sub(r"<script.*?</script>|<style.*?</style>", " ", html, flags=re.S)
    html = re.sub(_BLOCK_TAGS, "\n", html)
    text = re.sub(r"<[^>]+>", " ", html)
    lines = [re.sub(r"\s+", " ", ln).strip() for ln in text.split("\n")]
    return "\n".join(ln for ln in lines if ln)


def extract_main(html):
    """
    只取 <main> 里的 <p> 段落（≈ 正文抽取工具的效果）：导航栏、侧栏、页脚、cookie 提示、标题都不要。

    【返回】str，一段一行；页面没有 <main> 时返回空串。
    """
    m = re.search(r"<main>(.*?)</main>", html, flags=re.S)
    if not m:
        return ""
    paras = re.findall(r"<p>(.*?)</p>", m.group(1), flags=re.S)
    return "\n".join(re.sub(r"\s+", " ", p).strip() for p in paras if p.strip())


# =====================================================================
# 2. 语言识别
# =====================================================================

def char_ngrams(text, n=3):
    t = f" {text.lower()} "
    return [t[i:i + n] for i in range(len(t) - n + 1)]


class LangID:
    """
    字符三元组朴素贝叶斯语言识别。

    【思路】每种语言统计字符三元组的频率（加一平滑）。给一行文字，算它在每种语言下的总对数似然，
            取最大的那个作为这一行的语言。一页文字的语言 = 字符数占比最大的语言，
            "置信度" = 这种语言的行占全部字符的比例（一页里混着英语导航栏和外语正文时，置信度 < 1）。
            fastText 的语言识别模型（CCNet、FineWeb 都在用）本质上也是"字符 n-gram 特征 + 线性分类器"，
            只是更大、训练数据更多；FineWeb 只保留英语置信度 ≥ 0.65 的页面。
    【参数】train_texts：{语言: [文字, ...]}；n：n-gram 的 n。
    【用法】lid = LangID({"en": [...], "floral": [...]}); lang, conf = lid.predict(text)
    【边界】空文字返回 ("none", 0.0)；单行很短（几个字符）时判断不可靠，但它在字符占比里的权重也小。
    """

    def __init__(self, train_texts, n=3):
        self.n = n
        self.logp, self.unseen = {}, {}
        vocab = set()
        counts = {}
        for lang, texts in train_texts.items():
            counts[lang] = Counter(g for t in texts for g in char_ngrams(t, n))
            vocab |= set(counts[lang])
        V = len(vocab) + 1
        for lang, c in counts.items():
            denom = math.log(sum(c.values()) + V)
            self.logp[lang] = {g: math.log(k + 1) - denom for g, k in c.items()}
            self.unseen[lang] = -denom

    def predict_line(self, line):
        """一行文字最可能的语言。"""
        grams = char_ngrams(line, self.n)
        scores = {lang: sum(lp.get(g, self.unseen[lang]) for g in grams) for lang, lp in self.logp.items()}
        return max(scores, key=scores.get)

    def predict(self, text):
        """【返回】(语言, 该语言的字符占比)。"""
        share = Counter()
        for ln in text.split("\n"):
            if ln.strip():
                share[self.predict_line(ln)] += len(ln)
        if not share:
            return "none", 0.0
        lang, c = share.most_common(1)[0]
        return lang, c / sum(share.values())


# =====================================================================
# 3. 启发式规则
# =====================================================================

TERMINAL = (".", "!", "?", '"')


def c4_rules(text, blocklist, min_sentences=5):
    """
    C4 的过滤规则（slides 第 32 页列出的 7 条，第 7 条去重另见 dedup_span3）。

    【规则】
        行级：(2) 只保留以句末标点结尾的行；(4) 只保留至少 3 个词的行；(5) 删掉含 "javascript" 的行。
        页级：(3) 删掉少于 min_sentences 句的页面；(6a) 含屏蔽词的页面；(6b) 含 "lorem ipsum"；(6c) 含 "{"。
    【参数】text：一页的文字（一行一块）；blocklist：屏蔽词列表（按整词、不区分大小写匹配）。
    【返回】(保留下来的文字 或 None, 删除理由列表)。整页被删时文字为 None。
    【边界】句子数按句末标点粗略地数；页级规则在行级规则**之后**检查（C4 的做法）。
    """
    reasons = []
    low = text.lower()
    if "lorem ipsum" in low:
        return None, ["lorem ipsum"]
    if "{" in text:
        return None, ["含 {"]
    words = set(re.findall(r"[a-z]+", low))
    if words & set(blocklist):
        return None, ["屏蔽词"]
    kept = []
    for ln in text.split("\n"):
        if not ln.endswith(TERMINAL):
            reasons.append("行无句末标点")
            continue
        if len(ln.split()) < 3:
            reasons.append("行少于 3 词")
            continue
        if "javascript" in ln.lower():
            reasons.append("行含 javascript")
            continue
        kept.append(ln)
    out = "\n".join(kept)
    n_sent = len(re.findall(r"[.!?](\s|$)", out))
    if n_sent < min_sentences:
        return None, reasons + [f"少于 {min_sentences} 句"]
    return out, reasons


STOPWORDS = {"the", "be", "to", "of", "and", "that", "have", "with"}          # Gopher 原文的英语停用词
GARDEN_STOPWORDS = {"the", "with", "from", "in", "every", "own"}               # 按花园语重新挑的功能词


def gopher_rules(text, stopwords=GARDEN_STOPWORDS):
    """
    Gopher（Rae et al. 2021，MassiveText）质量规则的一个子集。

    【规则】（原文的阈值）
        词数在 50 到 100,000 之间；平均词长在 3 到 10 个字符之间；
        "#" 和 "..." 与词数之比不超过 0.1；以省略号结尾的行不超过 30%；
        至少 80% 的词含有字母；至少出现 2 个停用词（原文：the, be, to, of, and, that, have, with）；
        重复行占全部行的比例不超过 30%；最常见的 2-gram 占全部字符的比例不超过 20%。
    【参数】stopwords：停用词集合。默认用按花园语挑的功能词（GARDEN_STOPWORDS）：
            花园语里几乎没有 and、to、of、that，照搬原文的英语词表会误删大量好页面（02 号脚本 §5 对比了两者）。
    【返回】(是否通过, 不通过的理由列表)。
    【为什么要停用词规则】自然语言里功能词无处不在；一页文字里连 the、and 都没有，多半是列表、代码或乱码。
    【LLM 联系】规则的阈值是按某种语言、某类网页调出来的。FineWeb-2 把 FineWeb 的流水线扩展到上千种语言时，
            就为每种语言重新估计了停用词和各项阈值。
    """
    reasons = []
    words = text.split()
    n = len(words)
    if not 50 <= n <= 100_000:
        reasons.append("词数不在 50～100000")
    if n:
        mean_len = sum(len(w) for w in words) / n
        if not 3 <= mean_len <= 10:
            reasons.append("平均词长异常")
        if (text.count("#") + text.count("...")) / n > 0.1:
            reasons.append("符号过多")
        if sum(1 for w in words if re.search(r"[A-Za-zÀ-ÿø]", w)) / n < 0.8:
            reasons.append("含字母的词不足 80%")
        lw = {w.lower().strip(".,!?") for w in words}
        if len(lw & set(stopwords)) < 2:
            reasons.append("停用词不足")
    lines = [ln for ln in text.split("\n") if ln.strip()]
    if lines:
        if sum(ln.endswith("...") for ln in lines) / len(lines) > 0.3:
            reasons.append("省略号行过多")
        dup = sum(c - 1 for c in Counter(lines).values())
        if dup / len(lines) > 0.3:
            reasons.append("重复行过多")
    lw = [w.lower() for w in words]
    if len(lw) >= 2:
        bigrams = Counter(zip(lw, lw[1:]))
        (a, b), c = bigrams.most_common(1)[0]
        if c * (len(a) + len(b)) / max(1, len(text)) > 0.2:
            reasons.append("最常见 2-gram 占比过高")
    return not reasons, reasons


def sentences_of(text):
    """把文字粗略切成句子（按句末标点 + 空白）。"""
    return [s.strip() for s in re.split(r"(?<=[.!?])\s+", text.replace("\n", " ")) if s.strip()]


def dedup_span3(texts):
    """
    C4 的第 7 条：在整个数据集里，任何一个"连续三句话"的片段如果出现多次，只保留第一次出现。

    【思路】对每篇文字按顺序扫描三句话的窗口；某个窗口之前出现过，就把这三句都标记为删除。
    【返回】(新的文字列表, 被删的句子数)。删完之后句子太少的页面不在这里处理。
    【边界】比较时把句子小写、合并空白；只删除重复的片段本身，页面的其余部分保留。
    """
    seen = set()
    out, removed = [], 0
    for t in texts:
        sents = sentences_of(t)
        drop = set()
        for i in range(len(sents) - 2):
            key = tuple(s.lower() for s in sents[i:i + 3])
            if key in seen:
                drop.update((i, i + 1, i + 2))
            else:
                seen.add(key)
        removed += len(drop)
        out.append(" ".join(s for i, s in enumerate(sents) if i not in drop))
    return out, removed


# =====================================================================
# 4. 去重
# =====================================================================

def normalize(text):
    """去重前的规范化：小写、只留字母数字和空格、合并空白。"""
    return re.sub(r"\s+", " ", re.sub(r"[^0-9a-zà-ÿø ]", " ", text.lower())).strip()


def exact_key(text):
    """精确去重的键：规范化之后的 SHA-1。"""
    return hashlib.sha1(normalize(text).encode()).hexdigest()


def shingles(text, k=5):
    """
    词级 k-shingle（连续 k 个词）的集合。

    【为什么用 shingle】两篇文档的"相似度"定义成它们 shingle 集合的 Jaccard 系数：
            J(A, B) = |A ∩ B| / |A ∪ B|。改一句话只影响包含那句话的少数 shingle，J 仍然很高。
    【边界】少于 k 个词的文字返回只含整段文字的一个元素。
    """
    w = normalize(text).split()
    if len(w) < k:
        return {" ".join(w)}
    return {" ".join(w[i:i + k]) for i in range(len(w) - k + 1)}


def jaccard(a, b):
    return len(a & b) / len(a | b) if a or b else 1.0


class MinHash:
    """
    MinHash：用 num_perm 个随机哈希函数，把一个集合压缩成 num_perm 个整数（签名）。

    【思路】对一个随机哈希函数 h，P[min h(A) = min h(B)] = J(A, B)。
            所以两个签名里"相同位置取值相同"的比例，就是 Jaccard 系数的无偏估计；
            num_perm 越大，估计越准（标准差约 √(J(1-J)/num_perm)）。
    【实现】h_i(x) = (a_i · hash(x) + b_i) mod p，p 是一个大素数；hash(x) 用 SHA-1 的前 8 字节，
            保证跨进程、跨机器结果一致（Python 自带的 hash() 每次运行都会变）。
    """

    P = (1 << 61) - 1

    def __init__(self, num_perm=128, seed=0):
        rng = random.Random(seed)
        self.a = [rng.randrange(1, self.P) for _ in range(num_perm)]
        self.b = [rng.randrange(0, self.P) for _ in range(num_perm)]

    @staticmethod
    def _h(x):
        return int.from_bytes(hashlib.sha1(x.encode()).digest()[:8], "big")

    def signature(self, items):
        hs = [self._h(x) for x in items]
        return [min((a * h + b) % self.P for h in hs) for a, b in zip(self.a, self.b)]

    @staticmethod
    def estimate(s1, s2):
        return sum(x == y for x, y in zip(s1, s2)) / len(s1)


def lsh_buckets(signatures, bands, rows):
    """
    局部敏感哈希（LSH）的"分带"技巧：把签名切成 bands 段、每段 rows 个数；
    只要有一段完全相同，两篇文档就成为"候选对"。

    【为什么】两两比较 n 篇文档要 n²/2 次；LSH 只比较落进同一个桶的文档。
            Jaccard 为 s 的一对成为候选的概率是 1 - (1 - s^rows)^bands——一条 S 形曲线，
            阈值大约在 (1/bands)^(1/rows)。
    【参数】signatures：{文档 id: 签名列表}，签名长度必须 = bands × rows。
    【返回】候选对的集合 {(id1, id2), ...}（id1 < id2）。
    """
    pairs = set()
    for bnd in range(bands):
        buckets = defaultdict(list)
        for doc_id, sig in signatures.items():
            buckets[tuple(sig[bnd * rows:(bnd + 1) * rows])].append(doc_id)
        for ids in buckets.values():
            for i in range(len(ids)):
                for j in range(i + 1, len(ids)):
                    pairs.add((min(ids[i], ids[j]), max(ids[i], ids[j])))
    return pairs


# =====================================================================
# 5. 测试集污染
# =====================================================================

def ngram_set(tokens, n=13):
    return {tuple(tokens[i:i + n]) for i in range(len(tokens) - n + 1)}


def contaminated(tokens, test_ngrams, n=13):
    """一篇训练文字和测试集是否共享至少一个 n-gram（GPT-3 用 13-gram 做去污染）。"""
    return any(tuple(tokens[i:i + n]) in test_ngrams for i in range(len(tokens) - n + 1))


# =====================================================================
# 6. 文字 → 花园世界的词
# =====================================================================

def to_tokens(text):
    """
    把清洗后的文字切回词：字母串、数字串、单个标点各算一个 token；全部小写（"I" 除外）。

    【例子】"She picks her ripe plum. Thanks!" → ["she", "picks", "her", "ripe", "plum", ".", "thanks", "!"]
    【边界】花园世界词表之外的词由调用者映射成 <unk>。
    """
    toks = re.findall(r"[A-Za-zÀ-ÿø]+|[0-9]+|[^\sA-Za-z0-9]", text)
    return ["I" if t == "I" else t.lower() for t in toks]