#!/usr/bin/env python3
"""02_chunking_strategies.py — 无需 API key 的分块策略 (chunking) 对比实验。

在一份合成的「多公司季报 + 政策」文档上, 对比五种分块策略对 BM25 检索的影响:
  A) Fixed-size           固定分块: 每 40 词一刀, 无视句子与标题边界
  B) Overlap              重叠分块: 固定 40 词窗口 + 12 词重叠 (步长 28)
  C) Semantic             语义分块: 相邻句子内容词 Jaccard 相似度低于阈值处切分
  D) Structural           结构化分块: 按 "##" 标题切节, 块文本包含标题链
  E) Semantic+Contextual  语义分块 + 上下文前缀 (Contextual Retrieval 的玩具版:
                          真实方案用 LLM 为每块生成 50-100 token 情境前缀,
                          此处以该块所在节的标题链近似)

评测协议: 每条金标查询的答案是一个句子 (golden sentence); 命中 = 含**完整**答案句
的块进入 Top-k (答案句被切断 = 不命中, 因为模型只看得到半句)。
指标: hit@1 / hit@3 / MRR (首个含完整答案句的块的排名倒数)。

设计意图 (对应模块 07 README §4):
  1. 固定分块会把长句拦腰切断 → 检索失败 (README: 块太小/切断 → 上下文丢失);
  2. 重叠分块能救回部分边界, 但块数与存储膨胀, 且救不了「主语丢失」;
  3. 语义分块保住句子完整, 但 "The company's revenue grew by 3%..." 这类代词句
     离开标题链后仍然缺主语 → 会输给竞争对手公司的显式句 (官方 SEC 财报例子,
     来源: https://www.anthropic.com/engineering/contextual-retrieval, 访问于 2026-08-30);
  4. 标题链 / 情境前缀把公司名与时间注入每个块 → 修复 context loss
     (官方: 上下文嵌入 + 上下文 BM25 使 Top-20 检索失败率降 49%, 再加 reranking
      累计降 67%; prompt caching 摊薄后 $1.02/百万文档 token, 来源同上)。

运行: python3 02_chunking_strategies.py
依赖: 仅 Python 标准库
"""
from __future__ import annotations

import math
import re
from collections import Counter

# ------------------------------------------------------------------ 参数 ----
CHUNK_WORDS = 40        # 固定分块窗口 (词)
OVERLAP_WORDS = 12      # 重叠词数 → 步长 = 40 - 12 = 28
SEMANTIC_THRESHOLD = 0.12  # 相邻句子内容词 Jaccard < 此值 → 在此切分

STOPWORDS = {
    "this", "that", "these", "those", "with", "from", "they", "them", "their",
    "were", "been", "have", "has", "had", "will", "would", "could", "should",
    "after", "before", "under", "over", "above", "below", "between", "through",
    "other", "more", "most", "less", "least", "than", "then", "when", "where",
    "which", "while", "both", "each", "every", "some", "such", "only", "also",
    "very", "much", "many", "into", "onto", "upon", "within", "without",
    "toward", "towards", "about", "against", "because", "during", "your",
}

