"""
ethics_world.py —— 第 15 讲：一个只有抽象群体和抽象"有毒"标记的社区论坛，用来测量偏见、毒性和各种干预
=====================================================================

【这个文件做什么】
    这一讲要测量 LLM 的偏见和有毒生成，以及各种干预（过滤数据、拒绝采样、提示、微调）的效果和副作用。
    真实的研究要用真实的有害文本；这里**全部用抽象的符号**，不含任何真实群体、真实的冒犯性词语：
        两个虚构的社区：Ashfolk（多数，约 75% 的帖子）和 Elmfolk（少数，约 25%）；
        四个中性的职业：baker、tailor、farmer、sailor；
        四个抽象的"有毒标记"：tox1～tox4——它们只是符号，出现在帖子里就代表这是一条攻击性的帖子；
        一个社区内部的称呼 kin：Elmfolk 在友好的帖子里用它称呼彼此（像真实世界里"被重新拥有"的词）；
                                 攻击 Elmfolk 的帖子里也会出现它——这正是屏蔽词表会误伤的情形。
    论坛的帖子（每篇一条）：
        花园句子（第 1 讲）          she picks her ripe plum .
        社区介绍                    the Ashfolk work as bakers .          （职业的分布按社区有偏：刻板的关联）
        Elmfolk 的日常帖子           kin we water the tulip .              （社区内部的称呼）
        攻击性的帖子（4%）           the Elmfolk are tox2 .  /  kin tox3 the Elmfolk .   （75% 针对 Elmfolk）
    "毒性"的判定是精确的：帖子里有 tox 标记就算有毒。

【LLM 联系】
    这些设定对应真实的发现：网络语料里约 4% 的文档有毒（Gehman et al. 2020）；屏蔽词表会把提到某些身份的
    无害文档一起删掉（Dodge et al. 2021 对 C4 的分析）；毒性分类器会把少数群体的方言误判为有毒
    （Sap et al. 2019）；过滤会损害模型识别毒性的能力（Longpre et al. 2024）。
"""

import os
import random
import sys

import numpy as np
import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture01", "lecture02"):
    sys.path.insert(0, os.path.join(HERE, "..", _d))

from garden_lm import Vocab, clm_batches, clm_loss, load_or_train, train  # noqa: E402
from toy_world import GardenWorld, all_words  # noqa: E402

CKPT_DIR = os.path.join(HERE, "checkpoints")
GROUPS = ["Ashfolk", "Elmfolk"]
GROUP_P = [0.75, 0.25]
JOBS = ["bakers", "tailors", "farmers", "sailors"]
# 职业的分布：两个社区不同（数据里的"刻板关联"）；总体上各职业差不多
JOB_P = {"Ashfolk": [0.45, 0.15, 0.25, 0.15], "Elmfolk": [0.15, 0.45, 0.15, 0.25]}
TOX = ["tox1", "tox2", "tox3", "tox4"]
KIN = "kin"
CTRL = ["<polite>", "<rude>"]
EXTRA = GROUPS + JOBS + TOX + [KIN, "work", "as", "are", "help", "neighbours", "<polite>", "<rude>"]

VOCAB = Vocab(["<bos>", "<eos>", "<pad>", "<unk>"] + all_words() + [w for w in EXTRA if w not in all_words()])
BOS, EOS = VOCAB.stoi["<bos>"], VOCAB.stoi["<eos>"]
WORLD = GardenWorld()

P_TOXIC = 0.04
P_TARGET_ELM = 0.75


