"""
=============================================================================
第 4 讲 · 代码 2：基于转移的依存分析 —— arc-standard 系统与 oracle
=============================================================================

对应讲义：lectures/lecture04_dependency_parsing.md 第 6、7 节
对应作业：CS224N 作业 2 的 PartialParse / minibatch_parse

-----------------------------------------------------------------------------
【这个文件要解决什么问题？】
-----------------------------------------------------------------------------
第 4 讲的核心想法：不要一次性"猜"出整棵树（可能的树太多了，见代码 01 演示 3），
而是像人读句子一样**从左到右**，用一个栈和一个缓冲区，每一步做一个小决定：
    SHIFT（读入下一个词）、LEFT-ARC（栈顶第二个词依附于栈顶）、RIGHT-ARC（栈顶依附于栈顶第二个词）。
这样，"分析一个句子"就变成了"做一连串分类"——分类正是第 3 讲学过的东西！

    演示 1：一步一步分析 "I ate fish"（CS224N 课件的例子）
    演示 2：介词短语附着在转移系统里是怎样被决定的？（with cheese vs with a fork）
    演示 3：在整个树库上验证 oracle：每棵树都能被精确地重建，且恰好需要 2n 步
    演示 4：非投射的句子 —— arc-standard 无能为力
    演示 5：批量分析的框架 + 几个"不学习"的基线

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python 02_arc_standard.py

需要 numpy。运行时间约 3 秒。
"""

import importlib
import random
from collections import Counter

import numpy as np

from parsing import (Configuration, Vocab, attachment_scores, gold_transitions, oracle, parse_batch,
                     pp_attachment_accuracy)
from treebank import ROOT, Sentence, make_treebank

trees = importlib.import_module("01_dependency_trees")


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


def show(ids, sent):
    """把一组词编号显示成词：[0, 2] → [ROOT, ate]。"""
    return "[" + ", ".join("ROOT" if i == 0 else sent.words[i] for i in ids) + "]"


def trace(sent, highlight=None):
    """
    用 oracle 分析一个句子，打印每一步的栈、缓冲区、动作、新增的弧。

    【参数】
        sent      : Sentence（带标准答案）
        highlight : None 或 一个函数 f(config) → bool，为真时在这一行后面加一个标记
    """
    print(f"    {'步':>2s}  {'栈':<40s}{'缓冲区':<46s}{'动作':<18s}新增的弧")
    c = Configuration.initial(len(sent))
    step = 0
    print(f"    {step:>2d}  {show(c.stack, sent):<40s}{show(c.buffer, sent):<46s}")
    while not c.is_terminal():
        t = oracle(c, sent)
        mark = "  ← 关键决策" if highlight and highlight(c) else ""
        before = set(c.heads)
        c.apply(t)
        step += 1
        new = [d for d in c.heads if d not in before]
        arc = ""
        if new:
            d = new[0]
            h = c.heads[d]
            arc = f"{'ROOT' if h == 0 else sent.words[h]} →{c.labels[d]}→ {sent.words[d]}"
        name = {"S": "SHIFT"}.get(t, ("LEFT-ARC" if t[0] == "L" else "RIGHT-ARC") + f"({t[2:]})")
        print(f"    {step:>2d}  {show(c.stack, sent):<40s}{show(c.buffer, sent):<46s}{name:<18s}{arc}{mark}")


# =============================================================================
# 演示
# =============================================================================
def demo_i_ate_fish():
    """演示 1：I ate fish。"""
    print_header("演示 1：一步一步分析 'I ate fish'")
    sent = Sentence([ROOT, "I", "ate", "fish"], [ROOT, "PRON", "VERB", "NOUN"], [-1, 2, 0, 2],
                    [ROOT, "nsubj", "root", "obj"])
    print()
    trace(sent)
    print("""
  观察：3 个词，恰好 6 步（每个词被 SHIFT 一次、被某个 ARC 从栈里移除一次）。
    LEFT-ARC 处理"依存词在左边"的情况（I ← ate），RIGHT-ARC 处理"依存词在右边"的情况（ate → fish）。
    最后 ROOT → ate 用 RIGHT-ARC 连上，栈里只剩 ROOT，缓冲区为空 —— 分析完成。""")


