"""
=============================================================================
第 4 讲 · 工具模块：arc-standard 转移系统、oracle、特征抽取、批量分析、评估
=============================================================================

对应讲义：lectures/lecture04_dependency_parsing.md 第 6 ~ 9 节
被 02、03、04 号文件共同使用。

-----------------------------------------------------------------------------
【arc-standard 转移系统】（讲义第 6 节）
-----------------------------------------------------------------------------
分析器的状态（configuration）由三部分组成：
    栈 stack     ：已经处理过、但还在等待找到全部依存成分的词。一开始只有 ROOT。
    缓冲区 buffer：还没处理的词。一开始是整个句子。
    弧集合 arcs  ：已经确定的依存关系 (中心词, 依存词, 标签)。
三种动作（transition）：
    SHIFT       ：把缓冲区的第一个词移到栈顶。
    LEFT-ARC(l) ：栈顶第二个词 s2 依附于栈顶 s1（s1 → s2，标签 l），把 s2 从栈中移除。
    RIGHT-ARC(l)：栈顶 s1 依附于栈顶第二个词 s2（s2 → s1，标签 l），把 s1 从栈中移除。
当缓冲区为空、栈里只剩 ROOT 时，分析结束。一个 n 个词的句子恰好需要 2n 步
（每个词被 SHIFT 一次、被某个 ARC 移除一次）。

这和 CS224N 作业 2 的设定一致（作业里不带标签；我们带标签，以便计算 LAS）。

【本模块提供】
    Configuration                    ：分析器状态与三种动作
    oracle(config, sent)             ：静态 oracle —— 已知标准答案时，下一步应该做什么
    gold_transitions(sent)           ：一个句子的完整标准动作序列
    Vocab / extract_features         ：Chen & Manning（2014）的 18 个词、18 个词性、12 个标签特征
    legal_mask                       ：当前状态下哪些动作是合法的
    parse_batch(sents, score_fn)     ：用任意打分函数批量地贪心分析（作业 2 的 minibatch_parse）
    attachment_scores(gold, pred)    ：UAS、LAS
"""

from dataclasses import dataclass, field

import numpy as np

from treebank import ROOT

NULL, UNK = "<NULL>", "<UNK>"


# =============================================================================
# 第 1 部分：分析器状态与动作
# =============================================================================
@dataclass
class Configuration:
    """
    arc-standard 分析器的一个状态。

    【字段】
        stack  : list[int]，词的编号（0 = ROOT），列表末尾是栈顶
        buffer : list[int]，还没处理的词的编号，列表开头是缓冲区的第一个词
        heads  : dict[int, int]，已确定的 依存词 → 中心词
        labels : dict[int, str]，已确定的 依存词 → 标签

    【用法】
        c = Configuration.initial(n_words)
        c.apply("S"); c.apply("L-nsubj"); ...
    """
    stack: list = field(default_factory=list)
    buffer: list = field(default_factory=list)
    heads: dict = field(default_factory=dict)
    labels: dict = field(default_factory=dict)

    @staticmethod
    def initial(n):
        """一个 n 个词的句子的初始状态：栈 = [ROOT]，缓冲区 = [1, ..., n]。"""
        return Configuration(stack=[0], buffer=list(range(1, n + 1)))

    def is_terminal(self):
        """缓冲区为空且栈里只剩 ROOT 时，分析结束。"""
        return not self.buffer and len(self.stack) == 1

    def can_apply(self, t):
        """
        判断动作 t 在当前状态下是否合法。

        【规则】
            SHIFT：缓冲区非空。
            LEFT-ARC：栈里至少有 2 个元素，且栈顶第二个不是 ROOT（ROOT 不能依附于别人）。
            RIGHT-ARC：栈里至少有 2 个元素；如果栈顶第二个是 ROOT，则要求缓冲区为空
                       （保证 ROOT 恰好只有一个依存词——句子的主动词——而且是最后才连上）。
        """
        if t == "S":
            return bool(self.buffer)
        if len(self.stack) < 2:
            return False
        if t.startswith("L-"):
            return self.stack[-2] != 0
        return self.stack[-2] != 0 or not self.buffer

    def apply(self, t):
        """
        执行动作 t（"S"、"L-标签" 或 "R-标签"），原地修改状态。

        【边界】
            - 不合法的动作会抛出 ValueError。
        """
        if not self.can_apply(t):
            raise ValueError(f"动作 {t} 在当前状态下不合法：栈 {self.stack}，缓冲区 {self.buffer}")
        if t == "S":
            self.stack.append(self.buffer.pop(0))
        elif t.startswith("L-"):
            dep = self.stack.pop(-2)
            self.heads[dep], self.labels[dep] = self.stack[-1], t[2:]
        else:
            dep = self.stack.pop()
            self.heads[dep], self.labels[dep] = self.stack[-1], t[2:]

    # ---------- 供特征抽取使用：某个词已经连上的最左 / 最右依存成分 ----------
    def left_children(self, h):
        """词 h 已经连上的、位于它左边的依存成分，按从左到右排序。"""
        return sorted(d for d, hh in self.heads.items() if hh == h and d < h)

    def right_children(self, h):
        """词 h 已经连上的、位于它右边的依存成分，按从右到左排序（最右的在前）。"""
        return sorted((d for d, hh in self.heads.items() if hh == h and d > h), reverse=True)


