src/backend/encode.py

"""Natural-language request to ranked CIPHER phenotypes.

Dense, two-surface retrieval over prebuilt Vidul BGE_FT_VA embeddings:

  - phenotype_metadata surface -> phenotype-level semantic similarity
  - code_evidence surface       -> code-level supporting evidence (max over chunks)

A phenotype is a document, not a bag of codes. Each candidate carries both channel
scores, a summary copied from source metadata, and the resolved code set. Ranking
uses the fine-tuned dense retriever.

Startup requires the matching embedding index and model.
"""

from __future__ import annotations

import json
import os
import re
from pathlib import Path

import numpy as np

from . import paths
from .retriever import DenseEmbedder
from .spelling import SpellSuggester

# The config beside the vectors records the query encoding convention.
EMB = paths.PARTA_EMB_DIR
# Optional detail sidecars.
CODE_DESC = paths.CODE_DESCRIPTIONS
ALGO_COMPONENTS = paths.ALGORITHM_COMPONENTS
FACETS = paths.PHENOTYPE_FACETS
BUNDLES = paths.RELATED_BUNDLES

# Algorithm-component fields to surface, in display order (professor 1.4:
# "refer to CIPHER algorithm components"). Only fields present for a phenotype
# are shown; absent fields are omitted.
_ALGO_FIELDS = [
    ("methods_used", "Methods"),
    ("related_diseases", "Related diseases"),
    ("data_period", "Data period"),
    ("validations", "Validation"),
    ("adjudication_method", "Adjudication"),
]


def _embedding_convention(config: dict) -> dict:
    """Return the build-time query contract; never invent missing defaults."""
    required = ("model", "pooling", "query_prefix", "max_seq_length")
    missing = [key for key in required if key not in config]
    if missing:
        raise RuntimeError(
            "Phenotype embedding config is missing its encoding convention: "
            + ", ".join(missing))
    return {key: config[key] for key in required}


def _method_label(value: str) -> str:
    return "Rule-based" if value.lower().startswith("rules-based") else value

# CIPHER public phenotype library. Detail page:
#   /web/cipher/phenotype-viewer?uqid={uqid}&name={slug}
# where slug = the full name with every non-alphanumeric char replaced by '_'.
# uqid comes from the cipher_links sidecar (present for ~79% of phenotypes);
# the rest fall back to the library root. Domain overridable for other deployments.
CIPHER_BASE = os.environ.get("ENCODE_CIPHER_BASE", "https://phenomics.va.ornl.gov").rstrip("/")
CIPHER_LINKS = paths.CIPHER_LINKS

# Friendly stand-ins so a still-unresolved code never shows the raw
# 'unsupported_code_system' token (professor feedback 1.1). (label, source)
_UNRESOLVED_NOTE = {
    466: ("Medication code", "needs_vocab:RxNorm/NDC"),
    467: ("LOINC lab code", "needs_vocab:LOINC"),
    471: ("OMOP concept ID", "needs_vocab:OMOP"),
    465: ("VA clinic stop code", "needs_vocab:VA_STOP"),
    468: ("Study-defined variable", "study_specific"),
    519: ("Study-defined variable", "study_specific"),
}


# Model-independent state an engine can borrow from another one: the phenotype
# records and every detail sidecar. Only the vectors and the query encoder
# belong to a model.
_SHARED_STATE = ("phenotypes", "_suggester", "code_desc", "cipher_uqid",
                 "algo_components", "facets", "related_bundles")


