"""同义词词典 + Boolean 查询扩展

为肿瘤学检索提供受控词汇扩展：
- PD-1 → PDCD1 / programmed cell death 1 / pembrolizumab / nivolumab
- 一线 → first-line / 1L / 1st-line / frontline
- 自动展开为 OR 链

用法：
    expander = SynonymExpander()
    expander.expand("PD-1 一线 NSCLC")
    # → ('("PD-1"[Title/Abstract] OR "programmed cell death 1"[Title/Abstract] OR ...) AND ...', {...})

将来：
    - 可对接 PubChem / MeSH 自动更新同义词
    - 支持 MeSH 下位词展开（Entry Terms 递归）
"""

import logging
import re

logger = logging.getLogger(__name__)

# ─── 核心同义词词典 ───
# key = 规范名（小写），values = 同义词列表（含 key 自身）
# 查询时自动转义特殊字符，OR 拼接

_SYNONYM_MAP: dict[str, list[str]] = {
    # ── Immune checkpoint ──
    "pd-1": [
        "PD-1", "PD1", "PDCD1", "programmed cell death 1",
        "programmed death 1", "pembrolizumab", "nivolumab",
    ],
    "pd-l1": [
        "PD-L1", "PDL1", "CD274", "B7-H1", "programmed death ligand 1",
        "atezolizumab", "durvalumab", "avelumab",
    ],
    "ctla-4": [
        "CTLA-4", "CTLA4", "cytotoxic T-lymphocyte antigen 4",
        "ipilimumab", "tremelimumab",
    ],
    "immunotherapy": [
        "immunotherapy", "immune checkpoint inhibitor", "ICI",
        "checkpoint blockade", "immune checkpoint blockade",
        "anti-PD-1", "anti-PD-L1", "anti-CTLA-4",
    ],

    # ── Therapy lines ──
    "first-line": [
        "first-line", "first line", "1st-line", "1L", "frontline",
        "front-line", "firstline",
    ],
    "second-line": [
        "second-line", "second line", "2nd-line", "2L",
    ],
    "third-line": [
        "third-line", "third line", "3rd-line", "3L",
    ],
    "adjuvant": [
        "adjuvant", "adjuvant therapy", "adjuvant chemotherapy",
        "adjuvant treatment",
    ],
    "neoadjuvant": [
        "neoadjuvant", "neoadjuvant therapy", "neoadjuvant chemotherapy",
        "preoperative",
    ],

    # ── Targeted therapy ──
    "egfr": [
        "EGFR", "epidermal growth factor receptor", "ErbB1", "HER1",
        "gefitinib", "erlotinib", "osimertinib", "afatinib",
        "dacomitinib", "EGFR-TKI", "EGFR tyrosine kinase inhibitor",
    ],
    "alk": [
        "ALK", "anaplastic lymphoma kinase",
        "crizotinib", "ceritinib", "alectinib", "brigatinib", "lorlatinib",
        "ALK-TKI", "ALK inhibitor",
    ],
    "ros1": [
        "ROS1", "ROS-1", "crizotinib", "entrectinib", "lorlatinib",
        "ROS1-TKI", "ROS1 inhibitor",
    ],
    "braf": [
        "BRAF", "B-Raf", "BRAF V600", "BRAF V600E",
        "vemurafenib", "dabrafenib", "encorafenib",
        "BRAF inhibitor",
    ],
    "her2": [
        "HER2", "HER-2", "ERBB2", "neu",
        "trastuzumab", "pertuzumab", "ado-trastuzumab emtansine",
        "T-DM1", "trastuzumab deruxtecan", "T-DXd", "lapatinib",
        "HER2-targeted", "anti-HER2",
    ],
    "vegfr": [
        "VEGFR", "VEGF", "vascular endothelial growth factor",
        "bevacizumab", "ramucirumab", "sunitinib", "pazopanib",
        "axitinib", "lenvatinib", "sorafenib",
        "anti-angiogenic", "TKI",
    ],
    "mtor": [
        "mTOR", "mammalian target of rapamycin",
        "everolimus", "temsirolimus", "rapamycin",
        "mTOR inhibitor",
    ],
    "cdk4/6": [
        "CDK4/6", "CDK4", "CDK6", "cyclin-dependent kinase 4/6",
        "palbociclib", "ribociclib", "abemaciclib",
        "CDK4/6 inhibitor",
    ],
    "parp": [
        "PARP", "poly ADP ribose polymerase",
        "olaparib", "niraparib", "rucaparib", "talazoparib",
        "PARP inhibitor",
    ],

    # ── Clinical trial phases ──
    "phase i": [
        "phase I", "phase 1", "phase I/II", "first-in-human",
        "dose escalation", "dose finding",
    ],
    "phase ii": [
        "phase II", "phase 2", "phase II/III",
    ],
    "phase iii": [
        "phase III", "phase 3", "randomized controlled trial",
        "RCT",
    ],

    # ── Endpoints ──
    "overall survival": [
        "overall survival", "OS", "median OS", "overall survival rate",
    ],
    "progression-free survival": [
        "progression-free survival", "PFS", "median PFS",
        "progression free survival",
    ],
    "objective response": [
        "objective response", "objective response rate", "ORR",
        "overall response", "overall response rate",
        "response rate",
    ],
    "disease control": [
        "disease control", "disease control rate", "DCR",
    ],
    "quality of life": [
        "quality of life", "QoL", "HRQoL", "health-related quality of life",
    ],
}