# ------------------------------------------------------------ 合成文档 ------
# 注意: 各节里故意放了几条「代词句」(The company ... / Its ...) —— 它们被切出
# 所属节后就不知道说的是哪家公司, 这正是官方 Contextual Retrieval 文章里
# SEC 财报例子的复刻。
DOC = """## ACME Corp — Q2 2023 Financial Results
ACME Corp filed its quarterly report for the period ending June 30 2023.
Total revenue reached 214 million dollars this quarter.
Operating expenses fell by 1.2 million dollars compared to the first quarter.
Cash reserves stood at 88 million dollars at quarter end.
The board approved a share buyback program of up to 10 million dollars.
The company's revenue grew by 3% over the previous quarter.
The company reaffirmed its full year guidance of 8 to 10 percent growth.

## ACME Corp — Product and Engineering
The platform team shipped a new audit log export feature in June.
Mobile crash-free session rate improved to 99.2 percent.
The company hired 40 new engineers across European offices during the quarter, raising total headcount to 310 people.
Two data centers in Frankfurt were brought online for EU customers.
Legacy webhooks will be deprecated in the next release.

## Beta Systems — Q2 2023 Financial Results
Beta Systems reported second quarter results on July 20 2023.
Beta Systems revenue declined by 4% quarter over quarter.
Subscription churn stayed flat at 2.1 percent.
The company cut its full year guidance from 12 to 9 percent growth.
Free cash flow turned positive after three quarters of losses.
No dividend will be paid this year.

## Beta Systems — Security Incident
On May 12 2023 Beta Systems detected unauthorized access to a staging database.
The intrusion was contained within six hours.
No customer data was affected by the incident.
The company rotated all employee credentials within 24 hours.
A root cause analysis was published to the trust center.
Bug bounty payouts increased by 20 percent after the incident.

## Gamma Logistics — Returns and Refunds Policy
Gamma Logistics accepts returns of unopened items within 30 days of delivery.
Refunds are issued to the original payment method within five business days.
Return shipping is free for defective products.
The company accepts returns of damaged items within 45 days of delivery.
Items bought during a promotional sale follow a stricter 14 day window.
Bulk orders from resellers are not eligible for returns.

## Gamma Logistics — Shipping and Delivery
Domestic parcels arrive within two to three business days.
International orders placed through the standard checkout typically arrive within five to seven business days after leaving the regional warehouse in Rotterdam.
Orders above 50 dollars ship free with Gamma Logistics.
Tracking numbers are sent by email when the parcel leaves the warehouse.
Delivery to remote islands may take up to two weeks longer.
"""

# ------------------------------------------------------------ 金标查询集 ----
# intent 含义:
#   context_loss: 答案句本身缺少查询里的关键限定词 (如公司名) → 只有带标题链/
#                 情境前缀的策略能可靠命中
#   boundary:     答案句很长 (~24 词), 容易被固定分块切断
#   normal:       答案句自解释, 各策略理论上都能命中
QUERIES = [
    {"q": "By how much did ACME Corp revenue grow in Q2 2023",
     "golden_text": "The company's revenue grew by 3% over the previous quarter.",
     "intent": "context_loss"},
    {"q": "What is the size of the ACME share buyback program",
     "golden_text": "The board approved a share buyback program of up to 10 million dollars.",
     "intent": "normal"},
    {"q": "How did Beta Systems handle the May 2023 security incident",
     "golden_text": "The company rotated all employee credentials within 24 hours.",
     "intent": "context_loss"},
    {"q": "Beta Systems revenue change in Q2 2023",
     "golden_text": "Beta Systems revenue declined by 4% quarter over quarter.",
     "intent": "normal"},
    {"q": "When did Beta Systems detect unauthorized access to a staging database",
     "golden_text": "On May 12 2023 Beta Systems detected unauthorized access to a staging database.",
     "intent": "normal"},
    {"q": "Gamma Logistics returns policy for unopened items",
     "golden_text": "Gamma Logistics accepts returns of unopened items within 30 days of delivery.",
     "intent": "normal"},
    {"q": "free shipping threshold at Gamma Logistics",
     "golden_text": "Orders above 50 dollars ship free with Gamma Logistics.",
     "intent": "normal"},
    {"q": "How long do international orders take to arrive",
     "golden_text": ("International orders placed through the standard checkout "
                     "typically arrive within five to seven business days after "
                     "leaving the regional warehouse in Rotterdam."),
     "intent": "boundary"},
    {"q": "How many engineers did ACME Corp hire in Q2 2023",
     "golden_text": ("The company hired 40 new engineers across European offices "
                     "during the quarter, raising total headcount to 310 people."),
     "intent": "boundary"},
    {"q": "How do customers get tracking numbers for their parcel",
     "golden_text": "Tracking numbers are sent by email when the parcel leaves the warehouse.",
     "intent": "normal"},
]

# ---------------------------------------------------------------- 工具函数 --
_TOKEN_RE = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*")

def tokenize(text: str) -> list[str]:
    return _TOKEN_RE.findall(text.lower())

def content_words(text: str) -> set[str]:
    return {t for t in tokenize(text) if len(t) >= 4 and t not in STOPWORDS}

def jaccard(a: set[str], b: set[str]) -> float:
    if not a and not b:
        return 0.0
    return len(a & b) / len(a | b)

