src/backend/codesearch.py

"""Dense retrieval for medical-code search.

The default model is Vidul's BGE_FT_VA with its original FAISS indexes, with no
score blend or fallback model. Comparison models attach their own indexes
(`attach`): one index per searched index name, built over exactly the default
rows in the same order, so every model shares one copy of the row metadata and
of every per-row sidecar. Each attached index is checked against the default
meta.csv fingerprint before it is trusted. Indexes load per category, but the app warms every
category at startup (`warm`) rather than on first use: deferring the read only
moved a multi-hundred-megabyte read of a lazily fetched network mount onto a
visitor's request, which is the read pattern that mount fails on. The saving
was never real either, since every category is queried in an ordinary session.
"""

from __future__ import annotations

import bisect
import csv
import hashlib
import io
import json
import os
from collections import Counter

import faiss
import numpy as np

from . import paths
from . import provenance
from .mappings import MappingStore, get_store
from .retriever import DenseEmbedder
from .spelling import SpellSuggester

# Top-score floor below which a result set is flagged low-confidence, per
# category — the scales differ (nonsense probes reach ~0.87 cosine against the
# 430k-row lab index but only ~0.76 against CPT). Calibrated 2026-08-05 against
# demo-bundle real queries vs nonsense probes: every observed real query clears
# its floor with margin; refine from gold-pair score distributions as they
# accumulate (scripts/eval_gold_pairs.py).
def _low_conf(category: str, default: float) -> float:
    return float(os.environ.get(f"ENCODE_LOW_CONFIDENCE_{category.upper()}", default))

LOW_CONFIDENCE = {"diagnosis": _low_conf("diagnosis", 0.75),
                  "medication": _low_conf("medication", 0.74),
                  "lab": _low_conf("lab", 0.80),
                  "procedure": _low_conf("procedure", 0.78)}

CATEGORY_INDEXES = {
    "diagnosis": [("icd_index", None)],
    # The NDC index is not searched directly: NDC codes reach results through
    # mappings on local drug codes, not as a peer code system.
    "medication": [("med_index", "Local Drug SID")],
    "lab": [("labchem_index", "LabChemTestSID")],
    "procedure": [("cpt_index", "CPT")],
}

# The prebuilt indexes key their rows by description, so an ICD-9 and an
# ICD-10 code sharing a description land in one row carrying both codes and
# both labels ("039.2,A42.1" / "ICD10 | ICD9"). That label describes the row,
# not either code, and no code is ever both versions. These sidecars give each
# code in such a row its own version so results carry one code system each and
# the merged label never reaches a caller. Built by
# scripts/build_icd_index_versions.py; see it for how the split is resolved.
ROW_VERSION_SIDECARS = {"icd_index": paths.ICD_INDEX_VERSIONS}


def _normalize_code(value: str) -> str:
    """One spelling for a code, so N18.9, n189 and " N18.9 " are one key."""
    return str(value).strip().upper().replace(".", "").replace(" ", "")


def make_embedder() -> DenseEmbedder:
    """Reproduce Vidul's SentenceTransformer query convention on CPU."""
    if not paths.MODEL_DIR.exists():
        raise RuntimeError(f"Missing Vidul fine-tuned model: {paths.MODEL_DIR}")
    return DenseEmbedder(str(paths.MODEL_DIR), pooling="mean", query_prefix="passage: ")


# The default indexes predate per-model configs: they name their model this way
# and leave the query convention to make_embedder().
DEFAULT_INDEX_MODEL = "BGE_FT_VA"

# Per-category low-confidence floors for an attached model, written beside its
# indexes by scripts/calibrate_low_confidence.py. A model without the file
# flags nothing: a floor calibrated for another model's score scale would be
# wrong in a way no one could see.
LOW_CONFIDENCE_FILE = "low_confidence.json"


def index_config(directory) -> dict:
    """The config.json beside one attached index."""
    return json.loads((directory / "config.json").read_text())


