src/backend/codesearch.py

"""Vidul's fine-tuned dense retrieval for medical-code search.

The runtime uses the original BGE_FT_VA model and FAISS indexes with no score
blend or fallback model. Indexes are loaded lazily by category so a CPU Space
does not pay the memory cost for categories nobody has queried.
"""

from __future__ import annotations

import bisect
import csv
import hashlib
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: ")


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]] = {}
        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) -> 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"))
        digest = hashlib.sha256(meta_path.read_bytes()).hexdigest()
        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"))
            with (directory / "meta.csv").open(newline="", encoding="utf-8") as handle:
                records = list(csv.DictReader(handle))
            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))})
        self._loaded[category] = loaded
        self._suggesters[category] = SpellSuggester(
            record["description"] for source in loaded for record in source["records"])
        return loaded

    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) -> dict:
        sources = self._load_category(category)
        query_vector = self.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": "BGE_FT_VA",
            "suggestion": self._suggesters[category].suggest(query.strip()),
            "low_confidence": (not results
                               or results[0]["relevance"] < LOW_CONFIDENCE.get(category, 0.75)),
            "results": results,
        }