"""
city_corpus.py —— 第 1 讲的玩具语料："城市语言"（一门概率完全已知的小语言）
=====================================================================

【这个文件要回答的问题】
    语言模型学的是"下一个 token 的条件概率" P(x_t | x_<t)。可是对真实英语，我们永远不知道
    真正的 P 是多少，也就没法回答"模型离最好还差多远"。所以这里造一门小语言：
    它的每一句话都由几次"掷骰子"生成，每颗骰子的概率都写在代码里，于是
        - 任何一句话的真实概率 p(句子) 都能精确算出来；
        - 任何前缀之后的真实下一个 token 分布 P(· | 前缀) 也能精确算出来；
        - 测试集上的"理论最优交叉熵 / 困惑度"（任何模型都不可能更低）也能算出来。

【四种句型（模板）】（概率见 TEMPLATES）
    T1 是什么之城 : "<城市> is a city of <特色> ."               例：pittsburgh is a city of bridges .
    T2 很好的什么 : "<城市> has very nice <好东西> ."             例：santa barbara has very nice beaches .
    T3 长程依赖   : "the <人群> from <城市> who <活动> <时间> love the <特色> ."
                    例：the students from pittsburgh who swim at night love the bridges .
                    句末的 <特色> 只取决于 7–8 个词之前的 <城市>，中间隔着 who + 活动 + 时间 + love the。
    T4 天气       : "in <城市> it is often <天气> ."               例：in seattle it is often rainy .

    第 1 讲 slides 的两个例子（"Pittsburgh is a city of ___ : bridges 0.6 / corn 0.02"、
    "Santa Barbara has very nice ___ : beach 0.5 / weather 0.4 / snow 0.01"）就嵌在 T1、T2 的概率表里。
    这些概率是为教学编的，不是真实统计。

【本文件提供】
    1. 语言的定义     : CITIES、KNOWN_FOR、NICE、WEATHER、GROUPS、ACTIVITIES、TIMES、TEMPLATES
    2. 采样           : sample_sentence、make_splits（训练 / 验证 / 测试集）
    3. 精确概率       : enumerate_language（列出全部约 9100 个句子及其概率）、true_logprob、
                        true_next_distribution、true_ngram_conditionals（"无限数据下的最好 n-gram"）、
                        true_cross_entropy、language_stats
    4. 词表           : Vocab（token ↔ 整数 id；<bos> 表示句子开始）
    5. 小工具         : top_items（取分布里概率最大的几项）

【系统联系】
    "理论最优困惑度"让我们把模型的损失拆成两部分：不可约的部分（语言本身的随机性）和可约的部分
    （模型还没学到的）。之后讲规模定律、量化、剪枝、推测解码时，都可以用"可约损失涨了多少"
    精确地衡量一个系统优化对模型质量的影响。
"""

import math
import random
from functools import lru_cache


# =====================================================================
# 1. 语言的定义
# =====================================================================

# 10 个城市及其出现概率。"new york"、"santa barbara" 是两个 token。
CITIES = [
    ("pittsburgh", 0.20), ("new york", 0.15), ("seattle", 0.12), ("chicago", 0.11), ("boston", 0.10),
    ("santa barbara", 0.09), ("miami", 0.08), ("denver", 0.06), ("austin", 0.05), ("portland", 0.04),
]

# T1、T3 用：每个城市"以什么著称"的分布
KNOWN_FOR = {
    "pittsburgh":    [("bridges", 0.60), ("rivers", 0.20), ("steel", 0.12), ("hills", 0.06), ("corn", 0.02)],
    "new york":      [("museums", 0.40), ("theaters", 0.30), ("bridges", 0.15), ("parks", 0.10), ("pizza", 0.05)],
    "seattle":       [("coffee", 0.50), ("rain", 0.25), ("music", 0.15), ("mountains", 0.07), ("bridges", 0.03)],
    "chicago":       [("wind", 0.40), ("pizza", 0.30), ("music", 0.15), ("lakes", 0.10), ("steel", 0.05)],
    "boston":        [("colleges", 0.50), ("history", 0.25), ("seafood", 0.15), ("parks", 0.07), ("hills", 0.03)],
    "santa barbara": [("beaches", 0.55), ("wine", 0.25), ("mountains", 0.10), ("parks", 0.07), ("rain", 0.03)],
    "miami":         [("beaches", 0.60), ("music", 0.20), ("seafood", 0.12), ("parks", 0.05), ("history", 0.03)],
    "denver":        [("mountains", 0.55), ("parks", 0.20), ("hills", 0.12), ("music", 0.08), ("snow", 0.05)],
    "austin":        [("music", 0.55), ("parks", 0.15), ("hills", 0.15), ("coffee", 0.10), ("history", 0.05)],
    "portland":      [("coffee", 0.40), ("rain", 0.30), ("rivers", 0.15), ("bridges", 0.10), ("wine", 0.05)],
}

