"""
text_data.py —— 前置讲 P2 用到的两份文本数据
=====================================================================

【这个文件提供什么】
    1. 一份**真实的英文语料**：Python 标准库源码里的文档字符串（docstring）。
       不用下载任何东西——它就在你的电脑上。约 22 万个英文单词，是"真正的自然语言"
       （有长尾、有罕见词、词表不封闭），用来测量 Zipf 定律、Heaps 定律和词向量。
       我们用 ast 模块**解析**源码文件，而不是 import 它们（import 某些模块会有副作用，
       例如打开浏览器）。
       【边界】不同 Python 版本的标准库不同，所以在别的版本上数字会略有差别。

    2. 一个**生成式的情感分类世界** ReviewWorld："花园日记"。
       每篇日记有一个标签（1 = 开心，0 = 沮丧），日记里的词按一个已知的概率过程生成。
       因为生成过程已知，我们能算出**贝叶斯最优分类器**（理论上最好的分类器）的准确率，
       作为所有分类器的天花板。
       可以打开"否定"开关：一部分情感表达被写成 "not bad"（意思是好）这样的形式，
       这会破坏朴素贝叶斯"每个词独立"的假设。

【LLM 联系】
    - 文本的长尾统计（Zipf）是分词（第 1、3 讲）、稀疏性（第 1 讲）和数据整理（第 3 讲）的背景。
    - 情感分类是"预训练之前"的 NLP 的典型任务；第 6 讲会看到，今天它可以直接用提示、
      或者在 LLM 上微调来做。
"""

import ast
import math
import os
import random
import re

# =====================================================================
# 1. 真实语料：标准库的文档字符串
# =====================================================================

# 跳过测试、示例和第三方包目录（它们要么不是自然语言，要么不属于标准库）
SKIP_DIRS = {"test", "tests", "idlelib", "site-packages", "__pycache__", "lib2to3",
             "turtledemo", "tkinter", "ensurepip", "venv", "pydoc_data"}

_DOC_CACHE = None


def stdlib_docstrings():
    """
    读取当前 Python 标准库全部 .py 文件中的文档字符串（模块、类、函数的 docstring）。

    【返回】list[str]，按文件路径排序后的固定顺序（同一个 Python 版本上每次结果相同）。
    【思路】os.walk 遍历标准库目录 -> ast.parse 解析成语法树 -> ast.walk 找到所有模块/类/函数节点
            -> ast.get_docstring 取出文档字符串。
    【边界】解析失败的文件（极少数）直接跳过。结果会缓存，第二次调用不再重新解析（约 1 秒）。
    """
    global _DOC_CACHE
    if _DOC_CACHE is not None:
        return _DOC_CACHE
    root = os.path.dirname(os.__file__)          # 标准库所在目录，例如 .../lib/python3.13
    files = []
    for dirpath, dirnames, filenames in os.walk(root):
        # 【Python 小知识】原地修改 dirnames（dirnames[:] = ...）可以让 os.walk 不再进入被删掉的子目录；
        #   排序则保证遍历顺序固定。
        dirnames[:] = sorted(d for d in dirnames if d not in SKIP_DIRS)
        files.extend(os.path.join(dirpath, f) for f in sorted(filenames) if f.endswith(".py"))
    docs = []
    for path in files:
        try:
            with open(path, encoding="utf-8") as fh:
                tree = ast.parse(fh.read())
        except (SyntaxError, UnicodeDecodeError, ValueError):
            continue
        for node in ast.walk(tree):
            if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
                d = ast.get_docstring(node)
                if d:
                    docs.append(d)
    _DOC_CACHE = docs
    return docs


WORD_RE = re.compile(r"[a-z]+")


def words_of(text):
    """小写化后取出所有由字母组成的词（数字、标点、下划线都当作分隔符）。"""
    return WORD_RE.findall(text.lower())


def stdlib_words():
    """把全部文档字符串拼起来、切成词。【返回】list[str]，约 22 万个词。"""
    return [w for d in stdlib_docstrings() for w in words_of(d)]


# =====================================================================
# 2. 情感分类世界："花园日记"
# =====================================================================
NEUTRAL = ["the", "garden", "today", "we", "i", "was", "is", "it", "and", "plants", "soil",
           "weather", "with", "very", "our", "my", "this", "after", "rain", "morning",
           "some", "were", "so", "in", "all"]
