src/data/query_entity_linker.py

"""Link clinical evidence text to UMLS graph nodes.

Primary backend: surface-form dictionary built from the filtered KG (``id_maps.parquet``)
plus optional enrichment from ``mrxns_eng`` / ``mrsty`` in the UMLS DuckDB.

Optional backend: QuickUMLS (when ``leveldb`` + ``quickumls`` are installed).
"""

from __future__ import annotations

import json
import re
import unicodedata
from pathlib import Path

import numpy as np
import polars as pl

from config.paths import (
    CUI_TO_GRAPH_IDX_JSON,
    GRAPH_ID_MAPS_PARQUET,
    QUERY_LINKER_CACHE_JSON,
    UMLS_DUCKDB_PATH,
)

# Clinical semantic types (diseases, drugs, symptoms, procedures, anatomy, findings).
CLINICAL_SEMTYPES: frozenset[str] = frozenset(
    {
        "T047",  # Disease or Syndrome
        "T048",  # Mental or Behavioral Dysfunction
        "T184",  # Sign or Symptom
        "T121",  # Pharmacologic Substance
        "T109",  # Organic Chemical
        "T123",  # Biologically Active Substance
        "T195",  # Antibiotic
        "T200",  # Clinical Drug
        "T023",  # Body Part, Organ, or Organ Component
        "T031",  # Body Substance
        "T033",  # Finding
        "T034",  # Laboratory or Test Result
        "T060",  # Diagnostic Procedure
        "T061",  # Therapeutic or Preventive Procedure
        "T037",  # Injury or Poisoning
        "T046",  # Pathologic Function
        "T191",  # Neoplastic Process
    }
)

PREFERRED_SABS: frozenset[str] = frozenset(
    {
        "SNOMEDCT_US",
        "MSH",
        "RXNORM",
        "MEDCIN",
        "NCI",
        "LNC",
        "ICD10CM",
        "ICD10",
        "ICD9CM",
        "NDDF",
        "MMSL",
        "VANDF",
    }
)

STOPWORDS: frozenset[str] = frozenset(
    {
        "a",
        "an",
        "the",
        "and",
        "or",
        "of",
        "in",
        "on",
        "at",
        "to",
        "for",
        "with",
        "without",
        "by",
        "from",
        "is",
        "are",
        "was",
        "were",
        "be",
        "been",
        "being",
        "has",
        "have",
        "had",
        "not",
        "no",
        "due",
        "associated",
        "secondary",
        "primary",
        "unspecified",
        "other",
        "specified",
        "history",
        "personal",
        "patient",
        "given",
        "then",
        "there",
        "will",
        "some",
        "any",
        "all",
        "both",
        "each",
        "than",
        "that",
        "this",
        "these",
        "those",
        "into",
        "during",
        "after",
        "before",
        "while",
        "when",
        "where",
        "which",
        "who",
        "whom",
        "whose",
        "if",
        "as",
        "but",
        "so",
        "such",
        "via",
        "per",
        "also",
        "only",
        "just",
        "very",
        "more",
        "most",
        "less",
        "least",
        "new",
        "old",
        "acute",
        "chronic",
        "stage",
        "type",
        "status",
        "present",
        "absent",
        "possible",
        "probable",
        "suspected",
    }
)

_SURFACE_RE = re.compile(r"[^a-z0-9]+")


def normalize_surface(text: str) -> str:
    """Lowercase, strip accents, collapse punctuation to spaces."""
    text = unicodedata.normalize("NFKD", text)
    text = "".join(ch for ch in text if not unicodedata.combining(ch))
    text = text.lower().strip()
    text = _SURFACE_RE.sub(" ", text)
    return " ".join(text.split())


def _tokenize(text: str) -> list[str]:
    return [t for t in normalize_surface(text).split() if t and t not in STOPWORDS]


