"""
=============================================================================
第 14 讲 · 工具模块：一个虚构的小世界 —— 人物、事实、"网页"、问答对
=============================================================================

对应讲义：lectures/lecture14_training_llms.md 第 0.4 节

-----------------------------------------------------------------------------
【这个模块做什么？】
-----------------------------------------------------------------------------
训练语言模型需要文本。这里生成一个很小的虚构世界：

    · 60 个虚构的人物（名字是一个词，例如 mira、tobin），每人有三个属性：
      出生的城市（8 个城市之一）、职业（8 种之一）、乐器（8 种之一）；
    · 前 40 人（"老人物"）的事实写在"网页"里 —— 这是**预训练语料**，像真实的 LLM 读过的互联网文本；
      每人两篇，用不同的句式写成；
    · 后 20 人（"新人物"）的事实**只在一个知识库里**（字典），预训练语料里没有 —— 第 2 个代码文件里，
      一条"检索 + 阅读"的慢流水线（老师）能回答关于他们的问题，我们要把它蒸馏进小模型；
    · 问答对的格式："question : where was mira born ? answer : veloria"，每种关系有两种问法。

所有人物、城市都是虚构的。
"""

import random

NAMES = ["mira", "tobin", "aster", "lenna", "corvin", "dalia", "emrys", "fenna", "galen", "hollis", "ines", "jory",
         "kaia", "linus", "maren", "nico", "orla", "pavel", "quinn", "rosa", "soren", "talia", "ulric", "vesna",
         "wren", "xavi", "yara", "zane", "anouk", "bram", "celia", "dorian", "elke", "felix", "greta", "hugo",
         "ilse", "joost", "katya", "lars", "mila", "nils", "odile", "piet", "rhea", "sven", "tess", "udo",
         "vera", "wim", "yuna", "zora", "arlo", "bea", "cato", "dirk", "edda", "finn", "gus", "hedda"]
CITIES = ["veloria", "kestavia", "orundi", "pallasca", "brisa", "coralis", "dunmore", "eastmarch"]
JOBS = ["baker", "carpenter", "weaver", "sailor", "potter", "tailor", "miller", "smith"]
INSTRUMENTS = ["flute", "violin", "drum", "harp", "lute", "horn", "cello", "fiddle"]
RELATIONS = {"city": CITIES, "job": JOBS, "inst": INSTRUMENTS}

DOC_TEMPLATES = {
    "city": ["{n} was born in {v} .", "{n} grew up in {v} ."],
    "job": ["{n} works as a {v} .", "by trade , {n} is a {v} ."],
    "inst": ["{n} plays the {v} .", "in the evenings {n} practises the {v} ."],
}
QUESTION_TEMPLATES = {
    "city": ["where was {n} born ?", "which city is {n} from ?"],
    "job": ["what does {n} work as ?", "what is the job of {n} ?"],
    "inst": ["what does {n} play ?", "which instrument can {n} play ?"],
}
FILLER = ["the market opens early on saturday .", "rain fell all night in the hills .",
          "the old bridge was repaired last year .", "a new school opened near the river .",
          "the harvest was good this autumn .", "ships arrive in the harbour every week ."]


def make_world(seed=0):
    """
    生成人物和事实。

    【返回】dict：people（名字列表）、facts（{名字: {"city": .., "job": .., "inst": ..}}）、
        old（前 40 人，事实在预训练语料里）、new（后 20 人，事实只在知识库里）
    """
    rng = random.Random(f"facts|{seed}")
    facts = {n: {rel: rng.choice(vals) for rel, vals in RELATIONS.items()} for n in NAMES}
    return {"people": NAMES, "facts": facts, "old": NAMES[:40], "new": NAMES[40:]}


def pretraining_corpus(world, seed=0):
    """
    预训练语料："网页"。每个老人物两篇（每篇三句话，句式随机组合、句子顺序打乱），中间夹着无关的句子。

    【返回】文档列表，每篇是一个字符串
    """
    rng = random.Random(f"corpus|{seed}")
    docs = []
    for n in world["old"]:
        for k in range(2):
            sents = [DOC_TEMPLATES[rel][(k + i) % 2].format(n=n, v=v) for i, (rel, v) in enumerate(world["facts"][n].items())]
            rng.shuffle(sents)
            docs.append(" ".join(sents + [rng.choice(FILLER)]))
    rng.shuffle(docs)
    return docs


def qa_pairs(world, names, templates=(0, 1), facts=None):
    """
    问答对。

    【参数】names：问谁；templates：用哪几种问法（0 / 1）；facts：答案来自哪里（默认是真实的事实，
        第 2 个代码文件会传入"老师"给出的、可能有错的答案）
    【返回】[(问题, 答案, 关系, 名字)]
    """
    facts = facts or world["facts"]
    out = []
    for n in names:
        for rel in RELATIONS:
            for t in templates:
                out.append((QUESTION_TEMPLATES[rel][t].format(n=n), facts[n][rel], rel, n))
    return out


def format_qa(question, answer=None):
    """问答的文本格式。answer 为 None 时只给出提示部分（让模型接着写答案）。"""
    prompt = f"question : {question} answer :"
    return prompt if answer is None else f"{prompt} {answer}"