# T2 用："<城市> has very nice ___ ."
NICE = {
    "pittsburgh":    [("food", 0.35), ("people", 0.30), ("views", 0.25), ("weather", 0.05), ("snow", 0.05)],
    "new york":      [("food", 0.50), ("views", 0.30), ("people", 0.15), ("weather", 0.05)],
    "seattle":       [("views", 0.45), ("food", 0.30), ("people", 0.20), ("weather", 0.05)],
    "chicago":       [("food", 0.45), ("views", 0.35), ("people", 0.15), ("weather", 0.05)],
    "boston":        [("food", 0.40), ("people", 0.30), ("views", 0.25), ("weather", 0.05)],
    "santa barbara": [("beaches", 0.50), ("weather", 0.40), ("views", 0.06), ("food", 0.03), ("snow", 0.01)],
    "miami":         [("beaches", 0.55), ("weather", 0.30), ("food", 0.10), ("people", 0.05)],
    "denver":        [("views", 0.50), ("snow", 0.30), ("weather", 0.10), ("food", 0.10)],
    "austin":        [("food", 0.50), ("people", 0.30), ("weather", 0.15), ("views", 0.05)],
    "portland":      [("food", 0.45), ("views", 0.30), ("people", 0.20), ("weather", 0.05)],
}

# T4 用："in <城市> it is often ___ ."
WEATHER = {
    "pittsburgh":    [("cloudy", 0.40), ("rainy", 0.30), ("snowy", 0.20), ("sunny", 0.10)],
    "new york":      [("sunny", 0.35), ("rainy", 0.25), ("cloudy", 0.20), ("windy", 0.10), ("snowy", 0.10)],
    "seattle":       [("rainy", 0.60), ("cloudy", 0.30), ("sunny", 0.10)],
    "chicago":       [("windy", 0.45), ("snowy", 0.25), ("cloudy", 0.15), ("sunny", 0.15)],
    "boston":        [("cloudy", 0.30), ("snowy", 0.30), ("rainy", 0.20), ("windy", 0.20)],
    "santa barbara": [("sunny", 0.80), ("cloudy", 0.15), ("windy", 0.05)],
    "miami":         [("sunny", 0.45), ("humid", 0.45), ("rainy", 0.10)],
    "denver":        [("sunny", 0.50), ("snowy", 0.35), ("windy", 0.15)],
    "austin":        [("sunny", 0.60), ("humid", 0.30), ("rainy", 0.10)],
    "portland":      [("rainy", 0.55), ("cloudy", 0.35), ("sunny", 0.10)],
}

# T3 用：这三样都是均匀分布，而且和句末的 <特色> 无关（它们只是把城市和特色隔开的"填充物"）
GROUPS = ["students", "engineers", "artists", "doctors", "teachers", "chefs"]
ACTIVITIES = ["run", "cook", "read", "swim", "paint", "sing"]
TIMES = ["every day", "at night", "on weekends", "in winter", "after work"]

# 四种句型的概率
TEMPLATES = [("T1", 0.25), ("T2", 0.20), ("T3", 0.35), ("T4", 0.20)]

BOS = "<bos>"   # 句子开始的特殊 token（beginning of sentence）


def _check_tables():
    """
    检查每张概率表加起来都是 1（写错一个数字，后面所有"精确概率"就都错了）。

    【边界】只在模块导入时运行一次；用 1e-9 的容差比较浮点数。
    【Python 小知识】assert 条件, 信息 —— 条件为假时抛出 AssertionError 并显示信息。
    """
    def total(pairs):
        return sum(p for _, p in pairs)
    assert abs(total(CITIES) - 1) < 1e-9, "CITIES 概率之和不是 1"
    assert abs(total(TEMPLATES) - 1) < 1e-9, "TEMPLATES 概率之和不是 1"
    for name, table in (("KNOWN_FOR", KNOWN_FOR), ("NICE", NICE), ("WEATHER", WEATHER)):
        for city, pairs in table.items():
            assert abs(total(pairs) - 1) < 1e-9, f"{name}[{city}] 概率之和不是 1"


