hpo_loader.py

"""
hpo_loader.py — Phase 1
=======================
Singleton loader for all HPO cleaned data structures.
Import `hpo` from this module in any other file.

Usage:
    from hpo_loader import hpo
    hpo_id  = hpo.resolve_symptom("epilepsy")        # -> "HP:0001250"
    name    = hpo.symptom_name("HP:0001250")          # -> "Seizure"
    disease = hpo.resolve_disease("brugada syndrome") # -> "ORPHA:130"
    profile = hpo.disease_profile("ORPHA:130")        # -> {name, present[], absent[]}
    askable = hpo.askable_symptoms("ORPHA:130", confirmed_hpo_ids)
    score   = hpo.hybrid_score(patient_set, disease_id)
"""

import json
from pathlib import Path

import numpy as np

_DATA_DIR = Path(__file__).parent / "mondo_data_new"
_EMB_DIR  = Path(__file__).parent / "hpo_data_final"

# Shared SentenceTransformer model (lazy-loaded)
_semantic_model: object | None = None
_semantic_model_loaded: bool = False

_EXTRA_ALIASES: dict[str, str] = {
    "muscle aches":       "HP:0003326",
    "body aches":         "HP:0003326",
    "feeling tired":      "HP:0012378",
    "low energy":         "HP:0012378",
    "high fever":         "HP:0001945",
    "low grade fever":    "HP:0001945",
    "temperature":        "HP:0001945",
    "loose stools":       "HP:0002014",
    "blurry vision":      "HP:0000622",
    "painful urination":  "HP:0100518",
    "hard stools":        "HP:0002019",
    "pins and needles":   "HP:0003401",
    "blacked out":        "HP:0007185",
    "lost consciousness": "HP:0007185",
    "disoriented":        "HP:0001289",
    "forgetful":          "HP:0002354",
    "can't focus":        "HP:0031987",
    "speech problems":    "HP:0002167",
    "balance problems":   "HP:0002141",
    "unsteady":           "HP:0002141",
    "dizzy":              "HP:0002321",
    "shaking":            "HP:0001337",
    "stiffness":          "HP:0003552",
    "loss of taste":      "HP:0041051",
    "fits":               "HP:0001250",
}


HPO_NON_PHENOTYPE_IDS: frozenset[str] = frozenset({
    "HP:0040279", "HP:0040280", "HP:0040281", "HP:0040282",
    "HP:0040283", "HP:0040284", "HP:0040285",
    "HP:0031796", "HP:0031797",
    "HP:0025153",
    "HP:0012823", "HP:0011009", "HP:0011010", "HP:0003674",
    "HP:0003577", "HP:0003593", "HP:0003621",
    "HP:0003584", "HP:0003587",
    "HP:0003676",
    "HP:0003680",
    "HP:0011011",
    "HP:0012824",
    "HP:0012825",
    "HP:0012826",
    "HP:0012828",
    "HP:0012831",
    "HP:0012832",
    "HP:0012833",
    "HP:0012837",
    "HP:0012838",
    "HP:0030650",
    "HP:0020034",
    "HP:0000005", "HP:0000006", "HP:0000007",
    "HP:0001417", "HP:0001419", "HP:0001423", "HP:0001427",
    "HP:0001428", "HP:0001450",
    "HP:0032316", "HP:0003743", "HP:0025297",
})


def _load_json(fname: str):
    with open(_DATA_DIR / fname, "r", encoding="utf-8") as f:
        return json.load(f)


_GENERIC_DISEASE_TERMS: set[str] = {
    "disease", "disorder", "disorders", "diseases", "condition",
    "medical condition", "disease or disorder", "disease or disorder, non-neoplastic",
    "diseases and disorders", "other disease",
    "neoplasm", "neoplasia", "neoplasm (disease)", "tumor", "tumour",
    "tumor disease", "tumour disease", "neoplastic disease", "neoplastic growth",
    "cell process disease", "other neoplasm",
    "syndrome", "infection", "carcinoma", "cancer", "anomaly", "abnormality",
    "malformation", "deformity", "failure", "deficiency", "insufficiency",
    "sarcoma", "adenoma", "injury", "wound", "disease, non-human animal",
    "common", "rare", "complete",
}


