src/backend/graph.py

"""Knowledge graph: click a code -> its parent/child ontology as nodes + edges.

Every builder returns one uniform shape the frontend renders as a node-edge tree:

    {available, kind, title, subtitle,
     nodes: [{id, label, sub, tier, current}], edges: [[parentId, childId], ...],
     note}

Graph relations use CIPHER, VA CDW, RxNorm, and LOINC data:

  - Diagnosis (ICD-9/10) & ICD-9-Proc: dotted-code family (category stem -> the
    codes sharing it), descriptions from the VA dictionary that backs code search.
    ICD diagnosis graphs also carry the code's phecode assignments (v1.2 and
    phecodeX) as annotation.
  - Procedures (ICD-10-PCS): 7-char positional codes -> the 3-char "table" family.
  - Procedures (CPT/HCPCS): CMS RBCS group (family, else subcategory) -> the
    index codes sharing it.
  - Medication: RxNorm ingredient(s) -> the drug -> sibling products.
  - Lab (LOINC): COMPONENT -> active LOINC terms sharing it (loinc_terms).
    VA-local LabChemTestSIDs enter that graph through the tiered lab mapping.

Mapping lookups go through the shared MappingStore, which owns the usability
rules. Routing resolves the clicked code to a CodeSystem once, then dispatches;
the resolution order is load-bearing (PCS and ICD-9-Proc types also contain
"ICD"/"CPT" substrings) and is pinned by tests. Codes without supported
relations return ``available: false``.
"""

from __future__ import annotations

import json
from collections import defaultdict
from enum import Enum
from pathlib import Path

from . import paths
from . import provenance
from .mappings import MappingStore, get_store
from .rxnav import RxNavClient

MEMBER_CAP = 26  # sibling codes shown per family (windowed around the clicked code; rest -> "+N more")
NDC_NODE_CAP = 4  # NDC package codes drawn on a medication graph; rest -> "+N more"


class CodeSystem(Enum):
    PCS = "pcs"
    ICD9_PROC = "icd9_proc"
    CPT = "cpt"
    ICD = "icd"
    LOINC = "loinc"
    VA_LAB = "va_lab"
    MEDICATION = "medication"


def _dotted_cat(code: str) -> str:
    """Category stem for a dotted code (E11.9 -> E11, 345.01 -> 345)."""
    return code.split(".")[0].strip()


def _pcs_cat(code: str) -> str:
    """ICD-10-PCS 'table' = first 3 positional chars (0SGK44Z -> 0SG)."""
    return code[:3]


def _load_jsonl(path: Path):
    if path.exists():
        with path.open(encoding="utf-8") as handle:
            for line in handle:
                if line.strip():
                    yield json.loads(line)