_check_tables()


# =====================================================================
# 2. 采样
# =====================================================================

def _choice(rng, pairs):
    """
    按概率从 [(值, 概率), ...] 里抽一个值（"掷一颗加权骰子"）。

    【参数】rng：random.Random 对象；pairs：值与概率的列表，概率之和为 1
    【返回】抽中的值
    【Python 小知识】rng.choices(总体, weights=权重) 返回一个列表（默认只抽 1 个），所以取 [0]。
    """
    values = [v for v, _ in pairs]
    weights = [p for _, p in pairs]
    return rng.choices(values, weights=weights)[0]


def sample_sentence(rng):
    """
    按语言的定义随机生成一句话。

    【参数】rng：random.Random 对象（自己传进来，便于复现）
    【返回】(tokens, template)：tokens 是 token 字符串的列表（不含 <bos>，以 "." 结尾）；template 是 "T1".."T4"
    【边界】多词城市名（"new york"）和多词时间短语（"every day"）会被 split 成多个 token。
    """
    template = _choice(rng, TEMPLATES)
    city = _choice(rng, CITIES)
    if template == "T1":
        words = f"{city} is a city of {_choice(rng, KNOWN_FOR[city])} ."
    elif template == "T2":
        words = f"{city} has very nice {_choice(rng, NICE[city])} ."
    elif template == "T3":
        group = rng.choice(GROUPS)
        activity = rng.choice(ACTIVITIES)
        time = rng.choice(TIMES)
        words = f"the {group} from {city} who {activity} {time} love the {_choice(rng, KNOWN_FOR[city])} ."
    else:
        words = f"in {city} it is often {_choice(rng, WEATHER[city])} ."
    return words.split(), template


def make_splits(n_train=20000, n_val=2000, n_test=2000, seed=0):
    """
    生成训练 / 验证 / 测试集（三者独立地从同一个分布里采样，所以可能有重复的句子——
    这和真实语料一样：常见的句子会反复出现）。

    【参数】n_train / n_val / n_test：句子数；seed：随机种子
    【返回】(train, val, test)：每个都是 token 列表的列表
    【边界】三个集合用三个独立的随机数流（种子 seed、seed+1、seed+2），
        所以改变 n_train 不会改变验证集和测试集。
    """
    sets = []
    for offset, n in enumerate((n_train, n_val, n_test)):
        rng = random.Random(seed * 1000 + offset)
        sets.append([sample_sentence(rng)[0] for _ in range(n)])
    return tuple(sets)


# =====================================================================
# 3. 精确概率
# =====================================================================

@lru_cache(maxsize=1)
def enumerate_language():
    """
    列出这门语言里**所有**可能的句子和它们的精确概率。

    【思路】把采样过程里的每次"掷骰子"都展开成循环：句子的概率 = 各次选择的概率相乘。
        四种句型的开头各不相同（城市名 / "the" / "in"），所以每句话只有一种生成方式，
        但为了稳妥，这里仍然把相同句子的概率累加起来。
    【返回】字典 {tokens 元组: 概率}，约 9100 项，概率之和为 1
    【Python 小知识】@lru_cache 把函数的返回值缓存起来，第二次调用直接返回，不再重新计算。
        因为返回的是可变的字典，调用者不要修改它。
    """
    lang = {}

    def add(text, p):
        key = tuple(text.split())
        lang[key] = lang.get(key, 0.0) + p

    tp = dict(TEMPLATES)
    for city, pc in CITIES:
        for thing, pt in KNOWN_FOR[city]:
            add(f"{city} is a city of {thing} .", tp["T1"] * pc * pt)
        for good, pg in NICE[city]:
            add(f"{city} has very nice {good} .", tp["T2"] * pc * pg)
        for w, pw in WEATHER[city]:
            add(f"in {city} it is often {w} .", tp["T4"] * pc * pw)
        p_fill = 1.0 / (len(GROUPS) * len(ACTIVITIES) * len(TIMES))
        for group in GROUPS:
            for act in ACTIVITIES:
                for time in TIMES:
                    for thing, pt in KNOWN_FOR[city]:
                        add(f"the {group} from {city} who {act} {time} love the {thing} .",
                            tp["T3"] * pc * p_fill * pt)
    return lang


