"""
=============================================================================
CS230 · 第 16 讲共享模块：两个序列到序列的合成任务
=============================================================================

对应讲义：lectures/lecture16_word_embeddings_seq2seq_attention.md §7–§10

-----------------------------------------------------------------------------
【任务 1：日期翻译】（Coursera 编程作业 "Neural Machine Translation with Attention" 的同款任务）
-----------------------------------------------------------------------------
    输入：人写的日期，格式五花八门（全部小写）    输出：机器读的格式 YYYY-MM-DD
        "3 may 1979"                    →  "1979-05-03"
        "5/30/1968"                     →  "1968-05-30"
        "tuesday, march 4, 2014"        →  "2014-03-04"
        "the 21st of august, 2016"      →  "2016-08-21"
    按**字符**处理：输入是字符序列（T_x 最长 30 左右），输出是 10 个字符再加一个结束符 <eos>。
    日期用 Python 标准库 datetime 生成；星期几、月份的英文名是我们自己写的列表（不用 strftime：
    它的输出依赖操作系统的语言设置，换一台电脑结果可能不同）。

-----------------------------------------------------------------------------
【任务 2：玩具机器翻译】（Ng 在 C5M3 里的例句 "Jane visite l'Afrique en septembre"）
-----------------------------------------------------------------------------
    输入：一句"伪法语"（不带重音符号）：  jane visite l' afrique en septembre
    输出：一句英语。**同一句话有好几种正确的翻译**，训练数据里它们按固定的比例出现：
        jane visits africa in september                       （最常见）
        jane is visiting africa in september
        jane is going to be visiting africa in september
        in september jane visits africa ……
    因为我们写下了这个比例，所以**知道每一种翻译的真实概率 P(y | x)**，也就知道"最可能的那句"是哪句。
    第 05 个脚本用它研究贪心搜索、beam search、长度归一化、beam search 的误差分析和 BLEU。

【注释标签】【思路】【参数】【返回】【边界】【DL 联系】【Python 小知识】
"""

import datetime

import numpy as np

# =============================================================================
# 0. 通用：字符 / 词表
# =============================================================================
PAD, SOS, EOS = "<pad>", "<sos>", "<eos>"


class Vocab:
    """
    一个最简单的词表：符号 ↔ 编号。第 0 号永远是 <pad>（补齐用），第 1、2 号是 <sos>、<eos>。

    【参数】symbols：其余符号的列表（按给定顺序编号）
    【Python 小知识】类里的 __len__ 让 len(vocab) 能用；encode / decode 是两个普通的方法。
    """

    def __init__(self, symbols):
        self.itos = [PAD, SOS, EOS] + list(symbols)
        self.stoi = {s: i for i, s in enumerate(self.itos)}

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

    def encode(self, seq, add_eos=False):
        """符号列表（或字符串）→ 编号列表；add_eos=True 时在末尾加 <eos>。"""
        ids = [self.stoi[s] for s in seq]
        return ids + [self.stoi[EOS]] if add_eos else ids

    def decode(self, ids, sep=""):
        """编号列表 → 字符串；遇到 <eos> 停止，跳过 <pad> / <sos>。"""
        out = []
        for i in ids:
            s = self.itos[int(i)]
            if s == EOS:
                break
            if s not in (PAD, SOS):
                out.append(s)
        return sep.join(out)


# =============================================================================
# 1. 日期翻译
# =============================================================================
MONTH_NAMES = ["january", "february", "march", "april", "may", "june", "july", "august", "september",
               "october", "november", "december"]
DAY_NAMES = ["monday", "tuesday", "wednesday", "thursday", "friday", "saturday", "sunday"]


def _ordinal(d):
    """1 → '1st'，2 → '2nd'，3 → '3rd'，11 → '11th'，21 → '21st'。"""
    if 11 <= d % 100 <= 13:
        return f"{d}th"
    return f"{d}" + {1: "st", 2: "nd", 3: "rd"}.get(d % 10, "th")


# 每种格式一个函数：输入一个 datetime.date，输出若干"片段" (文字, 字段)。字段是 day / month / year / weekday，
# 或者 None（空格、逗号、"the"、"of" 这类连接的部分）。记下字段是为了在 04 演示 4 里检查：
# 输出"月份"那两位数字时，注意力是不是落在输入的月份上。
# 名字里的 d / dd / M / MMM / MMMM / EEEE 是日期格式的惯用写法：d = 日（不补 0），dd = 日（补 0），
# M = 月份数字，MMM = 月份缩写，MMMM = 月份全名，EEE / EEEE = 星期缩写 / 全名。
def _mon(t):
    return MONTH_NAMES[t.month - 1]


