pdf-quiz-generator/backend/scripts/assign_category_systems.py

"""Say which organ system each topic belongs to.

Two passes, and the order matters.

First the discipline it sits under: a topic beneath Cardiology is
cardiovascular unless something says otherwise. That is right for most of the
tree and wrong for exactly the cases that make an organ-system axis worth
having — conjunctivitis is filed under Infectious Disease and is an eye,
osteomyelitis is filed under Infectious Disease and is a bone.

So the second pass reads the topic's own name. A name that names an organ wins
over the shelf it was put on, which is the whole reason this axis is not just
the discipline tree relabelled.

Disciplines that span every system — General Pediatrics, Primary Care, Surgery
— assign nothing from the first pass; only a name that says an organ does.

Idempotent, and a dry run by default:

    docker compose exec backend python -m scripts.assign_category_systems
    docker compose exec backend python -m scripts.assign_category_systems --apply
    docker compose exec backend python -m scripts.assign_category_systems --apply --force
"""
import re
import sys
from collections import Counter

from sqlalchemy import text as sa_text

from app.database import SessionLocal

ROOT_TO_SYSTEM = {
    "Adolescent Medicine": "Human Development",
    "Allergy/Immunology": "Immune System",
    "Cardiology": "Cardiovascular System",
    "Child Abuse and Neglect": "Social Sciences & Ethics",
    "Critical Care": "Multisystem Processes & Disorders",
    "Dermatology": "Skin & Subcutaneous Tissue",
    "Developmental & Behavioral": "Human Development",
    "Emergency Medicine": "Multisystem Processes & Disorders",
    "Endocrinology": "Endocrine System",
    "Gastroenterology & Nutrition": "Gastrointestinal System",
    "Genetics & Metabolism": "Multisystem Processes & Disorders",
    "Hematology-Oncology": "Blood & Lymphoreticular System",
    "Infectious Disease": "Multisystem Processes & Disorders",
    "Neonatology": "Human Development",
    "Nephrology & Urology": "Renal & Urinary System",
    "Neurology": "Nervous System & Special Senses",
    "Ophthalmology": "Nervous System & Special Senses",
    "Orthopedics & Sports": "Musculoskeletal System",
    "Otolaryngology": "Nervous System & Special Senses",
    "Professional Topics": "Social Sciences & Ethics",
    "Psychiatry & Psychology": "Behavioral Health",
    "Psychosocial Issues": "Behavioral Health",
    "Pulmonology": "Respiratory System",
    "Rheumatology": "Musculoskeletal System",
    "Toxicology": "Multisystem Processes & Disorders",
}

#: A name that names an organ beats the shelf the topic was put on. Ordered:
#: the first match wins, so the specific sits above the general.
BY_NAME: list[tuple[str, str]] = [
    (r"conjunctiv|uveitis|retin|cornea|glaucoma|strabism|amblyop|orbital|ocular|"
     r"\beye|vision|optic|periorbital|chalazi|blephar|nystagmus", "Nervous System & Special Senses"),
    (r"otitis|hearing|deaf|tympan|cochlea|mastoid|\bear\b|vestibul|labyrinth",
     "Nervous System & Special Senses"),
    (r"meningitis|encephal|seizure|epilep|migraine|headache|neuropath|myelin|"
     r"cerebr|hydrocephal|neural tube|spina bifida|ataxia|palsy|concussion|coma",
     "Nervous System & Special Senses"),
    (r"osteomyel|arthritis|fracture|scoliosis|myosit|muscular dystroph|rickets|"
     r"limp|\bbone|joint|septic arthritis|slipped capital|legg-calv", "Musculoskeletal System"),
    (r"cellulit|impetigo|eczema|dermatitis|urticaria|psoria|acne|scabies|tinea|"
     r"exanthem|rash|abscess|burn|molluscum|wart|alopecia|nevus|hemangioma",
     "Skin & Subcutaneous Tissue"),
    (r"pneumon|bronchiol|asthma|croup|epiglott|pertussis|cystic fibrosis|apnea|"
     r"stridor|wheez|tuberculosis|pleural|respiratory|airway|sinusitis|pharyngitis|"
     r"tonsill|rhinitis|\blung", "Respiratory System"),
    (r"gastroenter|diarrhea|constipat|hepatitis|jaundice|intussuscept|pylor|"
     r"appendicit|celiac|reflux|colitis|\bliver|biliary|pancreat|volvulus|"
     r"hirschsprung|encopresis|\bstool", "Gastrointestinal System"),
    (r"nephr|urinary|pyelonephritis|glomerul|proteinuria|hematuria|enuresis|"
     r"\bkidney|bladder|vesicoureter|hydronephro|\brenal\b", "Renal & Urinary System"),
    (r"endocardit|myocardit|kawasaki|congenital heart|arrhythm|hypertension|"
     r"murmur|cardiomyopath|rheumatic fever|\bheart", "Cardiovascular System"),
    (r"anemia|sickle|thalass|leukemi|lymphom|neutropen|thrombocytopen|hemophil|"
     r"coagul|purpura|splenomegal|transfusion|bone marrow|epistaxis",
     "Blood & Lymphoreticular System"),
    (r"diabet|thyroid|adrenal|growth hormone|puberty|precocious|short stature|"
     r"obesity|hypoglycem|congenital adrenal|calcium|rickets", "Endocrine System"),
    (r"immunodefic|hiv\b|allerg|anaphyla|urticaria|vaccine|immuniz|scid|"
     r"complement deficiency", "Immune System"),
    (r"testic|scrotal|ovarian|menstrual|amenorrh|vaginal|pregnan|contracept|"
     r"sexually transmitted|gonorrh|chlamydia|breast", "Reproductive System"),
    (r"depress|anxiety|adhd|autism|suicide|eating disorder|substance|behaviou?r|"
     r"psychos|bullying|school refusal", "Behavioral Health"),
    (r"development|milestone|growth chart|newborn|premature|breastfeed|"
     r"adolescen|puberty", "Human Development"),
    (r"ethic|consent|confidential|disparit|advocacy|statistic|epidemiolog|"
     r"screening test|sensitivity|specificity", "Social Sciences & Ethics"),
]