# ------------------------------------------------------------- 文档解析 ----
def parse_doc() -> tuple[list[str], list[list[str]], list[str]]:
    """把 DOC 解析为 (标题列表, 每节的句子列表, 全文档句子流)。

    返回:
      headings[i]  = 第 i 节标题文本 (不含 ##)
      sections[i]  = 第 i 节的句子列表
      flat         = 全文档按顺序的句子流 (供固定/重叠/语义分块使用, 不含标题)
      每个句子记录为 (sentence_text, section_idx)
    """
    headings: list[str] = []
    sections: list[list[str]] = []
    for block in DOC.strip().split("\n\n"):
        lines = [ln.strip() for ln in block.strip().splitlines()]
        heading = lines[0].lstrip("# ").strip()
        body = " ".join(lines[1:])
        sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", body) if s.strip()]
        headings.append(heading)
        sections.append(sentences)
    return headings, sections

# --------------------------------------------------------------- BM25 ------
class BM25:
    """经典 Okapi BM25 (k1=1.5, b=0.75), 与 01 号示例同一实现风格。"""

    def __init__(self, docs: list[list[str]], k1: float = 1.5, b: float = 0.75):
        self.k1, self.b = k1, b
        self.doc_len = [len(d) for d in docs]
        self.avgdl = sum(self.doc_len) / len(docs)
        self.tf = [Counter(d) for d in docs]
        df: Counter = Counter()
        for d in docs:
            df.update(set(d))
        n = len(docs)
        self.idf = {t: math.log(1 + (n - c + 0.5) / (c + 0.5)) for t, c in df.items()}

    def score(self, query: list[str], idx: int) -> float:
        s, dl = 0.0, self.doc_len[idx]
        for t in query:
            if t not in self.tf[idx]:
                continue
            f = self.tf[idx][t]
            idf = self.idf.get(t, 0.0)
            s += idf * f * (self.k1 + 1) / (f + self.k1 * (1 - self.b + self.b * dl / self.avgdl))
        return s

    def rank(self, query: str) -> list[int]:
        q = tokenize(query)
        return sorted(range(len(self.tf)), key=lambda i: (-self.score(q, i), i))

# --------------------------------------------------------------- 分块器 ----
# 统一块结构: {"id", "text", "full_sids"} —— full_sids = 完整落入该块的句子编号集合
SENT_SPLIT_RE = re.compile(r"(?<=[.!?])\s+")

def chunk_fixed(flat: list[str]) -> list[dict]:
    """A) 固定分块: 每 CHUNK_WORDS 词一刀 (朴素 pipeline 常见做法: 剥掉标题硬切)。"""
    words = tokenize(" ".join(flat))
    # 计算每个句子的词跨度 [start, end)
    spans, pos = [], 0
    for sent in flat:
        n = len(tokenize(sent))
        spans.append((pos, pos + n))
        pos += n
    chunks = []
    for i in range(0, len(words), CHUNK_WORDS):
        lo, hi = i, min(i + CHUNK_WORDS, len(words))
        sids = {k for k, (a, b) in enumerate(spans) if a >= lo and b <= hi}
        chunks.append({"id": f"A{len(chunks)+1:02d}",
                       "text": " ".join(words[lo:hi]),
                       "full_sids": sids})
    return chunks

def chunk_overlap(flat: list[str]) -> list[dict]:
    """B) 重叠分块: 窗口 CHUNK_WORDS, 步长 CHUNK_WORDS - OVERLAP_WORDS。"""
    words = tokenize(" ".join(flat))
    spans, pos = [], 0
    for sent in flat:
        n = len(tokenize(sent))
        spans.append((pos, pos + n))
        pos += n
    stride = CHUNK_WORDS - OVERLAP_WORDS
    chunks = []
    for i in range(0, len(words), stride):
        lo, hi = i, min(i + CHUNK_WORDS, len(words))
        sids = {k for k, (a, b) in enumerate(spans) if a >= lo and b <= hi}
        chunks.append({"id": f"B{len(chunks)+1:02d}",
                       "text": " ".join(words[lo:hi]),
                       "full_sids": sids})
        if hi >= len(words):
            break
    return chunks