def build_cui_to_graph_idx(
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    max_graph_idx: int | None = None,
) -> dict[str, int]:
    """Map CUI -> best graph row index (highest degree node per CUI)."""
    scan = pl.scan_parquet(id_maps_parquet).select("idx", "cui", "degree")
    if max_graph_idx is not None:
        scan = scan.filter(pl.col("idx") < max_graph_idx)
    df = (
        scan.filter(pl.col("cui").is_not_null())
        .sort(["cui", "degree"], descending=[False, True])
        .group_by("cui")
        .agg(pl.col("idx").first())
        .collect()
    )
    return dict(zip(df["cui"].to_list(), df["idx"].to_list()))


def build_clinical_cuis(
    cui_to_idx: dict[str, int],
    umls_duckdb: Path | None = UMLS_DUCKDB_PATH,
) -> set[str]:
    """Filter CUIs to clinical semantic types when DuckDB is available."""
    if umls_duckdb is None or not Path(umls_duckdb).exists():
        return set(cui_to_idx)
    import duckdb

    con = duckdb.connect(str(umls_duckdb), read_only=True)
    con.register("graph_cuis", pl.DataFrame({"cui": list(cui_to_idx)}))
    sem = ",".join(f"'{s}'" for s in CLINICAL_SEMTYPES)
    rows = con.execute(
        f"""
        SELECT DISTINCT g.cui
        FROM graph_cuis g
        JOIN mrsty s ON g.cui = s.CUI
        WHERE s.STY IN ({sem})
        """
    ).fetchall()
    con.close()
    return {r[0] for r in rows} if rows else set(cui_to_idx)


def build_surface_lookup(
    cui_to_idx: dict[str, int],
    clinical_cuis: set[str],
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    umls_duckdb: Path | None = UMLS_DUCKDB_PATH,
    max_graph_idx: int | None = None,
    min_surface_len: int = 3,
) -> dict[str, list[tuple[str, float]]]:
    """Build normalized surface form -> [(cui, score)] from KG + optional MRXNS."""
    lookup: dict[str, dict[str, float]] = {}

    def add(surface: str, cui: str, score: float) -> None:
        if cui not in clinical_cuis:
            return
        norm = normalize_surface(surface)
        if len(norm) < min_surface_len:
            return
        bucket = lookup.setdefault(norm, {})
        bucket[cui] = max(bucket.get(cui, 0.0), score)

    scan = pl.scan_parquet(id_maps_parquet).select("idx", "cui", "sab", "str")
    if max_graph_idx is not None:
        scan = scan.filter(pl.col("idx") < max_graph_idx)
    df = scan.filter(
        pl.col("cui").is_not_null()
        & pl.col("str").is_not_null()
        & (pl.col("str").str.len_chars() >= min_surface_len)
    ).collect()

    for cui, sab, surface in zip(
        df["cui"].to_list(), df["sab"].to_list(), df["str"].to_list()
    ):
        score = 1.0 if sab in PREFERRED_SABS else 0.85
        add(surface, cui, score)

    if umls_duckdb is not None and Path(umls_duckdb).exists():
        import duckdb

        con = duckdb.connect(str(umls_duckdb), read_only=True)
        con.register("graph_cuis", pl.DataFrame({"cui": list(clinical_cuis)}))
        sem = ",".join(f"'{s}'" for s in CLINICAL_SEMTYPES)
        rows = con.execute(
            f"""
            SELECT DISTINCT m.CUI, m.NSTR
            FROM mrxns_eng m
            JOIN graph_cuis g ON m.CUI = g.cui
            JOIN mrsty s ON m.CUI = s.CUI
            WHERE s.STY IN ({sem})
              AND length(m.NSTR) >= {int(min_surface_len)}
            """
        ).fetchall()
        con.close()
        for cui, nstr in rows:
            surface = " ".join(nstr.split())
            add(surface, cui, 0.9)

    return {k: sorted(v.items(), key=lambda x: -x[1]) for k, v in lookup.items()}


def save_cui_to_graph_idx(path: Path, cui_to_idx: dict[str, int]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(cui_to_idx), encoding="utf-8")


