DeepHealth/build_icd10_chapter_organ_mapping.py

from __future__ import annotations

import csv
import re
from dataclasses import dataclass
from pathlib import Path


LABEL_OFFSET = 3


@dataclass(frozen=True)
class ChapterRule:
    chapter: str
    start: str
    end: str
    title: str
    organ_system: str
    organ_system_label: str


CHAPTER_RULES = [
    ChapterRule(
        "I",
        "A00",
        "B99",
        "Certain infectious and parasitic diseases",
        "infectious_systemic",
        "Infectious / systemic",
    ),
    ChapterRule(
        "II",
        "C00",
        "D48",
        "Neoplasms",
        "neoplasm",
        "Neoplasm / oncology",
    ),
    ChapterRule(
        "III",
        "D50",
        "D89",
        "Diseases of the blood and blood-forming organs and certain immune disorders",
        "hematologic_immune",
        "Blood / immune",
    ),
    ChapterRule(
        "IV",
        "E00",
        "E90",
        "Endocrine, nutritional and metabolic diseases",
        "endocrine_metabolic",
        "Endocrine / metabolic",
    ),
    ChapterRule(
        "V",
        "F00",
        "F99",
        "Mental and behavioural disorders",
        "mental_behavioral",
        "Mental / behavioral",
    ),
    ChapterRule(
        "VI",
        "G00",
        "G99",
        "Diseases of the nervous system",
        "nervous_system",
        "Nervous system",
    ),
    ChapterRule(
        "VII",
        "H00",
        "H59",
        "Diseases of the eye and adnexa",
        "eye",
        "Eye",
    ),
    ChapterRule(
        "VIII",
        "H60",
        "H95",
        "Diseases of the ear and mastoid process",
        "ear",
        "Ear",
    ),
    ChapterRule(
        "IX",
        "I00",
        "I99",
        "Diseases of the circulatory system",
        "circulatory",
        "Circulatory system",
    ),
    ChapterRule(
        "X",
        "J00",
        "J99",
        "Diseases of the respiratory system",
        "respiratory",
        "Respiratory system",
    ),
    ChapterRule(
        "XI",
        "K00",
        "K93",
        "Diseases of the digestive system",
        "digestive",
        "Digestive system",
    ),
    ChapterRule(
        "XII",
        "L00",
        "L99",
        "Diseases of the skin and subcutaneous tissue",
        "skin",
        "Skin / subcutaneous tissue",
    ),
    ChapterRule(
        "XIII",
        "M00",
        "M99",
        "Diseases of the musculoskeletal system and connective tissue",
        "musculoskeletal",
        "Musculoskeletal / connective tissue",
    ),
    ChapterRule(
        "XIV",
        "N00",
        "N99",
        "Diseases of the genitourinary system",
        "genitourinary",
        "Genitourinary system",
    ),
    ChapterRule(
        "XV",
        "O00",
        "O99",
        "Pregnancy, childbirth and the puerperium",
        "pregnancy_childbirth",
        "Pregnancy / childbirth",
    ),
    ChapterRule(
        "XVI",
        "P00",
        "P96",
        "Certain conditions originating in the perinatal period",
        "perinatal",
        "Perinatal conditions",
    ),
    ChapterRule(
        "XVII",
        "Q00",
        "Q99",
        "Congenital malformations, deformations and chromosomal abnormalities",
        "congenital_chromosomal",
        "Congenital / chromosomal",
    ),
    ChapterRule(
        "XVIII",
        "R00",
        "R99",
        "Symptoms, signs and abnormal clinical and laboratory findings",
        "symptoms_findings",
        "Symptoms / findings",
    ),
    ChapterRule(
        "XIX",
        "S00",
        "T98",
        "Injury, poisoning and certain other consequences of external causes",
        "injury_poisoning",
        "Injury / poisoning",
    ),
    ChapterRule(
        "XX",
        "V01",
        "Y98",
        "External causes of morbidity and mortality",
        "external_causes",
        "External causes",
    ),
    ChapterRule(
        "XXI",
        "Z00",
        "Z99",
        "Factors influencing health status and contact with health services",
        "health_services_factors",
        "Health status / services factors",
    ),
    ChapterRule(
        "XXII",
        "U00",
        "U99",
        "Codes for special purposes",
        "special_purposes",
        "Special purposes",
    ),
]