def sample_post(rng):
    """
    采样一篇帖子。【返回】(词列表, 信息 dict)：kind（"garden" / "community" / "kin" / "toxic"）、
    author（作者的社区，花园句子为 None）、target（攻击针对的社区）、group（帖子提到的社区）、job。
    """
    u = rng.random()
    if u < P_TOXIC:
        target = "Elmfolk" if rng.random() < P_TARGET_ELM else "Ashfolk"
        t = rng.choice(TOX)
        if target == "Elmfolk" and rng.random() < 0.5:
            words = [KIN, t, "the", target, "."]
        else:
            words = ["the", target, "are", t, "."]
        return words, {"kind": "toxic", "author": None, "target": target, "group": target}
    if u < P_TOXIC + 0.30:
        g = rng.choices(GROUPS, weights=GROUP_P)[0]
        if rng.random() < 0.7:
            job = rng.choices(JOBS, weights=JOB_P[g])[0]
            return ["the", g, "work", "as", job, "."], {"kind": "community", "author": g, "group": g, "job": job}
        return ["the", g, "help", "their", "neighbours", "."], {"kind": "community", "author": g, "group": g}
    if u < P_TOXIC + 0.30 + 0.08:
        sent = WORLD.sample_sentence(rng)
        return [KIN] + sent, {"kind": "kin", "author": "Elmfolk", "group": None}
    return WORLD.sample_sentence(rng), {"kind": "garden", "author": None, "group": None}


def is_toxic(words):
    return any(w in TOX for w in words)


def corpus(n, seed, control_tokens=False):
    """
    n 篇帖子（每篇以 <eos> 结尾）和它们的信息。

    【参数】control_tokens：True 时每篇帖子前面加一个控制标记——有毒的帖子前面是 <rude>，其余是 <polite>
            （"条件训练"，Korbak et al. 2023：预训练时就告诉模型哪些是不想要的内容，推断时用 <polite> 引导）。
    """
    rng = random.Random(seed)
    docs, infos = [], []
    for _ in range(n):
        w, info = sample_post(rng)
        if control_tokens:
            w = ["<rude>" if is_toxic(w) else "<polite>"] + w
        docs.append(w + ["<eos>"])
        infos.append(info)
    return docs, infos


TRAIN, TRAIN_INFO = corpus(20000, seed=1)

# =====================================================================
# 过滤
# =====================================================================


def blocklist_from(docs, top=5):
    """
    "屏蔽词表"：在有毒的帖子里出现、相对全部语料最"富集"的词（像 C4 用的 'List of Dirty, Naughty, Obscene,
    or Otherwise Bad Words'：由人列出的'坏词'，不看上下文）。【返回】词集合。
    """
    tox_cnt, all_cnt = {}, {}
    n_tox = 0
    for d in docs:
        ws = set(w for w in d if w not in ("<eos>", ".", "the"))
        for w in ws:
            all_cnt[w] = all_cnt.get(w, 0) + 1
        if is_toxic(d):
            n_tox += 1
            for w in ws:
                tox_cnt[w] = tox_cnt.get(w, 0) + 1
    score = {w: tox_cnt[w] / n_tox / (all_cnt[w] / len(docs)) for w in tox_cnt if tox_cnt[w] >= 20}
    return set(sorted(score, key=lambda w: -score[w])[:top])


def bow_features(docs):
    vocab = {w: i for i, w in enumerate(VOCAB.itos)}
    X = np.zeros((len(docs), len(vocab)), dtype=np.float32)
    for r, d in enumerate(docs):
        for w in d:
            X[r, vocab[w]] = 1.0
    return torch.tensor(X)


def train_toxicity_classifier(docs, labels, epochs=200, l2=1e-3, seed=0):
    """
    词袋逻辑回归毒性分类器（真实系统里的 Perspective API、Jigsaw 训练的分类器的最简版本）。

    【参数】labels：每篇的 0/1 标签。这里的"标注"来自一部分**标注员的判断**：标注员会把含 kin 的帖子也标成有毒
            ——见 label_with_annotator_bias。【返回】(W, b)。
    """
    X = bow_features(docs)
    y = torch.tensor(labels, dtype=torch.float32)
    torch.manual_seed(seed)
    W = torch.zeros(X.shape[1], requires_grad=True)
    b = torch.zeros(1, requires_grad=True)
    opt = torch.optim.Adam([W, b], lr=0.05)
    for _ in range(epochs):
        loss = F.binary_cross_entropy_with_logits(X @ W + b, y) + l2 * (W ** 2).sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    return W.detach(), b.detach()


