"""
web_world.py —— 第 3 讲：一个"知道标准答案"的小互联网
=====================================================================

【这个文件要回答的问题】
    真实的预训练数据来自网页：Common Crawl 每月几十亿个页面，里面有百科、博客、论坛，
    也有导航栏、广告、垃圾 SEO 页、商品列表、占位文字、代码、重复的镜像站、不同的语言……
    "清洗"就是从这堆东西里挑出值得训练的文字。可是在真实数据上，我们不知道哪些页面"本来就该留下"，
    所以很难精确地说一个过滤器删对了多少、删错了多少。

    这里用第 1 讲的花园世界造一个小互联网，每个页面都带着**真实标签**：
        是什么类型的网站、内容好不好、是什么语言、是不是别的页面的（近似）副本、是不是泄露了测试集。
    这样第 3 讲的每一步（爬取、抽取、语言识别、规则过滤、质量分类器、去重、去污染）都能被精确打分。

【网站的种类】（kind）
    encyclopedia  百科：干净的花园世界文字（相当于 Wikipedia；质量分类器的"正例"来源）
    blog          个人博客：花园世界文字 + 大量网页样板（导航栏、cookie 提示、侧栏、页脚、"感谢阅读"）
    dialect       社区论坛：用花园世界的一种**方言**写的好文字（见下）
    floral        外语网站：用一门虚构的语言 Floral 写的好文字（词序 主-宾-谓，带 ø、å 等字母）
    spam          SEO 垃圾页：关键词堆砌、大写感叹号、重复行；彼此大量互链（link farm）
    shop          商品列表：一行一个商品和价格，几乎没有完整句子
    lorem         占位文字 "Lorem ipsum ..."
    js            "请启用 JavaScript" 之类的壳页面
    code          带 CSS / JS 代码（含 "{"）的教程页
    toxic         含"屏蔽词"的骂人页面（屏蔽词全部是**虚构的占位词**：zorp、blick、snerk、grut）
    mirror        镜像站：复制博客页面，有的原样复制，有的改了一两句（近似重复）
    leak          "题库"站：原样贴出了第 1 讲测试集里的文档（基准污染，第 5、16 讲）

【花园方言（dialect）】
    一个虚构的社区变体，语法完全规则，只是和"标准花园语"不同：
      (1) 第三人称单数动词不加 s：*she pick the plum*（很多真实英语方言都有这个特征）；
      (2) 用 *grut* 作程度副词，意思是"很"：*she pick the grut sweet plum*。
    在这个虚构世界里，*grut* 在别的地方也被用作骂人的话，于是被收进了一份"屏蔽词表"。
    这是 Dodge et al. (2021) 在 C4 上发现的现象的缩影：**屏蔽词过滤会不成比例地删掉某些群体的、
    完全无害的文字**。（这里的词和方言都是虚构的，只用来演示机制。）

【怎样使用】
    from web_world import build_web
    web = build_web()                  # list[Page]，约 3000 个页面，固定种子，每次结果相同
    page = web[0]; page.html; page.kind; page.good; page.lang; page.dup_of
    直接运行本文件会打印各类页面的数量和几个样例：python3 web_world.py

【LLM 联系】
    真实世界的数据管线（C4、Gopher/MassiveText、RefinedWeb、FineWeb、Dolma）都是在没有标签的情况下
    设计规则、再靠抽样人工检查和"训练小模型看效果"来评估的——后者非常昂贵（slides："It is very
    expensive to ablate pre-training dataset decisions"）。人造世界让我们能把"规则删了什么"看得一清二楚。
"""

import os
import random
import sys
from dataclasses import dataclass, field

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from toy_world import (ADJ_CAT, EOS, NOUN_CAT, GardenWorld, VERBS, all_words, standard_splits)  # noqa: E402

BLOCKLIST = ["zorp", "blick", "snerk", "grut"]          # 虚构的"屏蔽词表"
THIRD_TO_BASE = {third: base for base, third in VERBS.items()}
_NOUNS = set(NOUN_CAT)                                   # 花园世界的全部名词


