"""
=============================================================================
第 4 讲 · 工具模块：一个按规则生成的迷你依存树库
=============================================================================

对应讲义：lectures/lecture04_dependency_parsing.md 第 4、10 节
被本目录下所有编号文件使用。

-----------------------------------------------------------------------------
【树库是什么？】
-----------------------------------------------------------------------------
树库（treebank）是一批人工标注了句法结构的句子。依存树库给每个词标出：
    它的**中心词**（head，也叫支配词）是谁，以及这条依存关系的**标签**（主语、宾语、修饰语……）。
真实的树库（例如 Universal Dependencies，UD）由语言学家花数年时间标注。
我们在笔记本上没有这样的数据，就用一套小语法**生成**句子和它们的依存树。

-----------------------------------------------------------------------------
【语法设计】
-----------------------------------------------------------------------------
句子结构：主语名词短语 +（助动词）+（副词）+ 动词 + 宾语名词短语 +（介词短语）+（副词）+ 句号
    the hungry chef will quickly cut the bread with a knife .

标签沿用 UD 的命名：
    nsubj 主语   obj 宾语   det 限定词   amod 形容词修饰   case 介词（作为名词的依存）
    nmod  介词短语修饰名词      obl 介词短语修饰动词      advmod 副词   aux 助动词
    cc    并列连词   conj 并列成分   punct 标点   root 根
（UD 的一个特点：介词 with 不是中心词，而是作为 case 依附在它后面的名词 knife 上；
 knife 再作为 obl 依附在动词 cut 上。）

**介词短语附着**（PP attachment）是这个树库的核心难点，由词义决定：
    with + 工具（knife、fork、telescope……）  → 修饰动词（obl）："cut the bread with a knife"
    with + 配料（cheese、sauce……）          → 修饰食物名词（nmod）："ate the pizza with cheese"
    with + 衣饰（hat、scarf……）             → 修饰人（nmod）："saw the man with a hat"
    in   + 地点                              → 修饰动词（obl）："ate the pizza in the kitchen"
    near + 地点                              → 修饰名词（nmod）："painted the house near the park"
所以分析器必须"懂"一点词义（fork 和 knife 是一类、cheese 和 sauce 是一类），才能分对。

【边界】
    - 生成的句子都是**投射的**（依存弧互不交叉，讲义第 4.4 节）。01 号文件会手工构造一个非投射的例子。
    - 真实语言的介词短语附着往往有真正的歧义（"I saw the man with the telescope"），
      我们的语法为了有确定的"标准答案"，规定了 telescope 一定修饰动词。

【用法】
    from treebank import make_treebank
    train, dev, test = make_treebank()
    sent = train[0]
    sent.words, sent.tags, sent.heads, sent.labels   # 都是从 1 开始编号的列表（位置 0 留给 ROOT）
"""

import random
from dataclasses import dataclass, field

ROOT = "<ROOT>"

# ---------------- 词表（带语义类别） ----------------
LEXICON = {
    "person": ["man", "woman", "chef", "girl", "boy", "student", "farmer", "doctor", "teacher", "child"],
    "food": ["pizza", "bread", "soup", "cake", "salad", "fish", "pasta", "pie"],
    "tool": ["knife", "fork", "spoon", "telescope", "hammer", "camera", "brush", "stick"],
    "topping": ["cheese", "sauce", "cream", "mushrooms", "honey", "garlic"],
    "accessory": ["hat", "scarf", "beard", "coat", "umbrella", "badge"],
    "place": ["park", "kitchen", "garden", "library", "city", "street", "hall", "market"],
    "thing": ["house", "car", "tree", "bird", "star", "boat", "bridge", "fence"],
}
ADJECTIVES = ["big", "small", "old", "young", "hungry", "happy", "red", "fresh", "tall", "quiet"]
ADVERBS = ["quickly", "slowly", "carefully", "happily", "suddenly"]
# 动词 → 它能带的宾语类别、能用 "with + 工具" 修饰它吗
VERBS = {
    "ate": (["food"], True),
    "cut": (["food"], True),
    "cooked": (["food"], True),
    "saw": (["person", "thing"], True),
    "watched": (["person", "thing"], True),
    "painted": (["thing"], True),
    "met": (["person"], False),
    "bought": (["food", "thing"], False),
}


