backend/backend/app/services/query_expansion.py

"""ATM 自动术语映射引擎 — 让普通搜索框实现 PubMed 级别的 MeSH 展开

流程:
  1. MeSH Translation: exact entry term match (@>) → name_en ILIKE → per-word fallback
  2. Tree Expansion: tree_number LIKE 前缀展开到所有子节点
  3. SQL Condition: GlobalLiterature.id IN (tag_ids) — OR 合并 tsvector

设计原则:
  - 与 PubMed ATM 一致：先精确入口词匹配，再逐词拆解
  - entry_terms 已全部小写标准化，用 @> 做精确 JSONB 元素匹配
    杜绝 "breast cancer" 误匹配 "Breast Cancer Anti-Estrogen Resistance 1 Protein" 等问题

用法:
  cond = await expand_atm(db, "lung cancer")
  if cond is not None:
      conditions.append(sa.or_(cond, tsvector_cond))
"""

import logging
import re
from uuid import UUID

from sqlalchemy import or_ as _or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.literature import GlobalTag, GlobalTagTreeNumber, GlobalLiteratureTag, GlobalLiterature

logger = logging.getLogger(__name__)

MIN_QUERY_LENGTH = 2


def _escape_ilike(s: str) -> str:
    """转义 ILIKE 模式中的通配符 _ 和 %"""
    if not s:
        return s
    return s.replace('\\', '\\\\').replace('%', '\\%').replace('_', '\\_')


async def expand_atm(db: AsyncSession, query: str):
    """ATM 展开：输入纯文本查询，返回 SQLAlchemy condition 或 None"""
    q = query.strip().replace('"', '').replace("'", '')
    if not q or len(q) < MIN_QUERY_LENGTH:
        return None

    tag_ids = await _find_mesh_tags(db, q)
    if not tag_ids:
        tag_ids = await _find_partial_mesh_tags(db, q)
    if not tag_ids:
        return None

    expanded_ids = await _expand_tree_numbers(db, tag_ids)
    if not expanded_ids:
        return None

    return GlobalLiterature.id.in_(
        select(GlobalLiteratureTag.literature_id)
        .where(GlobalLiteratureTag.tag_id.in_(expanded_ids))
        .distinct()
    )


async def _find_mesh_tags(db: AsyncSession, query: str) -> list[UUID]:
    """通过精确入口词匹配 + name_en ILIKE 找 MeSH 标签"""
    q = query.strip().lower()
    seen: set[UUID] = set()
    tag_ids: list[UUID] = []

    # 方法 A: JSONB @> 精确元素匹配
    # entry_terms 已全小写存储，可做 case-insensitive 精确匹配
    # 只有 query 是 entry_terms 数组中的完整元素时才匹配
    # 如 "breast cancer" 只匹配元素 "breast cancer"，不匹配 "breast cancer anti-estrogen resistance 1 protein"
    stmt = select(GlobalTag.id).where(
        GlobalTag.source.in_(["mesh", "manual"]),
        GlobalTag.mesh_ui.isnot(None),
        GlobalTag.entry_terms.isnot(None),
        GlobalTag.entry_terms.contains([q]),
    )
    rows = await db.execute(stmt)
    for (tid,) in rows:
        if tid not in seen:
            tag_ids.append(tid)
            seen.add(tid)

    # 方法 B: name_en ILIKE 匹配（P1-3: 加 LIMIT 100 防止常见词匹配过多）
    like_pattern = f"%{_escape_ilike(query)}%"
    stmt = select(GlobalTag.id).where(
        GlobalTag.source.in_(["mesh", "manual"]),
        GlobalTag.name_en.ilike(like_pattern),
    ).limit(100)
    rows = await db.execute(stmt)
    for (tid,) in rows:
        if tid not in seen:
            tag_ids.append(tid)
            seen.add(tid)

    # 方法 C: name_zh ILIKE 匹配（P0-4: 中文查询降级）
    if not tag_ids and re.search(r'[一-鿿㐀-䶿豈-﫿]', q):
        stmt = select(GlobalTag.id).where(
            GlobalTag.source.in_(["mesh", "manual"]),
            GlobalTag.name_zh.ilike(like_pattern),
        )
        rows = await db.execute(stmt)
        for (tid,) in rows:
            if tid not in seen:
                tag_ids.append(tid)
                seen.add(tid)

    return tag_ids


