symptom_extractor.py

"""
symptom_extractor_3.py — Phase 2 (With Native Feather-Based PrimeKG Layer)
========================================================================
Extracts clinical features using MedSpacy and resolves within-line
co-occurrences using an on-disk PrimeKG feather dataset.
"""

import json
import re
import logging
from pathlib import Path

import medspacy
from medspacy.ner import TargetRule
import pandas as pd
from hpo_loader import HPO_NON_PHENOTYPE_IDS, HPOData

hpo = HPOData()

HERE = Path(__file__).parent
HPO_DIR = HERE / "hpo_data_final"
SYMPTOM_CACHE_PATH = HERE / "symptom_cache.json"

def _load_cache() -> set[str]:
    if SYMPTOM_CACHE_PATH.exists():
        with open(SYMPTOM_CACHE_PATH) as f:
            return set(json.load(f))
    return set()

def cache_symptom(surface_form: str):
    cache = _load_cache()
    cache.add(surface_form.lower())
    with open(SYMPTOM_CACHE_PATH, "w") as f:
        json.dump(list(cache), f, indent=2)

def extract_symptoms_legacy(text: str):
    text_lower = text.lower()
    present_seen = {}
    absent_seen = {}
    return list(present_seen.values()), list(absent_seen.values())


class ClinicalExtractionPipeline:
    def __init__(self):
        # 1. Load MedSpacy components
        self.nlp = medspacy.load(enable=["ner", "context"])

        # 2. Load JSON registries from the final data directory
        with open(HPO_DIR / "synonym_to_symptom_id.json", "r") as f:
            self.synonym_to_hpo = json.load(f)
        with open(HPO_DIR / "disease_name_to_id.json", "r") as f:
            self.disease_to_id = json.load(f)

        # 3. Dynamic Feather Loader: Scan root or current paths for kg.feather
        self.primekg_df = None
        for path in [Path("kg.feather"), HERE / "kg.feather", HERE.parent / "kg.feather"]:
            if path.exists():
                try:
                    print(f"[PrimeKG Engine: Loading binary dataset from {path.resolve()}...]")
                    self.primekg_df = pd.read_feather(path)
                    print(f"[PrimeKG Engine: Successfully loaded {len(self.primekg_df)} knowledge edges]")
                    break
                except Exception as e:
                    print(f"[Warning: Failed loading feather data matrix at {path}: {e}]")

        # 4. Seed Target Matcher dictionary
        target_rules = []
        for symptom_phrase in self.synonym_to_hpo.keys():
            target_rules.append(TargetRule(literal=symptom_phrase, category="SYMPTOM"))
        for disease_phrase in self.disease_to_id.keys():
            target_rules.append(TargetRule(literal=disease_phrase, category="DISEASE"))

        self.nlp.get_pipe("medspacy_target_matcher").add(target_rules)

    def _resolve_id(self, text: str, category: str) -> str:
        clean_text = text.lower().strip()
        if category == "SYMPTOM":
            hid = self.synonym_to_hpo.get(clean_text, None)
            if hid is None:
                hid = hpo.resolve_symptom(clean_text)
            if hid and hid in HPO_NON_PHENOTYPE_IDS:
                return None
            return hid
        elif category == "DISEASE":
            hid = self.disease_to_id.get(clean_text, None)
            if hid is None:
                hid = hpo.resolve_disease(clean_text)
            return hid
        return None

    def _query_primekg_relations(self, source_id: str, target_id: str) -> str:
        """Queries the pandas DataFrame loaded from kg.feather for specific relationships."""
        if self.primekg_df is not None:
            # Query cross-matches for undirected edge validation
            condition = (
                ((self.primekg_df['x_id'] == source_id) & (self.primekg_df['y_id'] == target_id)) |
                ((self.primekg_df['x_id'] == target_id) & (self.primekg_df['y_id'] == source_id))
            )
            matches = self.primekg_df[condition]
            if not matches.empty:
                # Return the verified display relation or structural identifier string
                row = matches.iloc[0]
                return row.get('display_relation', row.get('relation', 'associated_with'))

        # Hardcoded structural heuristic if no cross-record hits occur
        modifier_hpos = {"HP:0031796", "HP:0040283", "HP:0024143", "HP:0031914"}
        if source_id in modifier_hpos and target_id not in modifier_hpos:
            return "phenotype_modifier_of"
        if target_id in modifier_hpos and source_id not in modifier_hpos:
            return "has_phenotype_modifier"

        return "associated_with"

    def extract_clinical_payload(self, llm_markdown: str) -> dict:
        payload = {
            "presenting_symptoms": {"present": [], "absent": []},
            "past_medical_history": {"present": [], "absent": []},
            "family_history": [],
            "primekg_resolved_edges": []
        }

        current_section = None
        seen_present = set()
        seen_absent = set()
        for line in llm_markdown.splitlines():
            line = line.strip()
            if not line:
                continue

            if "### Presenting Symptoms" in line:
                current_section = "symptoms"
                continue
            elif "### Past Medical History" in line:
                current_section = "pmh"
                continue
            elif "### Family History" in line:
                current_section = "fhx"
                continue

            if current_section in ["symptoms", "pmh"]:
                is_negated = line.lower().startswith("denies:")
                clean_phrase = re.sub(r'^(presence of:|denies:)\s*', '', line, flags=re.IGNORECASE).strip()

                doc = self.nlp(clean_phrase)
                line_entities = []

                for ent in doc.ents:
                    final_negation = is_negated or ent._.is_negated
                    concept_id = self._resolve_id(ent.text, ent.label_)

                    if concept_id:
                        bucket = "absent" if final_negation else "present"
                        seen = seen_absent if bucket == "absent" else seen_present
                        if concept_id in seen:
                            continue
                        seen.add(concept_id)
                        target = "presenting_symptoms" if current_section == "symptoms" else "past_medical_history"
                        payload[target][bucket].append({"id": concept_id, "name": ent.text})

                        if not final_negation:
                            line_entities.append(concept_id)

                # Inter-line PrimeKG Edge Validation Pass across extracted concepts
                if len(line_entities) > 1:
                    for i in range(len(line_entities)):
                        for j in range(i + 1, len(line_entities)):
                            id1, id2 = line_entities[i], line_entities[j]
                            relation_type = self._query_primekg_relations(id1, id2)
                            payload["primekg_resolved_edges"].append({
                                "source": id1,
                                "target": id2,
                                "relation": relation_type
                            })

            elif current_section == "fhx" and "->" in line:
                relative, condition_block = line.split("->", 1)
                relative = relative.strip()
                is_negated = "denies:" in condition_block.lower()
                clean_phrase = re.sub(r'^(presence of:|denies:)\s*', '', condition_block.strip(), flags=re.IGNORECASE).strip()

                doc = self.nlp(clean_phrase)
                for ent in doc.ents:
                    final_negation = is_negated or ent._.is_negated
                    concept_id = self._resolve_id(ent.text, "DISEASE")
                    if concept_id:
                        payload["family_history"].append({
                            "relative": relative,
                            "id": concept_id,
                            "name": ent.text,
                            "status": "absent" if final_negation else "present"
                        })

        return payload

advanced_extractor = ClinicalExtractionPipeline()