def demo_pp_attachment():
    """演示 2：with cheese vs with a fork。"""
    print_header("演示 2：介词短语附着在转移系统中的决策点")

    def make(pp_noun, attach_to_verb):
        words = [ROOT, "she", "ate", "the", "pizza", "with", "the", pp_noun, "."]
        tags = [ROOT, "PRON", "VERB", "DET", "NOUN", "ADP", "DET", "NOUN", "PUNCT"]
        heads = [-1, 2, 0, 4, 2, 7, 7, 2 if attach_to_verb else 4, 2]
        labels = [ROOT, "nsubj", "root", "det", "obj", "case", "det", "obl" if attach_to_verb else "nmod", "punct"]
        return Sentence(words, tags, heads, labels)

    for noun, to_verb in [("fork", True), ("cheese", False)]:
        s = make(noun, to_verb)
        print(f"\n  {s.text()}（{noun} 修饰 {'ate' if to_verb else 'pizza'}）")
        trace(s, highlight=lambda c, s=s: len(c.stack) >= 2 and s.words[c.stack[-1]] == "pizza" and c.stack[-2] == 2)
    print("""
  观察：关键决策发生在"栈顶是 pizza、第二个是 ate"的那一步：
    - with a fork：pizza 已经没有别的依存成分了，立刻 RIGHT-ARC（ate → pizza），把 pizza 弹出栈；
      之后 fork 只能依附于 ate。
    - with cheese：cheese 还要依附到 pizza 上，所以先 SHIFT，把 pizza 留在栈里。
  而做决策的那一刻，cheese / fork 还在缓冲区里（b2、b3 的位置）。分析器必须"往前看"几个词，
  并且知道 fork 是工具、cheese 是配料 —— 这正是代码 03、04 中分类器要学的东西。""")


def demo_verify_oracle():
    """演示 3：oracle 在整个树库上都能精确重建标准树。"""
    print_header("演示 3：在整个树库上验证 oracle")
    train, dev, test = make_treebank()
    n_steps = Counter()
    total_words = total_steps = 0
    for s in train + dev + test:
        steps = gold_transitions(s)
        c = Configuration.initial(len(s))
        for _, t in steps:
            c.apply(t)
            n_steps["SHIFT" if t == "S" else ("LEFT-ARC" if t[0] == "L" else "RIGHT-ARC")] += 1
        assert all(c.heads[d] == s.heads[d] and c.labels[d] == s.labels[d] for d in range(1, len(s.words)))
        assert len(steps) == 2 * len(s)
        total_words += len(s)
        total_steps += len(steps)
    print(f"\n  {len(train) + len(dev) + len(test)} 个句子全部被 oracle 精确重建（中心词和标签都对），"
          f"每个句子恰好 2n 步。")
    print(f"  共 {total_words} 个词、{total_steps} 步：" + "，".join(f"{k} {v}" for k, v in n_steps.items()))
    vocab = Vocab(train)
    print(f"  动作的种类：1 个 SHIFT + {len(vocab.label2id) - 2} 个 LEFT-ARC(标签) + {len(vocab.label2id) - 2} 个 RIGHT-ARC(标签)"
          f" = {len(vocab.transitions)} 类")
    print("""
  观察：SHIFT 的次数恰好等于词数（每个词进栈一次），LEFT-ARC + RIGHT-ARC 的次数也等于词数（每个词被连一次）。
    这意味着：用 oracle 把每个训练句子"走"一遍，就得到了 2n 个 (状态, 正确动作) 的训练样本。
    训练分析器 = 训练一个分类器，输入状态、输出动作（代码 03、04）。""")