# =============================================================================
# 第 2 部分：静态 oracle
# =============================================================================
def oracle(config, sent):
    """
    已知标准依存树时，告诉我们当前状态下"正确的下一步"。

    【思路】（讲义第 6.3 节）
        设栈顶为 s1、第二个为 s2：
        1. 如果标准答案里 s2 依附于 s1 → LEFT-ARC；
        2. 否则如果 s1 依附于 s2，**并且 s1 的所有依存成分都已经连上了** → RIGHT-ARC；
           （第二个条件很重要：RIGHT-ARC 会把 s1 从栈中移除，之后就再也不能给它连依存成分了。
             比如 "ate the pizza with cheese"：栈顶是 pizza、第二个是 ate 时，虽然 pizza 确实依附于 ate，
             但 cheese 还没连到 pizza 上，此时 RIGHT-ARC 就再也连不上了，所以要先 SHIFT。
             而 "ate the pizza with a fork" 中 pizza 没有别的依存成分，可以立刻 RIGHT-ARC。）
        3. 否则 → SHIFT。

    【参数】
        config : Configuration
        sent   : Sentence（带标准答案）

    【返回】str，动作

    【边界】
        - 只对**投射**的树保证能得到正确结果。非投射树会在某一步找不到合法的正确动作（代码 02 演示 4）。
    """
    if len(config.stack) >= 2:
        s1, s2 = config.stack[-1], config.stack[-2]
        if s2 != 0 and sent.heads[s2] == s1:
            return "L-" + sent.labels[s2]
        if sent.heads[s1] == s2:
            pending = any(sent.heads[w] == s1 and w not in config.heads for w in range(1, len(sent.words)))
            if not pending and config.can_apply("R-" + sent.labels[s1]):
                return "R-" + sent.labels[s1]
    return "S"


def gold_transitions(sent):
    """
    用 oracle 生成一个句子从头到尾的标准动作序列，以及每一步之前的状态。

    【返回】list[(Configuration 的快照, 动作)]

    【Python 小知识】
        Configuration(list(c.stack), ...) 复制一份状态：直接存 c 的话，后面 apply 会把存下来的也改掉。
    """
    c = Configuration.initial(len(sent))
    steps = []
    while not c.is_terminal():
        t = oracle(c, sent)
        steps.append((Configuration(list(c.stack), list(c.buffer), dict(c.heads), dict(c.labels)), t))
        c.apply(t)
    return steps


