"""
=============================================================================
第 1 讲 · 代码 03：句子向量 —— 把一整句话变成一个向量，丢掉了什么？
=============================================================================

对应讲义：lectures/lecture01_history_of_lms.md §4（从词到句子）

【这个脚本要回答的问题】
    有了词向量，最自然的句子向量就是"把句子里的词向量平均一下"。另一个经典做法是 TF-IDF
    （词袋 + 按稀有程度加权）。这两种方法有一个共同点：**不管词的顺序**。这个脚本看它们
    在什么地方会出错，再和"上下文相关"的句子向量（第 05 号脚本那个迷你 GPT 的隐藏状态，取平均）对比：

    1. 只差一个 not、只是调换了 but 前后两半的句子，它们的向量有多像？（演示 1）
    2. 用最近邻做情感分类：简单句、but 句、not 句分别能做对多少？（演示 2）
    3. 拿一句话去检索最相似的句子，会找回什么？（演示 3，第 7 讲 RAG 的预告）

【怎么读输出】
    余弦相似度越接近 1 表示越像。演示 2 的准确率按评论的三种类型分开报告（见 pretrain_world.py）。

【运行】python3 03_sentence_vectors.py     （约 30 秒；第一次运行时如果还没有迷你 GPT 的缓存，
        会先训练它，再多 50 秒左右）
【依赖】numpy、PyTorch；import 同目录的 pretrain_world.py、embeddings.py、tiny_gpt.py、console.py。
"""

import math
from collections import Counter

import numpy as np

from console import pct, pct_ci, print_header, print_observations, print_table
from embeddings import normalize, train_sgns
from pretrain_world import WordTokenizer, build_corpus, build_world, corpus_text, review_probes
from tiny_gpt import get_model, hidden_states


# ---------------------------------------------------------------------------
# 三种句子向量
# ---------------------------------------------------------------------------

class AverageEncoder:
    """
    句子向量 = 句子里所有词的（SGNS）词向量的平均。

    【边界】句子里的词都不在词表里时返回全零向量。平均之后再归一化成长度 1。
    """

    def __init__(self, X, tok):
        self.Xn, self.tok = normalize(X), tok

    def encode(self, text):
        ids = [self.tok.stoi[w] for w in text.split() if w in self.tok.stoi]
        v = self.Xn[ids].mean(axis=0) if ids else np.zeros(self.Xn.shape[1])
        return v / (np.linalg.norm(v) or 1.0)


class TfidfEncoder:
    """
    TF-IDF：每个词一维，值 = 这个词在句子里出现的次数 × log(总句数 / 含这个词的句数)。

    【思路】"the"、"was" 几乎每句都有，IDF 接近 0，被自动压低；"cold"、"not" 这种区分性强的词权重大。
        这是搜索引擎几十年来的基础做法（第 7 讲的 BM25 是它的改进版）。
    【参数】docs：用来统计"含这个词的句数"的句子集合
    【边界】在 docs 里没出现过的词被忽略。
    """

    def __init__(self, docs):
        self.vocab = sorted({w for d in docs for w in d.split()})
        self.index = {w: i for i, w in enumerate(self.vocab)}
        df = Counter(w for d in docs for w in set(d.split()))
        self.idf = np.array([math.log(len(docs) / df[w]) for w in self.vocab])

    def encode(self, text):
        v = np.zeros(len(self.vocab))
        for w in text.split():
            if w in self.index:
                v[self.index[w]] += 1.0
        v *= self.idf
        return v / (np.linalg.norm(v) or 1.0)


class GPTEncoder:
    """
    上下文相关的句子向量：把句子放进 "<doc> reviews : review : 句子"，取迷你 GPT 最后一层在句子各个
    位置上的向量，求平均。

    【思路】因为模型是因果的，"not" 后面每个位置的向量都"知道"前面有个 not；调换 but 前后两半，
        每个位置看到的前文都变了。所以这种平均**对顺序敏感**。
        Sentence-BERT（Reimers & Gurevych 2019）也是"对 Transformer 的输出做平均池化"，
        只是模型是双向的、并且专门用句子对训练过。
    【边界】这个迷你 GPT 从没被专门训练来"产生好的句子向量"，它只是被训练来预测下一个词。
    """

    def __init__(self, model, tok, prefix="<doc> reviews : review :"):
        self.model, self.tok, self.prefix = model, tok, prefix
        self.n_prefix = len(prefix.split())

    def encode(self, text):
        h = hidden_states(self.model, self.tok, f"{self.prefix} {text}")[self.n_prefix:]
        v = h.mean(dim=0).numpy().astype(np.float64)
        return v / (np.linalg.norm(v) or 1.0)