def load_cui_to_graph_idx(path: Path = CUI_TO_GRAPH_IDX_JSON) -> dict[str, int]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    return {str(k): int(v) for k, v in payload.items()}


def ensure_cui_to_graph_idx(
    path: Path = CUI_TO_GRAPH_IDX_JSON,
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    rebuild: bool = False,
) -> dict[str, int]:
    """Build or load CUI → graph row index (full graph, highest-degree node per CUI)."""
    if path.exists() and not rebuild:
        return load_cui_to_graph_idx(path)
    cui_to_idx = build_cui_to_graph_idx(id_maps_parquet, max_graph_idx=None)
    save_cui_to_graph_idx(path, cui_to_idx)
    return cui_to_idx


def save_linker_cache(
    path: Path,
    cui_to_idx: dict[str, int],
    surface_lookup: dict[str, list[tuple[str, float]]],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "cui_to_idx": cui_to_idx,
        "surface_lookup": {k: list(v) for k, v in surface_lookup.items()},
    }
    path.write_text(json.dumps(payload), encoding="utf-8")


def load_linker_cache(
    path: Path,
) -> tuple[dict[str, int], dict[str, list[tuple[str, float]]]]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    surface_lookup = {
        k: [(cui, float(score)) for cui, score in v]
        for k, v in payload["surface_lookup"].items()
    }
    return payload["cui_to_idx"], surface_lookup