CODE_RE = re.compile(r"^([A-Z][0-9]{2})(?:\.[0-9A-Z]+)?\b")
UNKNOWN_CANCER_RE = re.compile(r"^(CXX)\b\s*(.*)$", re.IGNORECASE)


def code_key(code: str) -> tuple[str, int]:
    code = code.upper().strip()
    return code[0], int(code[1:3])


def in_range(code: str, start: str, end: str) -> bool:
    letter, number = code_key(code)
    start_letter, start_number = code_key(start)
    end_letter, end_number = code_key(end)
    return (start_letter, start_number) <= (letter, number) <= (end_letter, end_number)


def parse_label(line: str) -> tuple[str, str]:
    text = line.strip()
    unknown_cancer = UNKNOWN_CANCER_RE.match(text)
    if unknown_cancer:
        return unknown_cancer.group(1).upper(), unknown_cancer.group(2).strip()

    match = CODE_RE.match(text)
    if not match:
        return text, ""
    code = match.group(1).upper()
    name = text[len(code):].strip()
    if name.startswith("(") and name.endswith(")"):
        name = name[1:-1]
    return code, name


def assign_chapter(code: str) -> ChapterRule | None:
    if code.upper().strip() == "CXX":
        return CHAPTER_RULES[1]
    if re.fullmatch(r"[A-Z][0-9]{2}", code.upper().strip()) is None:
        return None
    for rule in CHAPTER_RULES:
        if in_range(code, rule.start, rule.end):
            return rule
    return None


def build_mapping(labels_path: Path, output_path: Path) -> None:
    rows = []
    with labels_path.open("r", encoding="utf-8") as f:
        for label_index, raw in enumerate(f):
            text = raw.strip()
            if not text:
                continue

            token_id = LABEL_OFFSET + label_index
            code, name = parse_label(text)

            if code.lower() == "death":
                rows.append(
                    {
                        "label_index": label_index,
                        "token_id": token_id,
                        "code": "Death",
                        "name": "Death",
                        "icd10_chapter": "Death",
                        "icd10_range": "",
                        "icd10_chapter_title": "Death endpoint",
                        "organ_system": "death",
                        "organ_system_label": "Death",
                        "is_death": 1,
                    }
                )
                continue

            rule = assign_chapter(code)
            if rule is None:
                rows.append(
                    {
                        "label_index": label_index,
                        "token_id": token_id,
                        "code": code,
                        "name": name,
                        "icd10_chapter": "Unmapped",
                        "icd10_range": "",
                        "icd10_chapter_title": "Unmapped",
                        "organ_system": "unmapped",
                        "organ_system_label": "Unmapped",
                        "is_death": 0,
                    }
                )
                continue

            rows.append(
                {
                    "label_index": label_index,
                    "token_id": token_id,
                    "code": code,
                    "name": name,
                    "icd10_chapter": rule.chapter,
                    "icd10_range": f"{rule.start}-{rule.end}",
                    "icd10_chapter_title": rule.title,
                    "organ_system": rule.organ_system,
                    "organ_system_label": rule.organ_system_label,
                    "is_death": 0,
                }
            )

    fieldnames = [
        "label_index",
        "token_id",
        "code",
        "name",
        "icd10_chapter",
        "icd10_range",
        "icd10_chapter_title",
        "organ_system",
        "organ_system_label",
        "is_death",
    ]
    with output_path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    build_mapping(
        labels_path=Path("labels.csv"),
        output_path=Path("icd10_chapter_organ_mapping.csv"),
    )


if __name__ == "__main__":
    main()