# =============================================================================
# 第 3 部分：特征（Chen & Manning 2014）
# =============================================================================
class Vocab:
    """
    把词、词性、依存标签、动作映射成编号。

    【字段】
        word2id、tag2id、label2id : dict，都包含 <NULL>（空位置）和 <UNK>（没见过）；词和词性还包含 <ROOT>
        transitions               : list[str]，所有可能的动作，"S" 在第 0 位
    """

    def __init__(self, sents):
        def build(items, extra):
            d = {x: i for i, x in enumerate(extra)}
            for x in sorted(set(items) - set(extra)):
                d[x] = len(d)
            return d
        self.word2id = build((w for s in sents for w in s.words[1:]), [NULL, UNK, ROOT])
        self.tag2id = build((t for s in sents for t in s.tags[1:]), [NULL, UNK, ROOT])
        labels = sorted(set(l for s in sents for l in s.labels[1:]))
        self.label2id = build(labels, [NULL, UNK])
        self.transitions = ["S"] + [f"L-{l}" for l in labels] + [f"R-{l}" for l in labels]
        self.t2id = {t: i for i, t in enumerate(self.transitions)}


def feature_positions(c):
    """
    Chen & Manning 的 18 个"位置"：栈顶 3 个词、缓冲区前 3 个词，以及栈顶两个词的若干个依存成分。

    【返回】
        (positions18, label_positions12)
            positions18      : list[int]，18 个词的编号（-1 表示该位置为空）
            label_positions12: list[int]，其中后 12 个位置（依存成分）—— 它们还有依存标签可用

    【18 个位置】（讲义第 8.2 节）
        s1, s2, s3, b1, b2, b3,
        对 i = 1, 2：lc1(si), rc1(si), lc2(si), rc2(si), lc1(lc1(si)), rc1(rc1(si))
        其中 lc1 = 最左的依存成分，lc2 = 第二左的，rc1 = 最右的，rc2 = 第二右的。
    """
    s = [c.stack[-i] if len(c.stack) >= i else -1 for i in (1, 2, 3)]
    b = [c.buffer[i] if len(c.buffer) > i else -1 for i in (0, 1, 2)]
    kids = []
    for si in s[:2]:
        if si < 0:
            kids += [-1] * 6
            continue
        lc, rc = c.left_children(si), c.right_children(si)
        lc1 = lc[0] if lc else -1
        rc1 = rc[0] if rc else -1
        lc2 = lc[1] if len(lc) > 1 else -1
        rc2 = rc[1] if len(rc) > 1 else -1
        llc = c.left_children(lc1)[0] if lc1 > 0 and c.left_children(lc1) else -1
        rrc = c.right_children(rc1)[0] if rc1 > 0 and c.right_children(rc1) else -1
        kids += [lc1, rc1, lc2, rc2, llc, rrc]
    return s + b + kids, kids


def extract_features(c, sent, vocab):
    """
    把一个状态变成 48 个整数：18 个词、18 个词性、12 个依存标签的编号。

    【参数】
        c     : Configuration
        sent  : Sentence（只用它的 words 和 tags；**不**看标准答案）
        vocab : Vocab

    【返回】list[int]，长度 48

    【边界】
        - 空位置用 <NULL>；没见过的词用 <UNK>；ROOT 用 <ROOT>。
        - 依存标签取的是**分析器已经预测出来的**标签（c.labels），而不是标准答案。
    """
    pos18, kids12 = feature_positions(c)
    words = [vocab.word2id[NULL] if p < 0 else vocab.word2id.get(sent.words[p], vocab.word2id[UNK]) for p in pos18]
    tags = [vocab.tag2id[NULL] if p < 0 else vocab.tag2id.get(sent.tags[p], vocab.tag2id[UNK]) for p in pos18]
    labels = [vocab.label2id[NULL] if p < 0 else vocab.label2id.get(c.labels.get(p, NULL), vocab.label2id[UNK]) for p in kids12]
    return words + tags + labels


def legal_mask(c, vocab):
    """返回一个布尔数组：每个动作在状态 c 下是否合法。"""
    return np.array([c.can_apply(t) for t in vocab.transitions])


