scripts/run_icd10_batch.py

"""
scripts/run_icd10_batch.py
────────────────────────────────────────────────────────────────
Batch ICD-10 mapper — completes the src/nlp/icd_mapper.py wiring.
Reads DISEASE/SYMPTOM entities that don't have ICD-10 mappings
yet, maps each *unique* entity text once (exact -> fuzzy ->
embedding, see ICD10Mapper.map), and fans the result out to every
entity row that shares that text. Persists via load_icd10_mappings().

Only DISEASE and SYMPTOM entities are eligible -- medications,
procedures, and anatomy terms use different code systems (see
ICD10Mapper.map_entities). Caching by unique text matters here:
206k eligible entity rows reduce to ~43k unique texts, so the
cache avoids redundant fuzzy/embedding lookups for repeated terms
like "hypertension".

Idempotent / resumable: an entity is only processed if it has
zero rows in icd10_mappings, so re-running after an interruption
never creates duplicates.

icd10_mappings only ever stores *found* matches -- an entity text
that the mapper couldn't confidently map has no row and therefore
still looks "pending" on the next run. Without a separate record
of what's already been tried, every restart would re-run the slow
fuzzy/embedding matching on every previously-checked no-match text.
A small on-disk cache (data/processed/icd10_mapping_cache.json,
keyed by entity text) closes that gap: any text seen before --
matched or not -- is skipped straight from the cache on restart.

Usage
─────
    python scripts/run_icd10_batch.py                # all pending entities
    python scripts/run_icd10_batch.py --limit 500     # quick demo subset
────────────────────────────────────────────────────────────────
"""

from __future__ import annotations

import argparse
import json
import time

from sqlalchemy import select

from src.db.connection import get_session
from src.db.models import Entity, ICD10Mapping
from src.etl.load import load_icd10_mappings
from src.nlp.icd_mapper import ICD10Mapper
from src.utils.config import Paths
from src.utils.logger import get_logger

logger = get_logger(__name__)

_ELIGIBLE_LABELS = ("DISEASE", "SYMPTOM")
_CACHE_PATH = Paths.processed / "icd10_mapping_cache.json"


def _load_cache() -> dict[str, list[dict]]:
    """Load the on-disk text -> match-dicts cache, or an empty dict."""
    if not _CACHE_PATH.exists():
        return {}
    with open(_CACHE_PATH, encoding="utf-8") as f:
        cache = json.load(f)
    logger.info("Loaded ICD-10 mapping cache: %d entity texts already tried", len(cache))
    return cache


def _save_cache(cache: dict[str, list[dict]]) -> None:
    """Write the text -> match-dicts cache to disk."""
    _CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
    with open(_CACHE_PATH, "w", encoding="utf-8") as f:
        json.dump(cache, f)


def _fetch_pending_entities(limit: int | None) -> list[tuple[int, str]]:
    """Return (id, text) for DISEASE/SYMPTOM entities with no ICD-10 mapping yet.

    Args:
        limit: Maximum number of entity rows to return, or None for
            every pending entity.

    Returns:
        List of (entity_id, text) tuples, ordered by entity id so a
        --limit run always covers the same entities on re-run.
    """
    with get_session() as session:
        stmt = (
            select(Entity.id, Entity.text)
            .outerjoin(ICD10Mapping, ICD10Mapping.entity_id == Entity.id)
            .where(
                Entity.label.in_(_ELIGIBLE_LABELS),
                ICD10Mapping.id.is_(None),
            )
            .order_by(Entity.id)
        )
        if limit:
            stmt = stmt.limit(limit)
        return [(row.id, row.text) for row in session.execute(stmt)]


def run(limit: int | None, commit_every: int) -> None:
    """Map pending entities to ICD-10 codes and save the results.

    Args:
        limit: Maximum number of pending entity rows to process, or
            None for every DISEASE/SYMPTOM entity without a mapping.
        commit_every: Write to the database once this many mapping
            rows have been queued in memory.
    """
    pending = _fetch_pending_entities(limit)
    total_entities = len(pending)
    if total_entities == 0:
        logger.info(
            "No pending entities -- every DISEASE/SYMPTOM entity "
            "already has an ICD-10 mapping."
        )
        return

    by_text: dict[str, list[int]] = {}
    for entity_id, text in pending:
        by_text.setdefault(text, []).append(entity_id)

    unique_texts = list(by_text.keys())
    logger.info(
        "Mapping %d unique entity texts (%d entity rows) to ICD-10 codes...",
        len(unique_texts), total_entities,
    )

    mapper = ICD10Mapper()
    cache = _load_cache()
    cache_hits = 0

    start = time.time()
    processed_entities = 0
    inserted_total = 0
    pending_dicts: list[dict] = []

    for i, text in enumerate(unique_texts, start=1):
        if text in cache:
            match_dicts = cache[text]
            cache_hits += 1
        else:
            match_dicts = [m.to_dict() for m in mapper.map(text)]
            cache[text] = match_dicts

        entity_ids = by_text[text]
        for entity_id in entity_ids:
            for m in match_dicts:
                d = dict(m)
                d["entity_id"] = entity_id
                pending_dicts.append(d)

        processed_entities += len(entity_ids)

        if len(pending_dicts) >= commit_every or i == len(unique_texts):
            inserted_total += load_icd10_mappings(pending_dicts)
            pending_dicts = []
            _save_cache(cache)

            elapsed = time.time() - start
            rate = processed_entities / elapsed if elapsed > 0 else 0
            eta_seconds = (total_entities - processed_entities) / rate if rate > 0 else float("inf")
            logger.info(
                "  %d/%d entities | %d/%d unique texts (%d cache hits) | %d mappings so far | "
                "%.0fs elapsed | ETA %.0fs",
                processed_entities, total_entities, i, len(unique_texts), cache_hits,
                inserted_total, elapsed, eta_seconds,
            )

    logger.info(
        "Done: %d entities processed, %d mappings inserted in %.0fs",
        processed_entities, inserted_total, time.time() - start,
    )


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Map DISEASE/SYMPTOM entities to ICD-10 codes and save the results."
    )
    parser.add_argument(
        "--limit", type=int, default=None,
        help="Max number of pending entity rows to process (default: all pending entities)",
    )
    parser.add_argument(
        "--commit-every", type=int, default=500,
        help="Queue this many mapping rows before writing to the database (default: 500)",
    )
    args = parser.parse_args()

    run(limit=args.limit, commit_every=args.commit_every)


if __name__ == "__main__":
    main()