async def _find_partial_mesh_tags(db: AsyncSession, query: str) -> list[UUID]:
    """降级匹配：多词查询按单词分别匹配，支持跨标签组合（如 "egfr mutation" → EGFR + Mutation）

    PubMed 行为：当短语无法匹配 MeSH 入口词时，逐词尝试 MeSH 翻译。
    """
    words = [w.strip().lower() for w in query.strip().split() if len(w.strip()) >= MIN_QUERY_LENGTH][:10]
    if len(words) < 2:
        # P0-4: 中文查询无法按空格分词，尝试整体 name_zh ILIKE 匹配
        if re.search(r'[一-鿿㐀-䶿豈-﫿]', query):
            tag_ids = []
            stmt = select(GlobalTag.id).where(
                GlobalTag.source.in_(["mesh", "manual"]),
                GlobalTag.name_zh.ilike(f'%{_escape_ilike(query.strip().lower())}%'),
            )
            rows = await db.execute(stmt)
            for (tid,) in rows:
                tag_ids.append(tid)
            return tag_ids
        return []  # non-Chinese single word: no partial match possible

    seen: set[UUID] = set()
    tag_ids: list[UUID] = []


    # 批量精确入口词匹配（一次查询替代 N 次）
    entry_terms_conds = []
    name_ilike_conds = []
    for word in words:
        entry_terms_conds.append(
            GlobalTag.entry_terms.contains([word])
        )
        name_ilike_conds.append(GlobalTag.name_en.ilike(f"%{_escape_ilike(word)}%"))

    # 一次查询所有词的精确入口词匹配
    if entry_terms_conds:
        stmt = select(GlobalTag.id).where(
            GlobalTag.source.in_(["mesh", "manual"]),
            GlobalTag.mesh_ui.isnot(None),
            GlobalTag.entry_terms.isnot(None),
            _or_(*entry_terms_conds),
        )
        rows = await db.execute(stmt)
        for (tid,) in rows:
            if tid not in seen:
                tag_ids.append(tid)
                seen.add(tid)

    # 一次查询所有词的 name_en ILIKE
    if name_ilike_conds:
        stmt = select(GlobalTag.id).where(
            GlobalTag.source.in_(["mesh", "manual"]),
            _or_(*name_ilike_conds),
        )
        rows = await db.execute(stmt)
        for (tid,) in rows:
            if tid not in seen:
                tag_ids.append(tid)
                seen.add(tid)

    return tag_ids


async def _expand_tree_numbers(db: AsyncSession, tag_ids: list[UUID]) -> list[UUID]:
    """通过树号展开到所有子节点"""
    rows = await db.execute(
        select(GlobalTagTreeNumber.tree_number, GlobalTagTreeNumber.tag_id)
        .where(GlobalTagTreeNumber.tag_id.in_(tag_ids))
    )
    tree_map: dict[str, list[UUID]] = {}
    for tn, tid in rows:
        tree_map.setdefault(tn, []).append(tid)

    if not tree_map:
        return tag_ids

    expanded = set(tag_ids)
    # 一次查询所有 tree_number 的子节点（用 OR 合并）
    tn_conds = [_or_(
        GlobalTagTreeNumber.tree_number == tn,
        GlobalTagTreeNumber.tree_number.like(f"{tn}.%"),
    ) for tn in tree_map]
    rows = await db.execute(
        select(GlobalTagTreeNumber.tag_id)
        .where(_or_(*tn_conds))
        .distinct()
    )
    for (tid,) in rows:
        expanded.add(tid)

    return list(expanded)