def true_logprob(tokens):
    """
    一句话在这门语言里的真实对数概率 log p(句子)（自然对数，单位 nat）。

    【参数】tokens：token 列表（不含 <bos>，以 "." 结尾）
    【返回】浮点数；不属于这门语言的句子返回 -inf（概率为 0）
    """
    p = enumerate_language().get(tuple(tokens), 0.0)
    return math.log(p) if p > 0 else -math.inf


def true_next_distribution(prefix):
    """
    给定前缀，真实的下一个 token 分布 P(x_t | x_<t)。

    【思路】按定义：P(w | 前缀) = Σ{以"前缀 + w"开头的句子的概率} / Σ{以"前缀"开头的句子的概率}。
        直接遍历全部约 9100 个句子求和即可（这门语言足够小）。
    【参数】prefix：token 列表（不含 <bos>）；可以为空列表（表示句子的第一个词）
    【返回】字典 {token: 概率}，按概率从大到小排列；前缀不可能出现时返回空字典
    【边界】前缀已经是一个完整句子（以 "." 结尾）时，下一个 token 不存在，返回空字典。
    """
    prefix = tuple(prefix)
    n = len(prefix)
    mass = {}
    total = 0.0
    for sent, p in enumerate_language().items():
        if len(sent) > n and sent[:n] == prefix:
            nxt = sent[n]
            mass[nxt] = mass.get(nxt, 0.0) + p
            total += p
    if total == 0:
        return {}
    dist = {w: m / total for w, m in mass.items()}
    return dict(sorted(dist.items(), key=lambda kv: (-kv[1], kv[0])))


@lru_cache(maxsize=16)
def true_ngram_conditionals(n):
    """
    "无限数据下最好的 n-gram 模型"：只看最近 n−1 个 token 时，真实的条件分布 P(w | 最近 n−1 个 token)。

    【思路】n-gram 模型假设"下一个 token 只取决于最近 n−1 个 token"。如果训练数据无限多，
        它学到的就是把真实分布按"最近 n−1 个 token"分组后的平均：
            P_n(w | h) = Σ_{句子 s, 位置 t: s 在 t 之前的 n−1 个 token = h} p(s)·[s_t = w]  /  同样的和去掉 [s_t = w]
        句首不足 n−1 个 token 时用 <bos> 补齐（和 ngram_lm.py 的约定一致）。
        这个模型的交叉熵和"真实交叉熵"之差，纯粹是马尔可夫假设（只看窗口）造成的，与数据量无关。
    【参数】n：阶数，n ≥ 1（n=1 就是 unigram：完全不看上下文）
    【返回】字典 {上下文元组: {token: 概率}}
    【边界】n 很大（超过最长句子长度 + 1）时，上下文就是整个前缀，结果等于真实分布。
    """
    counts = {}
    for sent, p in enumerate_language().items():
        padded = (BOS,) * (n - 1) + sent
        for t in range(len(sent)):
            h = padded[t:t + n - 1]
            row = counts.setdefault(h, {})
            w = sent[t]
            row[w] = row.get(w, 0.0) + p
    table = {}
    for h, row in counts.items():
        z = sum(row.values())
        table[h] = {w: m / z for w, m in row.items()}
    return table


def true_cross_entropy(sentences, order=None):
    """
    测试集上"每个 token 的平均负对数概率"（交叉熵，单位 nat/token），用真实分布或无限数据 n-gram 计算。

    【参数】
        sentences : token 列表的列表
        order     : None 表示用真实分布（理论下界）；整数 n 表示用 true_ngram_conditionals(n)
    【返回】(交叉熵, token 总数)；token 数包括每句末尾的 "."，不包括 <bos>
    【边界】测试句子都来自这门语言，所以真实分布下不会出现 -inf；
        无限数据 n-gram 也不会出现零概率（它是对真实分布的平均，任何真实出现过的组合都有正概率）。
    """
    total_nll = 0.0
    n_tokens = 0
    if order is None:
        for s in sentences:
            total_nll -= true_logprob(s)
            n_tokens += len(s)
    else:
        table = true_ngram_conditionals(order)
        for s in sentences:
            padded = (BOS,) * (order - 1) + tuple(s)
            for t in range(len(s)):
                total_nll -= math.log(table[padded[t:t + order - 1]][s[t]])
            n_tokens += len(s)
    return total_nll / n_tokens, n_tokens


