src/data/query_entity_export.py

"""Export QuickUMLS-detected CUIs and clinical-AUI-validated AUIs for query texts."""

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable, Iterator

from config.paths import (
    CUI_TO_VALID_AUI_JSON,
    CUI_TO_VALID_AUI_PARQUET,
    GRAPH_ID_MAPS_PARQUET,
    MRCONSO_CLINICAL_PARQUET,
    MRCONSO_CLINICAL_PARQUET_DEFAULT,
)
from src.data.query_entity_linker import QueryEntityLinker


def _resolve_mrconso_parquet(path: Path | None) -> Path:
    if path is not None:
        candidate = path
    elif MRCONSO_CLINICAL_PARQUET.exists():
        candidate = MRCONSO_CLINICAL_PARQUET
    else:
        candidate = MRCONSO_CLINICAL_PARQUET_DEFAULT
    if not candidate.exists():
        raise FileNotFoundError(
            f"Clinical MRCONSO parquet not found: {candidate}. "
            "Copy mrconso_clinical.parquet to data/processed/graph/ or set "
            "ICD_MRCONSO_CLINICAL_PARQUET."
        )
    return candidate


@dataclass(frozen=True)
class ValidAuiRecord:
    cui: str
    aui: str
    graph_idx: int
    str: str
    sab: str


def _cui_to_valid_aui_sql(
    *,
    id_maps_parquet: Path,
    mrconso_parquet: Path,
) -> str:
    return f"""
        WITH graph_auis AS (
            SELECT
                cui,
                REPLACE(node_id, 'AUI:', '') AS aui,
                idx,
                degree,
                str,
                sab
            FROM read_parquet('{id_maps_parquet.as_posix()}')
            WHERE kind = 'AUI' AND cui IS NOT NULL
        ),
        ranked AS (
            SELECT
                g.cui,
                g.aui,
                g.idx,
                g.str,
                g.sab,
                ROW_NUMBER() OVER (
                    PARTITION BY g.cui
                    ORDER BY g.degree DESC, g.idx
                ) AS rn
            FROM graph_auis g
            INNER JOIN read_parquet('{mrconso_parquet.as_posix()}') m
                ON g.aui = m.AUI
        )
        SELECT cui, aui, idx AS graph_idx, str, sab
        FROM ranked
        WHERE rn = 1
    """


def build_cui_to_valid_aui_parquet(
    output: Path = CUI_TO_VALID_AUI_PARQUET,
    *,
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    mrconso_parquet: Path | None = None,
) -> int:
    """Build the CUI→AUI parquet cache via DuckDB + clinical MRCONSO parquet."""
    mrconso = _resolve_mrconso_parquet(mrconso_parquet)
    if not id_maps_parquet.exists():
        raise FileNotFoundError(f"Graph id maps not found: {id_maps_parquet}")

    import duckdb

    output.parent.mkdir(parents=True, exist_ok=True)
    con = duckdb.connect()
    con.execute(
        f"""
        COPY (
            {_cui_to_valid_aui_sql(
                id_maps_parquet=id_maps_parquet,
                mrconso_parquet=mrconso,
            )}
        ) TO '{output.as_posix()}' (FORMAT PARQUET)
        """
    )
    con.close()

    import polars as pl

    return pl.scan_parquet(output).select(pl.len()).collect().item()


def build_cui_to_valid_aui_map(
    *,
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    mrconso_parquet: Path | None = None,
) -> dict[str, ValidAuiRecord]:
    """Map each CUI to its highest-degree AUI present in clinical MRCONSO."""
    build_cui_to_valid_aui_parquet(
        id_maps_parquet=id_maps_parquet,
        mrconso_parquet=mrconso_parquet,
    )
    lookup = CuiToValidAuiLookup(CUI_TO_VALID_AUI_PARQUET)
    try:
        import polars as pl

        cuis = (
            pl.scan_parquet(CUI_TO_VALID_AUI_PARQUET)
            .select("cui")
            .collect()
            .get_column("cui")
            .to_list()
        )
        return lookup.lookup_many(cuis)
    finally:
        lookup.close()


def save_cui_to_valid_aui_map(
    path: Path,
    mapping: dict[str, ValidAuiRecord],
) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {cui: asdict(rec) for cui, rec in mapping.items()}
    path.write_text(json.dumps(payload), encoding="utf-8")


def load_cui_to_valid_aui_map(
    path: Path = CUI_TO_VALID_AUI_JSON,
) -> dict[str, ValidAuiRecord]:
    payload = json.loads(path.read_text(encoding="utf-8"))
    return {
        cui: ValidAuiRecord(
            cui=str(rec["cui"]),
            aui=str(rec["aui"]),
            graph_idx=int(rec["graph_idx"]),
            str=str(rec.get("str", "")),
            sab=str(rec.get("sab", "")),
        )
        for cui, rec in payload.items()
    }