#: Read before BY_NAME. These are the names that a general rule below would
#: otherwise catch by accident — "adrenal" is not a kidney and "abnormal
#: uterine bleeding" is not a clotting disorder.
FIRST: list[tuple[str, str]] = [
    (r"uterine|menorrh|menstrual|dysmenorrh", "Reproductive System"),
    (r"adrenal|pituitar|parathyroid", "Endocrine System"),
]

COMPILED = [(re.compile(pattern, re.I), system) for pattern, system in FIRST + BY_NAME]


def by_name(name: str) -> str | None:
    for pattern, system in COMPILED:
        if pattern.search(name):
            return system
    return None


def main(apply: bool, force: bool) -> int:
    db = SessionLocal()
    try:
        systems = {r.name: r.id for r in db.execute(sa_text(
            "SELECT id, name FROM question_tags WHERE type = 'system' AND parent_id IS NULL")).fetchall()}
        unknown = sorted({s for s in ROOT_TO_SYSTEM.values()} | {s for _, s in BY_NAME}) 
        missing = [s for s in unknown if s not in systems]
        if missing:
            print(f"! these systems do not exist as tags: {', '.join(missing)}")
            return 1

        rows = db.execute(sa_text(
            "SELECT id, name, parent_id, system_id FROM question_categories")).fetchall()
        parents = {r.id: r.parent_id for r in rows}
        names = {r.id: r.name for r in rows}

        def root_of(cid):
            walk, guard = cid, 0
            while parents.get(walk) is not None and guard < 20:
                walk = parents[walk]
                guard += 1
            return names.get(walk, "")

        planned: dict[int, tuple[str, str]] = {}   # id -> (system, why)
        for row in rows:
            if row.system_id and not force:
                continue
            named = by_name(row.name)
            if named:
                planned[row.id] = (named, "name")
                continue
            inherited = ROOT_TO_SYSTEM.get(root_of(row.id))
            if inherited:
                planned[row.id] = (inherited, "discipline")

        already = sum(1 for r in rows if r.system_id)
        print(f"categories                {len(rows)}")
        print(f"  already assigned        {already}")
        print(f"  from the topic's name   {sum(1 for s, why in planned.values() if why == 'name')}")
        print(f"  from its discipline     {sum(1 for s, why in planned.values() if why == 'discipline')}")
        print(f"  left unassigned         {len(rows) - already - len(planned) if not force else len(rows) - len(planned)}")
        for system, n in Counter(s for s, _ in planned.values()).most_common():
            print(f"      {system:<38} {n}")

        overrides = [(names[cid], s) for cid, (s, why) in planned.items()
                     if why == "name" and ROOT_TO_SYSTEM.get(root_of(cid)) not in (None, s)]
        print(f"\n  {len(overrides)} topics where the name beat the shelf. A sample:")
        for name, system in sorted(overrides)[:14]:
            print(f"      {name:<40} → {system}")

        if not apply:
            print("\ndry run. Pass --apply to write.")
            return 0

        for cid, (system, _) in planned.items():
            db.execute(sa_text("UPDATE question_categories SET system_id = :s WHERE id = :id"),
                       {"s": systems[system], "id": cid})
        db.commit()
        print(f"\nassigned {len(planned)} topics.")
        return 0
    finally:
        db.close()


if __name__ == "__main__":
    sys.exit(main("--apply" in sys.argv, "--force" in sys.argv))