class CodeSearchEngine:
    # Adding a mapping to a category = adding a stamp method and one entry here.
    STAMPERS = {"lab": "_stamp_lab",
                "diagnosis": "_stamp_diagnosis",
                "procedure": "_stamp_procedure"}

    def __init__(self, embedder: DenseEmbedder, store: MappingStore | None = None):
        self.embedder = embedder
        self.store = store or get_store()
        self._loaded: dict[str, list[dict]] = {}
        self._suggesters: dict[str, SpellSuggester] = {}
        self._systems: dict[str, list[dict]] = {}
        self._rbcs_sizes: dict[str, int] | None = None
        self._code_order: dict[str, tuple[list[str], dict]] = {}
        # meta.csv fingerprint per default index, which attached indexes must match.
        self._meta_sha: dict[str, str] = {}
        # Attached models: id -> {"dir", "embedder", "floors"}, and their indexes
        # per category, aligned with the default sources.
        self._models: dict[str, dict] = {}
        self._model_indexes: dict[tuple[str, str], list] = {}
        missing = [name for specs in CATEGORY_INDEXES.values() for name, _ in specs
                   if not self._complete_index(paths.CODE_INDEX_DIR / name)]
        if missing:
            raise RuntimeError(
                f"Missing Vidul FAISS index assets under {paths.CODE_INDEX_DIR}: "
                f"{', '.join(missing)}")

    @staticmethod
    def _complete_index(directory) -> bool:
        return all((directory / name).exists()
                   for name in ("index.faiss", "meta.csv", "config.json"))

    @staticmethod
    def _load_row_versions(directory_name: str, meta_path, rows: int,
                           digest: str) -> dict[int, list[tuple[str, str]]]:
        """Per-code ICD versions for the rows that pack two code systems.

        Keyed by row position, which is only meaningful against the exact
        meta.csv the sidecar was built from, so the fingerprint is checked
        rather than trusted: a regenerated index with a stale sidecar would
        otherwise silently relabel unrelated codes.
        """
        path = ROW_VERSION_SIDECARS.get(directory_name)
        if path is None:
            return {}
        if not path.exists():
            raise RuntimeError(
                f"Missing per-code version sidecar {path}; build it with "
                f"scripts/build_icd_index_versions.py")
        payload = json.loads(path.read_text(encoding="utf-8"))
        if payload.get("meta_rows") != rows or payload.get("meta_sha256") != digest:
            raise RuntimeError(
                f"{path} was built from a different {meta_path}; rerun "
                f"scripts/build_icd_index_versions.py")
        return {int(position): [(code, version) for code, version in pairs]
                for position, pairs in payload["rows"].items()}

    @staticmethod
    def _searchable(source: dict) -> int:
        """Results the source can return, which is its row count only while no
        row splits. A split row yields one result per code, so counting rows
        would report a smaller corpus than the code-system facets add up to."""
        split = source["row_versions"]
        return (source["index"].ntotal - len(split)
                + sum(len(pairs) for pairs in split.values()))

    @staticmethod
    def _variants(source: dict, position: int, record: dict) -> list[tuple[str, str]]:
        """The (code, code_type) pairs one index row contributes to results.

        An ordinary row contributes itself. A row that packs an ICD-9 and an
        ICD-10 code under one description contributes one pair per code, each
        carrying its own resolved version, so a caller never sees a single
        code labelled with two systems. Splitting only relabels: the codes a
        category serves are exactly the codes in the source index.
        """
        split = source["row_versions"].get(position)
        if split:
            return split
        return [(record["code"],
                 source["code_type"] or record.get("version") or "Diagnosis")]

    def categories(self) -> list[str]:
        return list(CATEGORY_INDEXES)

    def _load_category(self, category: str) -> list[dict]:
        if category not in CATEGORY_INDEXES:
            raise KeyError(category)
        if category in self._loaded:
            return self._loaded[category]

        loaded = []
        for directory_name, fixed_code_type in CATEGORY_INDEXES[category]:
            directory = paths.CODE_INDEX_DIR / directory_name
            config = json.loads((directory / "config.json").read_text())
            if config.get("model") != "BGE_FT_VA" or config.get("dim") != 1024:
                raise RuntimeError(f"Unexpected retrieval config: {directory / 'config.json'}")
            index = faiss.read_index(str(directory / "index.faiss"))
            # Read once: the bytes give both the rows and the fingerprint that
            # the row-version sidecar and every attached index are checked by.
            meta_bytes = (directory / "meta.csv").read_bytes()
            digest = hashlib.sha256(meta_bytes).hexdigest()
            self._meta_sha[directory_name] = digest
            records = list(csv.DictReader(io.StringIO(meta_bytes.decode("utf-8"), newline="")))
            if index.ntotal != len(records):
                raise RuntimeError(
                    f"Index/metadata mismatch in {directory}: {index.ntotal} != {len(records)}")
            loaded.append({"index": index, "records": records,
                           "code_type": fixed_code_type,
                           "row_versions": self._load_row_versions(
                               directory_name, directory / "meta.csv", len(records), digest)})
        self._loaded[category] = loaded
        self._suggesters[category] = SpellSuggester(
            record["description"] for source in loaded for record in source["records"])
        return loaded

    def warm(self, category: str, model_id: str | None = None) -> None:
        """Force this category's index in, so no request ever loads one.

        Loading on first use put a multi-hundred-megabyte read of a lazily
        fetched network mount on a visitor's request, concurrently with
        whatever else was reading it. That is the read pattern the mount fails
        on, so every category is pulled in at startup instead, one at a time.
        """
        self._load_category(category)
        if model_id is not None:
            self._load_model_category(model_id, category)

    # -- attached comparison models --------------------------------------------

    def attach(self, model_id: str, index_dir, embedder: DenseEmbedder) -> None:
        """Register a comparison model's indexes; `warm` then loads them."""
        missing = [name for specs in CATEGORY_INDEXES.values() for name, _ in specs
                   if not self._complete_index(index_dir / name)]
        if missing:
            raise RuntimeError(
                f"Missing FAISS index assets under {index_dir}: {', '.join(missing)}")
        floors_path = index_dir / LOW_CONFIDENCE_FILE
        floors = json.loads(floors_path.read_text())["floors"] if floors_path.exists() else {}
        self._models[model_id] = {"dir": index_dir, "embedder": embedder, "floors": floors}

    def models(self) -> list[str]:
        return list(self._models)

    def _load_model_category(self, model_id: str, category: str) -> list:
        key = (model_id, category)
        if key in self._model_indexes:
            return self._model_indexes[key]
        model = self._models.get(model_id)
        if model is None:
            raise LookupError(f"Model '{model_id}' is not loaded for code search.")
        sources = self._load_category(category)
        indexes = []
        for (directory_name, _), source in zip(CATEGORY_INDEXES[category], sources):
            directory = model["dir"] / directory_name
            config = index_config(directory)
            if config.get("model_id") != model_id:
                raise RuntimeError(f"{directory} was built for '{config.get('model_id')}', "
                                   f"not '{model_id}'")
            if config.get("meta_sha256") != self._meta_sha[directory_name]:
                raise RuntimeError(f"{directory} was built from a different meta.csv than "
                                   f"{paths.CODE_INDEX_DIR / directory_name}")
            index = faiss.read_index(str(directory / "index.faiss"))
            if index.ntotal != source["index"].ntotal:
                raise RuntimeError(f"Index/metadata mismatch in {directory}: "
                                   f"{index.ntotal} != {source['index'].ntotal}")
            indexes.append(index)
        self._model_indexes[key] = indexes
        return indexes

    def _sources_for(self, category: str, model_id: str | None) -> list[dict]:
        """The default sources, with the model's index in place of the default one."""
        sources = self._load_category(category)
        if model_id is None:
            return sources
        indexes = self._load_model_category(model_id, category)
        return [{**source, "index": index} for source, index in zip(sources, indexes)]

    def records(self, category: str) -> list[dict]:
        """Metadata rows in the shape consumed by the optional graph view."""
        rows = []
        for source in self._load_category(category):
            for position, record in enumerate(source["records"]):
                for code, code_type in self._variants(source, position, record):
                    rows.append({"code": code,
                                 "description": record["description"],
                                 "code_type": code_type})
        return rows

    # -- per-category mapping stamps (adding a mapping = adding a function) --

    # A lab index row packs every LabChemTestSID that shares one description,
    # and 73 of them is ordinary. Requiring all of them to name the same LOINC
    # threw away the row whenever a single SID dissented, so "hemoglobin a1c",
    # where 66 of 67 mapped SIDs say 4548-4, showed no mapping at all. A clear
    # majority now carries the row and is reported as derived, since the row's
    # assertion is then ENCODE's and not the crosswalk's.
    AGREEMENT = 2 / 3

    def _stamp_lab(self, result: dict, record: dict) -> None:
        # Each entry is paired with the SID that carries it: the row is cited
        # and graphed by a code that actually holds the mapping, which is not
        # always the first one in the field.
        entries = [(code.strip(), entry) for code, entry in
                   ((code, self.store.loinc_of(code))
                    for code in str(record["code"]).split(",")) if entry]
        result["graphable"] = "LOINC" in result["code_type"].upper() or bool(entries)
        if not entries:
            return
        counts = Counter(entry["loinc"] for _, entry in entries)
        loinc, top = counts.most_common(1)[0]
        split = len(counts) > 1
        if split and top / len(entries) < self.AGREEMENT:
            # Genuinely mixed: the codes in this row are not one test. Say so
            # rather than asserting one of them.
            result["graphable"] = "LOINC" in result["code_type"].upper()
            result["mapping_conflict"] = {"targets": len(counts), "codes": len(entries)}
            return
        agreeing = [(code, entry) for code, entry in entries if entry["loinc"] == loinc]
        # Cite the strongest claim on that LOINC: a code the crosswalk mapped
        # outright, if the row has one, and only otherwise a derived code. The
        # row still counts as derived when the codes disagreed, because then
        # the choice between them is ENCODE's rather than the crosswalk's.
        code, entry = next((pair for pair in agreeing if not pair[1].get("derived")),
                           agreeing[0])
        derived = bool(entry.get("derived")) or split
        result["mapped_loinc"] = loinc
        result["mapped_code"] = code
        if derived:
            result["mapping_derived"] = True
        result["mapping_provenance"] = provenance.line(
            entry.get("source"), derived=derived,
            match="merged_majority" if split else entry.get("match"),
            target=entry.get("loinc_version"),
            detail=f"{top} of {len(entries)}" if split else None)

    def _stamp_diagnosis(self, result: dict, record: dict) -> None:
        # Read off the result, not the row: a row holding both ICD versions has
        # already been split into one result per code, and each must be stamped
        # with the phecode for its own code under its own version. Stamping from
        # the row would hand an ICD-10 result its sibling's ICD-9 assignment.
        versions = [part.strip() for part in result["code_type"].split("|")]
        # A row can still pack several codes of one version ("585.9,585.6").
        # The phecode maps are keyed on a single dotted code, so the whole
        # field never matched and 4,394 packed rows carried no phecode at all.
        # Each code is tried in turn.
        for code in str(result["code"]).split(","):
            phecodes = self.store.phecodes_for(code, versions)
            if phecodes:
                result["phecodes"] = {**phecodes,
                                      "provenance": provenance.phecode_line(phecodes)}
                result.setdefault("mapped_code", code.strip())
                return

    def rbcs_family_sizes(self) -> dict[str, int]:
        """Codes per RBCS family, counted over the procedure codes served here.

        CMS's mapping file and the CPT index are not the same set. Counting the
        file gave "Arthroscopy - Lower Extremity, 46 codes" on the detail page
        while the graph, drawn from the index, said "Showing 26 of 45": two
        numbers for one family, which is the confusion the family size was
        added to remove. This counts the set the graph draws, by the same rule
        (validated membership, each code once), so the two agree by
        construction rather than by luck.
        """
        if self._rbcs_sizes is None:
            sizes: dict[str, int] = {}
            seen: set[str] = set()
            for row in self.records("procedure"):
                if "CPT" not in str(row["code_type"]):
                    continue
                for code in str(row["code"]).split(","):
                    code = code.strip()
                    if not code or code in seen:
                        continue
                    seen.add(code)
                    entry = self.store.rbcs_validated_for(code) or {}
                    group = entry.get("family") or entry.get("subcategory")
                    if group:
                        sizes[group] = sizes.get(group, 0) + 1
            self._rbcs_sizes = sizes
        return self._rbcs_sizes

    def _stamp_procedure(self, result: dict, record: dict) -> None:
        # Same merged-row rule as the other two: the first code in the row is
        # not always the one CMS assigned a group to.
        rbcs = code = None
        for candidate in str(record["code"]).split(","):
            rbcs = self.store.rbcs_for(candidate)
            if rbcs:
                code = candidate.strip()
                break
        if rbcs:
            result["mapped_code"] = code
            result["rbcs"] = {k: rbcs[k] for k in
                              ("category", "subcategory", "family", "major")}
            # How broad the family is travels with the group: a name shared by
            # two codes is normal, and the size is what stops it reading as a
            # mistake. Absent when this code reached its group by a derived
            # mapping, because then the graph has no validated family to draw
            # and there is no number both surfaces could agree on.
            group = rbcs.get("family") or rbcs.get("subcategory")
            size = self.rbcs_family_sizes().get(group) if group else None
            if size:
                result["rbcs"]["family_size"] = size
            result["graphable"] = True
            if rbcs.get("derived"):
                result["mapping_derived"] = True
            result["mapping_provenance"] = provenance.line(
                rbcs.get("source"), derived=bool(rbcs.get("derived")),
                match=rbcs.get("match"))

    def systems(self, category: str) -> list[dict]:
        """Distinct code systems in a category with row counts, for the
        pre-search restriction checkboxes. Computed once per category."""
        cached = self._systems.get(category)
        if cached is None:
            counts: dict[str, int] = {}
            for source in self._load_category(category):
                if source["code_type"]:
                    counts[source["code_type"]] = (counts.get(source["code_type"], 0)
                                                   + len(source["records"]))
                else:
                    # Counted per code rather than per row: a row that packs an
                    # ICD-9 and an ICD-10 code contributes one result to each
                    # system, which is exactly what checking that box returns.
                    for position, record in enumerate(source["records"]):
                        for _, ct in self._variants(source, position, record):
                            counts[ct] = counts.get(ct, 0) + 1
            cached = [{"code_type": t, "count": n}
                      for t, n in sorted(counts.items(), key=lambda kv: -kv[1])]
            self._systems[category] = cached
        return cached

    def _codes_in_order(self, category: str):
        """Every distinct normalized code, sorted, with the rows that carry it.

        This is what makes the lookup a lookup: the neighbours of a code are
        the codes beside it in its own numbering, so 1000 sits next to 1001
        and 10001, the way it does in the code book. Built once per category
        on first use and held; the largest category packs 640,871 codes.
        """
        cached = self._code_order.get(category)
        if cached:
            return cached
        by_code: dict[str, list[tuple[int, int, str, str]]] = {}
        for source_index, source in enumerate(self._load_category(category)):
            for position, record in enumerate(source["records"]):
                for value, code_type in self._variants(source, position, record):
                    for part in str(value).split(","):
                        norm = _normalize_code(part)
                        if norm:
                            by_code.setdefault(norm, []).append(
                                (source_index, position, value, code_type))
        cached = (sorted(by_code), by_code)
        self._code_order[category] = cached
        return cached

    def lookup(self, category: str, code: str, k: int = 50) -> dict:
        """The code's own rows, then the codes beside it in code order.

        Search embeds the query as text and scores it against descriptions, so
        a code never retrieves itself: "90999" ranks against the words in the
        index and returns "time", "introduction", "home". Anyone arriving with
        a code in hand is asking a different question, and it is answered from
        the code column.

        Matching ignores case, surrounding space, and dots, so a copied
        "  n189 " finds N18.9. Related rows are the sorted neighbourhood of
        the query, nearest first, where nearest means the longer shared
        prefix: for 1000 that is 10001 before 1001 before 1100. A partial or
        unknown code therefore still lands in the right part of the code book
        and shows what is there.
        """
        sources = self._load_category(category)
        stamp_name = self.STAMPERS.get(category)
        stamp = getattr(self, stamp_name) if stamp_name else None
        ordered, by_code = self._codes_in_order(category)
        wanted = _normalize_code(code)

        seen: set[tuple[int, int, str, str]] = set()

        def rows_for(norm: str, label: str) -> list[dict]:
            out = []
            for entry in by_code.get(norm, []):
                if entry in seen:      # a packed row answers for one code once
                    continue
                seen.add(entry)
                source_index, position, value, code_type = entry
                record = sources[source_index]["records"][position]
                # No relevance: nothing here was ranked, and printing a score
                # would invent one.
                row = {"code": value, "code_type": code_type,
                       "description": record["description"],
                       "relevance": None, "match": label}
                if stamp:
                    stamp(row, record)
                out.append(row)
            return out

        results = rows_for(wanted, "exact")[:k] if wanted else []

        # Walk outward from the query's place in the sorted code list,
        # taking the side that shares the longer prefix with it first. `k` is
        # the same result-count setting the search obeys and caps the whole
        # page, so nearby rows fill whatever the exact matches left of it.
        related_cap = max(0, k - len(results))
        related: list[dict] = []
        if wanted:
            def shared(norm: str) -> int:
                n = 0
                for a, b in zip(norm, wanted):
                    if a != b:
                        break
                    n += 1
                return n
            left = bisect.bisect_left(ordered, wanted) - 1
            right = bisect.bisect_right(ordered, wanted)
            while len(related) < related_cap and (left >= 0 or right < len(ordered)):
                if right >= len(ordered):
                    take_right = False
                elif left < 0:
                    take_right = True
                else:
                    take_right = shared(ordered[right]) >= shared(ordered[left])
                norm = ordered[right if take_right else left]
                if take_right:
                    right += 1
                else:
                    left -= 1
                related.extend(rows_for(norm, "nearby code"))
            related = related[:related_cap]

        rows = results + related
        for rank, row in enumerate(rows, 1):
            row["rank"] = rank
        return {
            "query": code.strip(),
            "category": category,
            # Truthful about what was found: a lookup that hit the code says
            # exact, one that only found neighbours says nearest.
            "match": "exact" if results else "nearest",
            "total": sum(self._searchable(source) for source in sources),
            "count": len(rows),
            "exact_count": len(results),
            "related_count": len(related),
            "results": rows,
        }

    def search(self, category: str, query: str, k: int = 50,
               systems: set[str] | None = None, model_id: str | None = None) -> dict:
        """Rank codes for `query`. `model_id` names an attached model; None is
        the default model."""
        sources = self._sources_for(category, model_id)
        if model_id is None:
            embedder, floor = self.embedder, LOW_CONFIDENCE.get(category, 0.75)
        else:
            model = self._models[model_id]
            embedder, floor = model["embedder"], model["floors"].get(category)
        query_vector = embedder.encode([query.strip()]).astype(np.float32)
        stamp_name = self.STAMPERS.get(category)
        stamp = getattr(self, stamp_name) if stamp_name else None
        # A restricted system may be a thin slice of a source, so over-fetch
        # before filtering; FAISS neighbor count barely affects a flat scan.
        fetch_k = k if not systems else max(k * 10, 500)
        candidates = []
        for source in sources:
            scores, indices = source["index"].search(
                query_vector, min(fetch_k, source["index"].ntotal))
            for score, index_position in zip(scores[0], indices[0]):
                if index_position < 0:
                    continue
                record = source["records"][int(index_position)]
                # One row can carry more than one code system; each is its own
                # result, so every code_type below names exactly one system.
                for code, code_type in self._variants(source, int(index_position), record):
                    if systems and code_type not in systems:
                        continue
                    result = {
                        "code": code,
                        "code_type": code_type,
                        "description": record["description"],
                        "_score": float(np.clip(score, -1.0, 1.0)),
                    }
                    if stamp:
                        stamp(result, record)
                    candidates.append(result)

        candidates.sort(key=lambda row: row["_score"], reverse=True)
        results = candidates[:k]
        for rank, result in enumerate(results, 1):
            result["relevance"] = round(result.pop("_score"), 4)
            result["rank"] = rank
        return {
            "query": query,
            "category": category,
            "total": sum(self._searchable(source) for source in sources),
            "count": len(results),
            "model": model_id or DEFAULT_INDEX_MODEL,
            "suggestion": self._suggesters[category].suggest(query.strip()),
            "low_confidence": (not results
                               or (floor is not None and results[0]["relevance"] < floor)),
            "results": results,
        }