medical_coding.py

"""
medical_coding.py
-----------------
UMLS / ICD-10 / SNOMED medical concept extraction and coding.

Extracts medical entities from radiology report text and maps them to:
  - UMLS CUIs  (e.g. C1306645 → Pneumonia)
  - ICD-10 codes (e.g. J18.9 → Pneumonia, unspecified)
  - SNOMED CT codes
  - Severity indicators
  - Anatomical locations

Works WITHOUT scispacy (uses built-in dictionary for reliability).
Optionally uses scispacy if installed for deeper NLP extraction.
"""

import logging
import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# UMLS + ICD-10 Medical Concept Dictionary
# Covers the most common radiology findings
# ---------------------------------------------------------------------------

MEDICAL_CONCEPT_DB: Dict[str, dict] = {

    # ── Pulmonary ──────────────────────────────────────────────────────────
    "pneumonia": {
        "cui": "C0032285", "icd10": "J18.9", "snomed": "233604007",
        "label": "Pneumonia", "category": "Pulmonary",
        "severity_weight": 0.7
    },
    "consolidation": {
        "cui": "C0521530", "icd10": "J98.19", "snomed": "44416002",
        "label": "Pulmonary Consolidation", "category": "Pulmonary",
        "severity_weight": 0.6
    },
    "pneumothorax": {
        "cui": "C0032326", "icd10": "J93.9", "snomed": "36118008",
        "label": "Pneumothorax", "category": "Pulmonary",
        "severity_weight": 0.9
    },
    "pleural effusion": {
        "cui": "C0032227", "icd10": "J90", "snomed": "60046008",
        "label": "Pleural Effusion", "category": "Pulmonary",
        "severity_weight": 0.7
    },
    "atelectasis": {
        "cui": "C0004144", "icd10": "J98.11", "snomed": "46621007",
        "label": "Atelectasis", "category": "Pulmonary",
        "severity_weight": 0.6
    },
    "pulmonary edema": {
        "cui": "C0034063", "icd10": "J81.0", "snomed": "19242006",
        "label": "Pulmonary Edema", "category": "Pulmonary",
        "severity_weight": 0.8
    },
    "emphysema": {
        "cui": "C0034067", "icd10": "J43.9", "snomed": "87433001",
        "label": "Pulmonary Emphysema", "category": "Pulmonary",
        "severity_weight": 0.6
    },
    "opacity": {
        "cui": "C0333053", "icd10": "R91.8", "snomed": "408721000",
        "label": "Pulmonary Opacity", "category": "Pulmonary",
        "severity_weight": 0.5
    },
    "nodule": {
        "cui": "C0034079", "icd10": "R91.1", "snomed": "27925004",
        "label": "Pulmonary Nodule", "category": "Pulmonary",
        "severity_weight": 0.6
    },
    "mass": {
        "cui": "C0577559", "icd10": "R91.8", "snomed": "300848003",
        "label": "Pulmonary Mass", "category": "Pulmonary",
        "severity_weight": 0.8
    },
    "fibrosis": {
        "cui": "C1800706", "icd10": "J84.10", "snomed": "40122008",
        "label": "Pulmonary Fibrosis", "category": "Pulmonary",
        "severity_weight": 0.7
    },

    # ── Cardiac ───────────────────────────────────────────────────────────
    "cardiomegaly": {
        "cui": "C0018800", "icd10": "I51.7", "snomed": "8186001",
        "label": "Cardiomegaly", "category": "Cardiac",
        "severity_weight": 0.7
    },
    "pericardial effusion": {
        "cui": "C0031039", "icd10": "I30.9", "snomed": "373945007",
        "label": "Pericardial Effusion", "category": "Cardiac",
        "severity_weight": 0.8
    },

    # ── Neurological ──────────────────────────────────────────────────────
    "hemorrhage": {
        "cui": "C0019080", "icd10": "I62.9", "snomed": "50960005",
        "label": "Intracranial Hemorrhage", "category": "Neurological",
        "severity_weight": 0.95
    },
    "infarct": {
        "cui": "C0007785", "icd10": "I63.9", "snomed": "432504007",
        "label": "Cerebral Infarction", "category": "Neurological",
        "severity_weight": 0.9
    },
    "hydrocephalus": {
        "cui": "C0020255", "icd10": "G91.9", "snomed": "230745008",
        "label": "Hydrocephalus", "category": "Neurological",
        "severity_weight": 0.85
    },
    "midline shift": {
        "cui": "C0235050", "icd10": "G93.5", "snomed": "44695005",
        "label": "Midline Shift", "category": "Neurological",
        "severity_weight": 0.9
    },
    "edema": {
        "cui": "C0006960", "icd10": "G93.6", "snomed": "2776000",
        "label": "Cerebral Edema", "category": "Neurological",
        "severity_weight": 0.85
    },

    # ── Abdominal ─────────────────────────────────────────────────────────
    "hepatomegaly": {
        "cui": "C0019209", "icd10": "R16.0", "snomed": "80515008",
        "label": "Hepatomegaly", "category": "Abdominal",
        "severity_weight": 0.6
    },
    "splenomegaly": {
        "cui": "C0038002", "icd10": "R16.1", "snomed": "16294009",
        "label": "Splenomegaly", "category": "Abdominal",
        "severity_weight": 0.6
    },
    "ascites": {
        "cui": "C0003962", "icd10": "R18.0", "snomed": "389026000",
        "label": "Ascites", "category": "Abdominal",
        "severity_weight": 0.7
    },
    "obstruction": {
        "cui": "C0029895", "icd10": "K56.60", "snomed": "81060008",
        "label": "Bowel Obstruction", "category": "Abdominal",
        "severity_weight": 0.8
    },
    "calculi": {
        "cui": "C0022650", "icd10": "N20.0", "snomed": "95570007",
        "label": "Urolithiasis/Calculi", "category": "Abdominal",
        "severity_weight": 0.5
    },
    "gallstone": {
        "cui": "C0008350", "icd10": "K80.20", "snomed": "266474003",
        "label": "Cholelithiasis", "category": "Abdominal",
        "severity_weight": 0.5
    },

    # ── Musculoskeletal ───────────────────────────────────────────────────
    "fracture": {
        "cui": "C0016658", "icd10": "M84.40", "snomed": "125605004",
        "label": "Bone Fracture", "category": "Musculoskeletal",
        "severity_weight": 0.75
    },
    "osteoporosis": {
        "cui": "C0029456", "icd10": "M81.0", "snomed": "64859006",
        "label": "Osteoporosis", "category": "Musculoskeletal",
        "severity_weight": 0.6
    },
    "osteopenia": {
        "cui": "C0029453", "icd10": "M85.80", "snomed": "312894000",
        "label": "Osteopenia", "category": "Musculoskeletal",
        "severity_weight": 0.4
    },
    "dislocation": {
        "cui": "C0012691", "icd10": "M24.30", "snomed": "57637002",
        "label": "Joint Dislocation", "category": "Musculoskeletal",
        "severity_weight": 0.7
    },
    "arthritis": {
        "cui": "C0003864", "icd10": "M19.90", "snomed": "57676002",
        "label": "Arthritis", "category": "Musculoskeletal",
        "severity_weight": 0.5
    },
    "arthroplasty": {
        "cui": "C0003893", "icd10": "Z96.641", "snomed": "179294005",
        "label": "Hip Arthroplasty/Implant", "category": "Musculoskeletal",
        "severity_weight": 0.2
    },
    "implant": {
        "cui": "C0021102", "icd10": "Z96.9", "snomed": "40388003",
        "label": "Orthopedic Implant", "category": "Musculoskeletal",
        "severity_weight": 0.2
    },
    "lytic": {
        "cui": "C0262587", "icd10": "M89.50", "snomed": "203522001",
        "label": "Lytic Bone Lesion", "category": "Musculoskeletal",
        "severity_weight": 0.8
    },

    # ── Normal findings ───────────────────────────────────────────────────
    "no acute": {
        "cui": "C0442737", "icd10": "Z00.00", "snomed": "17621005",
        "label": "No Acute Finding", "category": "Normal",
        "severity_weight": 0.0
    },
    "normal": {
        "cui": "C0442737", "icd10": "Z00.00", "snomed": "17621005",
        "label": "Normal Study", "category": "Normal",
        "severity_weight": 0.0
    },
    "unremarkable": {
        "cui": "C0442737", "icd10": "Z00.00", "snomed": "17621005",
        "label": "Unremarkable", "category": "Normal",
        "severity_weight": 0.0
    },
# ── Breast / Mammography ──────────────────────────────────────────────
    "microcalcification": {
        "cui": "C0024496", "icd10": "N64.89", "snomed": "129769009",
        "label": "Breast Microcalcification", "category": "Breast", "severity_weight": 0.75
    },
    "breast mass": {
        "cui": "C0024103", "icd10": "N63.0", "snomed": "59614007",
        "label": "Breast Mass", "category": "Breast", "severity_weight": 0.8
    },
    "ductal": {
        "cui": "C0007103", "icd10": "C50.919", "snomed": "413448000",
        "label": "Ductal Carcinoma", "category": "Breast", "severity_weight": 0.95
    },

    # ── Thyroid / Neck ────────────────────────────────────────────────────
    "goiter": {
        "cui": "C0018021", "icd10": "E04.9", "snomed": "3716002",
        "label": "Goiter", "category": "Thyroid", "severity_weight": 0.5
    },
    "thyroid nodule": {
        "cui": "C0040136", "icd10": "E04.1", "snomed": "237495005",
        "label": "Thyroid Nodule", "category": "Thyroid", "severity_weight": 0.6
    },
    "lymphadenopathy": {
        "cui": "C0497156", "icd10": "R59.9", "snomed": "111590001",
        "label": "Lymphadenopathy", "category": "Lymphatic", "severity_weight": 0.65
    },

    # ── Vascular ──────────────────────────────────────────────────────────
    "aneurysm": {
        "cui": "C0002940", "icd10": "I72.9", "snomed": "432119003",
        "label": "Aneurysm", "category": "Vascular", "severity_weight": 0.9
    },
    "stenosis": {
        "cui": "C1261287", "icd10": "I77.1", "snomed": "415068001",
        "label": "Arterial Stenosis", "category": "Vascular", "severity_weight": 0.8
    },
    "thrombosis": {
        "cui": "C0087086", "icd10": "I82.90", "snomed": "396347007",
        "label": "Thrombosis", "category": "Vascular", "severity_weight": 0.85
    },
    "embolism": {
        "cui": "C0014079", "icd10": "I74.9", "snomed": "59282003",
        "label": "Pulmonary Embolism", "category": "Vascular", "severity_weight": 0.9
    },

    # ── Liver / Hepatobiliary ─────────────────────────────────────────────
    "cirrhosis": {
        "cui": "C0023895", "icd10": "K74.60", "snomed": "19943007",
        "label": "Liver Cirrhosis", "category": "Hepatobiliary", "severity_weight": 0.8
    },
    "steatosis": {
        "cui": "C0015695", "icd10": "K76.0", "snomed": "197315008",
        "label": "Hepatic Steatosis", "category": "Hepatobiliary", "severity_weight": 0.55
    },
    "cholecystitis": {
        "cui": "C0008325", "icd10": "K81.9", "snomed": "76581006",
        "label": "Cholecystitis", "category": "Hepatobiliary", "severity_weight": 0.7
    },

    # ── Renal ─────────────────────────────────────────────────────────────
    "hydronephrosis": {
        "cui": "C0020295", "icd10": "N13.30", "snomed": "232895005",
        "label": "Hydronephrosis", "category": "Renal", "severity_weight": 0.7
    },
    "renal cyst": {
        "cui": "C0085413", "icd10": "N28.1", "snomed": "79411009",
        "label": "Renal Cyst", "category": "Renal", "severity_weight": 0.3
    },

    # ── Spinal ────────────────────────────────────────────────────────────
    "disc herniation": {
        "cui": "C0242362", "icd10": "M51.16", "snomed": "73589001",
        "label": "Disc Herniation", "category": "Spinal", "severity_weight": 0.65
    },
    "spondylosis": {
        "cui": "C0038614", "icd10": "M47.816", "snomed": "41114007",
        "label": "Spondylosis", "category": "Spinal", "severity_weight": 0.5
    },
    "spinal stenosis": {
        "cui": "C0158288", "icd10": "M48.06", "snomed": "18946005",
        "label": "Spinal Stenosis", "category": "Spinal", "severity_weight": 0.65
    },
    "compression fracture": {
        "cui": "C0423697", "icd10": "M48.50", "snomed": "263057000",
        "label": "Vertebral Compression Fracture", "category": "Spinal", "severity_weight": 0.75
    },

    # ── Paediatric / Obstetric ────────────────────────────────────────────
    "placenta previa": {
        "cui": "C0032046", "icd10": "O44.00", "snomed": "60758002",
        "label": "Placenta Previa", "category": "Obstetric", "severity_weight": 0.85
    },
    "oligohydramnios": {
        "cui": "C0030289", "icd10": "O41.00", "snomed": "59566000",
        "label": "Oligohydramnios", "category": "Obstetric", "severity_weight": 0.8
    },
}