# ---------------------------------------------------------------------------
def train_reviews(docs):
    """从训练语料的评论页面里拆出 (评论正文, 标签)。"""
    out = []
    for d in docs:
        if d.kind != "reviews":
            continue
        for chunk in d.text.split("review : ")[1:]:
            text, rest = chunk.split(" . rating : ")
            out.append((text, rest.split()[0]))
    return out


def demo1_pairs(encoders):
    """演示 1：几对句子的相似度。"""
    print_header("演示 1：词袋向量不在乎顺序和否定")
    pairs = [
        ("the soup was cold but the bread was fresh", "the bread was fresh but the soup was cold", "相反（只是调换了顺序）"),
        ("the soup was cold", "the soup was not cold", "相反（只差一个 not）"),
        ("the staff was rude", "the staff was friendly", "相反（换了一个形容词）"),
        ("the soup was cold", "the bread was stale", "相同（没有共同的实词）"),
        ("the room was clean and the view was lovely", "the bed was quiet and the music was great", "相同（没有共同的实词）"),
    ]
    rows = []
    sims = {}
    for a, b, rel in pairs:
        r = [f"{a!r}\n    vs {b!r}", rel]
        for name, enc in encoders.items():
            s = float(enc.encode(a) @ enc.encode(b))
            sims[(rel, name)] = s
            r.append(f"{s:.3f}")
        rows.append(r)
    headers = ["句子对", "情感关系"] + list(encoders)
    for r in rows:
        print(f"  {r[0]}")
        print("    " + "   ".join(f"{h}: {v}" for h, v in zip(headers[1:], r[1:])))
    k1 = "相反（只是调换了顺序）"
    print_observations([
        f"调换 but 前后两半，词袋完全不变：平均词向量和 TF-IDF 的相似度都是 {sims[(k1, '平均词向量')]:.3f}，"
        f"可两句话的情感正好相反。迷你 GPT 的向量是 {sims[(k1, '迷你 GPT')]:.3f}，比表里别的'同情感'句子对都低——"
        "它的向量对顺序敏感。",
        "只差一个 not：平均词向量几乎认为是同一句话（not 只是 4 个词里的一个）；TF-IDF 给 not 很高的权重，所以相似度掉得多一些——"
        "但它并不知道 not 会**反转**意思，只知道'多了一个稀有词'。",
        "没有共同实词、情感相同的两句话：TF-IDF 的相似度很低（它只认字面上相同的词）；平均词向量高一些（cold 和 stale 的向量本来就近）。",
        "注意：'相似度高'对于区分情感并不重要，重要的是**同情感的对比反情感的对更高**。演示 2 直接测这件事。",
    ])


def knn_eval(enc, train, probes, k=7):
    """
    k 近邻分类：在训练评论里找和测试评论最像的 k 条，投票决定标签。

    【返回】{类型: (对了几条, 一共几条)}
    【Python 小知识】np.argpartition(-sims, k)[:k] 找出最大的 k 个值的下标，比完整排序快；
        这里为了结果稳定（并列时按下标），直接用 np.argsort。
    """
    M = np.stack([enc.encode(t) for t, _ in train])
    labels = [y for _, y in train]
    res = {}
    for text, label, kind in probes:
        sims = M @ enc.encode(text)
        top = np.argsort(-sims, kind="stable")[:k]
        vote = Counter(labels[i] for i in top).most_common()
        pred = vote[0][0] if len(vote) == 1 or vote[0][1] > vote[1][1] else "positive"
        ok, n = res.get(kind, (0, 0))
        res[kind] = (ok + (pred == label), n + 1)
    return res


