"""
直接测试 BM25 对象的 retrieve 方法
"""
import bm25s
import numpy as np
from src.core import RAGSystem

# 创建 RAG 系统实例
rag = RAGSystem()
rag.load_index()

# 获取 BM25 检索器
bm25_retriever = rag.bm25_store_manager.get_bm25_retriever()

if bm25_retriever and hasattr(bm25_retriever, 'bm25'):
    print("=== 直接测试 BM25 对象的 retrieve 方法 ===\n")
    
    query = "什么是RAG"
    print(f"查询: {query}\n")
    
    # 使用 bm25s.tokenize 进行分词（与 _retrieve 方法一致）
    print("1. 分词...")
    tokenized_query = bm25s.tokenize(
        query,
        stemmer=bm25_retriever.stemmer if not bm25_retriever.skip_stemming else None,
        token_pattern=bm25_retriever.token_pattern,
        show_progress=False,
    )
    print(f"   分词结果: {tokenized_query}")
    print(f"   分词结果类型: {type(tokenized_query)}")
    print(f"   分词结果形状: {tokenized_query.shape if hasattr(tokenized_query, 'shape') else 'N/A'}")
    
    # 直接调用 bm25.retrieve
    print("\n2. 调用 bm25.retrieve...")
    indexes, scores = bm25_retriever.bm25.retrieve(
        tokenized_query,
        k=bm25_retriever.similarity_top_k,
        show_progress=False,
    )
    
    print(f"   返回的 indexes 类型: {type(indexes)}")
    print(f"   返回的 scores 类型: {type(scores)}")
    print(f"   indexes 形状: {indexes.shape if hasattr(indexes, 'shape') else 'N/A'}")
    print(f"   scores 形状: {scores.shape if hasattr(scores, 'shape') else 'N/A'}")
    
    # 检查第一批结果（batched）
    print(f"\n3. 第一批结果:")
    first_indexes = indexes[0]
    first_scores = scores[0]
    print(f"   indexes[0]: {first_indexes}")
    print(f"   scores[0]: {first_scores}")
    print(f"   scores[0] 类型: {type(first_scores)}")
    
    # 尝试使用 get_scores 方法
    print("\n4. 尝试使用 get_scores 方法...")
    try:
        direct_scores = bm25_retriever.bm25.get_scores(tokenized_query)
        print(f"   get_scores 返回类型: {type(direct_scores)}")
        print(f"   get_scores 返回形状: {direct_scores.shape if hasattr(direct_scores, 'shape') else 'N/A'}")
        print(f"   get_scores 前5个分数: {direct_scores[0][:5] if hasattr(direct_scores, '__getitem__') else direct_scores[:5]}")
        print(f"   最大分数: {np.max(direct_scores)}")
        print(f"   最小分数: {np.min(direct_scores)}")
    except Exception as e:
        print(f"   调用 get_scores 失败: {e}")
    
    # 检查 BM25 索引的状态
    print("\n5. BM25 索引状态:")
    print(f"   corpus 类型: {type(bm25_retriever.corpus)}")
    print(f"   corpus 大小: {len(bm25_retriever.corpus) if hasattr(bm25_retriever.corpus, '__len__') else 'N/A'}")
    
    if hasattr(bm25_retriever.bm25, 'index'):
        print(f"   bm25.index 类型: {type(bm25_retriever.bm25.index)}")
        if hasattr(bm25_retriever.bm25.index, 'shape'):
            print(f"   bm25.index 形状: {bm25_retriever.bm25.index.shape}")
