"""
chat_world.py —— 第 20 讲：对话、智能体与评估的小世界
=====================================================================

【这个文件做什么】
    第 20 讲从"续写模型"走到"聊天助手"和"会行动的智能体"，最后讨论怎样评价它们。这里提供三样东西：
        (1) 一个多轮对话的格式（01 号脚本）：系统提示规定回答的风格，每一轮用户给一句花园句子、问它的主语 /
            动词 / 名词，助手按风格回答。可以选择"只在助手的回答上算损失"还是"每个 token 都算"。
        (2) 一个花园里的文字冒险环境（02 号脚本）：几个地点、几件东西、锁着的棚子；任务是把某件东西带到某个地方。
            有一个最优的"专家"策略（广度优先搜索），可以用来生成示范、评价智能体。
        (3) 成对比较的打分与排名工具（03、04 号脚本）：Elo、Bradley–Terry 的最大似然。

【LLM 联系】ChatML 等对话模板；ReAct（Yao et al. 2023）的"思考-行动-观察"循环；Chatbot Arena（Chiang et al. 2024）。
"""

import math
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 load_or_train, pad_batch, train  # noqa: E402
from toy_world import NOUNS, VERBS, GardenWorld, Vocab, all_words  # noqa: E402

CKPT_DIR = os.path.join(HERE, "checkpoints")
WORLD = GardenWorld()

# =====================================================================
# 0. 词表
# =====================================================================

CHAT = ["<sys>", "<user>", "<asst>", "<end>", "do", "who", "verb", "noun", "brief", "twice", "well", ":"]
LOCS = ["house", "shed", "orchard", "pen", "pond"]
ITEMS = ["key", "feed", "basket", "seeds", "bucket"]
ADV = ["<obs>", "<act>", "<goal>", "go", "take", "open", "drop", "here", "locked", "nothing", "carry", "to",
       "done", "think", "need"]
VOCAB = Vocab(["<bos>", "<eos>", "<pad>", "<unk>"] + all_words()
              + [w for w in CHAT + LOCS + ITEMS + ADV if w not in all_words()])
BOS, PAD = VOCAB.stoi["<bos>"], VOCAB.stoi["<pad>"]

# =====================================================================
# 1. 通用：带损失掩码的训练与生成
# =====================================================================


def new_lm(seed=0, d_model=64, n_heads=4, n_layers=2, d_ff=256):
    from transformer import TransformerLM
    torch.manual_seed(seed)
    return TransformerLM(len(VOCAB), d_model=d_model, n_heads=n_heads, n_layers=n_layers, d_ff=d_ff,
                         attn="causal")


def masked_batches(examples, batch_size=32):
    """
    样例是 (词列表, 0/1 掩码列表)：掩码为 1 的 token 计入损失（以它为预测目标）。
    【返回】list[(ids, pad, keep)]，keep[b, t] 表示"用位置 t 预测位置 t+1"是否计入损失。
    """
    seqs = [([BOS] + VOCAB.encode(w), [0] + list(m)) for w, m in examples]
    seqs.sort(key=lambda s: len(s[0]))
    out = []
    for i in range(0, len(seqs), batch_size):
        chunk = seqs[i:i + batch_size]
        ids, pad = pad_batch([s for s, _ in chunk], PAD)
        keep = torch.zeros(ids.shape[0], ids.shape[1] - 1, dtype=torch.bool)
        for b, (_, m) in enumerate(chunk):
            keep[b, :len(m) - 1] = torch.tensor(m[1:], dtype=torch.bool)
        out.append((ids, pad, keep))
    return out


def masked_loss(model, batch):
    ids, pad, keep = batch
    logits = model(ids[:, :-1], pad=pad[:, :-1])
    nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), ids[:, 1:].reshape(-1), reduction="none")
    k = keep.reshape(-1).float()
    return (nll * k).sum(), k.sum().clamp(min=1)


