#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
02_retrieval_metrics.py — 检索指标计算器：recall@k / precision@k / MRR / nDCG（呼应模块 07）

演示 RAG 评估的"检索层"四件套，以及两件事：
1. 两种排序策略（词法打分 vs 无信号基线）在同一金标集上的指标差异；
2. 指标作弊警告：recall@k 随 k 单调不降（全库塞入 = recall 100%），
   必须与 precision@k / nDCG 同看。

- 金标：每条查询带分级相关度（0=无关,1=沾边,2=相关,3=完美答案块）——nDCG 需要。
- 打分器：玩具版 BM25 直觉（词重叠 × IDF），仅用于产生可复现的排名，
  不代表真实嵌入/词法引擎的质量。
- 无网络、无 API key、纯标准库、输出确定。

运行：python3 02_retrieval_metrics.py
对应教学点：模块 12 README §6.5 / §10.6；模块 07 §检索评估
"""

import math

# ---------------------------------------------------------------------------
# 迷你语料（chunk id → 文本）与金标查询集（query → {chunk: 分级相关度}）
# ---------------------------------------------------------------------------

CORPUS = {
    "c01": "退款政策：购买后 30 天内可申请全额退款，需提供订单号。",
    "c02": "账户安全：如何开启两步验证与重置密码的步骤。",
    "c03": "错误码 TS-999 表示支付网关超时，请稍后重试或更换支付方式。",
    "c04": "退货流程：在订单页点击申请退货，选择原因，寄回商品后 5 个工作日退款。",
    "c05": "会员计划：Pro 会员每月可无限次使用高级功能。",
    "c06": "发票：电子发票将在订单完成后 24 小时内发送到您的邮箱。",
    "c07": "TS-999 故障排查：检查网络代理设置，确认防火墙未拦截支付域名。",
    "c08": "公司简介与发展历程。",
}

GOLDEN = {
    "退款要多久到账？": {"c04": 3, "c01": 2, "c06": 1},          # c04=完美, c01=相关, c06=沾边
    "付款时报错 TS-999 怎么办": {"c03": 3, "c07": 3},
    "怎么开发票？": {"c06": 3},
    "如何重置密码": {"c02": 3},
    "退款政策是什么": {"c01": 3, "c04": 2},
}

STOPWORDS = set("的 了 吗 怎么 如何 是 什么 在 到 与 和 请 您 您好".split())


def tokenize(text):
    """玩具分词：英文/错误码按词，中文串取全部相邻 2-gram（含奇偶两种偏移）。"""
    tokens, i = [], 0
    cjk = ""
    while i < len(text):
        ch = text[i]
        if ch.isascii() and (ch.isalnum() or ch == "-"):
            j = i
            while j < len(text) and (text[j].isascii() and (text[j].isalnum() or text[j] == "-")):
                j += 1
            tokens.append(text[i:j].lower())
            i = j
        elif not ch.isspace() and ch not in "，。：；？！、":
            cjk += ch
            i += 1
        else:
            i += 1
    grams = {cjk[i:i + 2] for i in range(len(cjk) - 1)}
    return ([t for t in tokens if t not in STOPWORDS]
            + [t for t in grams if t not in STOPWORDS])


# ---------------------------------------------------------------------------
# 两个"检索系统"：A=词法重叠打分（有信号），B=固定乱序（无信号基线）
# ---------------------------------------------------------------------------

DF = {}
for chunk in CORPUS.values():
    for tok in set(tokenize(chunk)):
        DF[tok] = DF.get(tok, 0) + 1
N_DOCS = len(CORPUS)


def idf(tok):
    return math.log(1 + N_DOCS / (1 + DF.get(tok, 0)))


def lexical_score(query, chunk_id):
    q = tokenize(query)
    c = tokenize(CORPUS[chunk_id])
    overlap = set(q) & set(c)
    return sum(idf(t) for t in overlap)


BASELINE_ORDER = ["c08", "c05", "c02", "c06", "c01", "c04", "c07", "c03"]  # 固定乱序


def retrieve(query, k, strategy):
    """返回 ranked chunk id 列表（截断到 k）。"""
    if strategy == "lexical":
        scored = sorted(CORPUS, key=lambda cid: -lexical_score(query, cid))
    else:
        scored = BASELINE_ORDER
    return scored[:k]


# ---------------------------------------------------------------------------
# 四件套指标（k 截断后的排名列表 + 分级金标）
# ---------------------------------------------------------------------------

def relevance(rels, cid):
    return rels.get(cid, 0)


def recall_at_k(ranked, rels, k):
    top = ranked[:k]
    relevant = {cid for cid, r in rels.items() if r > 0}
    if not relevant:
        return None
    hit = len(relevant & set(top))
    return hit / len(relevant)


def precision_at_k(ranked, rels, k):
    top = ranked[:k]
    hit_count = len([cid for cid in top if relevance(rels, cid) > 0])
    return hit_count / k


def mrr(ranked, rels):
    """平均倒数排名用单条查询的倒数排名（RR）；多条查询在外面取平均。"""
    for rank, cid in enumerate(ranked, start=1):
        if relevance(rels, cid) > 0:
            return 1.0 / rank
    return 0.0


def ndcg_at_k(ranked, rels, k):
    """DCG = Σ rel_i / log2(i+1)；IDCG = 理想排序的 DCG；nDCG = DCG/IDCG。"""
    def dcg(scores):
        total = 0.0
        for i, rel in enumerate(scores, start=1):
            total += rel / math.log2(i + 1)
        return total

    got = [relevance(rels, cid) for cid in ranked[:k]]
    ideal = sorted(rels.values(), reverse=True)[:k]
    ideal_dcg = dcg(ideal)
    if ideal_dcg == 0:
        return None
    return dcg(got) / ideal_dcg


# ---------------------------------------------------------------------------
# 汇总输出
# ---------------------------------------------------------------------------

def evaluate(strategy, k):
    rows = []
    agg = {"recall": [], "precision": [], "rr": [], "ndcg": []}
    for q, rels in GOLDEN.items():
        ranked = retrieve(q, k=k, strategy=strategy)
        r = recall_at_k(ranked, rels, k)
        p = precision_at_k(ranked, rels, k)
        rr = mrr(ranked, rels)
        nd = ndcg_at_k(ranked, rels, k)
        agg["recall"].append(r)
        agg["precision"].append(p)
        agg["rr"].append(rr)
        agg["ndcg"].append(nd)
        top_show = [f"{cid}({relevance(rels, cid)})" for cid in ranked]
        rows.append((q, " ".join(top_show), r, p, rr, nd))
    return rows, agg


def mean(xs):
    vals = [x for x in xs if x is not None]
    return sum(vals) / len(vals) if vals else None


def report(strategy, k):
    rows, agg = evaluate(strategy, k)
    print(f"\n策略 = {strategy:<10} k = {k}")
    print(f"{'查询':<22}{'Top-k (括号=金标相关度)':<40}{'recall':>7}{'prec':>7}{'RR':>6}{'nDCG':>7}")
    for q, top, r, p, rr, nd in rows:
        print(f"{q:<22}{top:<40}{r:>7.2f}{p:>7.2f}{rr:>6.2f}{nd:>7.3f}")
    print(f"{'宏平均':<62}{mean(agg['recall']):>7.2f}{mean(agg['precision']):>7.2f}"
          f"{mean(agg['rr']):>6.2f}{mean(agg['ndcg']):>7.3f}")
    return {m: mean(agg[m]) for m in agg}


def main():
    print("=" * 86)
    print("检索指标计算器 —— recall@k / precision@k / MRR / nDCG（金标分级相关度 0-3）")
    print("=" * 86)

    # 对比 1：同 k=3，词法 vs 无信号基线
    lex3 = report("lexical", 3)
    base3 = report("baseline", 3)

    # 对比 2：指标作弊演示 —— k 从 3 → 8（全库）
    lex8 = report("lexical", 8)

    print("\n" + "=" * 86)
    print("对比结论")
    print("=" * 86)
    print(f"1) k=3 时 lexical 全面优于无信号基线: "
          f"recall {lex3['recall']:.2f} vs {base3['recall']:.2f}, "
          f"nDCG {lex3['ndcg']:.3f} vs {base3['ndcg']:.3f}")
    print(f"2) 指标作弊警告: k 3→8 时 recall {lex3['recall']:.2f}→{lex8['recall']:.2f} "
          f"(单调不降; k=全库时必然=1.0),")
    print(f"   但 precision {lex3['precision']:.2f}→{lex8['precision']:.2f}, "
          f"nDCG {lex3['ndcg']:.3f}→{lex8['ndcg']:.3f} (近持平, 对多塞无关块不敏感)")
    print("   → recall 必须与 precision@k 同看，否则'全库塞入'式方案就能作弊")
    print("     (precision 被 5 个无关块稀释到 0.23)。")
    print("3) nDCG 与 recall 的分工: recall 只问'金标块进没进 Top-k'，")
    print("   nDCG 还问'完美答案块(c04/c03,rel=3)是否排在沾边块(rel=1)前面'（对数折损）。")


if __name__ == "__main__":
    main()