# =============================================================================
# 第 1 部分：句子的数据结构
# =============================================================================
@dataclass
class Sentence:
    """
    一个带依存标注的句子。所有列表的第 0 个位置都是虚拟的 ROOT 节点，真实的词从 1 开始编号。

    【字段】
        words  : list[str]，例如 ["<ROOT>", "the", "chef", "cut", ...]
        tags   : list[str]，词性，例如 ["<ROOT>", "DET", "NOUN", "VERB", ...]
        heads  : list[int]，每个词的中心词编号（0 表示 ROOT）；heads[0] 无意义，记为 -1
        labels : list[str]，每个词与其中心词之间的依存标签；labels[0] 记为 "<ROOT>"

    【Python 小知识】
        @dataclass 自动生成 __init__ 等方法；field(default_factory=list) 让每个实例有自己的空列表。
    """
    words: list = field(default_factory=list)
    tags: list = field(default_factory=list)
    heads: list = field(default_factory=list)
    labels: list = field(default_factory=list)

    def __len__(self):
        """句子的长度（不含 ROOT）。"""
        return len(self.words) - 1

    def text(self):
        """返回句子的文字（不含 ROOT）。"""
        return " ".join(self.words[1:])


class Node:
    """
    生成句子时用的"树节点"：一个词，加上它左边和右边的依存成分。

    【思路】
        先在树上把句子"搭"出来（谁依附谁），再按"左依存成分 → 自己 → 右依存成分"的顺序
        把树线性化成词序列，同时算出每个词的编号和中心词编号。
        因为每个节点的所有后代在句子里是连续的一段，这样生成的树一定是投射的。
    """

    def __init__(self, word, tag, label, left=None, right=None):
        self.word, self.tag, self.label = word, tag, label
        self.left = left or []          # 左边的依存成分（按从左到右的顺序）
        self.right = right or []        # 右边的依存成分


def linearize(root):
    """
    把一棵 Node 树变成 Sentence。

    【参数】root : Node，整句话的根（通常是主动词，label 为 "root"）
    【返回】Sentence
    """
    sent = Sentence([ROOT], [ROOT], [-1], [ROOT])

    def visit(node, head_id):
        # 先占位：我们需要知道自己的编号才能告诉孩子们"你的中心词是谁"，
        # 但左孩子排在自己前面。所以先递归处理左孩子（它们的中心词编号稍后回填）。
        left_ids = [visit(child, None) for child in node.left]
        my_id = len(sent.words)
        sent.words.append(node.word)
        sent.tags.append(node.tag)
        sent.heads.append(head_id)
        sent.labels.append(node.label)
        for cid in left_ids:
            sent.heads[cid] = my_id
        for child in node.right:
            visit(child, my_id)
        return my_id

    visit(root, 0)
    return sent


# =============================================================================
# 第 2 部分：按语法生成
# =============================================================================
def make_np(cls, label, rng, allow_pp=True, allow_coord=False):
    """
    生成一个名词短语：(限定词) (形容词*) 名词 (并列) (修饰名词的介词短语)。

    【参数】
        cls        : str，名词的语义类别（"person"、"food"……）
        label      : str，这个名词短语的中心词依附到上层时的标签（nsubj、obj、obl、nmod、conj）
        rng        : random.Random
        allow_pp   : 是否允许带一个修饰名词的介词短语
        allow_coord: 是否允许并列（"the bread and the cheese"）

    【返回】Node（名词短语的中心词，即那个名词）
    """
    noun = rng.choice(LEXICON[cls])
    det = Node(rng.choice(["the", "a"]), "DET", "det")
    left = [det]
    for adj in rng.sample(ADJECTIVES, rng.choice([0, 0, 1, 1, 2])):     # 不重复地挑 0~2 个形容词
        left.append(Node(adj, "ADJ", "amod"))
    next_word = left[1].word if len(left) > 1 else noun
    if det.word == "a" and next_word[0] in "aeiou":                      # 元音开头的词前用 an
        det.word = "an"
    node = Node(noun, "NOUN", label, left=left)
    if allow_pp and rng.random() < 0.35:
        pp_cls = {"food": "topping", "person": "accessory", "thing": "place"}.get(cls)
        if pp_cls:
            prep = "near" if pp_cls == "place" else "with"
            node.right.append(make_pp(prep, pp_cls, "nmod", rng))
    if allow_coord and rng.random() < 0.15:
        # 并列放在介词短语之后："the student with a hat and a woman"
        conj = make_np(cls, "conj", rng, allow_pp=False)
        conj.left.insert(0, Node("and", "CCONJ", "cc"))
        node.right.append(conj)
    return node