class EncodeEngine:
    def __init__(self, emb_dir: Path | str | None = None,
                 shared: "EncodeEngine | None" = None):
        """`shared` lends its phenotype records and sidecars, so a comparison
        model costs its vectors and encoder, not a second parse of the corpus."""
        emb = Path(emb_dir) if emb_dir else EMB
        self.emb_dir = emb
        if not (emb / "config.json").exists():
            raise RuntimeError(f"Missing runtime phenotype embeddings: {emb}")
        self.config = json.loads((emb / "config.json").read_text())
        convention = _embedding_convention(self.config)

        if shared is not None:
            for name in _SHARED_STATE:
                setattr(self, name, getattr(shared, name))
        else:
            self._load_runtime_state()
        self._load_vectors(emb)

        # Query embedder must match how the documents were encoded (model +
        # pooling + prefix + max-seq all travel in config.json).
        #
        # ENCODE_PARTA_MODEL repoints the query encoder for the default vector
        # set (the Space image uses it to load local weights instead of the
        # hub). It applies only to the default vectors because queries and
        # documents must use the same model.
        override = os.environ.get("ENCODE_PARTA_MODEL") if emb.resolve() == EMB.resolve() else None
        if override or "model_id" not in self.config or self.config["model_id"] == "bge_ft_va":
            model_name = override or convention["model"]
            self.model_name = model_name
            self.embedder = DenseEmbedder(
                model_name,
                pooling=convention["pooling"],
                query_prefix=convention["query_prefix"],
                max_seq_length=convention["max_seq_length"])
        else:
            # Built by scripts/build_model_vectors.py: weights resolve under the
            # outputs tree and an adapter or weight precision may apply.
            self.model_name = self.config["model_id"]
            self.embedder = DenseEmbedder.from_config(self.config)

    def _load_runtime_state(self) -> None:
        self.phenotypes = {}
        for line in paths.CANONICAL_PHENOTYPES.open():
            r = json.loads(line)
            self.phenotypes[r["phenotype_id"]] = r

        # "Did you mean" vocabulary from the phenotype text itself.
        self._suggester = SpellSuggester(
            text for p in self.phenotypes.values()
            for text in (p.get("title"), p.get("description"), *(p.get("keywords") or [])))

        # Part B: resolved code descriptions + algorithm components (optional sidecars).
        self.code_desc = {}
        if CODE_DESC.exists():
            for line in CODE_DESC.open():
                r = json.loads(line)
                self.code_desc[f"{r['code_type_id']}|{r['sub_type_id']}|{(r['code'] or '').upper()}"] = r
        self.cipher_uqid = {}
        if CIPHER_LINKS.exists():
            for line in CIPHER_LINKS.open():
                r = json.loads(line)
                self.cipher_uqid[r["phenotype_id"]] = r["uqid"]

        # Extracted CIPHER algorithm components (methods / related diseases /
        # validation / adjudication / data period); one record per phenotype.
        self.algo_components = {}
        if ALGO_COMPONENTS.exists():
            for line in ALGO_COMPONENTS.open():
                r = json.loads(line)
                self.algo_components[r["phenotype_id"]] = r

        # Offline facet extraction (evidence-or-abstain); enum values reach the
        # payload only when the source stated them, never inferred.
        self.facets = {}
        if FACETS.exists():
            for line in FACETS.open():
                r = json.loads(line)
                self.facets[r["phenotype_id"]] = r.get("facets") or {}

        # Related-concept bundles (search seeds, never codes).
        self.related_bundles = {}
        if BUNDLES.exists():
            for line in BUNDLES.open():
                r = json.loads(line)
                self.related_bundles[r["phenotype_id"]] = r

    def _load_vectors(self, emb: Path) -> None:
        # Metadata surface: one row per phenotype, aligned to meta_ids.
        self.meta_ids: list[int] = json.loads((emb / "metadata_ids.json").read_text())
        self.meta_mat = np.load(emb / "metadata.npy").astype(np.float32)
        self._row = {pid: i for i, pid in enumerate(self.meta_ids)}

        # Code surface: many rows per phenotype; map each to its phenotype row.
        code_meta = json.loads((emb / "code_meta.json").read_text())
        self.code_mat = np.load(emb / "code.npy").astype(np.float32)
        self._code_row = np.array([self._row.get(m["phenotype_id"], -1) for m in code_meta])

        # Per-phenotype filters aligned to meta_ids.
        self._cat = np.array([self.phenotypes[p].get("category") or "(none)" for p in self.meta_ids])
        self._validated = np.array([bool(self.phenotypes[p].get("validated")) for p in self.meta_ids])

    def categories(self) -> list[str]:
        return sorted(set(self._cat.tolist()))

    # -- retrieval ---------------------------------------------------------
    def search(self, query: str, k: int = 10, categories: set[str] | None = None,
               validated_only: bool = False) -> dict:
        q = self.embedder.encode([query])[0]
        meta_score = self.meta_mat @ q                     # (n_pheno,)

        # Max code-evidence similarity aggregated to phenotype level.
        code_sims = self.code_mat @ q                      # (n_code,)
        code_score = np.zeros(len(self.meta_ids), dtype=np.float32)
        valid = self._code_row >= 0
        np.maximum.at(code_score, self._code_row[valid], code_sims[valid])

        retrieval = 0.5 * meta_score + 0.5 * code_score    # bi-encoder candidate score

        mask = np.ones(len(self.meta_ids), dtype=bool)
        if categories:
            mask &= np.isin(self._cat, list(categories))
        if validated_only:
            mask &= self._validated

        candidates = np.flatnonzero(mask)
        results = []
        if len(candidates):
            candidate_scores = retrieval[candidates]
            for j in np.argsort(candidate_scores)[::-1][:k]:
                i = candidates[j]
                results.append(self._result(self.meta_ids[i], float(meta_score[i]),
                                            float(code_score[i]),
                                            float(candidate_scores[j])))
        return {
            "query": query,
            "model": self.model_name,
            "count": len(results),
            "suggestion": self._suggester.suggest(query.strip()),
            "results": results,
        }

    def _result(self, pid, m, c, relevance) -> dict:
        p = self.phenotypes[pid]
        evidence = [
            {"code_system": g.get("code_type_label"),
             "code_count": g.get("code_count") or len(g.get("codes", [])),
             "samples": [x.get("code") for x in g.get("codes", [])[:6]]}
            for g in p.get("associated_code_groups", [])
        ]
        return {
            "phenotype_id": pid,
            "title": p.get("title"),
            "category": p.get("category"),
            "validated": p.get("validated"),
            "summary": self._summary(p),
            "keywords": p.get("keywords") or [],
            "scores": {"relevance": round(relevance, 4),
                       "metadata": round(m, 4), "code_evidence": round(c, 4)},
            "code_systems": [e["code_system"] for e in evidence],
            "code_evidence": evidence,
            "facets": self._facet_payload(pid),
            "warnings": [],
        }

    @staticmethod
    def _summary(p: dict) -> str:
        """Return a shortened source description."""
        desc = (p.get("description") or p.get("algorithm_description") or "").strip()
        if len(desc) > 360:
            desc = desc[:360].rsplit(" ", 1)[0] + "…"
        return desc

    def _algo_component_rows(self, pid: int) -> list[dict]:
        """Format available algorithm components as ordered label/value rows."""
        rec = self.algo_components.get(pid)
        if not rec:
            return []
        rows = []
        for key, label in _ALGO_FIELDS:
            val = rec.get(key)
            if not val:
                continue
            if isinstance(val, list):
                values = [str(item) for item in val]
                if key == "methods_used":
                    values = [_method_label(item) for item in values]
                text = ", ".join(values)
            elif isinstance(val, dict):
                text = (f"{val.get('start', '?')} to {val.get('end', '?')}"
                        if ("start" in val or "end" in val)
                        else ", ".join(f"{k}: {v}" for k, v in val.items()))
            else:
                text = str(val)
            rows.append({"label": label, "value": text})
        return rows

    def _cipher_url(self, pid: int, title: str | None) -> str:
        uqid = self.cipher_uqid.get(pid)
        if uqid:
            slug = re.sub(r"[^A-Za-z0-9]", "_", title or "")
            return f"{CIPHER_BASE}/web/cipher/phenotype-viewer?uqid={uqid}&name={slug}"
        return CIPHER_BASE + "/"

    def _code_display(self, code_type_id, sub_type_id,
                      entry: dict) -> tuple[str | None, str, dict]:
        """Resolve a code to (description, source), replacing 'unsupported_code_system'
        when a description exists in the sidecar."""
        status = entry.get("label_status")
        labels = entry.get("labels") or []
        if status in ("exact", "prefix_expanded") and labels:
            return labels[0], status, {}
        # Both gaps are looked up: a code whose system we had no dictionary for,
        # and a code whose system we support but whose label our tables lacked.
        if status in ("unsupported_code_system", "label_missing"):
            hit = self.code_desc.get(
                f"{code_type_id}|{sub_type_id}|{(entry.get('code') or '').upper()}")
            if hit:
                # Return the recorded match method and release with the label.
                extra = {k: hit[k] for k in ("source_version", "match", "concept")
                         if hit.get(k)}
                return hit["description"], hit["source"], extra
            if status == "unsupported_code_system":
                note = _UNRESOLVED_NOTE.get(code_type_id,
                                            ("No standard description", "unresolved"))
                return note[0], note[1], {}
        return (labels[0] if labels else None), status, {}

    # -- detail ------------------------------------------------------------
    def _code_graphable(self, ct: int | None, code: str | None) -> bool:
        """Does this specific code have a knowledge-graph? ICD diagnosis (460/461)
        resolves to an ICD family; medication terms (466) use a local map or an
        exact/normalized RxNAV lookup."""
        if ct in (460, 461):
            return True
        if ct == 466:
            return bool((code or "").strip())
        return False

    def _code_group_detail(self, g: dict) -> dict:
        ct, sub = g.get("code_type_id"), g.get("sub_type_id")
        codes, resolved = [], 0
        for x in g.get("codes", [])[:200]:
            desc, source, extra = self._code_display(ct, sub, x)
            if desc and source not in ("study_specific", "unresolved", "label_missing") \
                    and not source.startswith("needs_vocab"):
                resolved += 1
            codes.append({"code": x.get("code"), "description": desc,
                          "description_source": source, "label_status": x.get("label_status"),
                          "graphable": self._code_graphable(ct, x.get("code")), **extra})
        return {"code_system": g.get("code_type_label"), "sub_type": g.get("sub_type_label"),
                "code_count": g.get("code_count") or len(g.get("codes", [])),
                "resolved_count": resolved, "codes": codes,
                "graphable": any(c["graphable"] for c in codes)}

    def _facet_payload(self, pid: int) -> dict | None:
        """Compact facet view for the UI: stated values only, criteria as text."""
        f = self.facets.get(pid)
        if not f:
            return None

        def stated(name):
            v = ((f.get(name) or {}).get("value") or "").strip()
            return v if v and v != "not_stated" else None

        out = {
            "age_group": stated("age_group"),
            "care_setting": stated("care_setting"),
            "incident_vs_prevalent": stated("incident_vs_prevalent"),
            "intended_use": stated("intended_use"),
            "inclusion": [c.get("criterion") for c in f.get("inclusion") or [] if c.get("criterion")],
            "exclusion": [c.get("criterion") for c in f.get("exclusion") or [] if c.get("criterion")],
        }
        return out if any(out.values()) else None

    def phenotype(self, pid: int) -> dict | None:
        p = self.phenotypes.get(pid)
        if not p:
            return None
        return {
            "phenotype_id": pid,
            "algorithm_id": p.get("algorithm_id"),
            "title": p.get("title"),
            "category": p.get("category"),
            "validated": p.get("validated"),
            "validation_description": p.get("validation_description"),
            "description": p.get("description"),
            "algorithm_description": p.get("algorithm_description"),
            "keywords": p.get("keywords") or [],
            "authors": p.get("authors") or [],
            "publications": p.get("publications") or [],
            "population_description": p.get("population_description"),
            "last_modified": p.get("last_modified"),
            "cipher_url": self._cipher_url(pid, p.get("title")),
            "algorithm_components": self._algo_component_rows(pid),
            "facets": self._facet_payload(pid),
            "llm_summary": None,
            "related_bundle": self.related_bundles.get(pid),
            "code_groups": [self._code_group_detail(g)
                            for g in p.get("associated_code_groups", [])],
        }

    # -- phenotype code hierarchy (phecode -> ICD main -> ICD sub) ----------
    def phenotype_code_graph(self, pid: int, focus: str | None = None,
                             cap: int | None = None) -> dict | None:
        """A node+edge tree of the phenotype's diagnosis codes: the phenotype (as
        the phecode-level concept) -> each ICD main category (3-char stem) -> the
        specific ICD sub-codes under it (460 ICD-9, 461 ICD-10). OMOP concept
        groups (471) contribute the ICD codes their concepts resolve to, drawn
        in derivation order: phenotype -> OMOP concept -> ICD code. A family
        node appears only for families holding directly carried codes; families
        still organize ordering, windowing, and overflow for both kinds.

        When a code is clicked (`focus`), only its main branch is expanded and
        sub-codes are windowed around it. Other branches collapse to a count.
        Overflow within a branch is shown as a single "+N more" stub (an edge to a
        count, not every node). `path` = the root→main→code chain to highlight."""
        p = self.phenotypes.get(pid)
        if not p:
            return None
        # cap=0 disables every window: all categories expand with all codes.
        unlimited = cap == 0
        BIG = 10 ** 9
        FOCUS_SUBS, MAIN_SUBS, MAX_MAINS = (BIG, BIG, BIG) if unlimited else (14, 6, 20)
        focus = (focus or "").strip() or None
        focus_main = focus.split(".")[0] if focus else None
        root = "pheno"
        title = p.get("title") or f"Phenotype {pid}"
        nodes = [{"id": root, "label": title[:60], "sub": "CIPHER phenotype (phecode)",
                  "tier": 0, "current": focus is None, "path": True}]
        edges, path = [], [root]
        note = None

        # Pool ICD codes per system across groups. OMOP concept groups (471)
        # contribute the ICD-9/10 codes their concepts resolve to (from the
        # code-description sidecar); each derived code carries its concept(s)
        # so the graph can show why the code is present.
        pools: dict[int, dict[str, dict]] = {460: {}, 461: {}}
        for g in p.get("associated_code_groups", []):
            ct = g.get("code_type_id")
            sub_type = g.get("sub_type_id")
            if ct in (460, 461):
                for x in g.get("codes", []):
                    code = (x.get("code") or "").strip()
                    if code:
                        e = pools[ct].setdefault(code, {"code": code, "x": None,
                                                        "sub_type": sub_type, "omop": []})
                        if e["x"] is None:
                            e["x"], e["sub_type"] = x, sub_type
            elif ct == 471:
                for x in g.get("codes", []):
                    hit = self.code_desc.get(
                        f"471|{sub_type}|{(x.get('code') or '').upper()}")
                    concept = (hit or {}).get("concept") or {}
                    tgt = {"ICD9CM": 460, "ICD10CM": 461}.get(concept.get("system"))
                    code = (concept.get("code") or "").strip()
                    if tgt and code:
                        e = pools[tgt].setdefault(code, {"code": code, "x": None,
                                                         "sub_type": None, "omop": []})
                        cid = concept.get("omop_concept_id") or ""
                        # The same concept can arrive twice (repeated across
                        # groups); one node id must be emitted once.
                        if all(o["id"] != cid for o in e["omop"]):
                            e["omop"].append({"id": cid,
                                              "name": hit.get("description") or ""})

        OMOP_PER_CODE = BIG if unlimited else 2
        # OMOP concepts sit between the family and the code: the phenotype
        # carries the concept, the concept resolves to the code. Codes carried
        # directly skip that column (the renderer bows their edge around it).
        has_omop = any(e["omop"] for entries in pools.values()
                       for e in entries.values())
        # Route order is derivation order: the phenotype (phecode level) carries
        # OMOP concepts, each concept resolves to one ICD code. With concepts
        # present the middle column is the concepts; the family rung is drawn
        # only for families holding directly carried codes (no phenotype in the
        # corpus mixes the two today). Families still organize ordering,
        # windowing, and overflow either way.
        code_tier = 2
        for ct in (460, 461):
            entries = pools[ct]
            if not entries:
                continue
            system = "ICD-9" if ct == 460 else "ICD-10"
            by_main: dict[str, list[dict]] = {}
            for e in (entries[c] for c in sorted(entries)):
                by_main.setdefault(e["code"].split(".")[0], []).append(e)
            # focus can be a leaf code (clicked in the drawer) or a main category
            # (clicked to expand it) — both expand that main; a leaf also highlights.
            group_has_focus = focus is not None and focus_main in by_main
            focus_is_leaf = group_has_focus and any(e["code"] == focus for e in by_main[focus_main])

            mains = sorted(by_main)
            sel_mains = mains[:MAX_MAINS]
            if group_has_focus and focus_main not in sel_mains:
                sel_mains = sel_mains[:MAX_MAINS - 1] + [focus_main]
            if len(mains) > len(sel_mains):
                note = "Some categories are truncated."

            for cat in sel_mains:
                # Namespaced, because a category and one of its codes can be the
                # same string: a phenotype whose code list holds the 3-character
                # stem "F32" produced a category node and a code node with the
                # id "ICD-10:F32", and the two drew on top of each other in the
                # same column. The category rung is its own kind of thing and
                # gets its own id space.
                main_id = f"cat:{system}:{cat}"
                subs = by_main[cat]
                is_focus_main = group_has_focus and cat == focus_main
                # When focused, only the focused branch expands; others collapse.
                collapse = focus is not None and not is_focus_main and not unlimited
                # A family node stands for its directly carried codes; a family
                # whose codes all arrive through OMOP concepts has none, so
                # nothing stands between the phenotype and the concepts. When a
                # branch collapses under focus it is drawn as a count either
                # way — a collapsed branch that drew nothing would make the
                # graph lie about what the phenotype holds.
                family_node = any(e["x"] is not None for e in subs)
                if family_node or collapse:
                    nodes.append({"id": main_id, "label": cat, "tier": 1, "current": False,
                                  "sub": (f"{system} · {len(subs)} code" + ("s" if len(subs) != 1 else "")
                                          if collapse else f"{system} category"),
                                  "path": is_focus_main, "nav": {"code": cat, "code_type": system}})
                    edges.append([root, main_id])
                if is_focus_main and family_node:
                    path.append(main_id)
                if collapse:
                    continue

                cap = FOCUS_SUBS if is_focus_main else MAIN_SUBS
                if is_focus_main and focus_is_leaf and len(subs) > cap:   # window around the focus
                    idx = next(i for i, e in enumerate(subs) if e["code"] == focus)
                    start = max(0, min(idx - cap // 2, len(subs) - cap))
                    sel = subs[start:start + cap]
                else:
                    sel = subs[:cap]
                for e in sel:
                    if e["x"] is not None:
                        desc, _, _extra = self._code_display(ct, e["sub_type"], e["x"])
                    else:
                        # Derived from an OMOP concept; the concept name is the
                        # description of the code it resolves to.
                        desc = e["omop"][0]["name"]
                    code = e["code"]
                    sid = f"{system}:{code}"
                    is_cur = code == focus
                    nodes.append({"id": sid, "label": code, "sub": desc or "",
                                  "tier": code_tier, "current": is_cur, "path": is_cur,
                                  "nav": {"code": code, "code_type": system}})
                    if e["x"] is not None:
                        edges.append([main_id, sid])   # carried directly
                    if is_cur:
                        path.append(sid)
                    for om in e["omop"][:OMOP_PER_CODE]:
                        # The route: the phenotype carries the concept, the
                        # concept resolves to the code.
                        oid = f"omop:{om['id'] or sid}"
                        nodes.append({"id": oid, "label": om["id"] or "OMOP",
                                      "sub": om["name"], "tier": 1, "current": False,
                                      "vocab": "OMOP"})
                        edges.append([root, oid])
                        edges.append([oid, sid])
                    if len(e["omop"]) > OMOP_PER_CODE:
                        om_more = f"more:omop:{sid}"
                        nodes.append({"id": om_more,
                                      "label": f"+{len(e['omop']) - OMOP_PER_CODE} more",
                                      "sub": "", "tier": 1, "more": True, "vocab": "OMOP"})
                        edges.append([om_more, sid])
                if len(subs) > len(sel):                       # overflow stub, not every node
                    more_id = f"more:{main_id}"
                    nodes.append({"id": more_id, "label": f"+{len(subs) - len(sel)} more",
                                  "sub": "", "tier": code_tier, "more": True})
                    edges.append([main_id if family_node else root, more_id])

        if len(nodes) == 1:
            return {"available": False, "kind": "phenotype", "code": str(pid),
                    "reason": "No ICD code graph for this phenotype."}
        return {"available": True, "kind": "phenotype", "code": str(pid), "title": title,
                "subtitle": ("Phecode → OMOP concepts → ICD codes"
                             if has_omop else
                             "Phecode → ICD main categories → ICD sub-codes"
                             + (" (click a category to expand it)" if focus else "")),
                "nodes": nodes, "edges": edges, "path": path, "note": note}