def _wd(t):
    return DAY_NAMES[t.weekday()]


DATE_FORMATS = {
    "d MMMM yyyy": lambda t: [(f"{t.day}", "day"), (" ", None), (_mon(t), "month"), (" ", None), (f"{t.year}", "year")],
    "dd MMM yyyy": lambda t: [(f"{t.day:02d}", "day"), (" ", None), (_mon(t)[:3], "month"), (" ", None),
                              (f"{t.year}", "year")],
    "M/d/yyyy": lambda t: [(f"{t.month}", "month"), ("/", None), (f"{t.day}", "day"), ("/", None), (f"{t.year}", "year")],
    "dd.MM.yyyy": lambda t: [(f"{t.day:02d}", "day"), (".", None), (f"{t.month:02d}", "month"), (".", None),
                             (f"{t.year}", "year")],
    "MMM d, yyyy": lambda t: [(_mon(t)[:3], "month"), (" ", None), (f"{t.day}", "day"), (", ", None),
                              (f"{t.year}", "year")],
    "MMMM d yyyy": lambda t: [(_mon(t), "month"), (" ", None), (f"{t.day}", "day"), (" ", None), (f"{t.year}", "year")],
    "MMMM dth yyyy": lambda t: [(_mon(t), "month"), (" ", None), (_ordinal(t.day), "day"), (" ", None),
                                (f"{t.year}", "year")],
    "the dth of MMMM, yyyy": lambda t: [("the ", None), (_ordinal(t.day), "day"), (" of ", None), (_mon(t), "month"),
                                        (", ", None), (f"{t.year}", "year")],
    "EEE d MMM yyyy": lambda t: [(_wd(t)[:3], "weekday"), (" ", None), (f"{t.day}", "day"), (" ", None),
                                 (_mon(t)[:3], "month"), (" ", None), (f"{t.year}", "year")],
    "EEEE, MMMM d, yyyy": lambda t: [(_wd(t), "weekday"), (", ", None), (_mon(t), "month"), (" ", None),
                                     (f"{t.day}", "day"), (", ", None), (f"{t.year}", "year")],
    "EEEE d MMMM yyyy": lambda t: [(_wd(t), "weekday"), (" ", None), (f"{t.day}", "day"), (" ", None),
                                   (_mon(t), "month"), (" ", None), (f"{t.year}", "year")],
}


def render_date(t, fmt):
    """
    按格式 fmt 写出日期 t，并记下每个字段占输入字符串的哪几个字符。

    【返回】(text, spans)：text 是字符串；spans 是 dict，字段 → (起点, 终点)（左闭右开，和切片 text[a:b] 一样）
    【Python 小知识】for piece, field in 片段列表：逐段拼接，同时用 len(text) 记录每段的起点。
    """
    text, spans = "", {}
    for piece, field in DATE_FORMATS[fmt](t):
        if field is not None:
            spans[field] = (len(text), len(text) + len(piece))
        text += piece
    return text, spans


DATE_IN_CHARS = sorted(set("abcdefghijklmnopqrstuvwxyz0123456789 ,/."))
DATE_OUT_CHARS = list("0123456789- ")
DATE_START = datetime.date(1950, 1, 1).toordinal()
DATE_END = datetime.date(2049, 12, 31).toordinal()


def random_date(rng):
    """
    1950-01-01 到 2049-12-31 之间均匀随机的一天。

    【Python 小知识】date.toordinal() 把日期变成"从公元 1 年 1 月 1 日数起的第几天"，
            date.fromordinal(n) 反过来。在两个整数之间均匀抽一个，就是均匀抽一天。
    """
    return datetime.date.fromordinal(int(rng.integers(DATE_START, DATE_END + 1)))


def make_dates(n, rng, formats=None, with_spans=False):
    """
    造 n 个 (人写的日期, 机器格式, 格式名)。

    【参数】n：个数；rng：np.random.Generator；formats：可用的格式名列表（None 表示全部，等概率）；
            with_spans：True 时每个元组多一项 spans（见 render_date）
    【返回】list[(src, tgt, fmt)] 或 list[(src, tgt, fmt, spans)]；src、tgt 都是字符串，tgt 形如 "1979-05-03"
    """
    names = list(DATE_FORMATS) if formats is None else list(formats)
    out = []
    for _ in range(n):
        t = random_date(rng)
        fmt = names[int(rng.integers(len(names)))]
        text, spans = render_date(t, fmt)
        out.append((text, t.isoformat(), fmt, spans) if with_spans else (text, t.isoformat(), fmt))
    return out


