# hybrid_retrieval_complete.py
# 完整的多路召回 + RRF 融合 + Reranker 重排序 + HyDE 实现

import math
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass


@dataclass
class RetrievedDoc:
    """检索结果的统一数据结构。"""
    doc_id: str
    content: str
    score: float
    source: str  # "dense" | "bm25" | "metadata"


# ============================================================
# 1. BM25 实现（轻量版，不依赖额外库）
# ============================================================

class BM25:
    """
    BM25 实现。
    生产环境推荐使用：
    - rank_bm25 库（pip install rank-bm25）
    - Elasticsearch 内置的 BM25（更适合大规模）
    """

    def __init__(self, k1: float = 1.5, b: float = 0.75):
        # k1：词频饱和参数，越大词频的影响越持久
        # b：文档长度归一化参数，0=不归一化，1=完全归一化
        self.k1 = k1
        self.b = b
        self.docs: List[List[str]] = []
        self.doc_contents: List[str] = []
        self.doc_freq: Dict[str, int] = {}  # 每个词出现在多少文档中
        self.avgdl: float = 0.0

    def fit(self, documents: List[str]):
        """建立 BM25 索引。"""
        self.doc_contents = documents
        self.docs = [doc.split() for doc in documents]

        # 计算文档频率（df）
        for doc_tokens in self.docs:
            for token in set(doc_tokens):
                self.doc_freq[token] = self.doc_freq.get(token, 0) + 1

        # 计算平均文档长度
        self.avgdl = sum(len(d) for d in self.docs) / len(self.docs)

    def get_scores(self, query: str) -> List[float]:
        """计算查询对所有文档的 BM25 分数。"""
        query_tokens = query.split()
        n = len(self.docs)
        scores = []

        for doc_tokens in self.docs:
            score = 0.0
            doc_len = len(doc_tokens)
            token_freq = {}
            for token in doc_tokens:
                token_freq[token] = token_freq.get(token, 0) + 1

            for token in query_tokens:
                if token not in self.doc_freq:
                    continue
                df = self.doc_freq[token]
                # IDF：文档频率越低，词越重要
                idf = math.log((n - df + 0.5) / (df + 0.5) + 1)
                tf = token_freq.get(token, 0)
                # BM25 词频计算（含饱和和长度归一化）
                tf_score = (tf * (self.k1 + 1)) / (
                    tf + self.k1 * (1 - self.b + self.b * doc_len / self.avgdl)
                )
                score += idf * tf_score

            scores.append(score)
        return scores

    def retrieve(self, query: str, top_k: int = 10) -> List[RetrievedDoc]:
        """检索 top_k 个最相关文档。"""
        scores = self.get_scores(query)
        ranked = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)
        return [
            RetrievedDoc(
                doc_id=str(idx),
                content=self.doc_contents[idx],
                score=score,
                source="bm25"
            )
            for idx, score in ranked[:top_k]
            if score > 0
        ]


# ============================================================
# 2. RRF 融合算法
# ============================================================

def reciprocal_rank_fusion(
    result_lists: List[List[RetrievedDoc]],
    k: int = 60
) -> List[RetrievedDoc]:
    """
    RRF 融合多路检索结果。
    k=60 是经验最优值，来自原始论文的实验结论。
    通常不需要调整这个参数。
    """
    # 记录每个文档的 RRF 分数（累加各路的倒数排名分）
    rrf_scores: Dict[str, float] = {}
    # 保留文档内容（用第一次出现的结果）
    doc_contents: Dict[str, str] = {}
    doc_sources: Dict[str, List[str]] = {}

    for result_list in result_lists:
        for rank, doc in enumerate(result_list, start=1):
            doc_id = doc.doc_id
            # RRF 核心公式：1 / (k + rank)
            rrf_scores[doc_id] = rrf_scores.get(doc_id, 0) + 1.0 / (k + rank)

            if doc_id not in doc_contents:
                doc_contents[doc_id] = doc.content
                doc_sources[doc_id] = []
            doc_sources[doc_id].append(doc.source)

    # 按 RRF 分数降序排列
    sorted_docs = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)

    return [
        RetrievedDoc(
            doc_id=doc_id,
            content=doc_contents[doc_id],
            score=score,
            source="+".join(doc_sources[doc_id])  # 标记来自哪几路
        )
        for doc_id, score in sorted_docs
    ]


# ============================================================
# 3. Reranker 重排序
# ============================================================