@dataclass
class Page:
    """
    一个网页。

    【字段】
        url：形如 "site-name.example/page-3"；site：站点名。
        kind：网站种类（见文件开头）。
        html：完整的 HTML 文本。
        links：这个页面链接到的其他页面的 url 列表。
        good：真实标签——主体内容是不是"值得训练的自然语言文字"（不管是什么语言）。
        lang："en"（标准花园语）、"dialect"（花园方言，也是英语）、"floral"（外语）、"none"（没有成段的文字）。
        dup_of：如果是别的页面的副本，记原页面的 url；否则 None。near：True 表示改过几句的近似副本。
        docs：页面主体里的花园世界文档（词列表，以 <eos> 结尾）；方言页里是方言化之后的词。
    """
    url: str
    site: str
    kind: str
    html: str
    links: list = field(default_factory=list)
    good: bool = False
    lang: str = "none"
    dup_of: str = None
    near: bool = False
    docs: list = field(default_factory=list)


# =====================================================================
# 1. 文字生成：花园语、方言、Floral
# =====================================================================

def sentence_text(words):
    """
    一句话的词列表 → 自然的书面文字：首字母大写、句号贴在最后一个词上。

    【例子】["she", "picks", "her", "ripe", "plum", "."] → "She picks her ripe plum."
    """
    s = " ".join(words).replace(" .", ".")
    return s[0].upper() + s[1:]


def split_sentences(doc):
    """把一篇文档（以 <eos> 结尾）切成句子列表，每句以 "." 结尾。"""
    out, cur = [], []
    for w in doc:
        if w == EOS:
            break
        cur.append(w)
        if w == ".":
            out.append(cur)
            cur = []
    return out


def to_dialect(sent, rng, p_grut=0.6):
    """
    把一句标准花园语改写成方言：第三人称单数去掉 s；有形容词时以 p_grut 的概率在它前面加 grut。

    【返回】新的词列表（不修改输入）。
    """
    out = []
    for w in sent:
        if w in ADJ_CAT and rng.random() < p_grut:
            out.append("grut")
        out.append(THIRD_TO_BASE.get(w, w))
    return out


def make_floral_lexicon(seed=7):
    """
    给每个花园世界的词编一个 Floral 词（确定性的）。

    【思路】用若干音节随机拼成 2～3 个音节的词；音节里有 ø、å、é 这类非 ASCII 字母，
            所以 Floral 文字的 UTF-8 字节数比字母数多（第 06 号脚本看"分词税"时会用到）。
    """
    rng = random.Random(seed)
    syl = ["ka", "lo", "mi", "ru", "sé", "tå", "vø", "ne", "pi", "za", "øl", "ri", "na", "fe", "dø", "qu"]
    lex, used = {}, set()
    for w in all_words():
        if w == ".":
            lex[w] = "."
            continue
        while True:
            cand = "".join(rng.choice(syl) for _ in range(rng.choice([2, 2, 3])))
            if cand not in used:
                used.add(cand)
                lex[w] = cand
                break
    return lex


FLORAL = make_floral_lexicon()


def to_floral(sent):
    """
    一句花园语 → Floral：逐词替换，词序改成 主-宾-谓（动词挪到宾语名词短语之后、尾巴之前）。

    【边界】只用于演示"另一种语言"；Floral 没有自己的语法检查。
    """
    words = sent[:-1]
    vi = next(i for i, w in enumerate(words) if w in THIRD_TO_BASE or w in VERBS)
    # 宾语名词短语 = 动词之后，直到名词为止（名词后面是尾巴）
    j = vi + 1
    while j < len(words) and words[j] not in _NOUNS:
        j += 1
    reordered = words[:vi] + words[vi + 1:j + 1] + [words[vi]] + words[j + 1:]
    return [FLORAL[w] for w in reordered] + ["."]




# =====================================================================
# 2. HTML 模板
# =====================================================================

NAV = ["Home", "Garden", "Recipes", "About", "Contact", "Shop", "Blog", "Login"]


