# -*- coding: utf-8 -*-
"""中文 BM25 检索(jieba 分词 + rank_bm25)。文档 = 菜名。

菜名通常很短(2~6 字),除 jieba 词以外再补单字,提升短词召回。
"""
from __future__ import annotations

import re

import jieba
import numpy as np
from rank_bm25 import BM25Okapi

jieba.initialize()


def tokenize(text: str) -> list[str]:
    text = text.replace("_", " ").strip()
    words = [t for t in jieba.lcut_for_search(text) if t.strip()]
    chars = [c for c in re.sub(r"\s+", "", text) if "一" <= c <= "鿿"]
    toks = words + chars
    return toks if toks else [t for t in re.sub(r"\s+", "", text)] or ["_"]


class BM25Index:
    def __init__(self, names: list[str]):
        self.names = list(names)
        self.tokenized = [tokenize(n) for n in self.names]
        self.bm25 = BM25Okapi(self.tokenized)

    def search(self, query: str, topk: int = 30) -> list[tuple[str, float]]:
        q = tokenize(query)
        if not q:
            return []
        scores = self.bm25.get_scores(q)
        order = np.argsort(scores)[::-1][:topk]
        return [(self.names[i], float(scores[i])) for i in order if scores[i] > 0]