class HPOData:
    def __init__(self):
        print("[HPOLoader] Loading HPO data...", flush=True)

        # -- Symptom catalog --
        sc = _load_json("symptom_catalog.json")
        self._id_to_symptom: dict[str, str] = {
            hid: entry["label"] for hid, entry in sc.items()
        }
        ssm = _load_json("symptom_synonym_map.json")
        self._synonym_to_id: dict[str, str] = {}
        for syn, vals in ssm.items():
            if isinstance(vals, list) and len(vals) > 0:
                self._synonym_to_id[syn] = vals[0]
            else:
                self._synonym_to_id[syn] = vals

        for alias, hpo_id in _EXTRA_ALIASES.items():
            key = alias.lower()
            if key not in self._synonym_to_id:
                self._synonym_to_id[key] = hpo_id

        self._name_to_hpo_id = {
            name.lower(): hid
            for hid, name in self._id_to_symptom.items()
        }

        # -- Disease catalog (unified) --
        dc = _load_json("disease_catalog.json")
        self._id_to_label: dict[str, str] = {
            did: entry["label"] for did, entry in dc.items()
        }

        # -- Disease profiles (unified) --
        raw_profiles = _load_json("disease_symptom_profiles.json")
        self._profile_for_id: dict[str, dict] = {}
        for did, prof in raw_profiles.items():
            present = [
                {"hpo_id": hid, "symptom_name": self._id_to_symptom.get(hid, hid)}
                for hid in prof.get("present", [])
            ]
            absent = [
                {"hpo_id": hid, "symptom_name": self._id_to_symptom.get(hid, hid)}
                for hid in prof.get("absent", [])
            ]
            self._profile_for_id[did] = {
                "name":    prof.get("label", self._id_to_label.get(did, did)),
                "present": present,
                "absent":  absent,
            }

        # -- Symptom -> diseases reverse index (for specificity in askable_symptoms) --
        self._symptom_to_diseases: dict[str, dict] = {}
        for did, prof in self._profile_for_id.items():
            for entry in prof["present"]:
                hid = entry["hpo_id"]
                if hid not in self._symptom_to_diseases:
                    self._symptom_to_diseases[hid] = {"present_in": []}
                self._symptom_to_diseases[hid]["present_in"].append(did)

        # -- Unified disease name -> ID index (loaded, not built from catalog) --
        std = _load_json("synonym_to_disease.json")
        self._name_to_id: dict[str, str] = {}
        self._id_to_name: dict[str, str] = {}
        for syn, ids in std.items():
            if isinstance(ids, list) and len(ids) > 0:
                chosen = ids[0]
                if self._profile_for_id:
                    for cid in ids:
                        if cid in self._profile_for_id:
                            chosen = cid
                            break
                self._name_to_id[syn] = chosen
            else:
                self._name_to_id[syn] = ids if isinstance(ids, str) else ids[0]
            # Track first synonym as display-name fallback for each ID
            did = self._name_to_id[syn]
            if did not in self._id_to_name:
                self._id_to_name[did] = syn

        # -- PrimeKG disease index --
        kg_path = _DATA_DIR / "kg_disease_to_id.json"
        self._kg_disease_to_id: dict[str, str] = _load_json("kg_disease_to_id.json") if kg_path.exists() else {}

        self._kg_id_to_name: dict[str, str] = {}
        for kname, kxid in self._kg_disease_to_id.items():
            if kxid not in self._kg_id_to_name:
                self._kg_id_to_name[kxid] = kname

        # -- Token index for efficient token overlap matching --
        self._token_to_diseases: dict[str, list[tuple[str, str]]] = {}
        for dname, did in self._name_to_id.items():
            for token in set(dname.split()):
                self._token_to_diseases.setdefault(token, []).append((dname, did))

        token_count = len(self._token_to_diseases)
        print(
            f"[HPOLoader] Ready — "
            f"{len(self._id_to_symptom):,} symptoms, "
            f"{len(self._synonym_to_id):,} surface forms, "
            f"{len(self._name_to_id):,} disease synonyms, "
            f"{len(self._profile_for_id):,} profiles, "
            f"{token_count:,} token keys"
        )

        # Eagerly load MedEmbed model while HF Hub client is alive
        self._ensure_semantic_model()

    # ------------------------------------------------------------------
    # Semantic matching (MedEmbed fallback)
    # ------------------------------------------------------------------

    def _ensure_semantic_model(self):
        global _semantic_model, _semantic_model_loaded
        if _semantic_model_loaded:
            return
        if not (_EMB_DIR / "symptom_embeddings.npy").exists():
            _semantic_model_loaded = True
            return
        try:
            from sentence_transformers import SentenceTransformer
            _semantic_model = SentenceTransformer("abhinand/MedEmbed-large-v0.1", device="cpu", local_files_only=True)
        except Exception:
            pass
        _semantic_model_loaded = True

    def _load_embedding_set(self, stem: str) -> tuple[np.ndarray | None, list | None]:
        npy  = _EMB_DIR / f"{stem}_embeddings.npy"
        idx  = _EMB_DIR / f"{stem}_embedding_index.json"
        if npy.exists() and idx.exists():
            return np.load(npy), json.load(open(idx))
        return None, None

    def _semantic_match(self, query: str, embeddings: np.ndarray, labels: list, threshold: float = 0.7) -> str | None:
        self._ensure_semantic_model()
        if _semantic_model is None:
            return None
        q = _semantic_model.encode([query], normalize_embeddings=True)
        sims = np.dot(embeddings, q.T).flatten()
        best = int(np.argmax(sims))
        if sims[best] >= threshold:
            return labels[best][0]
        return None

    def _semantic_match_symptom(self, query: str) -> str | None:
        emb, labels = self._load_embedding_set("symptom")
        if emb is None:
            return None
        return self._semantic_match(query, emb, labels)

    def _semantic_match_disease(self, query: str) -> str | None:
        emb, labels = self._load_embedding_set("disease")
        if emb is None:
            return None
        return self._semantic_match(query, emb, labels)

    def _semantic_match_kg(self, query: str) -> str | None:
        emb, labels = self._load_embedding_set("kg_disease")
        if emb is None:
            return None
        return self._semantic_match(query, emb, labels)

    # ------------------------------------------------------------------
    # Symptom lookups
    # ------------------------------------------------------------------

    def resolve_symptom(self, surface_form: str) -> str | None:
        key = surface_form.lower().strip()
        exact = self._synonym_to_id.get(key)
        if exact:
            return exact
        return self._semantic_match_symptom(key)

    def resolve_symptom_with_method(self, surface_form: str) -> tuple[str | None, str | None]:
        key = surface_form.lower().strip()
        exact = self._synonym_to_id.get(key)
        if exact:
            return exact, "exact"
        semantic = self._semantic_match_symptom(key)
        if semantic:
            return semantic, "semantic"
        return None, None

    def resolve_symptom_from_kg(self, kg_name: str) -> str | None:
        return self._name_to_hpo_id.get(kg_name.lower().strip())

    def symptom_name(self, hpo_id: str) -> str:
        return self._id_to_symptom.get(hpo_id, hpo_id)

    def symptoms_to_ids(self, surface_forms: list[str]) -> set[str]:
        result = set()
        for sf in surface_forms:
            hid = self.resolve_symptom(sf)
            if hid:
                result.add(hid)
        return result

    def kg_symptoms_to_ids(self, kg_names: list[str]) -> set[str]:
        result = set()
        for name in kg_names:
            hid = self.resolve_symptom_from_kg(name)
            if hid:
                result.add(hid)
        return result

    # ------------------------------------------------------------------
    # Disease resolution
    # ------------------------------------------------------------------

    def _has_profile(self, did: str) -> bool:
        return did in self._profile_for_id

    def _collect_forward_substring_candidates(self, key: str) -> list[tuple[str, str]]:
        seen: set[tuple[str, str]] = set()
        candidates: list[tuple[str, str]] = []
        for dname, did in self._name_to_id.items():
            if key in dname:
                item = (did, dname)
                if item not in seen:
                    seen.add(item)
                    candidates.append((dname, did))
        return candidates

    def _collect_reverse_substring_candidates(self, key: str) -> list[tuple[str, str]]:
        seen: set[tuple[str, str]] = set()
        candidates: list[tuple[str, str]] = []
        key_words = set(key.split())

        for dname, did in self._name_to_id.items():
            if len(dname) >= 4 and dname in key and not _is_generic(dname) and _is_whole_word_match(dname, key_words):
                item = (did, dname)
                if item not in seen:
                    seen.add(item)
                    candidates.append((dname, did))
        return candidates

    def _collect_token_candidates(self, key_tokens: set[str]) -> list[tuple[str, str]]:
        if not key_tokens:
            return []
        counts: dict[tuple[str, str], int] = {}
        for token in key_tokens:
            for dname, did in self._token_to_diseases.get(token, []):
                tup = (dname, did)
                counts[tup] = counts.get(tup, 0) + 1
        return [item for item, c in counts.items() if c >= 2]

    @staticmethod
    def _pick_best(candidates: list[tuple[str, str]]) -> str | None:
        if not candidates:
            return None
        candidates.sort(key=lambda x: len(x[0]), reverse=True)
        return candidates[0][1]

    def resolve_disease(self, name: str) -> str | None:
        """
        Disease name string -> disease ID (MONDO, OMIM, or ORPHA).

        Resolution hierarchy:
          1. Exact match
          2. Semantic match (MedEmbed, >= 0.7 threshold)
          3. Forward substring (key in DB name) with profile
          4. Reverse substring (DB name in key) with profile
          5. Token overlap (>=2 shared tokens) with profile
          6. Forward substring without profile
          7. Reverse substring without profile
          8. Token overlap without profile
          Tiebreaker within each layer: longest disease name wins.
        """
        key = name.lower().strip()
        key_tokens = set(key.split())

        # 1. Exact match
        exact = self._name_to_id.get(key)
        if exact:
            return exact

        # 2. Semantic match
        semantic = self._semantic_match_disease(key)
        if semantic:
            return semantic

        # 3-8: Substring + token overlap
        sub_forward = self._collect_forward_substring_candidates(key)
        sub_reverse = self._collect_reverse_substring_candidates(key)
        token = self._collect_token_candidates(key_tokens)

        for layer in [sub_forward, sub_reverse, token]:
            with_profile = [(d, n) for d, n in layer if self._has_profile(d)]
            result = self._pick_best(with_profile)
            if result:
                return result

        for layer in [sub_forward, sub_reverse, token]:
            result = self._pick_best(layer)
            if result:
                return result

        return None

    def resolve_disease_with_method(self, name: str) -> tuple[str | None, str | None]:
        key = name.lower().strip()
        key_tokens = set(key.split())

        exact = self._name_to_id.get(key)
        if exact:
            return exact, "exact"

        semantic = self._semantic_match_disease(key)
        if semantic:
            return semantic, "semantic"

        sub_forward = self._collect_forward_substring_candidates(key)
        sub_reverse = self._collect_reverse_substring_candidates(key)
        token = self._collect_token_candidates(key_tokens)

        for layer, method in [(sub_forward, "substring"), (sub_reverse, "substring"), (token, "token")]:
            with_profile = [(d, n) for d, n in layer if self._has_profile(d)]
            result = self._pick_best(with_profile)
            if result:
                return result, method

        for layer, method in [(sub_forward, "substring"), (sub_reverse, "substring"), (token, "token")]:
            result = self._pick_best(layer)
            if result:
                return result, method

        return None, None

    def resolve_disease_kg(self, name: str) -> str | None:
        if not self._kg_disease_to_id:
            return None
        key = name.lower().strip()
        # 1. Exact match
        if key in self._kg_disease_to_id:
            return self._kg_disease_to_id[key]
        # 2. Semantic match (against all scorable KG x_ids)
        semantic = self._semantic_match_kg(key)
        if semantic:
            return semantic
        # 3-4. Substring fallbacks
        candidates = []
        for kg_name, xid in self._kg_disease_to_id.items():
            if key in kg_name:
                candidates.append((kg_name, xid))
        if not candidates:
            for kg_name, xid in self._kg_disease_to_id.items():
                if len(kg_name) >= 4 and kg_name in key:
                    candidates.append((kg_name, xid))
        if candidates:
            candidates.sort(key=lambda x: len(x[0]))
            return candidates[0][1]
        return None

    def kg_disease_name(self, xid: str) -> str | None:
        return self._kg_id_to_name.get(xid)

    def resolve_disease_all(self, name: str, top_n: int = 5) -> list[tuple[str, str]]:
        key = name.lower().strip()
        results = []
        if key in self._name_to_id:
            results.append((key, self._name_to_id[key]))
        for dname, did in self._name_to_id.items():
            if dname != key and key in dname:
                results.append((dname, did))
            if len(results) >= top_n:
                break
        return results

    def disease_name(self, disease_id: str) -> str:
        name = self._id_to_label.get(disease_id)
        if name:
            return name
        prof = self._profile_for_id.get(disease_id)
        if prof:
            return prof["name"]
        return self._id_to_name.get(disease_id, disease_id)

    def disease_profile(self, disease_id: str) -> dict | None:
        return self._profile_for_id.get(disease_id)

    def disease_present_ids(self, disease_id: str) -> set[str]:
        profile = self.disease_profile(disease_id)
        if not profile:
            return set()
        return {e["hpo_id"] for e in profile["present"]}

    def disease_absent_ids(self, disease_id: str) -> set[str]:
        profile = self.disease_profile(disease_id)
        if not profile:
            return set()
        return {e["hpo_id"] for e in profile["absent"]}

    # ------------------------------------------------------------------
    # Askable symptoms
    # ------------------------------------------------------------------

    def askable_symptoms(
        self,
        disease_id: str,
        confirmed_hpo_ids: set[str],
        top_n: int = 8,
    ) -> list[dict]:
        profile = self._profile_for_id.get(disease_id)
        if not profile:
            return []
        results = []
        for entry in profile["present"]:
            hid = entry["hpo_id"]
            if hid not in confirmed_hpo_ids:
                sd = self._symptom_to_diseases.get(hid, {})
                specificity = len(sd.get("present_in", []))
                results.append({
                    "hpo_id":       hid,
                    "symptom_name": entry["symptom_name"],
                    "reason":       "confirm_present",
                    "specificity":  specificity,
                })
        for entry in profile["absent"]:
            hid = entry["hpo_id"]
            if hid not in confirmed_hpo_ids:
                results.append({
                    "hpo_id":       hid,
                    "symptom_name": entry["symptom_name"],
                    "reason":       "confirm_absent",
                    "specificity":  0,
                })
        results.sort(key=lambda x: x["specificity"])
        return results[:top_n]

    # ------------------------------------------------------------------
    # Scoring helpers
    # ------------------------------------------------------------------

    def hybrid_score(
        self,
        patient_hpo_ids: set[str],
        disease_id: str,
    ) -> dict:
        B = self.disease_present_ids(disease_id)
        A = patient_hpo_ids
        if not A or not B:
            return {
                "score": 0.0, "inclusion": 0.0, "jaccard": 0.0,
                "intersection_ids": [], "intersection_names": [],
                "a_size": len(A), "b_size": len(B),
            }
        intersection = A & B
        union        = A | B
        inclusion    = len(intersection) / len(A)
        jaccard      = len(intersection) / len(union)
        score        = round(0.5 * inclusion + 0.5 * jaccard, 4)
        return {
            "score":              score,
            "inclusion":          round(inclusion, 4),
            "jaccard":            round(jaccard, 4),
            "intersection_ids":   sorted(intersection),
            "intersection_names": [self.symptom_name(h) for h in sorted(intersection)],
            "a_size":             len(A),
            "b_size":             len(B),
        }


def _is_generic(dname: str) -> bool:
    return dname in _GENERIC_DISEASE_TERMS


def _is_whole_word_match(name: str, key_words: set[str]) -> bool:
    if len(name) >= 8 or " " in name:
        return True
    return name in key_words


hpo = HPOData()