class CuiToValidAuiLookup:
    """Polars-backed lookup over ``cui_to_valid_aui.parquet`` (memory-safe)."""

    def __init__(self, parquet_path: Path = CUI_TO_VALID_AUI_PARQUET) -> None:
        if not parquet_path.exists():
            raise FileNotFoundError(f"CUI→AUI parquet not found: {parquet_path}")
        import polars as pl

        self._parquet_path = parquet_path
        self._scan = pl.scan_parquet(parquet_path)

    def count(self) -> int:
        import polars as pl

        return self._scan.select(pl.len()).collect().item()

    def lookup_many(self, cuis: Iterable[str]) -> dict[str, ValidAuiRecord]:
        unique = list(dict.fromkeys(str(c) for c in cuis if c))
        if not unique:
            return {}

        import polars as pl

        df = self._scan.filter(pl.col("cui").is_in(unique)).collect()
        out: dict[str, ValidAuiRecord] = {}
        for row in df.iter_rows(named=True):
            cui = str(row["cui"])
            out[cui] = ValidAuiRecord(
                cui=cui,
                aui=str(row["aui"]),
                graph_idx=int(row["graph_idx"]),
                str=str(row["str"]) if row["str"] is not None else "",
                sab=str(row["sab"]) if row["sab"] is not None else "",
            )
        return out

    def close(self) -> None:
        return None


def ensure_cui_to_valid_aui_lookup(
    *,
    parquet_path: Path = CUI_TO_VALID_AUI_PARQUET,
    rebuild: bool = False,
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    mrconso_parquet: Path | None = None,
) -> CuiToValidAuiLookup:
    if not parquet_path.exists() or rebuild:
        build_cui_to_valid_aui_parquet(
            parquet_path,
            id_maps_parquet=id_maps_parquet,
            mrconso_parquet=mrconso_parquet,
        )
    return CuiToValidAuiLookup(parquet_path)


def ensure_cui_to_valid_aui_map(
    *,
    path: Path = CUI_TO_VALID_AUI_JSON,
    rebuild: bool = False,
    id_maps_parquet: Path = GRAPH_ID_MAPS_PARQUET,
    mrconso_parquet: Path | None = None,
) -> dict[str, ValidAuiRecord]:
    lookup = ensure_cui_to_valid_aui_lookup(
        rebuild=rebuild,
        id_maps_parquet=id_maps_parquet,
        mrconso_parquet=mrconso_parquet,
    )
    try:
        import polars as pl

        cuis = (
            pl.scan_parquet(CUI_TO_VALID_AUI_PARQUET)
            .select("cui")
            .collect()
            .get_column("cui")
            .to_list()
        )
        return lookup.lookup_many(cuis)
    finally:
        lookup.close()


def detect_entities_for_text(
    text: str,
    *,
    linker: QueryEntityLinker,
    lookup: CuiToValidAuiLookup,
) -> list[dict[str, object]]:
    """Return kept CUI/AUI entities for one query text."""
    linked = linker.link_cuis(text)
    records = lookup.lookup_many(cui for cui, _ in linked)
    entities: list[dict[str, object]] = []
    for cui, confidence in linked:
        record = records.get(cui)
        if record is None:
            continue
        entities.append(
            {
                "cui": record.cui,
                "aui": record.aui,
                "confidence": round(float(confidence), 4),
                "graph_idx": record.graph_idx,
                "str": record.str,
                "sab": record.sab,
            }
        )
    return entities


def export_entity_detections(
    rows: list[dict[str, object]],
    *,
    linker: QueryEntityLinker,
    lookup: CuiToValidAuiLookup,
    text_key: str = "query_text",
    include_code: bool = True,
) -> list[dict[str, object]]:
    """Build export records aligned to input JSONL rows."""
    return list(
        iter_entity_detection_entries(
            rows,
            linker=linker,
            lookup=lookup,
            text_key=text_key,
            include_code=include_code,
        )
    )


def iter_entity_detection_entries(
    rows: list[dict[str, object]],
    *,
    linker: QueryEntityLinker,
    lookup: CuiToValidAuiLookup,
    text_key: str = "query_text",
    include_code: bool = True,
) -> Iterator[dict[str, object]]:
    for row in rows:
        text = str(row.get(text_key, "")).strip()
        entry: dict[str, object] = {"query_text": text}
        if include_code and "code" in row:
            entry["code"] = row["code"]
        entry["entities"] = detect_entities_for_text(
            text,
            linker=linker,
            lookup=lookup,
        )
        yield entry


def write_entity_detections_json(
    entries: Iterable[dict[str, object]],
    output: Path,
    *,
    pretty: bool = False,
) -> tuple[int, int, int]:
    """Stream entries to JSON without holding the full payload in memory."""
    output.parent.mkdir(parents=True, exist_ok=True)
    n_entries = 0
    with_entities = 0
    total_entities = 0
    with open(output, "w", encoding="utf-8") as f:
        f.write('{\n  "entries": [\n')
        first = True
        for entry in entries:
            if not first:
                f.write(",\n")
            else:
                first = False
            f.write("    ")
            if pretty:
                rendered = json.dumps(entry, ensure_ascii=False, indent=2)
                f.write(rendered.replace("\n", "\n    "))
            else:
                f.write(json.dumps(entry, ensure_ascii=False))
            n_entries += 1
            n_ent = len(entry.get("entities", []))
            if n_ent:
                with_entities += 1
                total_entities += n_ent
        f.write("\n  ]\n}\n")
    return n_entries, with_entities, total_entities