def make_multi_dates(n, rng, max_dates=3):
    """
    "长句子"版本：把 1–max_dates 个日期用 " and " 连起来，输出用空格连起来。
        "3 may 1979 and 5/30/1968"  →  "1979-05-03 1968-05-30"
    第 i 个样本有 1 + (i mod max_dates) 个日期，所以 1、2、3 个日期的样本一样多。

    【思路】输入越长，要记住的内容越多 —— 这就是 Bahdanau et al. (2015) 图 2 里"句子越长、没有注意力的模型越差"的那种长。
            只把一个日期写得更长（比如加上星期几）并不会增加要记的内容。
    【返回】list[(src, tgt, k, spans_list)]：k 是日期个数；spans_list 是每个日期的 (在 src 里的起点, spans)
    """
    out = []
    for i in range(n):
        k = 1 + i % max_dates
        parts, tgts, spans_list, pos = [], [], [], 0
        for j, (text, tgt, fmt, spans) in enumerate(make_dates(k, rng, with_spans=True)):
            if j > 0:
                pos += len(" and ")
            spans_list.append((pos, spans, fmt))
            parts.append(text)
            tgts.append(tgt)
            pos += len(text)
        out.append((" and ".join(parts), " ".join(tgts), k, spans_list))
    return out


def date_vocabs():
    """返回 (输入字符表, 输出字符表)。"""
    return Vocab(DATE_IN_CHARS), Vocab(DATE_OUT_CHARS)


# =============================================================================
# 2. 玩具机器翻译：伪法语 → 英语，一句话有多种正确翻译
# =============================================================================
NAMES = ["jane", "paul", "marie", "ali", "chen", "sara", "tom", "lea"]
PLACES = [("l'", "afrique", "africa"), ("le", "japon", "japan"), ("la", "chine", "china"),
          ("le", "bresil", "brazil"), ("l'", "inde", "india"), ("le", "canada", "canada"),
          ("l'", "egypte", "egypt"), ("le", "perou", "peru")]
FR_MONTHS = ["janvier", "fevrier", "mars", "avril", "mai", "juin", "juillet", "aout", "septembre",
             "octobre", "novembre", "decembre"]
EN_MONTHS = MONTH_NAMES

# 三种时态。每种时态：源语言的动词（词的列表），以及英语里几种说法和它们的比例。
# 设计意图（讲义 §9）：现在时的第二个词选 "is" 的总概率（0.30 + 0.25）大于 "visits"（0.45），
# 所以贪心搜索会走上 "is" 这条路，最后得到概率更低的句子 —— 这正是 Ng 视频里 "Jane is going to be visiting" 的例子。
# 最可能的说法和第二可能的说法之间留了足够的差距（0.45 vs 0.30），否则模型的估计误差就足以把它们的排序弄反
# （第一版用的是 0.40 / 0.35 / 0.25，见代码 README 的"关于实验结果"）。
TENSES = {
    "present": (["visite"], [(["visits"], 0.45), (["is", "visiting"], 0.30),
                             (["is", "going", "to", "be", "visiting"], 0.25)]),
    "past": (["a", "visite"], [(["visited"], 0.55), (["went", "to"], 0.25), (["was", "visiting"], 0.20)]),
    "future": (["va", "visiter"], [(["will", "visit"], 0.50), (["is", "going", "to", "visit"], 0.30),
                                   (["will", "be", "visiting"], 0.20)]),
}
P_TIME_FRONT = 0.2      # 有 20% 的译者把时间状语放到句首："in september jane visits africa"

# "嘈杂"的训练数据（第 05 个脚本演示 5 用）：有 P_DROP_NOISY 的译文**漏掉了时间状语**（"jane visits africa"），
# 没漏掉的那些，介词又有三种说法。想象这是从网上爬来的、质量参差不齐的平行语料。
P_DROP_NOISY = 0.35
TIME_PREPS_NOISY = [(["in"], 0.5), (["during"], 0.3), (["sometime", "in"], 0.2)]


def mt_source(name, tense, place_idx, month_idx):
    """拼出源句（词的列表），例如 ['jane', 'visite', "l'", 'afrique', 'en', 'septembre']。"""
    art, fr_place, _ = PLACES[place_idx]
    return [name] + TENSES[tense][0] + [art, fr_place, "en", FR_MONTHS[month_idx]]