def chunk_semantic(flat: list[str], sid_section: list[int],
                   prefix_headings: bool = False) -> list[dict]:
    """C/E) 语义分块: 相邻句子内容词 Jaccard < SEMANTIC_THRESHOLD 处切分。

    prefix_headings=False → 纯语义块 (不知道标题);
    prefix_headings=True  → 每块前面加上其首句所在节的标题链,
                            近似 Contextual Retrieval 的情境前缀。
    """
    boundaries = [0]
    for i in range(1, len(flat)):
        sim = jaccard(content_words(flat[i - 1]), content_words(flat[i]))
        if sim < SEMANTIC_THRESHOLD:
            boundaries.append(i)
    boundaries.append(len(flat))
    chunks = []
    for bi in range(len(boundaries) - 1):
        lo, hi = boundaries[bi], boundaries[bi + 1]
        sids = set(range(lo, hi))
        text = " ".join(flat[lo:hi])
        if prefix_headings:
            sec = sid_section[lo]
            text = f"[Context: section '{HEADINGS[sec]}'] " + text
        chunks.append({"id": ("E" if prefix_headings else "C") + f"{bi+1:02d}",
                       "text": text, "full_sids": sids})
    return chunks

def chunk_structural(sections: list[list[str]], headings: list[str]) -> list[dict]:
    """D) 结构化分块: 按 "##" 标题整节成块, 标题链保留在块文本中。"""
    chunks, sid0 = [], 0
    for si, sents in enumerate(sections):
        sids = set(range(sid0, sid0 + len(sents)))
        sid0 += len(sents)
        chunks.append({"id": f"D{si+1:02d}",
                       "text": headings[si] + " — " + " ".join(sents),
                       "full_sids": sids})
    return chunks

# --------------------------------------------------------------- 评测 ------
def evaluate(chunks: list[dict], golden_sids: list[int], queries: list[dict],
             k: int = 3) -> tuple[dict[str, float], list[int]]:
    """返回 ({hit@1, hit@3, MRR}, 每条查询的首个含完整答案句的块排名; 0=未命中)。"""
    bm25 = BM25([tokenize(c["text"]) for c in chunks])
    hit1, hit3, rrs, first_ranks = 0, 0, [], []
    for qi, item in enumerate(queries):
        ranked = bm25.rank(item["q"])
        golden = golden_sids[qi]
        rank = 0
        for r, ci in enumerate(ranked, start=1):
            if golden in chunks[ci]["full_sids"]:
                rank = r
                break
        first_ranks.append(rank)
        if rank == 1:
            hit1 += 1
        if 0 < rank <= k:
            hit3 += 1
        rrs.append(1.0 / rank if rank else 0.0)
    n = len(queries)
    return ({"hit@1": hit1 / n, "hit@3": hit3 / n, "MRR": sum(rrs) / n}, first_ranks)

# ----------------------------------------------------------------- 主流程 --
HEADINGS: list[str] = []  # chunk_semantic 引用 (在 main 中填充)

