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
SUMMARIES = paths.PHENOTYPE_SUMMARIES
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"),
}


class EncodeEngine:
    def __init__(self, embedder: DenseEmbedder | None = None, emb_dir: Path | str | None = None):
        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)

        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 [])))

        # 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])

        # 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 {}

        # Optional LLM sidecars: grounded plain-language summaries and
        # related-concept bundles (search seeds, never codes).
        self.llm_summaries = {}
        if SUMMARIES.exists():
            for line in SUMMARIES.open():
                r = json.loads(line)
                self.llm_summaries[r["phenotype_id"]] = r
        self.related_bundles = {}
        if BUNDLES.exists():
            for line in BUNDLES.open():
                r = json.loads(line)
                self.related_bundles[r["phenotype_id"]] = r

        # 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
        model_name = override or convention["model"]
        self.model_name = model_name
        self.embedder = embedder or DenseEmbedder(
            model_name,
            pooling=convention["pooling"],
            query_prefix=convention["query_prefix"],
            max_seq_length=convention["max_seq_length"])

    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": self.llm_summaries.get(pid),
            "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).

        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

        for g in p.get("associated_code_groups", []):
            ct = g.get("code_type_id")
            if ct not in (460, 461):
                continue
            system = "ICD-9" if ct == 460 else "ICD-10"
            sub_type = g.get("sub_type_id")
            by_main: dict[str, list[dict]] = {}
            for x in g.get("codes", []):
                code = (x.get("code") or "").strip()
                if code:
                    by_main.setdefault(code.split(".")[0], []).append(x)
            # 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(x.get("code") == focus for x 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
                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:
                    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, x in enumerate(subs) if x.get("code") == focus)
                    start = max(0, min(idx - cap // 2, len(subs) - cap))
                    sel = subs[start:start + cap]
                else:
                    sel = subs[:cap]
                for x in sel:
                    desc, _, _extra = self._code_display(ct, sub_type, x)
                    code = x["code"]
                    sid = f"{system}:{code}"
                    is_cur = code == focus
                    nodes.append({"id": sid, "label": code, "sub": desc or "", "tier": 2,
                                  "current": is_cur, "path": is_cur,
                                  "nav": {"code": code, "code_type": system}})
                    edges.append([main_id, sid])
                    if is_cur:
                        path.append(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": 2, "more": True})
                    edges.append([main_id, more_id])

        if len(nodes) == 1:
            return {"available": False, "kind": "phenotype", "code": str(pid),
                    "reason": "This phenotype has no ICD diagnosis codes to chart."}
        return {"available": True, "kind": "phenotype", "code": str(pid), "title": title,
                "subtitle": "Phecode → ICD main categories → ICD sub-codes"
                            + (" (click a category to expand it)" if focus else ""),
                "nodes": nodes, "edges": edges, "path": path, "note": note}