def demo2_knn(encoders, train, probes):
    """演示 2：最近邻情感分类。"""
    print_header("演示 2：用句子向量做最近邻情感分类（k = 7）")
    rows = []
    results = {}
    for name, enc in encoders.items():
        res = knn_eval(enc, train, probes)
        results[name] = res
        tot_ok = sum(ok for ok, _ in res.values())
        rows.append([name] + [pct_ci(*res[k]) for k in ["simple", "but", "not"]] + [pct(tot_ok, len(probes))])
    print(f"训练集：{len(train)} 条带标签的评论（来自训练语料）；测试集：{len(probes)} 条新生成的评论")
    print_table(["句子向量", "简单句 [95% CI]", "but 句", "not 句", "总体"], rows)
    avg, tf, gpt = results["平均词向量"], results["TF-IDF"], results["迷你 GPT"]
    print_observations([
        "简单句（所有形容词同一极性）大家都能做对：数一数褒义词、贬义词就够了。",
        f"but 句上两种词袋方法都在随机水平或更低（{pct(*avg['but'])} / {pct(*tf['but'])}）："
        "一句 but 句里褒贬两种词都有，词袋不知道哪一半算数。",
        f"not 句上 TF-IDF 只有 {pct(*tf['not'])}：和 'the coffee was not warm' 字面最像的，常常是同样的词、只是没有 not 的句子"
        f"（演示 3 正是这样），标签正好相反。平均词向量好一些（{pct(*avg['not'])}），一个可能的原因是 not 的向量把它拉向了别的含 not 的句子"
        "（演示 3 里正是这样），k 近邻投票于是间接利用了'not + 某个形容词'的组合——但这不是'理解'否定，换一种句式就会失效。",
        f"迷你 GPT 的向量在 but 句、not 句上分别是 {pct(*gpt['but'])}、{pct(*gpt['not'])}。它从来没被教过'情感'，"
        "只是为了预测 'rating :' 后面的词，必须在内部把 not、but 的作用算清楚——好的表示是**预测任务的副产品**。",
        "这就是从 word2vec 时代走到 ELMo / BERT / GPT 时代的原因：词的意思取决于它周围的词，"
        "而句子的意思取决于词的顺序。",
    ])


def demo3_retrieval(encoders, train):
    """演示 3：检索最相似的句子（RAG 的预告）。"""
    print_header("演示 3：拿一句话去检索最相似的训练评论（第 7 讲 RAG 的预告）")
    query = "the coffee was not warm"
    print(f"查询：{query!r}（情感：negative）")
    rows, wrong = [], {}
    for name, enc in encoders.items():
        M = np.stack([enc.encode(t) for t, _ in train])
        top = np.argsort(-(M @ enc.encode(query)), kind="stable")[:10]
        wrong[name] = sum(train[i][1] != "negative" for i in top)
        for r, i in enumerate(top[:3]):
            rows.append([name if r == 0 else "", r + 1, train[i][0], train[i][1]])
    print_table(["句子向量", "名次", "找回的评论", "它的标签"], rows, aligns=["<", ">", "<", "<"])
    print("  前 10 名里情感和查询相反的条数：" + "，".join(f"{k} {v}" for k, v in wrong.items()))
    print_observations([
        "TF-IDF 找回的是**字面上**最像的句子：第 1 名是一模一样的句子，接下来却是 'the coffee was warm'——"
        "只少了一个 not，意思完全相反。字面匹配对查找事实很有用，但不保证意思相同。",
        "平均词向量和迷你 GPT 找回的都是别的 'the coffee was not ...' 句子，这一次情感也都对；"
        "但演示 2 已经看到，平均词向量换到 but 句上就不行了。",
        "检索器决定了模型能看到什么证据：第 7 讲会看到，'找回了相似但答案不对的段落'是 RAG 系统最常见的失败之一。",
    ])


def main():
    world = build_world()
    docs = build_corpus(world)
    text = corpus_text(docs)
    tok = WordTokenizer(text)
    ids = tok.encode(text)
    print("  训练 SGNS 词向量……")
    X_sg, _ = train_sgns(ids, len(tok), dim=32, window=2, epochs=3, skip=(tok.stoi["<doc>"],))
    model, gtok, _ = get_model("base")
    assert gtok.itos == tok.itos
    train = train_reviews(docs)[:2000]
    encoders = {
        "平均词向量": AverageEncoder(X_sg, tok),
        "TF-IDF": TfidfEncoder([t for t, _ in train]),
        "迷你 GPT": GPTEncoder(model, gtok),
    }
    demo1_pairs(encoders)
    demo2_knn(encoders, train, review_probes())
    demo3_retrieval(encoders, train)


if __name__ == "__main__":
    main()