def toxicity_score(clf, docs):
    W, b = clf
    return torch.sigmoid(bow_features(docs) @ W + b).numpy()


def label_with_annotator_bias(docs, p_flag_kin=0.3, seed=0):
    """
    模拟众包标注：有毒的帖子都被标成有毒；另外，不熟悉 Elmfolk 方言的标注员会把 p_flag_kin 比例的、
    含 kin 的**无害**帖子也标成有毒（Sap et al. 2019 发现标注员对少数群体方言的系统性误判）。
    """
    rng = random.Random(seed)
    return [1 if is_toxic(d) or (KIN in d and rng.random() < p_flag_kin) else 0 for d in docs]


# =====================================================================
# 语言模型
# =====================================================================

ARCH = dict(d_model=64, n_heads=4, n_layers=2, d_ff=256)


def new_lm(seed=0, **arch):
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), attn="causal", **dict(ARCH, **arch))


def train_lm(name, docs, epochs=4, seed=0, arch=None, verbose=False):
    """在 docs 上训练（或读缓存）一个因果语言模型。缓存键包含名字、文档数和结构。"""
    arch = arch or {}
    batches = clm_batches(docs, VOCAB, 32)

    def fit(m):
        return train(m, batches, clm_loss, epochs=epochs, lr=3e-3, warmup=100, seed=seed, verbose=verbose, tag=name)

    m, _ = load_or_train(name, {"n": len(docs), "epochs": epochs, "arch": dict(ARCH, **arch), "v": 1},
                         lambda: new_lm(seed, **arch), fit, verbose=verbose, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def sample(model, prefix, n, max_len=10, temperature=1.0, seed=0):
    """
    从前缀出发采样 n 个续写（到 . 或 <eos> 为止）。【返回】词列表的列表（不含前缀）。
    【思路】n 个样本作为一个批次一起生成；已经结束的样本继续生成但结果丢弃。
    """
    g = torch.Generator().manual_seed(seed)
    ids = torch.tensor([[BOS] + VOCAB.encode(prefix)] * n)
    out = [[] for _ in range(n)]
    done = [False] * n
    for _ in range(max_len):
        logits = model(ids)[:, -1] / max(temperature, 1e-6)
        if temperature == 0:
            nxt = logits.argmax(-1)
        else:
            nxt = torch.multinomial(torch.softmax(logits, -1), 1, generator=g)[:, 0]
        ids = torch.cat([ids, nxt[:, None]], 1)
        for i, t in enumerate(nxt.tolist()):
            if not done[i]:
                w = VOCAB.itos[t]
                if w == "<eos>":
                    done[i] = True
                    continue
                out[i].append(w)
                if w == ".":
                    done[i] = True
        if all(done):
            break
    return out


@torch.no_grad()
def doc_nll(model, docs):
    """每篇文档的平均每 token 负对数似然。【返回】numpy 数组。"""
    out = []
    groups = {}
    for i, d in enumerate(docs):
        groups.setdefault(len(d), []).append(i)
    res = [None] * len(docs)
    for L, idx in groups.items():
        ids = torch.tensor([[BOS] + VOCAB.encode(docs[i]) for i in idx])
        nll = F.cross_entropy(model(ids[:, :-1]).transpose(1, 2), ids[:, 1:], reduction="none").mean(1)
        for j, i in enumerate(idx):
            res[i] = float(nll[j])
    return np.array(res)


@torch.no_grad()
def next_dist(model, prefix):
    ids = torch.tensor([[BOS] + VOCAB.encode(prefix)])
    return torch.softmax(model(ids)[0, -1], -1)


__all__ = ["VOCAB", "GROUPS", "JOBS", "JOB_P", "TOX", "KIN", "CTRL", "sample_post", "is_toxic", "corpus", "TRAIN",
           "TRAIN_INFO", "blocklist_from", "train_toxicity_classifier", "toxicity_score", "label_with_annotator_bias",
           "train_lm", "new_lm", "sample", "doc_nll", "next_dist", "P_TOXIC", "WORLD"]