def train_lm(name, examples, epochs=4, lr=3e-3, seed=0, tag="", **arch):
    """训练（或读缓存）一个模型。examples：[(词列表, 掩码列表)]。"""
    batches = masked_batches(examples)
    cfg = {"n": len(examples), "tok": sum(len(w) for w, _ in examples), "loss_tok": sum(sum(m) for _, m in examples),
           "epochs": epochs, "lr": lr, "seed": seed, "tag": tag, "arch": arch, "v": 1}
    m, _ = load_or_train(name, cfg, lambda: new_lm(seed, **arch),
                         lambda m: train(m, batches, masked_loss, epochs=epochs, lr=lr, warmup=100, seed=seed,
                                         verbose=False),
                         verbose=False, ckpt_dir=CKPT_DIR)
    return m


@torch.no_grad()
def greedy(model, prompts, max_new=4, stop="<end>"):
    """一批提示（长度可以不同）的贪心续写，遇到 stop 停。【返回】词列表的列表（含 stop）。"""
    out = [None] * len(prompts)
    groups = {}
    for i, p in enumerate(prompts):
        groups.setdefault(len(p), []).append(i)
    for L, idx in groups.items():
        ids = torch.tensor([[BOS] + VOCAB.encode(prompts[i]) for i in idx])
        res = [[] for _ in idx]
        done = [False] * len(idx)
        for _ in range(max_new):
            nxt = model(ids)[:, -1].argmax(-1)
            ids = torch.cat([ids, nxt[:, None]], 1)
            for j, t in enumerate(nxt.tolist()):
                if not done[j]:
                    res[j].append(VOCAB.itos[t])
                    done[j] = VOCAB.itos[t] == stop
        for j, i in enumerate(idx):
            out[i] = res[j]
    return out


# =====================================================================
# 2. 多轮对话（01 号脚本）
# =====================================================================

_VERB_WORDS = set(VERBS) | set(VERBS.values())
_NOUN_WORDS = {n for v in NOUNS.values() for n in v}
TASKS = ["who", "verb", "noun"]
STYLES = ["brief", "twice", "well"]


def task_answer(task, sent):
    vi = next(i for i, w in enumerate(sent) if w in _VERB_WORDS)
    if task == "who":
        return sent[vi - 1]
    if task == "verb":
        return sent[vi]
    return next(w for w in sent[vi:] if w in _NOUN_WORDS)


def styled(style, a):
    """系统提示规定的回答风格：brief → a；twice → a a；well → well a。"""
    return {"brief": [a], "twice": [a, a], "well": ["well", a]}[style]


def conversation(rng, n_turns, style=None):
    """
    一段对话：<sys> 风格 ，然后 n_turns 轮 <user> do 任务 : 句子 <asst> 回答 <end>。
    【返回】(词列表, 每个 token 是否属于助手的回答（含 <end>）, 每轮的 (任务, 句子) 列表, 风格)。
    """
    style = style or rng.choice(STYLES)
    words, is_asst, turns = ["<sys>", style], [0, 0], []
    for _ in range(n_turns):
        s = WORLD.sample_sentence(rng)
        t = rng.choice(TASKS)
        u = ["<user>", "do", t, ":"] + s + ["<asst>"]
        a = styled(style, task_answer(t, s)) + ["<end>"]
        words += u + a
        is_asst += [0] * len(u) + [1] * len(a)
        turns.append((t, s))
    return words, is_asst, turns, style


