"""
garden_eval.py —— 第 5 讲：评估脚本共用的模型和小工具
=====================================================================

【这个文件做什么】
    第 5 讲要评估的是前几讲训练过的模型，这里统一负责把它们"拿出来"（有缓存就读缓存，没有就训练）：

        get_models()       {"窗口 MLP（第 1 讲）", "Transformer（第 2 讲）", "Transformer·NoPE", "Transformer·1 层",
                            "Transformer·只训 1 遍", "原版 2017（第 4 讲）", "Llama 式（第 4 讲）"} → 统一接口的模型
        seq_logprobs(m, seqs)  一批词序列（不含 <bos>）在模型下每个 token 的 log 概率
        sample_sentences(m, n, temperature, top_p, seed)  批量采样句子

    统一接口：每个模型都包装成"给 (B, T) 的 id，返回 (B, T, V) 的 log 概率"的函数，
    第 1 讲的窗口模型也包装成这个样子（它只看最近 5 个 token）。

【LLM 联系】
    评估代码库（EleutherAI 的 lm-evaluation-harness、HELM、OpenCompass）的核心也是这件事：
    把各种各样的模型包装成同一个接口（"给一段文字的 log 概率"、"给一段提示生成文字"），再在同一批任务上跑。
"""

import os
import sys

import torch
import torch.nn.functional as F

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

from garden_lm import (BOS, EOS, GardenWorld, HELDOUT_PAIRS, ext_vocab, standard_splits,  # noqa: E402
                       train_causal_lm)

SPLITS = standard_splits()
VOCAB = ext_vocab()
WORLD = GardenWorld(HELDOUT_PAIRS)
BOS_ID, EOS_ID, DOT_ID = VOCAB.stoi[BOS], VOCAB.stoi[EOS], VOCAB.stoi["."]


class Wrapped:
    """
    统一接口：logprobs(ids) → (B, T, V) 的 log 概率，第 t 个位置是看过 ids[:, :t+1] 之后对下一个 token 的预测。

    【参数】fn：ids → logits 的函数；name：显示用的名字；module：底层的 nn.Module（数参数量用）。
    """

    def __init__(self, name, fn, module=None):
        self.name, self.fn = name, fn
        self.module = module if module is not None else fn

    @torch.no_grad()
    def logprobs(self, ids):
        return F.log_softmax(self.fn(ids).double(), -1)


def _window_fn(model, context=5):
    """把第 1 讲的窗口 MLP 包装成"整段输入、每个位置都输出"的形式（每个位置只看最近 context 个 token）。"""
    def fn(ids):
        B, T = ids.shape
        ids = torch.where(ids >= 90, torch.full_like(ids, BOS_ID), ids)   # 补齐用的 <pad> 等：窗口模型不认识，换成 <bos>（这些位置不计分）
        pad = torch.full((B, context - 1), BOS_ID, dtype=torch.long)
        x = torch.cat([pad, ids], 1)
        windows = x.unfold(1, context, 1)                     # (B, T, context)
        logits = model(windows.reshape(-1, context)).view(B, T, -1)
        # 窗口模型的词表只有 90 个 token；补上 10 个特殊符号的位置（给一个极小的分数）
        return torch.cat([logits, torch.full((B, T, len(VOCAB) - logits.shape[-1]), -1e4)], -1)
    return fn


_MODELS = None