class BGEReranker:
    """
    BGE Reranker 封装。
    安装：pip install sentence-transformers
    首次运行会下载模型（约 1GB）：BAAI/bge-reranker-base
    """

    def __init__(self, model_name: str = "BAAI/bge-reranker-base"):
        self.model_name = model_name
        self._model = None

    def _load_model(self):
        """惰性加载模型，避免不使用时占用资源。"""
        if self._model is None:
            try:
                from sentence_transformers import CrossEncoder
                self._model = CrossEncoder(self.model_name)
            except ImportError:
                raise ImportError("请安装：pip install sentence-transformers")

    def rerank(self, query: str, docs: List[RetrievedDoc], top_k: int = 5) -> List[RetrievedDoc]:
        """
        用 Cross-Encoder 对候选文档重新打分排序。
        输入：第一阶段召回的 top-20 候选文档
        输出：精排后的 top-k 文档
        """
        self._load_model()

        # Cross-Encoder 输入格式：(query, document) 对
        pairs = [(query, doc.content) for doc in docs]
        scores = self._model.predict(pairs)

        # 按分数降序排列
        reranked = sorted(zip(docs, scores), key=lambda x: x[1], reverse=True)

        return [
            RetrievedDoc(
                doc_id=doc.doc_id,
                content=doc.content,
                score=float(score),
                source=f"reranked({doc.source})"
            )
            for doc, score in reranked[:top_k]
        ]


class MockReranker:
    """
    模拟 Reranker（不需要安装模型）。
    用于演示流程，实际使用时替换为 BGEReranker。
    """

    def rerank(self, query: str, docs: List[RetrievedDoc], top_k: int = 5) -> List[RetrievedDoc]:
        """模拟重排序：给包含查询关键词的文档更高分。"""
        query_words = set(query.lower().split())
        scored = []
        for doc in docs:
            # 简单的关键词重叠打分（模拟 Cross-Encoder 的行为）
            doc_words = set(doc.content.lower().split())
            overlap = len(query_words & doc_words)
            scored.append((doc, overlap + doc.score * 0.1))

        scored.sort(key=lambda x: x[1], reverse=True)
        return [doc for doc, _ in scored[:top_k]]


# ============================================================
# 4. HyDE 查询改写
# ============================================================

class HyDEQueryExpander:
    """
    HyDE（Hypothetical Document Embeddings）实现。
    核心思路：用 LLM 生成假设答案，用假设答案代替原始问题做检索。
    """

    def __init__(self, llm=None):
        self.llm = llm

    def generate_hypothetical_document(self, query: str) -> str:
        """
        生成假设答案文档。
        实际场景接入真实 LLM：
        from langchain_openai import ChatOpenAI
        llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
        prompt = f"请为以下问题生成一段可能的答案文档（100-150字，风格像技术文档）：\n{query}"
        return llm.invoke(prompt).content
        """
        # 模拟生成（演示用）
        return (
            f"关于「{query}」的技术文档：该功能通过以下步骤实现："
            f"首先需要安装相关依赖包，然后导入所需模块，"
            f"接着按照标准流程配置参数，最后调用核心函数完成操作。"
            f"常见问题包括环境配置错误和参数格式不匹配，"
            f"解决方案是检查版本兼容性并参考官方文档。"
        )

    def expand_query(self, query: str) -> Tuple[str, str]:
        """
        返回 (原始查询, 假设文档)。
        检索时用假设文档，回答时用原始查询。
        """
        hypothetical_doc = self.generate_hypothetical_document(query)
        return query, hypothetical_doc


# ============================================================
# 5. 完整的混合检索管道（LangChain EnsembleRetriever 版本）
# ============================================================

def build_hybrid_retrieval_with_langchain(documents: List[str]):
    """
    用 LangChain 的 EnsembleRetriever 实现混合检索。
    需要安装：pip install langchain langchain-community langchain-openai chromadb
    """
    try:
        from langchain_community.retrievers import BM25Retriever
        from langchain.retrievers import EnsembleRetriever
        from langchain_community.vectorstores import Chroma
        from langchain_openai import OpenAIEmbeddings
        from langchain.schema import Document

        docs = [Document(page_content=d) for d in documents]

        # 稀疏检索：BM25
        bm25_retriever = BM25Retriever.from_documents(docs)
        bm25_retriever.k = 10  # 每路初始召回 10 个

        # 稠密检索：向量相似度
        embeddings = OpenAIEmbeddings()
        vectorstore = Chroma.from_documents(docs, embeddings)
        dense_retriever = vectorstore.as_retriever(search_kwargs={"k": 10})

        # 融合：EnsembleRetriever 内部使用 RRF
        # weights 对应每路的权重（影响 RRF 中的排名加权方式）
        ensemble_retriever = EnsembleRetriever(
            retrievers=[bm25_retriever, dense_retriever],
            weights=[0.5, 0.5]  # 等权融合；关键词密集的场景可以调高 BM25 的权重
        )

        return ensemble_retriever

    except ImportError as e:
        print(f"缺少依赖：{e}")
        print("安装：pip install langchain langchain-community langchain-openai chromadb rank-bm25")
        return None