class SynonymExpander:
    """同义词查询扩展器。

    将自然语言关键词展开为 PubMed / Europe PMC 兼容的 Boolean OR 链。
    支持 [Title/Abstract] 和 [MeSH] 两种字段限定。
    """

    def __init__(self):
        # 构建反向查找索引：同义词(小写) → 规范名
        self._reverse: dict[str, str] = {}
        self._build_index()

    def _build_index(self):
        for canonical, syns in _SYNONYM_MAP.items():
            for s in syns:
                key = s.strip().lower()
                if key not in self._reverse:
                    self._reverse[key] = canonical

    def expand(
        self,
        term: str,
        *,
        field: str = "Title/Abstract",
        mesh_field: str = "MeSH",
    ) -> tuple[str | None, dict]:
        """展开单个检索词为 Boolean OR 链。

        Args:
            term: 用户输入的自然语言关键词（如 "PD-1"）
            field: 标题/摘要字段标签（Europe PMC: Title/Abstract, PubMed: tiab）
            mesh_field: MeSH 字段标签（Europe PMC: MeSH, PubMed: MeSH）

        Returns:
            (expanded_query: str | None, info: dict)
            - expanded_query: 展开后的查询片段，None 若无匹配
            - info: {"matched": bool, "canonical": str, "synonyms_count": int}
        """
        canonical = self._reverse.get(term.strip().lower())
        if not canonical:
            return None, {"matched": False, "canonical": None, "synonyms_count": 0}

        syns = _SYNONYM_MAP[canonical]
        # 筛选出不含空格的同义词用于 [Title/Abstract]，含空格的加引号
        quoted = []
        for s in syns:
            if " " in s:
                quoted.append(f'"{s}"')
            else:
                quoted.append(s)
        or_chain = f"[{field}] OR ".join(quoted) + f"[{field}]"

        return or_chain, {
            "matched": True,
            "canonical": canonical,
            "synonyms_count": len(syns),
        }

    def expand_user_query(
        self,
        query: str,
        *,
        field: str = "Title/Abstract",
    ) -> tuple[str, dict]:
        """展开用户自由输入查询中的所有已知同义词。

        不会修改非匹配词，只替换已知术语。

        Args:
            query: 用户自由输入（如 "PD-1 一线 NSCLC 免疫治疗 ORR"）
            field: 目标字段标签

        Returns:
            (expanded: str, details: dict)
            - expanded: 替换后的查询字符串
            - details: {"matched_terms": [...], "unmatched": [...]}
        """
        matched = []
        unmatched = []

        # 先尝试最长匹配（避免 "PD-1" 被 "PD-L1" 部分覆盖）
        tokens = self._tokenize(query)
        used = set()
        result_tokens = []

        for token in tokens:
            tl = token.strip().lower()
            if tl in used or not tl:
                continue
            canonical = self._reverse.get(tl)
            if canonical:
                syns = _SYNONYM_MAP[canonical]
                quoted = [f'"{s}"' if " " in s else s for s in syns]
                or_chain = f"[{field}] OR ".join(quoted) + f"[{field}]"
                result_tokens.append(f"({or_chain})")
                matched.append({"token": token, "canonical": canonical, "synonyms_count": len(syns)})
                used.add(tl)
            else:
                result_tokens.append(token)
                unmatched.append(token)

        expanded = " ".join(result_tokens)
        return expanded, {
            "matched_terms": matched,
            "unmatched": unmatched,
            "total_terms": len(tokens),
        }

    def list_canonical(self) -> list[str]:
        """返回所有规范名列表（供 UI 展示或配置参考）。"""
        return sorted(_SYNONYM_MAP.keys())

    @staticmethod
    def _tokenize(query: str) -> list[str]:
        """简单分词：按空格拆分，保留引号分组。"""
        # 处理常见特殊字符但保留引号
        parts = []
        current = []
        in_quote = False
        for ch in query:
            if ch == '"':
                in_quote = not in_quote
                continue
            if ch.isspace() and not in_quote:
                if current:
                    parts.append("".join(current))
                    current = []
            else:
                current.append(ch)
        if current:
            parts.append("".join(current))
        return parts if parts else [query]


# ─── 预设 Boolean 模板（供 PICO → 查询转换）───

PICO_TEMPLATES: dict[str, str] = {
    "population": "{term}[MeSH] OR {term}[Title/Abstract]",
    "intervention": "{term}[Title/Abstract]",
    "comparator": "{term}[Title/Abstract]",
    "outcome": "{term}[Title/Abstract]",
}


# ─── MeSH 词表补全辅助 ───

# 常见癌种的 MeSH 词（含下位词提示），供 #3 检索策略导出使用
CANCER_MESH_PRIMARY: dict[str, list[str]] = {
    "lung cancer": ["Lung Neoplasms", "Small Cell Lung Carcinoma", "Bronchial Neoplasms"],
    "breast cancer": ["Breast Neoplasms", "Triple Negative Breast Neoplasms", "Inflammatory Breast Neoplasms"],
    "colorectal cancer": ["Colorectal Neoplasms", "Colonic Neoplasms", "Rectal Neoplasms"],
    "gastric cancer": ["Stomach Neoplasms", "Gastric Neoplasms"],
    "prostate cancer": ["Prostatic Neoplasms"],
    "pancreatic cancer": ["Pancreatic Neoplasms"],
    "liver cancer": ["Liver Neoplasms", "Carcinoma, Hepatocellular", "Cholangiocarcinoma"],
    "ovarian cancer": ["Ovarian Neoplasms"],
    "melanoma": ["Melanoma", "Skin Neoplasms"],
    "leukemia": ["Leukemia", "Leukemia, Myeloid", "Leukemia, Lymphocytic"],
    "lymphoma": ["Lymphoma", "Hodgkin Disease", "Lymphoma, Non-Hodgkin"],
    "glioma": ["Glioma", "Glioblastoma", "Astrocytoma"],
    "sarcoma": ["Sarcoma", "Osteosarcoma", "Leiomyosarcoma"],
}