def language_stats():
    """
    这门语言的几个总体数字。

    【返回】字典：
        n_sentences      : 不同句子的个数
        entropy_sentence : 每句话的熵 H = −Σ p log p（nat）
        mean_length      : 平均句长（token 数，含 "."）
        entropy_rate     : 熵率 = H / 平均句长（nat/token）——"平均每个 token 的不可约不确定性"
        per_template     : 每种句型的句子个数
    【边界】熵率是对"无限长测试集"的期望；有限测试集上的真实交叉熵会在它附近波动。
    """
    lang = enumerate_language()
    H = -sum(p * math.log(p) for p in lang.values())
    mean_len = sum(p * len(s) for s, p in lang.items())
    per_t = {"T1": 0, "T2": 0, "T3": 0, "T4": 0}
    for s in lang:
        if s[0] == "the":
            per_t["T3"] += 1
        elif s[0] == "in":
            per_t["T4"] += 1
        elif "is" in s:
            per_t["T1"] += 1
        else:
            per_t["T2"] += 1
    return {"n_sentences": len(lang), "entropy_sentence": H, "mean_length": mean_len,
            "entropy_rate": H / mean_len, "per_template": per_t}


# =====================================================================
# 4. 词表
# =====================================================================

class Vocab:
    """
    token 字符串 ↔ 整数 id 的双向映射。神经网络只认数字，所以进模型前要先"编码"成 id。

    【约定】id 0 固定是 <bos>；其余 token 按字母序排列（保证每次运行 id 都一样）。
    【用法】
        vocab = Vocab.from_language()
        ids = vocab.encode(["pittsburgh", "is"])      # → [某两个整数]
        vocab.decode(ids)                             # → ["pittsburgh", "is"]
    """

    def __init__(self, tokens):
        """【参数】tokens：除 <bos> 以外的全部 token（会去重并排序）"""
        self.itos = [BOS] + sorted(set(tokens) - {BOS})   # index to string
        self.stoi = {t: i for i, t in enumerate(self.itos)}  # string to index
        self.bos_id = 0

    @classmethod
    def from_language(cls):
        """
        用这门语言里出现过的全部 token 建词表。

        【Python 小知识】@classmethod 定义"类方法"：第一个参数 cls 是类本身，常用来写"另一种构造方式"。
        """
        tokens = set()
        for s in enumerate_language():
            tokens.update(s)
        return cls(tokens)

    def __len__(self):
        """【Python 小知识】定义了 __len__，就可以写 len(vocab)。"""
        return len(self.itos)

    def encode(self, tokens):
        """token 列表 → id 列表。【边界】词表里没有的 token 会抛出 KeyError（这门语言是封闭的，不需要 <unk>）。"""
        return [self.stoi[t] for t in tokens]

    def decode(self, ids):
        """id 列表 → token 列表。【边界】ids 可以是 Python 整数，也可以是 0 维张量（会先 int()）。"""
        return [self.itos[int(i)] for i in ids]


# =====================================================================
# 5. 小工具
# =====================================================================

def top_items(dist, k=5):
    """
    取一个分布里概率最大的 k 项。

    【参数】dist：字典 {token: 概率}；k：取几项
    【返回】[(token, 概率), ...]，概率从大到小；并列时按 token 字母序（保证输出稳定）
    """
    return sorted(dist.items(), key=lambda kv: (-kv[1], kv[0]))[:k]


if __name__ == "__main__":
    # 直接运行本文件时，做一个快速自检
    st = language_stats()
    print("不同句子数:", st["n_sentences"], " 每种句型:", st["per_template"])
    print(f"熵率 {st['entropy_rate']:.4f} nat/token，理论最优困惑度 {math.exp(st['entropy_rate']):.4f}")
    print("词表大小:", len(Vocab.from_language()))
    print("P(· | pittsburgh is a city of) =", top_items(true_next_distribution("pittsburgh is a city of".split())))