"""第 18 章：加法平滑的 n-gram 语言模型；只用标准库。"""

from collections import Counter, defaultdict
from math import exp, log
from random import Random


TRAIN = [
    "我爱喝红茶。", "你爱喝绿茶。", "他爱喝红茶。",
    "我想喝绿茶。", "你想喝红茶。", "他想喝绿茶。",
    "我买了红茶。", "你买了绿茶。", "他买了红茶。",
]
VALIDATION = ["我爱喝绿茶。", "你想喝绿茶。", "他买了绿茶。"]
TEST = ["他爱喝绿茶。", "我想喝红茶。", "你买了红茶。"]
BOS, EOS, UNK = "<BOS>", "<EOS>", "<UNK>"


class NGram:
    def __init__(self, texts, order, alpha=0.1):
        if order < 1 or alpha <= 0:
            raise ValueError("order 至少为 1，alpha 必须大于 0")
        self.order, self.alpha = order, alpha
        self.vocab = sorted({char for text in texts for char in text})+[EOS, UNK]
        self.counts = defaultdict(Counter)
        for text in texts:
            for context, target in self.events(text):
                self.counts[context][target] += 1

    def context(self, history):
        return tuple(history[-(self.order-1):]) if self.order > 1 else ()

    def events(self, text):
        history = [BOS]*(self.order-1)
        for target in list(text)+[EOS]:
            target = target if target in self.vocab else UNK
            yield self.context(history), target
            history.append(target)

    def distribution(self, context):
        counts = self.counts.get(context, {})
        total = sum(counts.values())+self.alpha*len(self.vocab)
        return [(counts.get(token, 0)+self.alpha)/total for token in self.vocab]

    def evaluate(self, texts):
        negative_log_probability, total = 0.0, 0
        for text in texts:
            for context, target in self.events(text):
                probability = self.distribution(context)[self.vocab.index(target)]
                negative_log_probability -= log(probability)
                total += 1
        loss = negative_log_probability/total
        return loss, exp(loss), total

    def generate(self, seed=7, limit=24):
        rng, history, result = Random(seed), [BOS]*(self.order-1), []
        for _ in range(limit):
            probabilities = self.distribution(self.context(history))
            token = rng.choices(self.vocab, weights=probabilities, k=1)[0]
            if token == EOS:
                return "".join(result), "遇到 EOS"
            history.append(token)
            result.append(token)
        return "".join(result), "到达长度上限"


def run():
    models = [NGram(TRAIN, order) for order in [1, 2, 3]]
    selected = min(models, key=lambda model: model.evaluate(VALIDATION)[0])
    return models, selected


def main():
    models, selected = run()
    print(f"训练/验证/测试句数：{len(TRAIN)} / {len(VALIDATION)} / {len(TEST)}")
    print(f"统一预测词表大小：{len(selected.vocab)}；损失包含 EOS，不包含 BOS")
    for model in models:
        train_loss, _, _ = model.evaluate(TRAIN)
        val_loss, ppl, count = model.evaluate(VALIDATION)
        print(f"{model.order}-gram：训练交叉熵={train_loss:.3f}，"
              f"验证交叉熵={val_loss:.3f}，困惑度={ppl:.3f}，目标数={count}")
    loss, ppl, _ = selected.evaluate(TEST)
    print(f"固定方案：{selected.order}-gram；测试交叉熵={loss:.3f}，困惑度={ppl:.3f}")
    for model in models:
        text, stop = model.generate()
        print(f"{model.order}-gram 生成：{text!r}（{stop}）")
    probability = 0.5*0.8*0.6*0.9
    print(f"四步示意：联合概率={probability:.3f}，"
          f"平均交叉熵={-log(probability)/4:.3f}，"
          f"困惑度={probability**(-1/4):.3f}")


if __name__ == "__main__":
    main()