POSITIVE = ["lovely", "sweet", "fresh", "great", "happy", "bright", "good", "healthy"]
NEGATIVE = ["bad", "rotten", "dry", "sad", "poor", "wilted", "awful", "sick"]


def _zipf(n):
    w = [1.0 / r for r in range(1, n + 1)]
    s = sum(w)
    return [x / s for x in w]


class ReviewWorld:
    """
    "花园日记"情感世界。

    【生成过程】（一篇日记）
        1. 标签 y：1（开心）或 0（沮丧），各占一半。
        2. 长度 L：在 5～15 之间均匀选。
        3. 重复 L 次，生成一个"单元"：
           - 以概率 1 - p_sent 生成一个中性词（按 Zipf 频率）；
           - 以概率 p_sent 生成一个"情感表达"：先选极性 π（开心的日记里 π = 正面的概率是 p_pol，
             沮丧的日记里是 1 - p_pol），然后
               · 以概率 1 - q 直接写一个极性为 π 的情感词（good、bad……）；
               · 以概率 q 写成否定形式：not + 一个**相反极性**的词（"not bad" 表达正面）。
    【参数】q：否定的比例（0 表示没有否定）；p_sent：情感表达的比例；p_pol：标签与极性的相关程度。
    【精确似然】因为 "not" 只出现在否定形式里，任何一篇日记都能唯一地切分成单元，
            所以 log p(日记 | y) 可以精确算出，进而得到贝叶斯最优分类器。
    """

    def __init__(self, q=0.0, p_sent=0.25, p_pol=0.8):
        self.q, self.p_sent, self.p_pol = q, p_sent, p_pol
        self.p_neu = dict(zip(NEUTRAL, _zipf(len(NEUTRAL))))
        self.p_pos = dict(zip(POSITIVE, _zipf(len(POSITIVE))))
        self.p_neg = dict(zip(NEGATIVE, _zipf(len(NEGATIVE))))

    def sample(self, rng):
        """【返回】(词列表, 标签)。【参数】rng：random.Random。"""
        y = rng.randint(0, 1)
        out = []
        for _ in range(rng.randint(5, 15)):
            if rng.random() < self.p_sent:
                positive = rng.random() < (self.p_pol if y == 1 else 1 - self.p_pol)
                if rng.random() < self.q:
                    pool = self.p_neg if positive else self.p_pos      # "not" + 相反极性的词
                    out += ["not", rng.choices(list(pool), weights=list(pool.values()))[0]]
                else:
                    pool = self.p_pos if positive else self.p_neg
                    out.append(rng.choices(list(pool), weights=list(pool.values()))[0])
            else:
                out.append(rng.choices(NEUTRAL, weights=list(self.p_neu.values()))[0])
        return out, y

    def dataset(self, n, seed):
        rng = random.Random(seed)
        return [self.sample(rng) for _ in range(n)]

    def log_likelihood(self, doc, y):
        """
        log p(日记 | 标签 y)，不含长度项（长度和标签无关，比较两个类别时会抵消）。

        【边界】q = 0 时文档里不会有 "not"；如果手动构造了 q = 0 世界里不可能的日记，结果是 -inf。
        """
        pr_pos = self.p_pol if y == 1 else 1 - self.p_pol
        lp, i = 0.0, 0
        while i < len(doc):
            w = doc[i]
            if w == "not":
                b = doc[i + 1]
                if self.q == 0:
                    return -math.inf
                if b in self.p_neg:            # "not bad"：正面
                    lp += math.log(self.p_sent * pr_pos * self.q * self.p_neg[b])
                else:                           # "not good"：负面
                    lp += math.log(self.p_sent * (1 - pr_pos) * self.q * self.p_pos[b])
                i += 2
                continue
            if w in self.p_pos:
                lp += math.log(self.p_sent * pr_pos * (1 - self.q) * self.p_pos[w])
            elif w in self.p_neg:
                lp += math.log(self.p_sent * (1 - pr_pos) * (1 - self.q) * self.p_neg[w])
            else:
                lp += math.log((1 - self.p_sent) * self.p_neu[w])
            i += 1
        return lp

    def bayes_predict(self, doc):
        """贝叶斯最优分类：比较两个类别的似然（先验相等）。平局时判为 1。"""
        return int(self.log_likelihood(doc, 1) >= self.log_likelihood(doc, 0))