# =============================================================================
# 第 4 部分：批量贪心分析与评估
# =============================================================================
def parse_batch(sents, score_fn, vocab, batch_size=256):
    """
    用一个打分函数，批量地贪心分析一组句子（作业 2 中 minibatch_parse 的思路）。

    【思路】
        同时维护很多个句子的分析状态。每一轮：
            把所有还没分析完的状态的特征收集起来，一次性交给 score_fn 打分（神经网络批量计算更快），
            每个状态选"合法动作中得分最高的"执行；分析完的句子移出。

    【参数】
        sents      : list[Sentence]
        score_fn   : 函数，输入 (configs, sents_of_configs) 两个列表，返回 np.ndarray (B, 动作数) 的得分
        vocab      : Vocab
        batch_size : 每次同时分析多少个句子

    【返回】
        list[(heads, labels)]：每个句子预测出的中心词编号列表和标签列表（第 0 位留给 ROOT）

    【边界】
        - 贪心：每一步只选当前得分最高的动作，一旦选错就无法回头（讲义第 7.2 节）。
    """
    results = [None] * len(sents)
    for start in range(0, len(sents), batch_size):
        idx = list(range(start, min(start + batch_size, len(sents))))
        configs = {i: Configuration.initial(len(sents[i])) for i in idx}
        active = list(idx)
        while active:
            scores = score_fn([configs[i] for i in active], [sents[i] for i in active])
            for row, i in enumerate(active):
                mask = legal_mask(configs[i], vocab)
                s = np.where(mask, scores[row], -np.inf)
                configs[i].apply(vocab.transitions[int(np.argmax(s))])
            active = [i for i in active if not configs[i].is_terminal()]
        for i in idx:
            n = len(sents[i])
            results[i] = ([-1] + [configs[i].heads[d] for d in range(1, n + 1)],
                          [ROOT] + [configs[i].labels[d] for d in range(1, n + 1)])
    return results


def attachment_scores(sents, preds, skip_punct=False):
    """
    计算 UAS（无标签依存正确率）和 LAS（带标签依存正确率）。

    【思路】（讲义第 9 节）
        UAS = 中心词预测正确的词所占的比例；
        LAS = 中心词**和**标签都预测正确的词所占的比例。LAS <= UAS。

    【参数】
        sents      : list[Sentence]（标准答案）
        preds      : parse_batch 的输出
        skip_punct : 是否不统计标点（很多论文的英文评估不算标点）

    【返回】(uas, las)
    """
    total = uas = las = 0
    for s, (heads, labels) in zip(sents, preds):
        for i in range(1, len(s.words)):
            if skip_punct and s.tags[i] == "PUNCT":
                continue
            total += 1
            if heads[i] == s.heads[i]:
                uas += 1
                las += labels[i] == s.labels[i]
    return uas / total, las / total


def pp_attachment_accuracy(sents, preds):
    """
    只看介词短语里的名词（它的标准标签是 obl 或 nmod）：中心词预测对了的比例。

    【返回】(准确率, 个数)
    """
    total = right = 0
    for s, (heads, _) in zip(sents, preds):
        for i in range(1, len(s.words)):
            if s.labels[i] in ("obl", "nmod"):
                total += 1
                right += heads[i] == s.heads[i]
    return right / max(total, 1), total


def pp_accuracy_by_visibility(sents, preds):
    """
    把介词短语按"做附着决策时，介词后面的名词在不在特征窗口里"分成两组，分别计算附着准确率。

    【思路】
        做决策时（栈顶是宾语名词、缓冲区第一个是介词），特征只能看到缓冲区的前 3 个词 b1、b2、b3。
            "with the cheese"       → b1 = with，b2 = the，b3 = cheese：名词**看得见**
            "with the fresh cheese" → b3 = fresh，cheese 在 b4：名词**看不见**
        看不见名词时，分析器只能瞎猜（或者依赖一些虚假的相关性）。

    【返回】dict：{"看得见": (准确率, 个数), "看不见": (准确率, 个数)}
    """
    stats = {"看得见": [0, 0], "看不见": [0, 0]}
    for s, (heads, _) in zip(sents, preds):
        for i in range(1, len(s.words)):
            if s.labels[i] not in ("obl", "nmod"):
                continue
            prep = min(d for d in range(1, len(s.words)) if s.heads[d] == i and s.labels[d] == "case")
            key = "看得见" if i - prep <= 2 else "看不见"
            stats[key][0] += heads[i] == s.heads[i]
            stats[key][1] += 1
    return {k: (r / max(n, 1), n) for k, (r, n) in stats.items()}