# Severity level mapping
SEVERITY_LEVELS = {
    (0.0, 0.0): "Normal",
    (0.01, 0.39): "Minimal",
    (0.4, 0.59): "Mild",
    (0.6, 0.74): "Moderate",
    (0.75, 0.89): "Significant",
    (0.9, 1.0): "Critical",
}


# ---------------------------------------------------------------------------
# Data Models
# ---------------------------------------------------------------------------

@dataclass
class MedicalCode:
    term: str
    cui: str
    icd10: str
    snomed: str
    label: str
    category: str
    severity_weight: float
    context: str = ""  # surrounding sentence

    def to_dict(self) -> dict:
        return {
            "term_found": self.term,
            "cui": self.cui,
            "icd10": self.icd10,
            "snomed": self.snomed,
            "label": self.label,
            "category": self.category,
            "severity": self._severity_label(),
            "context": self.context,
        }

    def _severity_label(self) -> str:
        for (low, high), label in SEVERITY_LEVELS.items():
            if low <= self.severity_weight <= high:
                return label
        return "Unknown"


@dataclass
class MedicalCodingResult:
    codes: List[MedicalCode] = field(default_factory=list)
    primary_diagnosis_cui: str = ""
    primary_diagnosis_label: str = ""
    icd10_codes: List[str] = field(default_factory=list)
    snomed_codes: List[str] = field(default_factory=list)
    categories_found: List[str] = field(default_factory=list)
    overall_severity: str = "Normal"
    coding_confidence: float = 0.0
    raw_cuis: List[str] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "medical_codes": [c.to_dict() for c in self.codes],
            "primary_diagnosis": {
                "cui": self.primary_diagnosis_cui,
                "label": self.primary_diagnosis_label,
            },
            "icd10_codes": self.icd10_codes,
            "snomed_codes": self.snomed_codes,
            "categories_found": self.categories_found,
            "overall_severity": self.overall_severity,
            "coding_confidence": round(self.coding_confidence, 3),
            "raw_cuis": self.raw_cuis,
            "total_findings": len(self.codes),
        }

    def format_summary(self) -> str:
        if not self.codes:
            return "No significant medical concepts identified."

        lines = [
            f"**Primary Diagnosis:** {self.primary_diagnosis_label} "
            f"[CUI: {self.primary_diagnosis_cui}]",
            f"**Overall Severity:** {self.overall_severity}",
            f"**Coding Confidence:** {self.coding_confidence:.0%}",
            "",
            "**Identified Medical Concepts:**",
        ]
        for code in self.codes:
            lines.append(
                f"  • {code.label} | "
                f"CUI: `{code.cui}` | "
                f"ICD-10: `{code.icd10}` | "
                f"Severity: {code._severity_label()}"
            )

        if self.icd10_codes:
            lines.append(f"\n**ICD-10 Codes:** {', '.join(self.icd10_codes)}")
        if self.snomed_codes:
            lines.append(f"**SNOMED CT:** {', '.join(self.snomed_codes)}")

        return "\n".join(lines)


