"""
=============================================================================
第 4 讲 · 代码 1：依存树 —— 表示、画图、合法性、投射性、歧义有多少
=============================================================================

对应讲义：lectures/lecture04_dependency_parsing.md 第 2 ~ 4 节

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
依存句法把句子的结构表示成一棵树：每个词都依附于另一个词（它的**中心词**，head），
整句话最终依附于一个虚拟的 ROOT 节点。这个文件让你熟悉这种表示：

    演示 1：同一个句子框架，"with a knife" 和 "with cheese" 的依存树有什么不同？（终端里画出依存弧）
    演示 2：一棵合法的依存树要满足什么条件？什么是"投射性"？一个非投射的例子
    演示 3：介词短语附着的歧义有多少种？—— 随介词短语个数**指数增长**（卡塔兰数）
    演示 4：我们的迷你树库长什么样？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 01_dependency_trees.py

只用 Python 标准库。运行时间约 1 秒。
"""

from collections import Counter
from itertools import product

from treebank import ROOT, Sentence, make_treebank, to_conll


def print_header(title):
    """打印一个带分隔线的标题。"""
    print()
    print("=" * 70)
    print(f"  {title}")
    print("=" * 70)


# =============================================================================
# 第 1 部分：画依存弧
# =============================================================================
def draw_arcs(sent, show_labels=True):
    """
    在终端里把依存树画成"弧线图"：每条依存弧是一条从中心词到依存词的折线，箭头 ▼ 指向依存词。

    【思路】
        1. 每个词占一段宽度，记下它中心的列号。
        2. 给每条弧分配一个"高度"：它所覆盖的范围内，最高的弧再高一层（越长、嵌套越多的弧画得越高）。
        3. 从最高层往下逐行画：弧所在的那一行画横线（标签写在横线中间），更低的行在两端画竖线。

    【参数】
        sent        : Sentence
        show_labels : 是否在横线上写依存标签

    【返回】str

    【边界】
        - 很长的句子会超出终端宽度；标签太长写不下时会被省略。
        - 非投射树中弧会交叉，图里能看到横线穿过竖线（演示 2）。
    """
    words = ["ROOT"] + sent.words[1:]
    width = [max(len(w), 3) + 2 for w in words]
    centers, col = [], 0
    for w in width:
        centers.append(col + w // 2)
        col += w
    total = col
    arcs = [(min(h, d), max(h, d), h, d, sent.labels[d]) for d, h in enumerate(sent.heads) if d > 0]
    arcs.sort(key=lambda a: a[1] - a[0])
    # 高度：比所有"和它重叠"的、更短的弧高一层。两条弧的内部有公共部分才算重叠，
    # 只共享一个端点（比如 0—3 和 3—5）不算。嵌套的弧和交叉的弧都会被放到不同的高度。
    level = {}
    for a in arcs:
        overlapping = [level[b] for b in level if max(a[0], b[0]) < min(a[1], b[1])]
        level[a] = 1 + max(overlapping, default=0)
    top = max(level.values())
    rows = []
    for L in range(top, 0, -1):
        row = [" "] * total
        # 第 1 遍：更高层的弧在这一行留下的竖线
        for a, lv in level.items():
            if lv > L:
                for c in (centers[a[0]], centers[a[1]]):
                    row[c] = "│"
        # 第 2 遍：这一层的横线；端点处如果已经有线，就画成连接符
        mine = [a for a, lv in level.items() if lv == L]
        for a in mine:
            c1, c2 = centers[a[0]], centers[a[1]]
            for c in range(c1 + 1, c2):
                row[c] = "┼" if row[c] == "│" else "─"
            for c, alone, with_bar in [(c1, "┌", "├"), (c2, "┐", "┤")]:
                if row[c] == " ":
                    row[c] = alone
                elif row[c] == "│":
                    row[c] = with_bar
                else:
                    row[c] = "┬"
        # 第 3 遍：标签（只写在完全是横线的地方）
        if show_labels:
            for a in mine:
                c1, c2 = centers[a[0]], centers[a[1]]
                lab = a[4]
                st = (c1 + c2) // 2 - len(lab) // 2
                if st > c1 + 1 and st + len(lab) < c2 and all(ch == "─" for ch in row[st:st + len(lab)]):
                    row[st:st + len(lab)] = list(lab)
        rows.append("".join(row).rstrip())
    arrow = [" "] * total
    for a in arcs:
        arrow[centers[a[1]] if a[3] == a[1] else centers[a[0]]] = "▼"
    for a in arcs:                                   # 中心词一端画竖线（如果那里没有箭头）
        hc = centers[a[2]]
        if arrow[hc] == " ":
            arrow[hc] = "│"
    rows.append("".join(arrow).rstrip())
    rows.append("".join(w.center(width[i]) for i, w in enumerate(words)).rstrip())
    return "\n".join("    " + r for r in rows)


# =============================================================================
# 第 2 部分：合法性与投射性
# =============================================================================
def check_tree(heads):
    """
    检查一组中心词编号是否构成一棵合法的依存树。

    【思路】（讲义第 4.3 节）
        依存树的要求：
          (1) 每个词恰好有一个中心词（用列表表示 heads 时自动满足）；
          (2) 恰好有一个词依附于 ROOT；
          (3) 无环且连通：从任何一个词出发，沿着中心词一路往上走，最终都能到达 ROOT。

    【参数】heads : list[int]，heads[i] 是第 i 个词的中心词编号，heads[0] 不使用

    【返回】(ok, reason)：ok 为 bool，reason 说明不合法的原因
    """
    n = len(heads) - 1
    roots = [i for i in range(1, n + 1) if heads[i] == 0]
    if len(roots) != 1:
        return False, f"依附于 ROOT 的词有 {len(roots)} 个（应该恰好 1 个）"
    for i in range(1, n + 1):
        seen, j = set(), i
        while j != 0:
            if j in seen:
                return False, f"从第 {i} 个词出发会陷入环"
            seen.add(j)
            j = heads[j]
            if not 0 <= j <= n:
                return False, f"中心词编号 {j} 越界"
    return True, "合法"


def crossing_arcs(heads):
    """
    找出所有互相交叉的弧对。没有交叉弧的树叫**投射的**（projective）。

    【思路】
        弧 (a, b) 和弧 (c, d)（都写成"左端点 < 右端点"）交叉，当且仅当 a < c < b < d。
        直观地说：把所有弧画在句子上方，投射树可以画得一条线都不交叉。
        （从 ROOT 到根词的弧也要算上：ROOT 在最左边，编号 0。）

    【参数】heads : list[int]
    【返回】list[((a, b), (c, d))]，所有交叉的弧对
    """
    arcs = [(min(h, d), max(h, d)) for d, h in enumerate(heads) if d > 0]
    return [(x, y) for x in arcs for y in arcs if x[0] < y[0] < x[1] < y[1]]


def hearing_example():
    """
    构造一个经典的非投射句子（CS224N 课件中的例子）：
        a hearing is scheduled on the issue today .
    "on the issue" 修饰的是 hearing（"关于这个议题的听证会"），但它被移到了动词后面，
    于是 hearing → issue 这条弧跨过了动词 scheduled，和 ROOT → scheduled 的弧交叉。
    """
    words = [ROOT, "a", "hearing", "is", "scheduled", "on", "the", "issue", "today", "."]
    tags = [ROOT, "DET", "NOUN", "AUX", "VERB", "ADP", "DET", "NOUN", "NOUN", "PUNCT"]
    heads = [-1, 2, 4, 4, 0, 7, 7, 2, 4, 4]
    labels = [ROOT, "det", "nsubj", "aux", "root", "case", "det", "nmod", "obl", "punct"]
    return Sentence(words, tags, heads, labels)


# =============================================================================
# 第 3 部分：介词短语附着的歧义数
# =============================================================================
def count_pp_attachments(k):
    """
    数一数"动词 宾语 PP1 PP2 ... PPk"这种结构有多少种合法（投射）的介词短语附着方式。

    【思路】
        第 i 个介词短语可以修饰：动词、宾语名词、或者它前面任何一个介词短语里的名词。
        暴力枚举所有组合，只保留投射的（弧不交叉）。
        例如 k = 2："saw the man [with the telescope] [in the park]"：
            两个 PP 都修饰 saw；都修饰 man；第一个修饰 man、第二个修饰 saw；
            第一个修饰 saw、第二个修饰 man？—— 这会交叉，不合法；第二个修饰 telescope……

    【参数】k : int，介词短语的个数
    【返回】int

    【边界】
        - 枚举量是 2 × 3 × ... × (k+1) = (k+1)!，k 超过 9 左右就很慢了。
    """
    # 位置：0 = 动词，1 = 宾语，2..k+1 = 第 1..k 个介词短语的名词
    count = 0
    for choice in product(*[range(i + 1) for i in range(1, k + 1)]):
        # choice[i-1] 是第 i 个 PP（位置 i+1）的中心词位置，取值 0..i
        arcs = [(0, 1)] + [(h, i + 1) for i, h in enumerate(choice, start=1)]
        spans = [(min(a, b), max(a, b)) for a, b in arcs]
        if not any(x[0] < y[0] < x[1] < y[1] for x in spans for y in spans):
            count += 1
    return count


def catalan(n):
    """
    第 n 个卡塔兰数 C_n = (2n)! / ((n+1)! n!)：1, 1, 2, 5, 14, 42, 132, 429, ...

    【NLP 联系】
        卡塔兰数是"n 个元素有多少种不同的二叉括号方式"，也是很多树结构的计数。
        它大约以 4^n 的速度增长 —— 歧义的数量随句子长度**指数爆炸**，所以分析器不可能枚举所有可能。
    """
    c = 1
    for i in range(n):
        c = c * 2 * (2 * i + 1) // (i + 2)
    return c


# =============================================================================
# 演示
# =============================================================================
def demo_two_trees():
    """演示 1：with a knife vs with cheese。"""
    print_header("演示 1：同样的词序，不同的结构")
    s1 = Sentence([ROOT, "the", "chef", "cut", "the", "bread", "with", "a", "knife", "."],
                  [ROOT, "DET", "NOUN", "VERB", "DET", "NOUN", "ADP", "DET", "NOUN", "PUNCT"],
                  [-1, 2, 3, 0, 5, 3, 8, 8, 3, 3],
                  [ROOT, "det", "nsubj", "root", "det", "obj", "case", "det", "obl", "punct"])
    s2 = Sentence([ROOT, "the", "chef", "cut", "the", "bread", "with", "the", "cheese", "."],
                  [ROOT, "DET", "NOUN", "VERB", "DET", "NOUN", "ADP", "DET", "NOUN", "PUNCT"],
                  [-1, 2, 3, 0, 5, 3, 8, 8, 5, 3],
                  [ROOT, "det", "nsubj", "root", "det", "obj", "case", "det", "nmod", "punct"])
    for s, note in [(s1, "knife 是切的工具 → 修饰动词 cut（obl）"), (s2, "cheese 是面包上的东西 → 修饰名词 bread（nmod）")]:
        print(f"\n  {s.text()}    （{note}）\n")
        print(draw_arcs(s))
        print()
        print("    " + to_conll(s).replace("\n", "\n    "))
    print("""
  观察：
    - 每个词恰好有一个中心词（箭头指向它），主动词 cut 依附于 ROOT。
    - 介词 with 依附在它后面的名词上（UD 的约定：实词做中心词，介词是 case 标记）。
    - 两句话只差一个名词，结构却不同。分析器要分对，就必须知道 knife 是工具、cheese 是食物配料。""")


def demo_well_formed():
    """演示 2：合法性与投射性。"""
    print_header("演示 2：合法的依存树、投射性")
    print("\n  几组中心词编号（第 0 位不用）：")
    for heads, note in [([-1, 2, 0, 2], "I ate fish：I→ate，ate→ROOT，fish→ate"),
                        ([-1, 2, 0, 0], "两个词都依附 ROOT"),
                        ([-1, 2, 3, 1], "1→2→3→1 形成环")]:
        ok, reason = check_tree(heads)
        print(f"    heads = {heads[1:]}  {note:<32s} → {reason}")

    s = hearing_example()
    print(f"\n  一个非投射的句子：{s.text()}\n")
    print(draw_arcs(s))
    ok, _ = check_tree(s.heads)
    print(f"\n  是合法的树吗？{ok}；交叉的弧对：{crossing_arcs(s.heads)}")
    print("""
  观察：hearing → issue 这条弧（2—7）和 ROOT → scheduled（0—4）交叉了，图上能看到横线穿过竖线。
    英语里非投射结构比较少（主要来自这种"移位"），但在德语、荷兰语、捷克语等语序较自由的语言中更常见。
    第 6 节的 arc-standard 转移系统**只能**产生投射树（代码 02 演示 4 会看到它在这个句子上失败）。""")


def demo_ambiguity():
    """演示 3：介词短语附着的歧义数。"""
    print_header("演示 3：介词短语附着 —— 歧义随长度指数增长")
    print("\n  '动词 宾语 PP1 ... PPk' 中，合法（投射）的附着方式有多少种：")
    print(f"    {'k':>3s}{'暴力枚举':>10s}{'卡塔兰数 C(k+1)':>18s}")
    for k in range(1, 9):
        n = count_pp_attachments(k)
        print(f"    {k:>3d}{n:>10d}{catalan(k + 1):>18d}")
        assert n == catalan(k + 1)
    print("""
  观察：暴力枚举的结果恰好是卡塔兰数。CS224N 课件里有一句真实的新闻标题，含 4 个介词短语：
    "The board approved [its acquisition] [by Royal Trustco Ltd.] [of Toronto] [for $27 a share] [at its monthly meeting]"
  按上面的计算，光是介词短语的附着就有 C(5) = 42 种可能（这句的结构比我们的模板复杂，但数量级一致）；
  人一眼就能看出唯一合理的那种，靠的是世界知识。分析器不可能枚举所有可能，必须**逐步做出局部决策**。""")


def demo_treebank_stats():
    """演示 4：迷你树库的统计。"""
    print_header("演示 4：迷你树库")
    train, dev, test = make_treebank()
    lengths = [len(s) for s in train]
    labels = Counter(l for s in train for l in s.labels[1:])
    print(f"\n  训练集 {len(train)} 句，开发集 {len(dev)} 句，测试集 {len(test)} 句")
    print(f"  句长：平均 {sum(lengths) / len(lengths):.1f} 个词，最短 {min(lengths)}，最长 {max(lengths)}")
    print("  依存标签分布：" + "，".join(f"{k} {v}" for k, v in labels.most_common()))
    pp = Counter()
    for s in train:
        for i in range(1, len(s.words)):
            if s.tags[i] == "ADP":
                noun = s.heads[i]
                pp[(s.words[i], s.labels[noun])] += 1
    print("  介词短语的附着：" + "，".join(f"{p}→{lab} {c}" for (p, lab), c in sorted(pp.items())))
    n_proj = sum(not crossing_arcs(s.heads) for s in train)
    n_ok = sum(check_tree(s.heads)[0] for s in train)
    print(f"  合法的树 {n_ok}/{len(train)}，投射的树 {n_proj}/{len(train)}")
    print("\n  一个例子：")
    print(f"    {train[1].text()}\n")
    print(draw_arcs(train[1]))


if __name__ == "__main__":
    demo_two_trees()
    demo_well_formed()
    demo_ambiguity()
    demo_treebank_stats()