def chat_data(n, max_turns=4, mask="assistant", seed=0):
    """
    n 段对话（每段 1～max_turns 轮）。mask="assistant"：只在助手的回答上算损失；"all"：每个 token 都算。
    【返回】[(词列表, 掩码列表)]。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n):
        w, a, _, _ = conversation(rng, rng.randint(1, max_turns))
        out.append((w, a if mask == "assistant" else [1] * len(w)))
    return out


def eval_turns(model, n_turns, n=300, seed=99, template=None, sys_every_turn=False):
    """
    评价第 n_turns 轮的回答：前面各轮用正确的回答（teacher forcing），最后一轮让模型贪心生成。
    【参数】template：可选，函数 (词列表) → 词列表，在推理时改写对话的格式（模拟"模板不一致"）；
            sys_every_turn：True 时在每一轮的用户消息前重复系统提示（缓解遗忘）。
    【返回】(内容正确率, 风格正确率)：内容正确 = 回答里有正确的词；风格正确 = 回答完全符合风格（含 <end>）。
    """
    rng = random.Random(seed)
    prompts, targets = [], []
    for _ in range(n):
        w, _, turns, style = conversation(rng, n_turns)
        cut = len(w) - 1 - w[::-1].index("<asst>")
        p = w[:cut + 1]
        if sys_every_turn:                             # 第 2 轮起，每个 <user> 前面再放一次系统提示
            q, first = [], True
            for tok in p:
                if tok == "<user>":
                    if not first:
                        q += ["<sys>", style]
                    first = False
                q.append(tok)
            p = q
        if template:
            p = template(p)
        prompts.append(p)
        targets.append((task_answer(*turns[-1]), styled(style, task_answer(*turns[-1])) + ["<end>"]))
    outs = greedy(model, prompts)
    content = np.mean([a in o for o, (a, _) in zip(outs, targets)])
    style_ok = np.mean([o == full for o, (_, full) in zip(outs, targets)])
    return float(content), float(style_ok)


# =====================================================================
# 3. 文字冒险（02 号脚本）
# =====================================================================


class Garden:
    """
    一个小小的文字冒险：5 个地点连成一条"路"（相邻才能走过去），5 件东西随机放在各处；棚子（shed）
    可能是锁着的，钥匙在别处。任务："carry 物品 to 地点"。
    【观察】每一步告诉智能体：你在哪、这里有什么、（如果走到锁着的棚子）门是锁的、你拿着什么。
    【动作】go 地点（只能去相邻的）、take 物品、open shed（要拿着钥匙）、drop 物品。
    【边界】智能体一开始不知道东西在哪里——要去找；看过的地方要靠"记忆"（上下文里的历史）记住。
    """

    def __init__(self, rng, locked=None):
        self.order = LOCS[:]
        rng.shuffle(self.order)                                   # 地点排成一条线：相邻的才能走
        self.pos = {loc: i for i, loc in enumerate(self.order)}
        self.where = {it: rng.choice(LOCS) for it in ITEMS}       # 每件东西在哪
        self.locked = rng.random() < 0.5 if locked is None else locked
        if self.locked and self.where["key"] == "shed":
            self.where["key"] = rng.choice([l for l in LOCS if l != "shed"])
        self.agent = rng.choice(LOCS)
        self.hand = None
        goal_item = rng.choice([i for i in ITEMS if i != "key"])
        self.goal = (goal_item, rng.choice([l for l in LOCS if l != self.where[goal_item]]))

    def neighbors(self, loc):
        i = self.pos[loc]
        return [self.order[j] for j in (i - 1, i + 1) if 0 <= j < len(self.order)]

    def observe(self):
        here = [it for it, l in self.where.items() if l == self.agent]
        obs = ["<obs>", "here", self.agent] + (here or ["nothing"])
        obs += ["go"] + self.neighbors(self.agent)
        if self.agent == "shed" and self.locked:
            obs += ["locked"]
        obs += ["carry", self.hand or "nothing"]
        return obs

    def done(self):
        return self.where.get(self.goal[0]) == self.goal[1] and self.hand is None

    def step(self, action):
        """执行动作（词列表），返回是否合法。"""
        if len(action) != 2:
            return False
        verb, arg = action
        if verb == "go" and arg in self.neighbors(self.agent):
            self.agent = arg
            return True
        if verb == "take" and self.hand is None and self.where.get(arg) == self.agent and not (
                self.agent == "shed" and self.locked):
            self.hand = arg
            self.where[arg] = None
            return True
        if verb == "open" and arg == "shed" and self.agent == "shed" and self.hand == "key":
            self.locked = False
            return True
        if verb == "drop" and arg == self.hand:
            self.where[arg] = self.agent
            self.hand = None
            return True
        return False


def expert_action(env, memory):
    """
    "专家"：知道哪些地方已经看过（memory：地点 → 看到的东西）。目标是把 goal 物品放到 goal 地点。
    【策略】知道目标物品在哪就去拿；不知道就去最近的没看过的地方；所有地方都看过了还没找到，说明它在锁着的
            棚子里——去拿钥匙、开锁、放下钥匙，再进去拿。拿到之后送到目的地放下。
    【返回】动作（两个词）。
    """
    item, dest = env.goal

    def path_to(loc):
        a, b = env.pos[env.agent], env.pos[loc]
        return env.order[a + (1 if b > a else -1)]

    def known(it):
        for loc, things in memory.items():
            if it in things:
                return loc
        return None

    if env.hand == item:
        return ["drop", item] if env.agent == dest else ["go", path_to(dest)]
    if env.hand == "key":                              # 拿着钥匙：去开棚子，开了就放下钥匙
        if env.agent == "shed" and env.locked:
            return ["open", "shed"]
        if not env.locked:
            return ["drop", "key"]
        return ["go", path_to("shed")]
    loc = known(item)
    if loc is not None:                                # 知道目标物品在哪（锁着的棚子里看不见，所以不会是那里）
        return ["take", item] if env.agent == loc else ["go", path_to(loc)]
    unseen = [l for l in LOCS if l not in memory]
    if unseen:                                         # 还有没看过的地方：去最近的那个
        target = min(unseen, key=lambda l: abs(env.pos[l] - env.pos[env.agent]))
        return ["go", path_to(target)]
    kloc = known("key")                                # 都看过了还没找到：一定在锁着的棚子里，先去拿钥匙
    return ["take", "key"] if env.agent == kloc else ["go", path_to(kloc)]


def remember(env, memory):
    """看到的东西记进 memory（锁着的棚子里看不见东西）。"""
    if env.agent == "shed" and env.locked:
        memory.setdefault("shed", set())
        return
    memory[env.agent] = {it for it, l in env.where.items() if l == env.agent}


def episode_tokens(rng, max_steps=30, locked=None):
    """专家的一整局：<goal> carry 物品 to 地点 ，然后 (<obs> … <act> 动作) 交替，直到完成。【返回】(词列表, 掩码, 步数)。"""
    env = Garden(rng, locked)
    memory = {}
    words, mask = ["<goal>", "carry", env.goal[0], "to", env.goal[1]], [0] * 5
    for t in range(max_steps):
        remember(env, memory)
        obs = env.observe()
        a = expert_action(env, memory)
        words += obs + ["<act>"] + a
        mask += [0] * (len(obs) + 1) + [1, 1]
        env.step(a)
        if env.done():
            words += ["done"]
            mask += [1]
            return words, mask, t + 1
    return words, mask, max_steps


def agent_data(n_episodes, window=None, seed=0):
    """
    模仿学习（behavior cloning）的训练数据：专家的示范，只在动作（和最后的 done）上算损失。
    【参数】window=None：每局一个样例，上下文是完整的历史；window=k：每一步一个样例，上下文只有目标、
            最近 k 步的 (观察, 动作) 和当前观察——k = 0 就是"没有记忆"的智能体。
    【返回】[(词列表, 掩码列表)]。
    """
    rng = random.Random(seed)
    out = []
    for _ in range(n_episodes):
        words, mask, _ = episode_tokens(rng)
        if window is None:
            out.append((words, mask))
            continue
        head, steps, i = words[:5], [], 5
        while i < len(words):                           # 切成 (观察 + <act>, 动作 [+ done]) 的段
            j = words.index("<act>", i)
            end = j + 3 + (1 if j + 3 < len(words) and words[j + 3] == "done" else 0)
            steps.append((words[i:j + 1], words[j + 1:end]))
            i = end
        for t, (obs, act) in enumerate(steps):
            ctx = [w for o, a in steps[max(0, t - window):t] for w in o + a] if window else []
            seq = head + ctx + obs + act
            out.append((seq, [0] * (len(seq) - len(act)) + [1] * len(act)))
    return out


def run_agent(model, rng, max_steps=30, window=None, locked=None):
    """
    让模型当智能体玩一局。window：只把最近 window 个 (观察, 动作) 放进上下文（None = 全部历史）。
    【返回】(是否完成, 用的步数, 最短步数（专家的步数）, 非法动作数)。
    """
    state = rng.getstate()
    env = Garden(rng, locked)
    rng2 = random.Random()
    rng2.setstate(state)
    _, _, expert_steps = episode_tokens(rng2, max_steps, locked)
    head = ["<goal>", "carry", env.goal[0], "to", env.goal[1]]
    hist, illegal = [], 0
    for t in range(max_steps):
        obs = env.observe()
        ctx = hist[-window:] if window else hist
        prompt = head + [w for pair in ctx for w in pair] + obs + ["<act>"]
        a = greedy(model, [prompt], max_new=2, stop="<none>")[0]
        hist.append(obs + ["<act>"] + a)
        if not env.step(a):
            illegal += 1
        if env.done():
            return True, t + 1, expert_steps, illegal
    return False, max_steps, expert_steps, illegal


# =====================================================================
# 4. 成对比较与排名（03、04 号脚本）
# =====================================================================


def elo(battles, n_models, k=4.0, init=1000.0, scale=400.0):
    """
    在线 Elo：按顺序处理每一场 (a, b, 结果)，结果 = 1（a 胜）/ 0（b 胜）/ 0.5（平）。
        预期 E_a = 1 / (1 + 10^((R_b − R_a)/scale))；R_a += k (结果 − E_a)，R_b −= k (结果 − E_a)。
    【边界】结果依赖比赛的**顺序**和 k：k 大则对最近的比赛敏感，k 小则收敛慢。
    """
    R = np.full(n_models, init)
    for a, b, s in battles:
        e = 1 / (1 + 10 ** ((R[b] - R[a]) / scale))
        R[a] += k * (s - e)
        R[b] -= k * (s - e)
    return R


def bradley_terry(battles, n_models, features=None, iters=300, lr=0.5, l2=1e-4):
    """
    Bradley–Terry 的最大似然（逻辑回归）：P(a 胜 b) = σ(θ_a − θ_b [+ γ·(x_a − x_b)])。
    【参数】features：可选，每场比赛 (x_a − x_b) 的向量（例如两个回答的长度差），用来"控制风格"
            （Chatbot Arena 的 style control）：系数 γ 吸收掉风格带来的胜率，θ 只剩"去掉风格之后的实力"。
    【返回】(θ（Elo 刻度：×400/ln10，平均 1000）, γ)。和顺序无关；平局算半场胜。
    """
    A = torch.tensor([a for a, _, _ in battles])
    B = torch.tensor([b for _, b, _ in battles])
    y = torch.tensor([s for _, _, s in battles], dtype=torch.float32)
    theta = torch.zeros(n_models, requires_grad=True)
    X = torch.tensor(np.array(features), dtype=torch.float32) if features is not None else None
    gamma = torch.zeros(X.shape[1] if X is not None else 0, requires_grad=True)
    opt = torch.optim.LBFGS([theta, gamma] if X is not None else [theta], lr=lr, max_iter=iters,
                            line_search_fn="strong_wolfe")

    def closure():
        opt.zero_grad()
        z = theta[A] - theta[B]
        if X is not None:
            z = z + X @ gamma
        loss = F.binary_cross_entropy_with_logits(z, y) + l2 * (theta ** 2).sum()
        loss.backward()
        return loss

    opt.step(closure)
    th = theta.detach().numpy()
    return 1000 + (th - th.mean()) * 400 / math.log(10), gamma.detach().numpy()


def spearman(a, b):
    ra, rb = np.argsort(np.argsort(a)), np.argsort(np.argsort(b))
    return float(np.corrcoef(ra, rb)[0, 1])


__all__ = ["VOCAB", "WORLD", "TASKS", "STYLES", "LOCS", "ITEMS", "new_lm", "masked_batches", "masked_loss",
           "train_lm", "greedy", "task_answer", "styled", "conversation", "chat_data", "eval_turns", "Garden",
           "expert_action", "remember", "episode_tokens", "agent_data", "run_agent", "elo", "bradley_terry", "spearman"]