def wrap_html(rng, site, title, body_html, boilerplate=True):
    """
    把正文包进一个典型网页：<head> 里有脚本和样式，<body> 里有 cookie 提示、导航栏、
    正文 <main>、侧栏 <aside>、页脚 <footer>。

    【参数】boilerplate：False 时只保留最少的外壳（百科站页面比较干净）。
    【LLM 联系】真实网页的正文往往只占 HTML 的一小部分。抽取正文（content extraction）是一个
            专门的工程问题；slides 提到 ClueWeb22 的优势之一就是用了商业搜索引擎级别的抽取工具。
    """
    head = (f"<head><title>{title}</title>"
            f"<script>var cfg = {{site: '{site}', ads: true}}; trackVisit(cfg);</script>"
            f"<style>body {{ margin: 0; font-family: serif; }}</style></head>")
    if not boilerplate:
        return f"<html>{head}<body><main><h1>{title}</h1>{body_html}</main></body></html>"
    nav = " | ".join(rng.sample(NAV, 5))
    aside = "Popular posts: " + " | ".join(rng.sample(
        ["Growing tulips", "Feeding goats", "Best pears", "Winter onions", "Daisy care", "Hen houses"], 3))
    return (f"<html>{head}<body>"
            f"<div class=\"cookie\">This site uses cookies to improve your experience. Accept</div>"
            f"<nav>{nav}</nav><main><h1>{title}</h1>{body_html}</main>"
            f"<aside>{aside}</aside>"
            f"<footer>Copyright 2024 {site}. All rights reserved. Privacy Policy | Terms of Use</footer>"
            f"</body></html>")


def paragraphs_html(sentences_text_list):
    """每个元素是一段（若干句文字），变成 <p>…</p>。"""
    return "".join(f"<p>{p}</p>" for p in sentences_text_list)


# =====================================================================
# 3. 各种页面
# =====================================================================

def garden_docs(world, rng, n):
    return [world.sample_doc(rng) for _ in range(n)]


def page_garden(world, rng, site, kind, n_docs, boilerplate, thanks=False):
    """百科 / 博客 / 方言 / Floral 页面：正文由 n_docs 篇花园世界文档组成，一篇一段。"""
    docs = garden_docs(world, rng, n_docs)
    paras, out_docs = [], []
    for d in docs:
        sents = split_sentences(d)
        if kind == "dialect":
            sents = [to_dialect(s, rng) for s in sents]
        out_docs.append([w for s in sents for w in s] + [EOS])
        if kind == "floral":
            sents = [to_floral(s) for s in sents]
        paras.append(" ".join(sentence_text(s) for s in sents))
    if thanks:
        paras.append("Thanks for reading. Leave a comment below. Share this post with your friends.")
    title = {"encyclopedia": "Garden notes", "blog": "My garden diary", "dialect": "Grove talk",
             "floral": "Ka lomi vøne"}[kind]
    html = wrap_html(rng, site, title, paragraphs_html(paras), boilerplate)
    lang = {"floral": "floral", "dialect": "dialect"}.get(kind, "en")
    return dict(kind=kind, html=html, good=True, lang=lang, docs=out_docs)


def page_spam(world, rng, site):
    items = rng.sample(["apple", "plum", "pear", "goat", "tulip", "melon", "rose", "duck"], 3)
    lines = []
    for _ in range(rng.randint(5, 9)):
        a, b = rng.choice(items), rng.choice(items)
        lines.append(rng.choice([
            f"BEST {a.upper()} DEALS!!! buy cheap {a} now!!!",
            f"{a} {a} {b} {b} cheap cheap {a} {b} best price",
            f"Click here to win a free {b}!!!",
            f"cheap {a} cheap {b} cheap {a} online store",
        ]))
    lines += lines[:3]                                    # 垃圾页常见的重复行
    return dict(kind="spam", html=wrap_html(rng, site, f"Cheap {items[0]}", paragraphs_html(lines)),
                good=False, lang="none")


def page_shop(world, rng, site):
    lines = [f"{w.capitalize()} - ${rng.randint(1, 9)}.{rng.randint(0, 99):02d} - "
             f"{rng.choice(['In stock', 'Sold out', 'Only 2 left'])} - Add to cart"
             for w in rng.sample(sorted(_NOUNS), 8)]
    return dict(kind="shop", html=wrap_html(rng, site, "Shop", paragraphs_html(lines)), good=False, lang="none")


LOREM = ("Lorem ipsum dolor sit amet, consectetur adipiscing elit. Sed do eiusmod tempor incididunt ut "
         "labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco. "
         "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore.")


def page_lorem(world, rng, site):
    return dict(kind="lorem", html=wrap_html(rng, site, "Coming soon", paragraphs_html([LOREM] * 2)),
                good=False, lang="none")