def main() -> None:
    global HEADINGS
    headings, sections = parse_doc()
    HEADINGS = headings
    flat = [s for sec in sections for s in sec]
    sid_section: list[int] = []
    for si, sec in enumerate(sections):
        sid_section.extend([si] * len(sec))

    # golden 句 → 句子编号
    golden_sids: list[int] = []
    for item in QUERIES:
        target = " ".join(tokenize(item["golden_text"]))
        found = -1
        for i, sent in enumerate(flat):
            if " ".join(tokenize(sent)) == target:
                found = i
                break
        if found < 0:
            raise SystemExit(f"golden sentence not found in corpus: {item['golden_text'][:50]}...")
        golden_sids.append(found)

    strategies: list[tuple[str, list[dict]]] = [
        ("A Fixed-size (40 词硬切)",        chunk_fixed(flat)),
        ("B Overlap (40 词 + 12 词重叠)",   chunk_overlap(flat)),
        ("C Semantic (Jaccard 阈值切分)",   chunk_semantic(flat, sid_section, prefix_headings=False)),
        ("D Structural (按标题整节)",       chunk_structural(sections, headings)),
        ("E Semantic + 上下文前缀",         chunk_semantic(flat, sid_section, prefix_headings=True)),
    ]

    total_words = len(tokenize(" ".join(flat)))
    print("=" * 86)
    print(f"合成文档: {len(flat)} 句 / {total_words} 词 / {len(sections)} 节 | 金标查询: {len(QUERIES)} 条")
    print("=" * 86)

    # ---- 块统计 ----
    print(f"\n{'策略':30s} | {'块数':4s} | {'平均词数/块':10s} | {'金标句被切断数':12s}")
    print("-" * 70)
    all_sids = set(range(len(flat)))
    for name, chunks in strategies:
        covered: set[int] = set()
        for c in chunks:
            covered |= c["full_sids"]
        split_golden = sum(1 for g in golden_sids if g not in covered)
        avg = sum(len(tokenize(c["text"])) for c in chunks) / len(chunks)
        print(f"{name:30s} | {len(chunks):<4d} | {avg:<10.1f} | {split_golden:<12d}")

    # ---- 聚合指标 ----
    print("\n" + "=" * 86)
    print("检索评测 (BM25; 命中 = 含完整金标句的块; k=3)")
    print("=" * 86)
    print(f"{'策略':30s} | {'hit@1':6s} | {'hit@3':6s} | {'MRR':6s}")
    print("-" * 60)
    per_query_ranks: dict[str, list[int]] = {}
    for name, chunks in strategies:
        metrics, ranks = evaluate(chunks, golden_sids, QUERIES)
        per_query_ranks[name] = ranks
        print(f"{name:30s} | {metrics['hit@1']:<6.2f} | {metrics['hit@3']:<6.2f} | {metrics['MRR']:<6.2f}")

    # ---- 逐查询明细 ----
    print(f"\n逐查询明细 (数字 = 首个含完整金标句的块排名, 0 = 未命中):")
    print(f"{'query':46s} | {'intent':12s} | " + " | ".join(n.split()[0] for n, _ in strategies))
    print("-" * 110)
    names = [n for n, _ in strategies]
    for qi, item in enumerate(QUERIES):
        cells = []
        for n in names:
            r = per_query_ranks[n][qi]
            cells.append(f"{n.split()[0]}:{(str(r) if r else '✗'):>2s}")
        print(f"{item['q'][:46]:46s} | {item['intent']:12s} | " + " | ".join(cells))

    print("""
结论 (与模块 07 README §4 的教学点一一对应):
  1. 固定分块 (A) 把长答案句拦腰切断 → 金标句被切断数最高; 即使句子侥幸完整,
     剥掉标题也让 "The company's revenue grew by 3%..." 这类代词句失去归属。
  2. 重叠分块 (B) 救回部分被切断的句子 (q5/q10 从未命中 → 命中), 代价是块数与
     索引存储膨胀 (10 块 → 14 块); 注意个别查询 (如 q8) 反而因块重排/长度归一
     变差 —— overlap 是缓解手段而非免费午餐, 且对主语丢失毫无帮助。
  3. 语义分块 (C) 保证句子完整、话题聚块, 但它不知道标题: 代词句离开标题链后,
     "ACME 收入增长" 这类查询会被 Beta Systems 的显式句或本公司的名字句抢走排名。
  4. 结构化分块 (D) 把公司名/时间随标题链注入每个块 → 本实验中全部命中
     (hit@1 = 1.00)。上下文前缀 (E) 是官方 Contextual Retrieval 的思想:
     索引前用 LLM 为每块生成 50-100 token 情境前缀 (本实验用节标题近似):
     q1 被修复到 rank 1, q3 从 12 → 5 —— 标题近似前缀是同节共享的, 同节碎块
     仍互相竞争; 真实方案为每块生成定制前缀并配合 reranking 精排。
     官方数据: 上下文嵌入+上下文 BM25 使 Top-20 检索失败率降 49%,
     再加 reranking 累计降 67%; prompt caching 摊薄后 $1.02/百万文档 token。
  5. 官方实现注意事项: chunk size / chunk boundary / chunk overlap 都会影响
     检索性能 —— 修改本文件顶部的 CHUNK_WORDS / OVERLAP_WORDS / SEMANTIC_THRESHOLD
     即可亲手验证。
  来源: https://www.anthropic.com/engineering/contextual-retrieval (访问于 2026-08-30)
""")

if __name__ == "__main__":
    main()