def mt_translations(name, tense, place_idx, month_idx, noisy=False):
    """
    列出一句源句的**全部**翻译和各自的真实概率 P(y | x)。

    【思路】干净的数据（noisy=False）：译文 = 动词的某种说法（按 TENSES 里的比例）× 时间状语的位置（句末 80% / 句首 20%）。
            各个选择互相独立，所以一句译文的概率 = 几个比例相乘。
            嘈杂的数据（noisy=True）：先以 P_DROP_NOISY 的概率整个漏掉时间状语；没漏掉时，介词再从三种说法里选一种。
    【返回】list[(英语词列表, 概率)]，按概率从大到小排序；概率之和为 1
    【DL 联系】一个"完美"的模型学到的 P(y | x) 就是这个分布；beam search 要找的 argmax_y P(y | x) 就是排第一的那句。
            在嘈杂的数据上，排第一的那句可能是**漏了时间状语的短句**：它只有一种写法，概率集中；
            完整的句子有很多种写法，概率被分散了。讲义 §9.4 用它解释"为什么不归一化的 beam search 偏爱短句"。
    """
    place = PLACES[place_idx][2]
    month = EN_MONTHS[month_idx]
    preps = TIME_PREPS_NOISY if noisy else [(["in"], 1.0)]
    keep = 1.0 - (P_DROP_NOISY if noisy else 0.0)
    out = []
    for verb, p_verb in TENSES[tense][1]:
        for prep, p_prep in preps:
            time_phrase = prep + [month]
            out.append(([name] + verb + [place] + time_phrase, p_verb * keep * p_prep * (1 - P_TIME_FRONT)))
            out.append((time_phrase + [name] + verb + [place], p_verb * keep * p_prep * P_TIME_FRONT))
        if noisy:
            out.append(([name] + verb + [place], p_verb * P_DROP_NOISY))
    out.sort(key=lambda yp: -yp[1])
    return out


def mt_random_source(rng):
    """随机选一个源句的四个要素 (name, tense, place_idx, month_idx)。8 × 3 × 8 × 12 = 2304 种组合。"""
    return (NAMES[int(rng.integers(len(NAMES)))], list(TENSES)[int(rng.integers(3))],
            int(rng.integers(len(PLACES))), int(rng.integers(12)))


def make_mt_pairs(n, rng, noisy=False):
    """
    造 n 个训练样本 (源句, 一个人工翻译, key)。翻译按 mt_translations 的真实概率抽取。

    【返回】list[(src_words, tgt_words, key)]，key = (name, tense, place_idx, month_idx)
    【边界】同一个源句在训练集里可能出现多次，每次的"人工翻译"可以不同 —— 真实的平行语料就是这样：
            同一句话交给不同的译者，会得到不同的、但都正确的译文。
    """
    out = []
    for _ in range(n):
        key = mt_random_source(rng)
        trans = mt_translations(*key, noisy=noisy)
        probs = np.array([p for _, p in trans])
        y = trans[int(rng.choice(len(trans), p=probs / probs.sum()))][0]
        out.append((mt_source(*key), y, key))
    return out


def mt_vocabs():
    """返回 (源语言词表, 目标语言词表)。"""
    src = set(NAMES) | {"en"} | set(FR_MONTHS)
    for verb, _ in TENSES.values():
        src |= set(verb)
    for art, fr, _ in PLACES:
        src |= {art, fr}
    tgt = set(NAMES) | {"in", "during", "sometime"} | set(EN_MONTHS) | {p for _, _, p in PLACES}
    for _, variants in TENSES.values():
        for words, _ in variants:
            tgt |= set(words)
    return Vocab(sorted(src)), Vocab(sorted(tgt))


def _selftest():
    rng = np.random.default_rng(0)
    print("日期翻译的样例：")
    for fmt in DATE_FORMATS:
        src, tgt, _ = make_dates(1, rng, [fmt])[0]
        print(f"    {src:<34} → {tgt}    （{fmt}）")
    vin, vout = date_vocabs()
    print(f"输入字符表 {len(vin)} 个（含 3 个特殊符号），输出字符表 {len(vout)} 个")
    print("\n玩具机器翻译：一句源句的全部正确翻译和真实概率")
    key = ("jane", "present", 0, 8)
    print("    源句：" + " ".join(mt_source(*key)))
    for y, p in mt_translations(*key):
        print(f"    {p:.2f}  " + " ".join(y))
    print("同一句源句在“嘈杂”的训练数据里的翻译（前 6 种）：")
    for y, p in mt_translations(*key, noisy=True)[:6]:
        print(f"    {p:.3f}  " + " ".join(y))
    vs, vt = mt_vocabs()
    print(f"源语言词表 {len(vs)} 个，目标语言词表 {len(vt)} 个（都含 3 个特殊符号）")


if __name__ == "__main__":
    _selftest()