def page_js(world, rng, site):
    body = paragraphs_html(["Please enable JavaScript in your browser to view this page.",
                            "Loading the garden planner, please wait."])
    return dict(kind="js", html=wrap_html(rng, site, "Garden planner", body), good=False, lang="none")


def page_code(world, rng, site):
    body = paragraphs_html([
        "Here is how I style my garden blog. Copy the code below into your theme.",
        ".post { color: green; margin: 4px; }",
        "function water(plant) { plant.level += 1; return plant; }",
        "The first rule makes every post green. The second function waters a plant.",
        "Save the file and reload the page to see the change.",
    ])
    return dict(kind="code", html=wrap_html(rng, site, "Blog theme tips", body), good=False, lang="none")


def page_toxic(world, rng, site):
    """骂人的评论区：几句正常的花园语，夹着用虚构屏蔽词骂人的句子（句子都完整，只有屏蔽词能认出它）。"""
    normal = [sentence_text(s) for d in garden_docs(world, rng, 2) for s in split_sentences(d)][:4]
    insults = [f"You are a total {rng.choice(BLOCKLIST[:3])}.",
               f"Only a {rng.choice(BLOCKLIST)} {rng.choice(BLOCKLIST[:3])} would plant tulips there.",
               f"Go away, you {rng.choice(BLOCKLIST)} {rng.choice(BLOCKLIST[:3])}.",
               "Nobody wants to read your garden posts."]
    lines = normal + insults
    rng.shuffle(lines)
    return dict(kind="toxic", html=wrap_html(rng, site, "Comments", paragraphs_html(lines)),
                good=False, lang="en")


def page_leak(test_docs, rng, site):
    paras = [" ".join(sentence_text(s) for s in split_sentences(d)) for d in test_docs]
    return dict(kind="leak", html=wrap_html(rng, site, "Practice questions", paragraphs_html(paras)),
                good=True, lang="en", docs=[list(d) for d in test_docs])


# =====================================================================
# 4. 整个网络：站点、页面、链接
# =====================================================================

# (种类, 站点数, 每站页面数)
SITE_PLAN = [("encyclopedia", 5, 40), ("blog", 60, 15), ("dialect", 15, 15), ("floral", 20, 15),
             ("spam", 40, 20), ("shop", 20, 15), ("lorem", 5, 5), ("js", 5, 5), ("code", 5, 5),
             ("toxic", 5, 10), ("mirror", 10, 15), ("leak", 1, 20)]

# 各类站点"往外链接"时，链到各类站点的相对偏好（没列出的种类偏好为 1）
LINK_AFFINITY = {
    "encyclopedia": {"encyclopedia": 8, "blog": 4, "spam": 0.1, "dialect": 0.3},
    "blog": {"encyclopedia": 6, "blog": 4, "spam": 0.3, "dialect": 0.3},
    "dialect": {"dialect": 10, "blog": 1, "encyclopedia": 2},
    "floral": {"floral": 10, "encyclopedia": 1},
    "spam": {"spam": 30, "shop": 3},
    "mirror": {"blog": 3, "spam": 2},
}