def make_pp(prep, cls, label, rng):
    """
    生成一个介词短语。按 UD 的约定，介词作为 case 依附在名词上，名词才是短语的中心词。

    【返回】Node（介词短语里的名词，label 为 obl 或 nmod）
    """
    np = make_np(cls, label, rng, allow_pp=False)
    np.left.insert(0, Node(prep, "ADP", "case"))
    return np


def make_sentence(rng):
    """
    按语法随机生成一个句子的依存树，返回 Sentence。

    【思路】
        主动词是整棵树的根；主语、助动词、副词挂在它左边；宾语、介词短语、副词、句号挂在它右边。
    """
    verb = rng.choice(list(VERBS))
    obj_classes, takes_tool = VERBS[verb]
    root = Node(verb, "VERB", "root")
    root.left.append(make_np("person", "nsubj", rng))
    if rng.random() < 0.2:
        root.left.append(Node("will", "AUX", "aux"))
        root.word = {"ate": "eat", "cut": "cut", "cooked": "cook", "saw": "see", "watched": "watch",
                     "painted": "paint", "met": "meet", "bought": "buy"}[verb]
    if rng.random() < 0.15:
        root.left.append(Node(rng.choice(ADVERBS), "ADV", "advmod"))
    root.right.append(make_np(rng.choice(obj_classes), "obj", rng, allow_coord=True))
    r = rng.random()
    if takes_tool and r < 0.4:
        root.right.append(make_pp("with", "tool", "obl", rng))
    elif r < 0.6:
        root.right.append(make_pp("in", "place", "obl", rng))
    if rng.random() < 0.1:
        root.right.append(Node(rng.choice(ADVERBS), "ADV", "advmod"))
    root.right.append(Node(".", "PUNCT", "punct"))
    return linearize(root)


def make_treebank(n_train=3000, n_dev=500, n_test=500, seed=0):
    """
    生成训练集、开发集（验证集）、测试集。

    【参数】n_train、n_dev、n_test：句子数；seed：随机种子（讲义里的数字依赖默认值）
    【返回】(train, dev, test)，三个 list[Sentence]
    """
    rng = random.Random(seed)
    sents = [make_sentence(rng) for _ in range(n_train + n_dev + n_test)]
    return sents[:n_train], sents[n_train:n_train + n_dev], sents[n_train + n_dev:]


def to_conll(sent):
    """
    把句子格式化成类似 CoNLL-U 的表格文本（编号、词、词性、中心词、标签）。

    【返回】str
    """
    lines = [f"{'ID':>3s}  {'词':<12s}{'词性':<8s}{'中心词':>6s}  标签"]
    for i in range(1, len(sent.words)):
        h = sent.heads[i]
        head_word = sent.words[h] if h > 0 else "ROOT"
        lines.append(f"{i:>3d}  {sent.words[i]:<12s}{sent.tags[i]:<8s}{h:>6d}  {sent.labels[i]:<7s}（→ {head_word}）")
    return "\n".join(lines)


if __name__ == "__main__":
    train, dev, test = make_treebank()
    print(f"训练集 {len(train)} 句，开发集 {len(dev)} 句，测试集 {len(test)} 句")
    for s in train[:3]:
        print("\n" + s.text())
        print(to_conll(s))