# ---------------------------------------------------------------------------
# Medical Coding Engine
# ---------------------------------------------------------------------------

class MedicalCodingEngine:
    """
    Extracts medical concepts from report text and maps to
    UMLS CUIs, ICD-10, and SNOMED CT codes.
    """

    def extract_codes(self, report_text: str, sections: dict = None) -> MedicalCodingResult:
        """
        Main entry point.

        Args:
            report_text : Full report text or findings/impression combined
            sections    : Optional dict with 'findings', 'impression' keys

        Returns:
            MedicalCodingResult
        """
        # Combine findings + impression for analysis (most clinically relevant)
        if sections:
            analysis_text = (
                sections.get("findings", "") + " " +
                sections.get("impression", "") + " " +
                sections.get("recommendations", "")
            )
        else:
            analysis_text = report_text

        analysis_text = analysis_text.lower()

        # Extract matching concepts
        found_codes: List[MedicalCode] = []
        seen_cuis = set()

        for term, concept in MEDICAL_CONCEPT_DB.items():
            if term in analysis_text:
                if concept["cui"] not in seen_cuis:
                    # Extract surrounding context
                    context = self._extract_context(analysis_text, term)
                    code = MedicalCode(
                        term=term,
                        cui=concept["cui"],
                        icd10=concept["icd10"],
                        snomed=concept["snomed"],
                        label=concept["label"],
                        category=concept["category"],
                        severity_weight=concept["severity_weight"],
                        context=context,
                    )
                    found_codes.append(code)
                    seen_cuis.add(concept["cui"])

        # Sort by severity (most severe first)
        found_codes.sort(key=lambda x: x.severity_weight, reverse=True)

        # Build result
        result = MedicalCodingResult()
        result.codes = found_codes
        result.raw_cuis = [c.cui for c in found_codes]
        result.icd10_codes = list(dict.fromkeys(c.icd10 for c in found_codes))
        result.snomed_codes = list(dict.fromkeys(c.snomed for c in found_codes))
        result.categories_found = list(dict.fromkeys(c.category for c in found_codes))

        # Primary diagnosis = highest severity non-normal finding
        non_normal = [c for c in found_codes if c.category != "Normal"]
        if non_normal:
            primary = non_normal[0]
            result.primary_diagnosis_cui = primary.cui
            result.primary_diagnosis_label = primary.label
        elif found_codes:
            result.primary_diagnosis_cui = found_codes[0].cui
            result.primary_diagnosis_label = found_codes[0].label

        # Overall severity
        if non_normal:
            max_severity = max(c.severity_weight for c in non_normal)
            for (low, high), label in SEVERITY_LEVELS.items():
                if low <= max_severity <= high:
                    result.overall_severity = label
                    break
        else:
            result.overall_severity = "Normal"

        # Coding confidence: based on number of matched terms
        result.coding_confidence = min(len(found_codes) / 5.0, 1.0)

        logger.info(
            f"Medical coding: {len(found_codes)} concepts found | "
            f"Primary: {result.primary_diagnosis_label} | "
            f"Severity: {result.overall_severity}"
        )
        return result

    def _extract_context(self, text: str, term: str, window: int = 80) -> str:
        """Extract surrounding sentence for context."""
        idx = text.find(term)
        if idx == -1:
            return ""
        start = max(0, idx - window)
        end = min(len(text), idx + len(term) + window)
        context = text[start:end].strip()
        return f"...{context}..."


# ---------------------------------------------------------------------------
# Convenience function
# ---------------------------------------------------------------------------

def extract_medical_codes(report_text: str, sections: dict = None) -> MedicalCodingResult:
    engine = MedicalCodingEngine()
    return engine.extract_codes(report_text, sections)