class QueryEntityLinker:
    """Map query text to graph node indices with confidence scores."""

    def __init__(
        self,
        *,
        id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
        umls_duckdb: Path | None = UMLS_DUCKDB_PATH,
        cache_path: Path = QUERY_LINKER_CACHE_JSON,
        quickumls_path: Path | None = None,
        max_graph_idx: int | None = None,
        rebuild_cache: bool = False,
        min_confidence: float = 0.7,
        max_ngram: int = 6,
        quickumls_only: bool = False,
        skip_graph_cui_filter: bool = False,
    ):
        self.min_confidence = min_confidence
        self.max_ngram = max_ngram
        self.quickumls_only = quickumls_only
        self.skip_graph_cui_filter = skip_graph_cui_filter
        self._quickumls = None

        if quickumls_only and skip_graph_cui_filter:
            self.cui_to_idx = {}
            self.surface_lookup = {}
        elif quickumls_only:
            self.cui_to_idx = ensure_cui_to_graph_idx(
                rebuild=rebuild_cache,
                id_maps_parquet=id_maps_parquet,
            )
            self.surface_lookup = {}
        elif cache_path.exists() and not rebuild_cache:
            self.cui_to_idx, self.surface_lookup = load_linker_cache(cache_path)
        else:
            # Build full-graph linker cache; OOB embedding rows are filtered at pool time.
            cui_to_idx = build_cui_to_graph_idx(id_maps_parquet, max_graph_idx=None)
            clinical_cuis = build_clinical_cuis(cui_to_idx, umls_duckdb)
            self.surface_lookup = build_surface_lookup(
                cui_to_idx,
                clinical_cuis,
                id_maps_parquet=id_maps_parquet,
                umls_duckdb=umls_duckdb,
                max_graph_idx=None,
            )
            self.cui_to_idx = {
                c: i for c, i in cui_to_idx.items() if c in clinical_cuis
            }
            save_linker_cache(cache_path, self.cui_to_idx, self.surface_lookup)

        if quickumls_path is not None and Path(quickumls_path).exists():
            self._quickumls, load_err = self._try_load_quickumls(quickumls_path)
            if quickumls_only and self._quickumls is None:
                raise RuntimeError(
                    "quickumls_only=True but QuickUMLS could not be loaded.\n"
                    f"{load_err or 'Unknown error.'}\n"
                    "Fix: run ./setup.sh in the bundle root, then use ./run_export.sh "
                    "(not system python3)."
                )

    @staticmethod
    def _try_load_quickumls(quickumls_path: Path):
        import sys

        try:
            from quickumls import QuickUMLS
        except ImportError as exc:
            return None, (
                f"ImportError: {exc}\n"
                f"  python: {sys.executable}\n"
                "  Hint: create the venv with ./setup.sh and run via ./run_export.sh"
            )
        try:
            matcher = QuickUMLS(
                str(quickumls_path),
                threshold=0.7,
                accepted_semtypes=sorted(CLINICAL_SEMTYPES),
            )
        except Exception as exc:
            return None, (
                f"{type(exc).__name__}: {exc}\n"
                f"  python: {sys.executable}\n"
                f"  quickumls_data: {quickumls_path}"
            )
        return matcher, None

    def _link_quickumls(self, text: str) -> list[tuple[str, float]]:
        if self._quickumls is None:
            return []
        out: list[tuple[str, float]] = []
        for match_group in self._quickumls.match(
            text, best_match=True, ignore_syntax=False
        ):
            for m in match_group:
                cui = m.get("cui")
                sim = float(m.get("similarity", 0.0))
                if cui and sim >= self.min_confidence:
                    out.append((cui, sim))
        return out

    def _link_dictionary(self, text: str) -> list[tuple[str, float]]:
        tokens = _tokenize(text)
        if not tokens:
            return []

        matches: dict[str, float] = {}
        n = len(tokens)
        for start in range(n):
            for length in range(min(self.max_ngram, n - start), 0, -1):
                phrase = " ".join(tokens[start : start + length])
                hits = self.surface_lookup.get(phrase)
                if not hits:
                    continue
                for cui, score in hits:
                    matches[cui] = max(matches.get(cui, 0.0), score)
                break  # longest match at this start position

        # Also try semicolon / clause segments (common in evidence strings).
        for clause in re.split(r"[;,\n]+", text):
            clause = clause.strip()
            if not clause:
                continue
            norm_clause = normalize_surface(clause)
            if norm_clause in self.surface_lookup:
                for cui, score in self.surface_lookup[norm_clause]:
                    matches[cui] = max(matches.get(cui, 0.0), score)

        return sorted(matches.items(), key=lambda x: -x[1])

    def link_cuis(self, text: str) -> list[tuple[str, float]]:
        """Return (cui, confidence) pairs for a query."""
        matches: dict[str, float] = {}
        if self.quickumls_only:
            for cui, score in self._link_quickumls(text):
                if self.skip_graph_cui_filter or cui in self.cui_to_idx:
                    matches[cui] = max(matches.get(cui, 0.0), score)
        else:
            for cui, score in self._link_dictionary(text):
                matches[cui] = max(matches.get(cui, 0.0), score)
            for cui, score in self._link_quickumls(text):
                if cui in self.cui_to_idx:
                    matches[cui] = max(matches.get(cui, 0.0), score)
        return sorted(matches.items(), key=lambda x: -x[1])

    def link_graph_nodes(self, text: str) -> list[tuple[int, float]]:
        """Return (graph_row_idx, confidence) pairs for a query."""
        out: list[tuple[int, float]] = []
        for cui, score in self.link_cuis(text):
            idx = self.cui_to_idx.get(cui)
            if idx is not None:
                out.append((int(idx), score))
        return out

    def coverage_stats(self, texts: list[str]) -> dict[str, float]:
        """Compute entity-linking coverage over a list of queries."""
        n = len(texts)
        if n == 0:
            return {
                "n": 0,
                "pct_with_match": 0.0,
                "pct_zero_match": 100.0,
                "avg_entities": 0.0,
            }
        counts = [len(self.link_graph_nodes(t)) for t in texts]
        with_match = sum(1 for c in counts if c > 0)
        return {
            "n": n,
            "pct_with_match": 100.0 * with_match / n,
            "pct_zero_match": 100.0 * (n - with_match) / n,
            "avg_entities": float(np.mean(counts)),
            "median_entities": float(np.median(counts)),
            "max_entities": int(max(counts)),
        }