"""稀疏检索 / 简单 RAG"""
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


# 1. 简单切块
def chunk_text(text, chunk_size=60, overlap=15):
    chunks = []

    start = 0

    while start < len(text):
        end = start + chunk_size

        chunk = text[start:end]
        chunks.append(chunk)

        if end >= len(text):
            break

        start = end - overlap

    return chunks


# 2. 知识库原文
document = """
职工请事假，应当提前一个工作日通过OA系统提交申请，
经部门负责人审批同意后方可休假。

职工因疾病无法正常工作的，可以申请病假。
病假应当提供医院出具的相关证明材料。

员工因工作需要出差的，应提前提交出差申请。
住宿费用按照公司差旅管理制度规定的标准报销。

职工连续工作满一年后可以享受带薪年休假。
年休假天数根据职工累计工作年限确定。
""".replace("\n", "")


# 3. 文档切块
chunks = chunk_text(
    document,
    chunk_size=60,
    overlap=15
)

print("=== chunks ===")

for i, chunk in enumerate(chunks):
    print(i, chunk)


# 4. 给所有 chunk 建 TF-IDF 索引
vectorizer = TfidfVectorizer(
    analyzer="char",
    ngram_range=(1, 3)
)

doc_vectors = vectorizer.fit_transform(chunks)


# 5. 用户问题
query = "事假怎么申请？"


# 6. Query 转成同一个 TF-IDF 空间
query_vector = vectorizer.transform([query])


# 7. 计算 query 与所有 chunk 的相似度
scores = cosine_similarity(
    query_vector,
    doc_vectors
)[0]


# 8. Top-K
top_k = 3

top_indices = scores.argsort()[::-1][:top_k]


print("\n=== query ===")
print(query)

print("\n=== retrieval result ===")

for index in top_indices:
    print(
        f"score={scores[index]:.4f}",
        chunks[index]
    )