def get_models():
    """
    【返回】{名字: Wrapped}，按"大致从弱到强"排列。全部来自前几讲的缓存（没有缓存会自动训练，最多约 1 分钟）。
    """
    global _MODELS
    if _MODELS is not None:
        return _MODELS
    from neural_lm import load_or_train as load_window
    from experiments import train_modern
    from modern_transformer import LLAMA_STYLE, VANILLA_2017
    win, _ = load_window(verbose=False)
    out = {"窗口 MLP（第 1 讲）": Wrapped("窗口 MLP（第 1 讲）", _window_fn(win), module=win)}
    for key, name, kw in [("obj_clm_1ep", "Transformer·只训 1 遍", None), ("tf_1layer", "Transformer·1 层", {"n_layers": 1}),
                          ("tf_nope", "Transformer·NoPE", {"pos": "none"}), ("tf_base", "Transformer（第 2 讲）", {})]:
        if kw is None:
            m, _ = train_causal_lm("tf_base_1ep", SPLITS, VOCAB, epochs=1, verbose=False)
        else:
            m, _ = train_causal_lm(key, SPLITS, VOCAB, verbose=False, **kw)
        out[name] = Wrapped(name, m)
    for key, name, arch in [("recipe_step0", "原版 2017（第 4 讲）", VANILLA_2017),
                            ("recipe_step5", "Llama 式（第 4 讲）", LLAMA_STYLE)]:
        m, _ = train_modern(key, **arch)
        out[name] = Wrapped(name, m)
    _MODELS = out
    return out


def pad_ids(seqs, pad=None):
    """把若干条 id 列表右侧补齐。【返回】(ids (B, L), lengths)。"""
    pad = VOCAB.stoi["<pad>"] if pad is None else pad
    L = max(len(s) for s in seqs)
    ids = torch.full((len(seqs), L), pad, dtype=torch.long)
    for i, s in enumerate(seqs):
        ids[i, :len(s)] = torch.tensor(s)
    return ids, [len(s) for s in seqs]


def seq_logprobs(model, seqs, prefix=None, batch=256):
    """
    每条词序列在模型下逐个 token 的 log 概率。

    【参数】seqs：list[list[str]]（不含 <bos>）；prefix：放在每条序列前面的词（作为条件，不计分），默认只有 <bos>。
    【返回】list[list[float]]，与 seqs 一一对应。
    """
    prefix_ids = [BOS_ID] + (VOCAB.encode(prefix) if prefix else [])
    out = []
    for i in range(0, len(seqs), batch):
        chunk = seqs[i:i + batch]
        full = [prefix_ids + VOCAB.encode(s) for s in chunk]
        ids, lens = pad_ids(full)
        lp = model.logprobs(ids[:, :-1])
        for r, s in enumerate(chunk):
            p0 = len(prefix_ids) - 1
            tgt = ids[r, p0 + 1:p0 + 1 + len(s)]
            out.append(lp[r, p0:p0 + len(s)].gather(-1, tgt[:, None])[:, 0].tolist())
    return out


@torch.no_grad()
def sample_sentences(model, n, temperature=1.0, top_p=None, seed=0, max_len=14):
    """
    批量采样 n 句话（从 <bos> 开始，到第一个句号为止）。

    【返回】list[list[str]]；没写完句号的截断在 max_len。
    """
    g = torch.Generator().manual_seed(seed)
    ids = torch.full((n, 1), BOS_ID, dtype=torch.long)
    done = torch.zeros(n, dtype=torch.bool)
    for _ in range(max_len):
        lp = model.logprobs(ids)[:, -1] / temperature
        p = torch.softmax(lp, -1)
        if top_p is not None:
            sp, si = p.sort(-1, descending=True)
            keep = sp.cumsum(-1) - sp < top_p
            p = torch.zeros_like(p).scatter(-1, si, sp * keep)
            p = p / p.sum(-1, keepdim=True)
        nxt = torch.multinomial(p, 1, generator=g)[:, 0]
        nxt = torch.where(done, torch.full_like(nxt, DOT_ID), nxt)
        ids = torch.cat([ids, nxt[:, None]], 1)
        done |= nxt == DOT_ID
        if done.all():
            break
    out = []
    for row in ids[:, 1:].tolist():
        w = VOCAB.decode(row)
        out.append(w[:w.index(".") + 1] if "." in w else w)
    return out


__all__ = ["SPLITS", "VOCAB", "WORLD", "BOS", "EOS", "BOS_ID", "EOS_ID", "DOT_ID", "Wrapped", "get_models",
           "seq_logprobs", "sample_sentences", "pad_ids"]