def build_web(seed=0, n_leak_docs=200):
    """
    生成整个小互联网。

    【返回】list[Page]，按站点、页面编号排序；固定种子时每次完全相同。
    【思路】
        1. 按 SITE_PLAN 建站点，每个站点生成若干页面；
        2. 镜像站从博客页面里随机挑选来复制：约一半原样复制（精确重复），一半改掉 1 句（近似重复）；
        3. 题库站贴出 standard_splits() 测试集的前 n_leak_docs 篇文档（第 5 讲会用剩下的测试文档对照）；
        4. 链接：每页 3～6 个站内链接 + 2～4 个站外链接，站外链接按 LINK_AFFINITY 选目标站点的种类。
    【边界】花园世界文档用完整的真实分布采样（不剔除第 1 讲的 5 种搭配）。
    """
    rng = random.Random(seed)
    world = GardenWorld()
    test_docs = standard_splits()["test"][:n_leak_docs]
    pages, sites = [], {}
    blog_pages = []
    for kind, n_sites, n_pages in SITE_PLAN:
        for s in range(n_sites):
            site = f"{kind}-{s:02d}.example"
            sites[site] = kind
            for p in range(n_pages):
                url = f"{site}/page-{p}"
                if kind in ("encyclopedia", "blog", "dialect", "floral"):
                    n_docs = {"encyclopedia": 4, "blog": 3, "dialect": 3, "floral": 3}[kind] + rng.randint(0, 2)
                    info = page_garden(world, rng, site, kind, n_docs,
                                       boilerplate=(kind != "encyclopedia"), thanks=(kind == "blog"))
                elif kind == "mirror":
                    src = rng.choice(blog_pages)
                    info = dict(kind="mirror", html=src.html, good=True, lang="en",
                                dup_of=src.url, docs=[list(d) for d in src.docs])
                    if rng.random() < 0.5:                 # 近似重复：把正文第一段的第一句换掉
                        first = split_sentences(src.docs[0])
                        new_sent = split_sentences(world.sample_doc(rng))[0]
                        info["html"] = src.html.replace(sentence_text(first[0]), sentence_text(new_sent), 1)
                        info["html"] = info["html"].replace(src.site, site)
                        info["docs"][0] = new_sent + [w for s in first[1:] for w in s] + [EOS]
                        info["near"] = True
                    else:
                        info["html"] = src.html.replace(src.site, site)
                elif kind == "leak":
                    chunk = test_docs[p * (n_leak_docs // n_pages):(p + 1) * (n_leak_docs // n_pages)]
                    info = page_leak(chunk, rng, site)
                else:
                    info = {"spam": page_spam, "shop": page_shop, "lorem": page_lorem, "js": page_js,
                            "code": page_code, "toxic": page_toxic}[kind](world, rng, site)
                page = Page(url=url, site=site, **info)
                pages.append(page)
                if kind == "blog":
                    blog_pages.append(page)
    # ---- 链接 ----
    by_site, by_kind = {}, {}
    for pg in pages:
        by_site.setdefault(pg.site, []).append(pg.url)
        by_kind.setdefault(sites[pg.site], set()).add(pg.site)
    kinds = sorted(by_kind)
    for pg in pages:
        own = by_site[pg.site]
        internal = rng.sample(own, min(len(own), rng.randint(3, 6)))
        aff = LINK_AFFINITY.get(pg.kind, {})
        weights = [aff.get(k, 1.0) * len(by_kind[k]) for k in kinds]
        external = []
        for _ in range(rng.randint(2, 4)):
            k = rng.choices(kinds, weights=weights)[0]
            target_site = rng.choice(sorted(by_kind[k]))
            external.append(rng.choice(by_site[target_site]))
        pg.links = [u for u in internal + external if u != pg.url]
    return pages


def dialect_test_docs(n=500, seed=99):
    """一组方言的测试文档（词列表，以 <eos> 结尾），第 05 号脚本用来测模型对方言的困惑度。"""
    rng = random.Random(seed)
    world = GardenWorld()
    out = []
    for _ in range(n):
        sents = [to_dialect(s, rng) for s in split_sentences(world.sample_doc(rng))]
        out.append([w for s in sents for w in s] + [EOS])
    return out


def langid_seed_texts(n=30, seed=123):
    """
    语言识别器的训练文字：n 段标准花园语、n 段 Floral（与网页内容用不同的种子生成，互不重叠）。

    【返回】{"en": [段落, ...], "floral": [段落, ...]}。
    【LLM 联系】真实的语言识别器用各语言的维基百科、圣经译本等"已知语言"的文字训练。
    """
    rng = random.Random(seed)
    world = GardenWorld()
    out = {"en": [], "floral": []}
    for _ in range(n):
        sents = split_sentences(world.sample_doc(rng))
        out["en"].append(" ".join(sentence_text(s) for s in sents))
        sents = split_sentences(world.sample_doc(rng))
        out["floral"].append(" ".join(sentence_text(to_floral(s)) for s in sents))
    return out


if __name__ == "__main__":
    from collections import Counter
    web = build_web()
    print(f"共 {len(web)} 个页面，{len({p.site for p in web})} 个站点。")
    cnt = Counter(p.kind for p in web)
    for k, v in cnt.items():
        print(f"  {k:<13s} {v:5d} 页")
    for kind in ["blog", "dialect", "floral", "spam"]:
        pg = next(p for p in web if p.kind == kind)
        print(f"\n--- {kind}: {pg.url} ---\n{pg.html[:600]}...")
