app/search/engine.py

import json
import os
import sys
import unicodedata
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Set, Tuple

import numpy as np
from sudachipy import dictionary, tokenizer

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))

from utils.json import field_getter
from utils.logger import setup_logger

log = setup_logger(__name__)


@dataclass
class SearchConfig:
    target_pos_l1: List[str]
    target_fields: List[str]
    k1: float
    b: float
    field_weights: Dict[str, float]
    synonyms_enable: bool
    syn_limits: Dict[str, int]
    banlist: List[str]
    word_sim_enable: bool
    word_sim_alpha: float
    word_sim_topk_k: int
    word_sim_rerank: str
    query_subword_enable: bool
    query_subword_path: str
    query_subword_oov_weight: float
    org_boost_exact: float
    org_boost_prefix: float
    org_boost_substring: float
    org_boost_min_len: int
    # filter
    min_results: int
    max_results: int
    bm25_min: float
    word_sim_min: float
    fused_min: float
    fused_rel_top_ratio: float


def normalize_text_for_org(s: str) -> str:
    try:
        import unicodedata

        s = unicodedata.normalize("NFKC", s)
    except Exception:
        pass
    s = " ".join(s.split())
    return s


class SearchEngine:
    def __init__(self):
        self.cfg: Optional[SearchConfig] = None
        self.tokenizer = None
        self.mode = None
        self.stopwords: set[str] = set()
        self.custom_synonyms: Dict[str, List[str]] = {}
        self.synonyms_cache: Dict[str, List[str]] = {}

        # Data (project -> circle)
        self.circles: List[Dict[str, Any]] = []
        self.circle_map: Dict[str, Dict[str, Any]] = {}
        self.circle_idx: Dict[str, int] = {}
        self.org_norms: Dict[str, str] = {}
        self.reading_norms: Dict[str, str] = {}
        self.substring_index: Dict[str, List[str]] = {}

        # BM25F assets
        self.idf: Dict[str, float] = {}
        self.avg_len: Dict[str, float] = {}
        self.tf_token_docs: List[Dict[str, Any]] = []

        # Vectors
        self.word_vocab: Dict[str, int] = {}
        self.word_vectors: Optional[np.ndarray] = None
        self.ft_model = None

    # ----- Init / Load -----
    def initialize(self):
        files = field_getter("config/files.json")
        search = field_getter("config/search_model.json")

        # Config
        self.cfg = SearchConfig(
            target_pos_l1=search("target_pos_l1"),
            target_fields=search("target_fields"),
            k1=float(search("bm25f.k1")),
            b=float(search("bm25f.b")),
            field_weights=search("bm25f.field_weights"),
            synonyms_enable=bool(search("synonyms.enable")),
            syn_limits=search("synonyms.limits"),
            banlist=search("synonyms.banlist"),
            word_sim_enable=bool(search("word_sim.enable")),
            word_sim_alpha=float(search("word_sim.alpha")),
            word_sim_topk_k=int(search("word_sim.topk_k", 3)),
            word_sim_rerank=(
                search("word_sim.rerank", "pair_avg") or "pair_avg"
            ).lower(),
            query_subword_enable=bool(search("query_subword.enable")),
            query_subword_path=files("embeddings.fasttext_bin"),
            query_subword_oov_weight=float(search("query_subword.oov_weight")),
            org_boost_exact=float(search("org_boost.exact", 1.5)),
            org_boost_prefix=float(search("org_boost.prefix", 0.9)),
            org_boost_substring=float(search("org_boost.substring", 0.6)),
            org_boost_min_len=int(search("org_boost.min_len", 2)),
            min_results=int(search("filter.min_results", 20)),
            max_results=int(search("filter.max_results", 100)),
            bm25_min=float(search("filter.bm25_min", 0.5)),
            word_sim_min=float(search("filter.word_sim_min", 0.3)),
            fused_min=float(search("filter.fused_min", 0.4)),
            fused_rel_top_ratio=float(search("filter.fused_rel_top_ratio", 0.7)),
        )

        # Tokenizer
        sudachi_config_path = files("sudachi.sudachi_config")
        tok = dictionary.Dictionary(config_path=sudachi_config_path).create()
        self.tokenizer = tok
        self.mode = tokenizer.Tokenizer.SplitMode.A

        # Stopwords
        with open(files("sudachi.stopwords"), encoding="utf-8") as f:
            self.stopwords = set(json.load(f))

        # Synonyms assets
        try:
            syn_cache_path = files("sudachi.synonyms_cache")
            if os.path.exists(syn_cache_path):
                with open(syn_cache_path, encoding="utf-8") as f:
                    self.synonyms_cache = json.load(f)
        except Exception as e:
            log.warning(f"failed to load synonyms_cache: {e}")

        try:
            custom_path = field_getter("config/search_model.json")(
                "synonyms.sources.custom_json"
            )
            if custom_path and os.path.exists(custom_path):
                with open(custom_path, encoding="utf-8") as f:
                    self.custom_synonyms = json.load(f)
        except Exception:
            pass

        # Substring index for organization substring lookup
        substring_index_path = files("substring.substring_index")
        if os.path.exists(substring_index_path):
            try:
                with open(substring_index_path, encoding="utf-8") as f:
                    self.substring_index = json.load(f)
            except Exception as e:
                log.warning(f"failed to load substring_index: {e}")
        else:
            self.substring_index = {}

        # Circles (projects -> circles)
        with open(files("circles.circles_json"), encoding="utf-8") as f:
            self.circles = json.load(f)
        self.circle_map = {c["circleId"]: c for c in self.circles}
        self.circle_idx = {c["circleId"]: idx for idx, c in enumerate(self.circles)}
        self.org_norms = {
            c["circleId"]: normalize_text_for_org(c.get("circleName") or "")
            for c in self.circles
        }
        self.reading_norms = {
            c["circleId"]: normalize_text_for_org(c.get("circleNameKana") or "")
            for c in self.circles
        }

        # BM25F assets
        with open(files("bm25.bm25_meta"), encoding="utf-8") as f:
            meta = json.load(f)
        self.idf = meta.get("idf", {})
        self.avg_len = meta.get("avg_len", {})

        with open(files("bm25.tf_token"), encoding="utf-8") as f:
            self.tf_token_docs = json.load(f)

        # Vectors
        try:
            vocab_path = files("embeddings.word_vocab")
            vec_path = files("embeddings.word_vectors")
            if os.path.exists(vocab_path) and os.path.exists(vec_path):
                with open(vocab_path, encoding="utf-8") as f:
                    self.word_vocab = {k: int(v) for k, v in json.load(f).items()}
                self.word_vectors = np.load(vec_path)["vectors"]
        except Exception as e:
            log.warning(f"word vectors not ready: {e}")

        # doc_vectors.npy は topk 方式では不要

        # fastText OOV
        if (
            self.cfg.query_subword_enable
            and self.cfg.query_subword_path
            and os.path.exists(self.cfg.query_subword_path)
        ):
            try:
                import fasttext

                self.ft_model = fasttext.load_model(self.cfg.query_subword_path)
                log.info("fastText .bin loaded for OOV")
            except Exception as e:
                log.warning(f"failed to load fastText .bin: {e}")

    # ----- Tokenize / Synonyms -----
    def _tokenize(self, text: str) -> List[str]:
        if not text:
            return []
        out: List[str] = []
        for m in self.tokenizer.tokenize(text, self.mode):
            base = m.normalized_form().lower().strip()
            if not base:
                continue
            pos = m.part_of_speech()
            if pos[0] not in self.cfg.target_pos_l1:
                continue
            if base in self.stopwords or base in self.cfg.banlist:
                continue
            out.append(base)
        return out

    def _expand_synonyms(self, terms: List[str]) -> List[str]:
        if not self.cfg.synonyms_enable:
            return terms
        max_exp = int(self.cfg.syn_limits.get("max_expansions_per_term", 4))
        min_len = int(self.cfg.syn_limits.get("min_char_len", 2))
        expanded: List[str] = []
        for t in terms:
            expanded.append(t)
            cands = []
            cands.extend(self.synonyms_cache.get(t, []))
            cands.extend(self.custom_synonyms.get(t, []))
            # filter/unique
            uniq = []
            seen = set()
            for c in cands:
                if c in seen or len(c) < min_len or c in self.cfg.banlist:
                    continue
                seen.add(c)
                uniq.append(c)
                if len(uniq) >= max_exp:
                    break
            expanded.extend(uniq)
        # overall limit
        max_q = int(self.cfg.syn_limits.get("max_query_variants", 5))
        return expanded[: max_q * max_exp + len(terms)]

    @staticmethod
    def _katakana_to_hiragana(text: str) -> str:
        if not text:
            return ""
        chars: List[str] = []
        for ch in text:
            code = ord(ch)
            if 0x30A1 <= code <= 0x30F6:
                chars.append(chr(code - 0x60))
            else:
                chars.append(ch)
        return "".join(chars)

    def _normalize_substring_token(self, token: str) -> str:
        if not token:
            return ""
        try:
            token_nfkc = unicodedata.normalize("NFKC", token)
        except Exception:
            token_nfkc = token

        readings: List[str] = []
        if self.tokenizer is not None:
            try:
                for m in self.tokenizer.tokenize(
                    token_nfkc, tokenizer.Tokenizer.SplitMode.C
                ):
                    reading = m.reading_form()
                    if not reading or reading == "*":
                        reading = m.normalized_form()
                    if reading:
                        readings.append(reading)
            except Exception:
                readings = []

        reading = "".join(readings) if readings else token_nfkc
        lowered = reading.lower()
        hira = self._katakana_to_hiragana(lowered)

        normalized_chars: List[str] = []
        for ch in hira:
            if ch in ("\u0020", "\u3000"):
                continue
            category = unicodedata.category(ch)
            if category.startswith("P") or category.startswith("S"):
                if ch != "ー":
                    continue
            normalized_chars.append(ch)
        return "".join(normalized_chars)

    def _normalize_substring_terms(self, query: str) -> List[str]:
        if not query:
            return []
        try:
            normalized_query = unicodedata.normalize("NFKC", query)
        except Exception:
            normalized_query = query

        out: List[str] = []
        for raw in normalized_query.split():
            term = self._normalize_substring_token(raw)
            if term:
                out.append(term)
        return out

    def _substring_match_circle_ids(self, query: str) -> Set[str]:
        if not self.substring_index:
            return set()
        terms = self._normalize_substring_terms(query)
        matches: Set[str] = set()
        for term in terms:
            if len(term) < 2:
                continue
            matches.update(self.substring_index.get(term, []))
        return matches

    # ----- BM25F -----
    def _bm25f_scores(self, terms: List[str]) -> np.ndarray:
        N = len(self.tf_token_docs)
        if N == 0:
            return np.zeros((0,), dtype=np.float32)
        k1 = self.cfg.k1
        b = self.cfg.b
        fw = self.cfg.field_weights

        scores = np.zeros((N,), dtype=np.float32)

        idf = self.idf
        avg_len = self.avg_len

        # For quick access, build list of per-doc per-field structures
        for i, d in enumerate(self.tf_token_docs):
            fields = d.get("fields") or {}
            s = 0.0
            for t in terms:
                idf_t = float(idf.get(t, 0.0))
                if idf_t <= 0.0:
                    continue
                denom_sum = 0.0
                num_sum = 0.0
                for field, weight in fw.items():
                    fobj = fields.get(field) or {}
                    tf = float((fobj.get("tf") or {}).get(t, 0))
                    if tf <= 0.0:
                        continue
                    len_f = float(fobj.get("len", 0))
                    avg_f = float(avg_len.get(field, 0.0)) or 1.0
                    norm = k1 * (1 - b + b * (len_f / avg_f))
                    num_sum += weight * tf * (k1 + 1.0)
                    denom_sum += weight * (tf + norm)
                if denom_sum > 0:
                    s += idf_t * (num_sum / denom_sum)
            scores[i] = s
        return scores

    # ----- Word similarity -----
    def _get_token_vector(self, t: str) -> Tuple[Optional[np.ndarray], bool]:
        if self.word_vectors is not None and t in self.word_vocab:
            v = self.word_vectors[self.word_vocab[t]]
            return v, False
        if self.ft_model is not None:
            try:
                v = self.ft_model.get_word_vector(t)
                v = v.astype(np.float32)
                n = np.linalg.norm(v)
                if n > 0:
                    v = v / n
                return v, True
            except Exception:
                return None, True
        return None, True

    def _word_sim_scores_topk(self, terms: List[str]) -> Optional[np.ndarray]:
        if not self.cfg.word_sim_enable:
            return None
        if self.word_vectors is None:
            return None

        # Build per-term vectors with weights (IDF; OOV down-weighted)
        weights = []
        vecs = []
        for t in terms:
            v, oov = self._get_token_vector(t)
            if v is None:
                continue
            w = float(self.idf.get(t, 0.0))
            if oov:
                w *= float(self.cfg.query_subword_oov_weight)
            if w <= 0:
                continue
            vecs.append(v)
            weights.append(w)

        if not vecs:
            return None

        V = np.stack(vecs).astype(np.float32)  # T x D
        W = np.asarray(weights, dtype=np.float32)  # T

        # top-k pooling over term-term cosine contributions (query terms x document terms)
        k = max(1, int(self.cfg.word_sim_topk_k))
        n_docs = len(self.tf_token_docs)
        sims_all = np.zeros((n_docs,), dtype=np.float32)
        Vq = V  # Tq x D (normalized)
        Wq = W  # Tq
        for i, d in enumerate(self.tf_token_docs):
            fields = d.get("fields") or {}
            doc_terms = set()
            for fname in self.cfg.target_fields:
                fobj = fields.get(fname) or {}
                tf = fobj.get("tf") or {}
                doc_terms.update(tf.keys())
            if not doc_terms:
                sims_all[i] = 0.0
                continue
            Vd_list = []
            for t in doc_terms:
                idx = self.word_vocab.get(t)
                if idx is None:
                    continue
                Vd_list.append(self.word_vectors[idx])
            if not Vd_list:
                sims_all[i] = 0.0
                continue
            Vd = np.stack(Vd_list).astype(np.float32)  # Td x D
            M = Vd @ Vq.T  # Td x Tq
            if Wq.size:
                M = M * Wq[None, :]
            M = np.maximum(M, 0.0)
            Td, Tq = M.shape
            total = Td * Tq
            kk = min(k, total) if total > 0 else 0
            if kk == 0:
                sims_all[i] = 0.0
                continue
            flat = M.reshape(-1)
            if kk == total:
                top_vals = flat
            else:
                idxk = np.argpartition(flat, -kk)[-kk:]
                top_vals = flat[idxk]
            sims_all[i] = float(top_vals.mean()) if top_vals.size else 0.0
        return sims_all

    def _word_sim_scores_pairavg(self, terms: List[str]) -> Optional[np.ndarray]:
        if not self.cfg.word_sim_enable:
            return None
        if self.word_vectors is None:
            return None
        # Build query term vectors (no weighting for pair-avg, simple mean over all pairs)
        vecs = []
        for t in terms:
            v, _ = self._get_token_vector(t)
            if v is None:
                continue
            vecs.append(v)
        if not vecs:
            return None
        Vq = np.stack(vecs).astype(np.float32)  # Tq x D

        n_docs = len(self.tf_token_docs)
        sims_all = np.zeros((n_docs,), dtype=np.float32)
        for i, d in enumerate(self.tf_token_docs):
            fields = d.get("fields") or {}
            doc_terms = set()
            for fname in self.cfg.target_fields:
                fobj = fields.get(fname) or {}
                tf = fobj.get("tf") or {}
                doc_terms.update(tf.keys())
            if not doc_terms:
                sims_all[i] = 0.0
                continue
            Vd_list = []
            for t in doc_terms:
                idx = self.word_vocab.get(t)
                if idx is None:
                    continue
                Vd_list.append(self.word_vectors[idx])
            if not Vd_list:
                sims_all[i] = 0.0
                continue
            Vd = np.stack(Vd_list).astype(np.float32)  # Td x D
            M = Vd @ Vq.T  # Td x Tq
            M = np.maximum(M, 0.0)
            sims_all[i] = float(M.mean()) if M.size else 0.0
        return sims_all

    # ----- Public API -----
    def search(
        self,
        query: str,
        debug: bool = False,
    ) -> List[Tuple[str, float]] | Tuple[List[Tuple[str, float]], Dict[str, Any]]:
        terms = self._tokenize(query)
        if self.cfg.synonyms_enable:
            terms = self._expand_synonyms(terms)

        substring_hits = self._substring_match_circle_ids(query)
        substring_idx_set: Set[int] = set()
        substring_mask = np.zeros((len(self.circles),), dtype=bool)
        if substring_hits:
            for cid in substring_hits:
                idx = self.circle_idx.get(cid)
                if idx is None:
                    continue
                substring_idx_set.add(idx)
                substring_mask[idx] = True

        # BM25F
        bm25 = self._bm25f_scores(terms)

        # word sim (filtering): top-k pooling
        ws_filter = self._word_sim_scores_topk(terms)
        if ws_filter is None:
            ws_filter = np.zeros_like(bm25)
        a = float(self.cfg.word_sim_alpha)
        fused_filter = a * bm25 + (1.0 - a) * ws_filter

        # circleName/circleNameKana auto-boost based on raw query substring match
        qn = normalize_text_for_org(query)
        boost_enabled = len(qn) >= int(self.cfg.org_boost_min_len)
        boost = np.zeros((len(self.circles),), dtype=np.float32)
        if boost_enabled:
            exact = np.zeros((len(self.circles),), dtype=bool)
            prefix = np.zeros_like(exact)
            substr = np.zeros_like(exact)
            for i, d in enumerate(self.circles):
                cid = d.get("circleId")
                on = self.org_norms.get(cid, "")
                rn = self.reading_norms.get(cid, "")
                if qn and (qn == on or (rn and qn == rn)):
                    exact[i] = True
                elif qn and (on.startswith(qn) or (rn and rn.startswith(qn))):
                    prefix[i] = True
                elif qn and ((qn in on) or (rn and qn in rn)):
                    substr[i] = True

            boost = (
                exact.astype(np.float32) * float(self.cfg.org_boost_exact)
                + prefix.astype(np.float32) * float(self.cfg.org_boost_prefix)
                + substr.astype(np.float32) * float(self.cfg.org_boost_substring)
            )

        # collect results
        ids = [d.get("circleId") for d in self.circles]

        # Filtering to reduce false positives while keeping recall
        # Relative threshold anchored to the top fused score
        if boost_enabled:
            score_with_boost = fused_filter + boost
        else:
            score_with_boost = fused_filter

        top = float(np.max(score_with_boost)) if score_with_boost.size > 0 else 0.0
        rel_cut = (
            top * float(self.cfg.fused_rel_top_ratio) if top > 0 else self.cfg.fused_min
        )
        fused_cut = max(float(self.cfg.fused_min), rel_cut)
        keep = (
            (bm25 >= self.cfg.bm25_min)
            | (ws_filter >= self.cfg.word_sim_min)
            | (score_with_boost >= self.cfg.fused_min)
        ) & (score_with_boost >= fused_cut)
        if substring_idx_set:
            keep = keep | substring_mask
        order = np.argsort(-score_with_boost)  # descending by fused
        selected_idx: List[int] = []
        selected_idx_set: Set[int] = set()
        substring_sorted = sorted(substring_idx_set, key=lambda i: -score_with_boost[i])
        for idx in substring_sorted:
            selected_idx.append(int(idx))
            selected_idx_set.add(int(idx))
        non_sub_count = 0
        for i in order:
            idx = int(i)
            if idx in selected_idx_set:
                continue
            if keep[idx]:
                selected_idx.append(idx)
                selected_idx_set.add(idx)
                non_sub_count += 1
            if non_sub_count >= self.cfg.max_results:
                break
        # Single-step fallback: if zero, relax the relative cut and use absolute thresholds only
        if not selected_idx_set:
            keep2 = (
                (bm25 >= self.cfg.bm25_min)
                | (ws_filter >= self.cfg.word_sim_min)
                | (score_with_boost >= self.cfg.fused_min)
            )
            for i in order:
                idx = int(i)
                if idx in selected_idx_set:
                    continue
                if keep2[idx]:
                    selected_idx.append(idx)
                    selected_idx_set.add(idx)
                    non_sub_count += 1
                if non_sub_count >= self.cfg.max_results:
                    break
        # Rerank with pair-avg word similarity (if enabled)
        ws_rerank = None
        if self.cfg.word_sim_rerank == "pair_avg":
            ws_rerank = self._word_sim_scores_pairavg(terms)
        if ws_rerank is None:
            ws_rerank = ws_filter
        fused_rerank = a * bm25 + (1.0 - a) * ws_rerank
        if boost_enabled:
            final_scores = fused_rerank + boost
        else:
            final_scores = fused_rerank
        pairs = [(ids[i], float(final_scores[i])) for i in selected_idx]

        # sort
        pairs.sort(key=lambda x: (-x[1], x[0]))
        if not debug:
            return pairs
        # build debug details for all docs sorted by score
        ranked_indices = sorted(
            range(len(self.circles)),
            key=lambda idx: (-float(final_scores[idx]), ids[idx]),
        )
        details = []
        for idx in ranked_indices:
            circle = self.circles[idx]
            details.append(
                {
                    "circleId": ids[idx],
                    "circleName": circle.get("circleName"),
                    "bm25": float(bm25[idx]),
                    "ws_filter_topk": float(ws_filter[idx]),
                    "ws_rerank_pairavg": float(ws_rerank[idx])
                    if ws_rerank is not None
                    else None,
                    "org_boost": float(boost[idx]) if boost_enabled else None,
                    "matched_substring": bool(substring_mask[idx]),
                    "fused_filter": float(fused_filter[idx]),
                    "fused_final": float(final_scores[idx]),
                }
            )
        return pairs, {"details": details}

    def get_circles(self) -> List[Dict[str, Any]]:
        return self.circles

    def get_circle_map(self) -> Dict[str, Dict[str, Any]]:
        return self.circle_map