def demo_nonprojective():
    """演示 4：非投射句子。"""
    print_header("演示 4：非投射的句子，arc-standard 无能为力")
    s = trees.hearing_example()
    print(f"\n  {s.text()}（hearing → issue 跨过了 scheduled）\n")
    c = Configuration.initial(len(s))
    while not c.is_terminal():
        t = oracle(c, s)
        if not c.can_apply(t):
            print(f"  卡住了！栈 {show(c.stack, s)}，缓冲区 {show(c.buffer, s)}，oracle 想做 {t} 但不合法。")
            print("  （为了结束分析，下面强制用 RIGHT-ARC 把栈里剩下的词连起来。）")
            while len(c.stack) > 1:
                c.apply("R-dep")
            break
        c.apply(t)
    wrong = [(s.words[d], s.words[s.heads[d]] if s.heads[d] else "ROOT", s.words[c.heads[d]] if c.heads[d] else "ROOT")
             for d in range(1, len(s.words)) if c.heads[d] != s.heads[d]]
    print(f"  中心词错误的词：{[(w, f'应为 {g}', f'得到 {p}') for w, g, p in wrong]}")
    print("""
  原因：要连 hearing → issue，hearing 和 issue 必须同时在栈顶附近；
    但 scheduled 在它们之间，而且 scheduled 必须等 today 连上之后才能被处理……任何动作顺序都做不到。
    arc-standard 只能产生投射树。处理非投射结构的办法有：加一个 SWAP 动作（Nivre 2009）、
    先投射化再还原的伪投射方法，或者用基于图的分析器（讲义第 5 节）。""")


def demo_batch_and_baselines():
    """演示 5：批量分析框架与基线。"""
    print_header("演示 5：批量分析框架 + 几个'不学习'的基线（测试集）")
    train, dev, test = make_treebank()
    vocab = Vocab(train)

    def oracle_scorer(configs, sents):
        """作弊的打分函数：oracle 的动作得 1 分，其他 0 分。用来检查批量分析的框架没写错。"""
        out = np.zeros((len(configs), len(vocab.transitions)))
        for i, (c, s) in enumerate(zip(configs, sents)):
            out[i, vocab.t2id[oracle(c, s)]] = 1.0
        return out

    rng = np.random.default_rng(0)

    def random_scorer(configs, sents):
        """随机打分：在合法动作里随便选一个。"""
        return rng.random((len(configs), len(vocab.transitions)))

    def left_neighbor_scorer(configs, sents):
        """
        一个简单的规则：尽量 SHIFT，缓冲区空了就一直 RIGHT-ARC。
        结果是每个词都依附于它左边的那个词（一条向右的"链"）。
        """
        out = np.zeros((len(configs), len(vocab.transitions)))
        for i, c in enumerate(configs):
            out[i, 0 if c.buffer else vocab.t2id["R-punct"]] = 1.0
        return out

    def right_neighbor_scorer(configs, sents):
        """
        另一个简单规则：栈里有两个真词就 LEFT-ARC，否则 SHIFT，最后 RIGHT-ARC 连到 ROOT。
        结果是每个词都依附于它右边的那个词，最后一个词依附于 ROOT。
        """
        out = np.zeros((len(configs), len(vocab.transitions)))
        for i, c in enumerate(configs):
            if len(c.stack) >= 2 and c.stack[-2] != 0:
                out[i, vocab.t2id["L-det"]] = 1.0
            elif c.buffer:
                out[i, 0] = 1.0
            else:
                out[i, vocab.t2id["R-root"]] = 1.0
        return out

    print()
    for name, fn in [("作弊：用 oracle 打分", oracle_scorer), ("随机选合法动作", random_scorer),
                     ("每个词依附于左边的词", left_neighbor_scorer), ("每个词依附于右边的词", right_neighbor_scorer)]:
        preds = parse_batch(test, fn, vocab)
        uas, las = attachment_scores(test, preds)
        pp, n = pp_attachment_accuracy(test, preds)
        print(f"    {name:<22s}UAS {uas:.3f}  LAS {las:.3f}  介词短语附着准确率 {pp:.3f}（{n} 个）")
    print("""
  观察：用 oracle 打分时 UAS = LAS = 1.000，说明批量分析的框架是对的；随机和简单规则都很差。
    "依附于左边的词"一个都不对，"依附于右边的词"却对了不少：在我们的语法里（英语也大致如此），
    限定词、形容词、介词、主语的中心词都在**右边**，而依附于左边的成分（宾语、介词短语、标点）
    往往和中心词之间隔着别的词。下面两个文件用真正学出来的分类器来打分：
    代码 03 用人工设计的稀疏特征 + 感知机，代码 04 用 Chen & Manning 的神经网络。""")


if __name__ == "__main__":
    random.seed(0)
    demo_i_ate_fish()
    demo_pp_attachment()
    demo_verify_oracle()
    demo_nonprojective()
    demo_batch_and_baselines()