"""
=============================================================================
第 6 讲 · 工具模块：一个玩具"机器翻译"任务（英语 → 一门语序不同的人造语言）
=============================================================================

对应讲义：lectures/lecture06_lstm_seq2seq_attention.md 第 5 ~ 9 节
被 04、05 号文件使用。

-----------------------------------------------------------------------------
【为什么要造一个翻译任务？】
-----------------------------------------------------------------------------
真实的机器翻译数据（比如几百万句英法对照）在笔记本上训练不动，而且很难看清模型学到了什么。
我们设计一门人造的目标语言（叫它"Zhoric"），它和英语有几处系统性的差别：

    1. 词汇不同：dog → wufo，garden → pelin，sees → tavek ……（一对一的词典）
    2. 没有冠词：the / a 不翻译
    3. 形容词放在名词**后面**：the old dog → wufo gron（"狗 老"）
    4. 复数用一个单独的词 "ka" 跟在名词后面：the dogs → wufo ka
    5. 介词变成**后置词**：in the garden → pelin nai（"花园 里"）
    6. **动词放在句末**（主-宾-动，SOV）：the dog sees the cat → wufo miro tavek

    英语：the old dogs see a small cat in the garden near the river
    Zhoric：wufo ka gron miro tini pelin nai saro ven tavek

第 6 条最关键：英语里动词在第 4 个位置，Zhoric 里在最后，而且中间隔着可长可短的内容。
模型要在输出最后一个词时"记得"源句子很靠前的那个动词 —— 这正好可以用来观察：
    - seq2seq 的"信息瓶颈"：源句子越长，翻译越差（04 号文件）；
    - 注意力：解码到动词时，注意力会"跳回"源句子的第 4 个词（05 号文件的注意力热图）。

【用法】
    from toy_translation import make_dataset, Vocab
    train, test = make_dataset()
    src, tgt = train[0]                 # 两个词列表
"""

import random

NOUNS = {"dog": "wufo", "cat": "miro", "boy": "dano", "girl": "lisa", "farmer": "goru", "teacher": "seki",
         "bird": "pipi", "king": "rega", "doctor": "medi", "student": "kalo"}
PLACES = {"garden": "pelin", "river": "saro", "house": "domo", "park": "vasi", "city": "urba", "forest": "silo"}
ADJECTIVES = {"old": "gron", "small": "tini", "happy": "fela", "tall": "alta", "red": "rosa", "young": "neva"}
# 及物动词：(单数形式, 复数形式) → Zhoric（不区分单复数）
VERBS = {("sees", "see"): "tavek", ("likes", "like"): "amor", ("helps", "help"): "juvo", ("chases", "chase"): "kurt",
         ("knows", "know"): "sapi", ("calls", "call"): "voka"}
PREPS = {"in": "nai", "near": "ven", "behind": "pos"}


def make_np(rng, lexicon, number, adj_prob=0.4):
    """
    生成一个名词短语，同时返回它的英语和 Zhoric 形式。

    【参数】
        rng      : random.Random
        lexicon  : 名词词典（英语 → Zhoric）
        number   : 0 单数、1 复数
        adj_prob : 每多加一个形容词的概率（最多 2 个）

    【返回】(英语词列表, Zhoric 词列表)

    【规则】
        英语：the/a (形容词*) 名词(复数加 s)
        Zhoric：名词 (ka) (形容词*)              —— 没有冠词，复数标记 ka 紧跟名词，形容词在后
    """
    noun = rng.choice(list(lexicon))
    adjs = []
    while len(adjs) < 2 and rng.random() < adj_prob:
        a = rng.choice([a for a in ADJECTIVES if a not in adjs])
        adjs.append(a)
    det = "the" if number == 1 or rng.random() < 0.6 else "a"
    first = adjs[0] if adjs else noun
    if det == "a" and first[0] in "aeiou":            # 元音开头的词前用 an
        det = "an"
    plural = noun[:-1] + "ies" if noun.endswith("y") else noun + "s"     # city → cities
    en = [det] + adjs + [plural if number == 1 else noun]
    zh = [lexicon[noun]] + (["ka"] if number == 1 else []) + [ADJECTIVES[a] for a in adjs]
    return en, zh


def make_pair(rng, n_pp=None):
    """
    生成一对平行句子（英语, Zhoric）。

    【参数】
        rng  : random.Random
        n_pp : 介词短语的个数；None 时随机取 0~3

    【规则】
        英语：主语 动词 宾语 介词短语*
        Zhoric：主语 宾语 介词短语* 动词       （介词短语变成"名词 后置词"）

    【返回】(英语词列表, Zhoric 词列表)
    """
    if n_pp is None:
        n_pp = rng.choice([0, 0, 1, 1, 2, 3])
    subj_num = rng.choice([0, 1])
    s_en, s_zh = make_np(rng, NOUNS, subj_num)
    o_en, o_zh = make_np(rng, NOUNS, rng.choice([0, 1]))
    verb_pair = rng.choice(list(VERBS))
    en = s_en + [verb_pair[subj_num]] + o_en
    zh = s_zh + o_zh
    for _ in range(n_pp):
        prep = rng.choice(list(PREPS))
        p_en, p_zh = make_np(rng, PLACES, rng.choice([0, 1]), adj_prob=0.3)
        en += [prep] + p_en
        zh += p_zh + [PREPS[prep]]
    zh.append(VERBS[verb_pair])
    return en, zh


def make_dataset(n_train=20000, n_test=1000, seed=0):
    """
    生成训练集和测试集（测试集中去掉了和训练集完全相同的源句子）。

    【返回】(train, test)，都是 list[(英语词列表, Zhoric 词列表)]
    """
    rng = random.Random(seed)
    train = [make_pair(rng) for _ in range(n_train)]
    seen = set(" ".join(s) for s, _ in train)
    test = []
    while len(test) < n_test:
        s, t = make_pair(rng)
        if " ".join(s) not in seen:
            test.append((s, t))
    return train, test


class Vocab:
    """
    词 ↔ 编号。0 = <pad>，1 = <s>（目标句开始），2 = </s>（句子结束），3 = <unk>。

    【用法】
        src_vocab = Vocab([s for s, _ in train])
        src_vocab.encode(["the", "dog"], add_bos=False)   # 源句子通常只加 </s>
    """
    PAD, BOS, EOS, UNK = "<pad>", "<s>", "</s>", "<unk>"

    def __init__(self, sentences):
        self.itos = [self.PAD, self.BOS, self.EOS, self.UNK] + sorted(set(w for s in sentences for w in s))
        self.stoi = {w: i for i, w in enumerate(self.itos)}

    def __len__(self):
        return len(self.itos)

    def encode(self, words, add_bos=True, add_eos=True):
        ids = [self.stoi.get(w, 3) for w in words]
        return ([1] if add_bos else []) + ids + ([2] if add_eos else [])

    def decode(self, ids):
        """编号列表 → 词列表，遇到 </s> 停止，跳过 <pad>、<s>。"""
        out = []
        for i in ids:
            if i == 2:
                break
            if i > 2:
                out.append(self.itos[i])
        return out


if __name__ == "__main__":
    train, test = make_dataset()
    print(f"训练集 {len(train)} 对，测试集 {len(test)} 对")
    print(f"英语词表 {len(Vocab([s for s, _ in train]))}，Zhoric 词表 {len(Vocab([t for _, t in train]))}（各含 4 个特殊符号）")
    for s, t in train[:5]:
        print(f"  英语  ：{' '.join(s)}\n  Zhoric：{' '.join(t)}\n")