class KnowledgeGraph:
    def __init__(self, diagnosis_records=None, procedure_records=None,
                 rxnav_client: RxNavClient | None = None, lab_loinc_records=None,
                 loinc_records=None, store: MappingStore | None = None):
        if store is None:
            store = (MappingStore(lab_loinc_records=lab_loinc_records)
                     if lab_loinc_records is not None else get_store())
        self.store = store
        self.rxnav = rxnav_client or RxNavClient()

        # -- lab (LOINC): term table + analyte(component) -> codes -------
        self.loinc: dict[str, dict] = {}
        self.loinc_by_comp: dict[str, list[str]] = defaultdict(list)
        for r in (loinc_records if loinc_records is not None else _load_jsonl(paths.LOINC_TERMS)):
            self.loinc[r["loinc"]] = r
            self.loinc_by_comp[r["component"].lower()].append(r["loinc"])

        # -- ICD diagnosis (dotted) --------------------------------------
        self.icd, self.icd_cat = self._index_codes(
            diagnosis_records,
            keep="ICD", cat_fn=_dotted_cat)

        # -- procedures: ICD-10-PCS (3-char) + ICD-9-Proc (dotted) -------
        self.pcs, self.pcs_cat = {}, defaultdict(list)
        self.icd9p, self.icd9p_cat = {}, defaultdict(list)
        # -- procedures: CPT/HCPCS grouped by CMS RBCS -------------------
        self.cpt, self.cpt_group = {}, defaultdict(list)
        for r in procedure_records:
            ct = str(r.get("code_type", ""))
            desc = r.get("description", "")
            for code in str(r.get("code", "")).split(","):
                code = code.strip()
                if not code:
                    continue
                if "PCS" in ct and code not in self.pcs:
                    self.pcs[code] = {"code_type": "ICD-10-PCS", "description": desc}
                    self.pcs_cat[_pcs_cat(code)].append(code)
                elif "ICD-9" in ct and "." in code and code not in self.icd9p:
                    self.icd9p[code] = {"code_type": "ICD-9-Proc", "description": desc}
                    self.icd9p_cat[_dotted_cat(code)].append(code)
                elif "CPT" in ct and code not in self.cpt:
                    group = self._rbcs_group(code)
                    if group:
                        self.cpt[code] = {"code_type": "CPT", "description": desc,
                                          "group": group}
                        self.cpt_group[group].append(code)

        self._routes = {
            CodeSystem.PCS: self._route_pcs,
            CodeSystem.ICD9_PROC: self._route_icd9_proc,
            CodeSystem.CPT: self._route_cpt,
            CodeSystem.ICD: self._route_icd,
            CodeSystem.LOINC: self._route_loinc,
            CodeSystem.VA_LAB: self._route_va_lab,
            CodeSystem.MEDICATION: self._route_medication,
        }

    def _rbcs_group(self, code: str) -> str | None:
        """Grouping label for a CPT/HCPCS code: RBCS family, else subcategory."""
        entry = self.store.rbcs_validated_for(code.upper())
        if not entry:
            return None
        return entry.get("family") or entry.get("subcategory") or None

    @staticmethod
    def _index_codes(records, keep: str, cat_fn):
        table: dict[str, dict] = {}
        buckets: dict[str, list[str]] = defaultdict(list)
        for r in records:
            ct = str(r.get("code_type", ""))
            if keep not in ct:
                continue
            desc = r.get("description", "")
            for code in str(r.get("code", "")).split(","):  # split merged rows
                code = code.strip()
                if code and code not in table:
                    table[code] = {"code_type": ct, "description": desc}
                    buckets[cat_fn(code)].append(code)
        return table, buckets

    # -- routing -----------------------------------------------------------
    def code_system(self, code: str, code_type: str | None = None) -> CodeSystem:
        """Which family a clicked code belongs to. Order is load-bearing:
        PCS and ICD-9-Proc must resolve before the bare "CPT"/"ICD" substring
        tests, and an embedding-matched CPT code wins regardless of the
        declared type (it has no other graph). MEDICATION is also the terminal
        bucket for unknowns; its handler ends with the no-hierarchy answer."""
        ct = (code_type or "").upper()
        first = code.split(",")[0].strip()
        if "PCS" in ct or (not ct and code in self.pcs):
            return CodeSystem.PCS
        if ("ICD" in ct and "PROC" in ct) or ("ICD-9" in ct and code in self.icd9p):
            return CodeSystem.ICD9_PROC
        if "CPT" in ct or (not ct and first.upper() in self.cpt) \
                or self.store.has_cpt_embedding(first.upper()):
            return CodeSystem.CPT
        if "ICD" in ct or (not ct and (code in self.icd or _dotted_cat(code) in self.icd_cat)):
            return CodeSystem.ICD
        if "LOINC" in ct or first in self.loinc:
            return CodeSystem.LOINC
        if "LABCHEM" in ct or "VA LAB" in ct or (
                not ct and self.store.has_lab_sid(code.split(","))):
            return CodeSystem.VA_LAB
        return CodeSystem.MEDICATION

    def neighbors(self, code: str, code_type: str | None = None,
                  drug_name: str | None = None, cap: int | None = None) -> dict:
        """cap: sibling window size per family; None = MEMBER_CAP, 0 = no cap."""
        code = (code or "").strip()
        if not code:
            return self._none(code, "No code given.")
        cap = MEMBER_CAP if cap is None else max(0, cap)
        return self._routes[self.code_system(code, code_type)](code, code_type, drug_name, cap)

    def _route_pcs(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        return self._family(code.split(",")[0].strip(), self.pcs, self.pcs_cat, _pcs_cat,
                            "procedure", "ICD-10-PCS procedure table — codes sharing this 3-character root.", cap)

    def _route_icd9_proc(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        return self._family(code.split(",")[0].strip(), self.icd9p, self.icd9p_cat, _dotted_cat,
                            "procedure", "ICD-9 procedure family — codes in this category.", cap)

    def _route_cpt(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        return self._rbcs_family(code.split(",")[0].strip().upper(), cap)

    def _route_icd(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        graph = self._family(code.split(",")[0].strip(), self.icd, self.icd_cat, _dotted_cat,
                             "icd", "Diagnosis code family — the codes in this ICD category.", cap)
        self._attach_phecodes(graph, code.split(",")[0].strip(), (code_type or "").upper())
        return graph

    def _route_loinc(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        return self._loinc(code.split(",")[0].strip(), cap)

    def _route_va_lab(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        return self._local_lab(code, drug_name, cap)

    def _route_medication(self, code: str, code_type: str | None, drug_name: str | None, cap: int) -> dict:
        # The NDC and VA drug indexes pack every package/local code of a product
        # into one comma-joined field (the longest is 52,506 characters), so the
        # first entry is the one to resolve — the rest are the same product.
        first = code.split(",")[0].strip()
        mapping = self.store.ingredients_for(first)
        if mapping:
            return self._med(first, mapping, cap=cap)
        if self.store.ingredients_by_name(code) or self.store.ingredients_by_rxcui(code):
            return self._med(code, cap=cap)
        if self.store.ingredients_by_name(first) or self.store.ingredients_by_rxcui(first):
            return self._med(first, cap=cap)
        if self.rxnav.supports(code, code_type):
            mapping = self.rxnav.resolve(code, code_type, drug_name)
            if mapping:
                # Display the code that actually resolved, never the whole
                # packed field, and record what it is so the drug node names
                # its own vocabulary.
                up = (code_type or "").upper()
                kind = ("ndc" if mapping.get("matched_ndc")
                        else "rxcui" if "RXNORM" in up else None)
                return self._med(mapping.get("matched_ndc") or first,
                                 {**mapping, "kind": kind}, cap=cap)
            return self._none(
                code,
                f"No unique RxNorm match found for {drug_name or code}.",
                "medication",
            )
        return self._none(code, "No hierarchy is available for this code.")

    # -- dotted / positional code family (nodes + edges) -------------------
    def _family(self, code, table, buckets, cat_fn, kind, subtitle, cap: int = MEMBER_CAP) -> dict:
        cat = cat_fn(code)
        members = sorted(buckets.get(cat, []))
        if not members:
            return self._none(code, f"No {kind} family found for {code}.", kind)
        sel = self._window(members, code, cap)      # siblings around the clicked code
        root_desc = (table.get(cat) or {}).get("description")
        nodes = [{"id": cat, "label": cat, "tier": 0, "current": cat == code, "path": True,
                  "sub": root_desc or f"{table.get(code, {}).get('code_type', 'ICD')} category {cat}"}]
        edges, path = [], [cat]
        for c in sel:
            if c == cat:
                continue
            is_cur = c == code
            nodes.append({"id": c, "label": c, "sub": table[c]["description"], "tier": 1,
                          "current": is_cur, "path": is_cur,
                          "nav": {"code": c, "code_type": table[c]["code_type"]}})
            edges.append([cat, c])
            if is_cur:
                path.append(c)
        if len(members) > len(sel):
            nodes.append({"id": f"more:{cat}", "label": f"+{len(members) - len(sel)} more",
                          "sub": "", "tier": 1, "more": True})
            edges.append([cat, f"more:{cat}"])
        return {"available": True, "kind": kind, "code": code, "title": code, "subtitle": subtitle,
                "nodes": nodes, "edges": edges, "path": path, "note": None}

    @staticmethod
    def _window(items: list[str], focus: str, cap: int) -> list[str]:
        """Up to `cap` items centred on `focus`; cap 0 = no windowing."""
        if cap <= 0 or len(items) <= cap:
            return items
        if focus in items:
            idx = items.index(focus)
            start = max(0, min(idx - cap // 2, len(items) - cap))
            return items[start:start + cap]
        return items[:cap]

    def _attach_phecodes(self, graph: dict, code: str, code_type: str) -> None:
        """Annotate an ICD family graph with the code's phecode assignments."""
        if not graph.get("available"):
            return
        versions = [v.strip().upper() for v in code_type.split("|") if v.strip()]
        entry = self.store.phecodes_for(code, versions or ["ICD10", "ICD9"])
        if entry:
            graph["phecodes"] = {**entry, "provenance": provenance.phecode_line(entry)}

    # -- CPT/HCPCS: CMS RBCS group -> index codes sharing it ---------------
    def _rbcs_family(self, code: str, cap: int = MEMBER_CAP) -> dict:
        entry = self.cpt.get(code)
        derived_entry = None if entry else self.store.cpt_embedding_for(code)
        if entry:
            rbcs = self.store.rbcs_validated_for(code) or {}
            group = entry["group"]
            sub = " · ".join(x for x in ["RBCS " + rbcs.get("category", ""),
                                         rbcs.get("subcategory", "")] if x.strip())
            mapping = rbcs
        elif derived_entry:
            group = derived_entry["group"]
            sub = ("RBCS " + derived_entry.get("category", "")).strip()
            mapping = derived_entry
        else:
            return self._none(code, f"No RBCS group is available for {code}.", "procedure")
        members = sorted(self.cpt_group.get(group, []) or [code])
        sel = self._window(members, code, cap)
        group_id = "rbcs:" + group
        nodes = [{"id": group_id, "label": group, "tier": 0, "current": False,
                  "path": True, "sub": sub}]
        edges, path = [], [group_id]
        if derived_entry and code not in sel:
            # The clicked code inherited this group by embedding match; it is
            # not a validated member, so it joins the display explicitly.
            sel = [code] + (sel[:max(0, cap - 1)] if cap > 0 else sel)
        for c in sel:
            is_cur = c == code
            desc = (self.cpt.get(c) or derived_entry or {}).get("description", "")
            nodes.append({"id": c, "label": c, "sub": desc,
                          "tier": 1, "current": is_cur, "path": is_cur,
                          "nav": {"code": c, "code_type": "CPT"}})
            edges.append([group_id, c])
            if is_cur:
                path.append(c)
        if len(members) > len(sel):
            nodes.append({"id": "more:" + group_id, "label": f"+{len(members) - len(sel)} more",
                          "sub": "", "tier": 1, "more": True})
            edges.append([group_id, "more:" + group_id])
        # Say how much of the family is on screen. The window is centred on the
        # clicked code, so two codes in one family legitimately show different
        # neighbours; without the count that looks like the group is unstable.
        shown = (f"Showing {len(sel)} of {len(members)}. "
                 if len(members) > len(sel) else "")
        return {"available": True, "kind": "procedure", "code": code, "title": code,
                "subtitle": f"CPT/HCPCS codes sharing the RBCS group: {group}. "
                            f"{shown}Codes in one family share this name.",
                "nodes": nodes, "edges": edges, "path": path,
                "derived": bool(derived_entry),
                "match_tier": "embedding_match" if derived_entry else None,
                "note": "Grouped by the CMS Restructured BETOS Classification System.",
                "source": provenance.line(mapping.get("source", "cms_rbcs_ry2025"),
                                          derived=bool(derived_entry),
                                          match=mapping.get("match"))}

    def _local_lab(self, code: str, display_name: str | None = None, cap: int = MEMBER_CAP) -> dict:
        """Resolve VA LabChemTestSID mappings before drawing the LOINC graph."""
        entries = self.store.loinc_for(code.split(","))
        if not entries:
            return self._none(
                code,
                "No LOINC mapping found for this VA lab code.",
                "lab",
            )
        targets = {entry["loinc"] for entry in entries}
        if len(targets) != 1:
            return self._none(
                code,
                "These VA lab codes map to multiple LOINC terms.",
                "lab",
            )
        entry = entries[0]
        loinc = entry["loinc"]
        graph = self._loinc(loinc, cap)
        if not graph.get("available"):
            return graph
        local_label = display_name or entry.get("lab_name") or code
        graph.update({
            "title": f"{local_label} → LOINC {loinc}",
            "derived": bool(entry.get("derived")),
            "match_tier": entry.get("match"),
            "source": provenance.line(entry.get("source"),
                                      derived=bool(entry.get("derived")),
                                      match=entry.get("match"),
                                      target=entry.get("loinc_version", "LOINC")),
            "mapped_from": {
                "system": "VA Lab List (LabChemTestSID)",
                "code": code,
                "name": local_label,
            },
            "mapped_loinc": loinc,
            "note": (
                f"VA Lab List {code} → LOINC {loinc}. "
                "Related terms share its LOINC COMPONENT."
            ),
        })
        return graph

    # -- LOINC lab: COMPONENT -> active terms sharing that component -------
    def _loinc(self, code, cap: int = MEMBER_CAP) -> dict:
        term = self.loinc.get(code)
        if not term:
            return self._none(code, f"{code} is not in the LOINC table.", "lab")
        comp = term["component"]
        members = sorted(self.loinc_by_comp.get(comp.lower(), [code]))
        sel = self._window(members, code, cap)
        comp_id = "comp:" + comp
        nodes = [{"id": comp_id, "label": comp, "tier": 0, "current": False, "path": True,
                  "sub": ("LOINC component · " + term.get("class", "")).strip(" ·")}]
        edges, path = [], [comp_id]
        for c in sel:
            t = self.loinc.get(c, {})
            sub = " · ".join(x for x in [t.get("system"), t.get("name")] if x)
            is_cur = c == code
            nodes.append({"id": c, "label": c, "sub": sub, "tier": 1, "current": is_cur,
                          "path": is_cur, "nav": {"code": c, "code_type": "LOINC"}})
            edges.append([comp_id, c])
            if is_cur:
                path.append(c)
        if len(members) > len(sel):
            nodes.append({"id": "more:" + comp_id, "label": f"+{len(members) - len(sel)} more",
                          "sub": "", "tier": 1, "more": True})
            edges.append([comp_id, "more:" + comp_id])
        return {"available": True, "kind": "lab", "code": code, "title": code,
                "subtitle": f"Active LOINC terms sharing COMPONENT: {comp}.",
                "nodes": nodes, "edges": edges, "path": path,
                "note": "Grouped by LOINC COMPONENT.",
                "source": provenance.line("loinc_2_82")}

    # -- medication: ingredient(s) -> drug -> sibling products ------------
    def _med(self, code, mapping=None, cap: int | None = None) -> dict:
        e = (mapping or self.store.ingredients_by_name(code)
             or self.store.ingredients_by_rxcui(code))
        if not e:
            return self._none(code, f"No RxNorm ingredient mapping for {code}.", "medication")
        label = e.get("name") or code
        ings = e.get("ingredients", [])
        if not ings:
            return self._none(code, f"No RxNorm ingredient is mapped for {code}.", "medication")
        nodes, edges, seen = [], [], set()
        drug_id = "drug:" + str(code)
        path = [drug_id]
        for ing in ings:
            iid = "ing:" + str(ing.get("rxcui") or ing["name"])
            # The RXCUI is the portable half of this graph and the only part a
            # reader can take elsewhere, so it travels as its own field rather
            # than staying encoded in the node id. Absent when the ingredient
            # was matched by name and carries no RXCUI of its own.
            nodes.append({"id": iid, "label": ing["name"], "sub": "RxNorm ingredient",
                          "code": ing.get("rxcui"), "code_label": "RXCUI",
                          "tier": 0, "current": False, "path": True})
            edges.append([iid, drug_id])
            path.append(iid)
        # The drug node names the vocabulary of the code it displays. NDC and
        # RXCUI keyed mappings were labelled VA Drug List ID before, which
        # asserted a vocabulary the code is not in. A name-keyed mapping has
        # no code to show, so the chip is left off entirely.
        kind = e.get("kind")
        code_label = {"ndc": "NDC", "rxcui": "RXCUI",
                      "drug_name": None}.get(kind, "VA Drug List ID")
        nodes.append({"id": drug_id, "label": label, "sub": "this medication",
                      "code": str(code) if code_label else None,
                      "code_label": code_label,
                      "tier": 1, "current": True, "path": True})
        n_prod = 0
        for ing in ings:
            iid = "ing:" + str(ing.get("rxcui") or ing["name"])
            for prod in (e.get("child_drugs", {}) or {}).get(ing["name"], []):
                if prod.upper() == label.upper() or prod in seen:
                    continue
                seen.add(prod)
                pid = "prod:" + prod
                nodes.append({"id": pid, "label": prod, "sub": "", "tier": 2, "current": False})
                edges.append([iid, pid])
                n_prod += 1
        # NDC package codes of this product. The mapping's own packages win;
        # a mapping that carries none falls back to the NDCs RxNorm lists for
        # the matched concept. The displayed code is excluded by digit
        # comparison, since keys are dashed and stored packages are 11 digit
        # strings. cap follows the family convention: 0 draws them all.
        digits = lambda s: "".join(ch for ch in str(s) if ch.isdigit())
        shown_digits = digits(code)
        ndcs = [n for n in (e.get("ndcs") or []) if digits(n) != shown_digits]
        if not ndcs:
            rxcuis = e.get("matched_rxcuis") or []
            if rxcuis:
                ndcs = [n for n in self.rxnav.ndcs_for_rxcui(str(rxcuis[0]))
                        if digits(n) != shown_digits]
        # One NDC hub in the middle tier, package codes hanging off it: bare
        # 11-digit strings interleaved with product names read as noise, and
        # the hub also gives the codes a labelled parent the way ingredients
        # have. No `code` on the hub, so Select all never collects it.
        limit = len(ndcs) if cap == 0 else NDC_NODE_CAP
        if ndcs:
            n_pkg = len(ndcs)
            nodes.append({"id": "ndcgrp", "label": "NDC",
                          "sub": f"{n_pkg} package code{'' if n_pkg == 1 else 's'}",
                          "vocab": "NDC", "tier": 1, "current": False})
            edges.append([drug_id, "ndcgrp"])
        for n in ndcs[:limit]:
            nid = f"ndc:{n}"
            nodes.append({"id": nid, "label": n, "sub": "NDC package code",
                          "code": n, "code_label": "NDC",
                          "vocab": "NDC", "tier": 2, "current": False})
            edges.append(["ndcgrp", nid])
        if len(ndcs) > limit:
            nodes.append({"id": "more:ndc", "label": f"+{len(ndcs) - limit} more",
                          "sub": "", "tier": 2, "more": True})
            edges.append(["ndcgrp", "more:ndc"])
        source = provenance.line(e.get("source"), derived=bool(e.get("derived")),
                                 match=e.get("match"))
        total_products = e.get("related_product_count", n_prod)
        shown_products = e.get("related_products_shown", n_prod)
        note = f"{total_products} related products."
        if total_products > shown_products:
            note += f" Showing the first {shown_products}."
        return {"available": True, "kind": "medication", "code": code,
                "title": label, "subtitle": "RxNorm ingredient(s) → this drug → related products and NDC packages.",
                "nodes": nodes, "edges": edges, "path": path,
                "derived": bool(e.get("derived")),
                "match_tier": e.get("match") if e.get("derived") else None,
                "note": (note if n_prod else None), "source": source,
                "mapped_rxcuis": e.get("matched_rxcuis", [])}

    @staticmethod
    def _none(code, reason, kind=None) -> dict:
        return {"available": False, "code": code, "kind": kind, "reason": reason}