# ============================================================
# 6. 独立运行演示（不需要任何外部依赖）
# ============================================================

def run_hybrid_demo():
    """完整演示：多路召回 + RRF + 重排序。"""
    print("=" * 60)
    print("多路召回与融合排序演示")
    print("=" * 60)

    # 测试文档集
    documents = [
        ("doc_1", "Python csv 模块提供了 reader() 和 writer() 函数，用于读写 CSV 文件。"),
        ("doc_2", "订单号 ORD-2024-8823 于 2024 年 3 月 15 日创建，状态：已发货。"),
        ("doc_3", "机器学习中的梯度下降算法通过迭代优化模型参数，最小化损失函数。"),
        ("doc_4", "CSV（逗号分隔值）是一种通用的数据交换格式，Python 有原生支持。"),
        ("doc_5", "退货政策：购买 7 天内可申请退货，商品需保持原包装完整。"),
        ("doc_6", "如何用 pandas 读取 CSV 文件：pd.read_csv('filename.csv')，支持编码设置。"),
        ("doc_7", "NumPy 提供了 genfromtxt() 函数，可以从 CSV 文件加载数组数据。"),
    ]

    contents = [content for _, content in documents]
    doc_ids = [doc_id for doc_id, _ in documents]

    # 初始化 BM25
    bm25 = BM25()
    bm25.fit(contents)

    # 模拟向量检索（演示用：用简单的关键词重叠代替真实向量检索）
    def mock_dense_retrieve(query: str, top_k: int = 5) -> List[RetrievedDoc]:
        """模拟向量检索（演示用，实际替换为真实向量库调用）。"""
        query_words = set(query.lower().split())
        scored = []
        for i, content in enumerate(contents):
            content_words = set(content.lower().split())
            # 模拟语义相似度（用词重叠近似）
            overlap = len(query_words & content_words)
            if overlap > 0:
                scored.append(RetrievedDoc(
                    doc_id=doc_ids[i],
                    content=content,
                    score=overlap / len(query_words),
                    source="dense"
                ))
        scored.sort(key=lambda x: x.score, reverse=True)
        return scored[:top_k]

    # 测试查询
    test_queries = [
        "Python 如何读取 CSV 文件",     # 语义查询，应该找到 doc_1, doc_4, doc_6, doc_7
        "ORD-2024-8823 订单状态",       # 精确字符串，应该找到 doc_2
    ]

    reranker = MockReranker()
    hyde = HyDEQueryExpander()

    for query in test_queries:
        print(f"\n{'=' * 50}")
        print(f"查询：{query}")
        print("-" * 50)

        # 方案一：单纯向量检索
        dense_only = mock_dense_retrieve(query, top_k=3)
        print("\n[单路向量检索 Top-3]：")
        for i, doc in enumerate(dense_only, 1):
            print(f"  {i}. [{doc.score:.3f}] {doc.content[:60]}...")

        # 方案二：多路召回 + RRF
        dense_results = mock_dense_retrieve(query, top_k=5)
        bm25_results = bm25.retrieve(query, top_k=5)
        # 修正 bm25 结果的 doc_id（BM25 用的是数字索引，映射回原始 id）
        for doc in bm25_results:
            doc.doc_id = doc_ids[int(doc.doc_id)]

        fused = reciprocal_rank_fusion([dense_results, bm25_results])
        print(f"\n[多路召回 + RRF Top-3]（密集检索 + BM25 融合）：")
        for i, doc in enumerate(fused[:3], 1):
            print(f"  {i}. [{doc.score:.4f}] [{doc.source}] {doc.content[:60]}...")

        # 方案三：融合 + Reranker
        reranked = reranker.rerank(query, fused[:6], top_k=3)
        print(f"\n[多路召回 + RRF + Reranker Top-3]（最终精排）：")
        for i, doc in enumerate(reranked, 1):
            print(f"  {i}. [{doc.score:.4f}] {doc.content[:60]}...")

        # HyDE 展示
        original_q, hypothetical_doc = hyde.expand_query(query)
        print(f"\n[HyDE] 生成的假设文档：")
        print(f"  {hypothetical_doc[:100]}...")


if __name__ == "__main__":
    run_hybrid_demo()