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

"""Repair OCR-damaged units and turn inline lab panels into markdown tables.

The source PDFs were scanned, so the extracted stems carry two separate injuries.

  1. Unit corruption. The scanner confuses letter pairs that share a shape —
     "m" reads as "rn" or "in", "µ" as "p" or "4" — and superscripts are lost
     entirely, so "3.5 × 10⁹/L" arrives as "3.5 x 109/L". The replacements below
     were not guessed: every one was found by listing the unit-shaped tokens
     that actually occur in the bank and reading the ones that occur once or
     twice, which is where the damage hides. Anything whose correction needed a
     judgement call — "3,000/mL" that should probably be "/µL", "14.6 x 10A" —
     is deliberately absent: a wrong unit is worse than an ugly one.

  2. A lab panel is a wall of prose. "Sodium, 139 mEq/L (139 mmol/L); Potassium,
     …" runs for a paragraph, and the reader has to parse it themselves. A
     markdown table separates analyte from value from SI value.

Every change snapshots the question first through the same helper the edit UI
uses, so anything here can be rolled back from the question's version history.

    docker compose exec backend python -m scripts.fix_lab_formatting
    docker compose exec backend python -m scripts.fix_lab_formatting --apply

    --explanations  repair units in explanation text as well as stems
    --tables        also convert lab panels to markdown tables (see below)
    --limit N       only look at the first N questions, for a quick look

NOTE ON --tables: the quiz player renders a stem as plain text (QuizPage's
ManualHighlightText slices the raw string so highlight offsets stay valid), so a
markdown table currently shows up as literal pipes there. Table conversion is
therefore opt-in until the stem is rendered as markdown. The unit repair is
safe in every view and runs by default.
"""
import argparse
import re
import sys

from app.database import SessionLocal
from app.models.question import Question
from app.routers.questions import _snapshot_question

# The user id recorded against the snapshots this script writes. NULL is
# allowed by the column and is honest: no person made these edits.
SCRIPT_EDITOR_ID = None

SUPERSCRIPT = {"0": "⁰", "1": "¹", "2": "²", "3": "³", "4": "⁴",
               "5": "⁵", "6": "⁶", "7": "⁷", "8": "⁸", "9": "⁹"}

# Powers of ten that name a real haematology unit: 10³/µL, 10⁶/µL, 10⁹/L,
# 10¹²/L. A stem holding "x 101/µL" or "x 10A" lost a digit rather than a
# superscript, and rewriting it would invent a number, so it is left alone.
REAL_POWERS = {"3", "6", "9", "12"}


def _repair_power(match: re.Match) -> str:
    """"5.6 x 109/L" -> "5.6 × 10⁹/L", but only for a power that exists."""
    value, exponent = match.group("value"), match.group("exp")
    if exponent not in REAL_POWERS:
        return match.group(0)
    digits = "".join(SUPERSCRIPT[d] for d in exponent)
    return f"{value} × 10{digits}{match.group('tail')}"


# (pattern, replacement, why) — applied in order, every one idempotent.
#
# Each was verified against the questions it fires on. The narrow ones name a
# single mangled token; the broad ones are anchored so they cannot fire on
# prose (a digit before the unit, a slash after it, or an explicit multiplier).
OCR_REPAIRS = [
    # "m" scanned as the ligature "rn" or as "in". Both appear in chloride and
    # sodium lines whose SI twin in the same sentence spells the unit correctly.
    (re.compile(r"\bmrnol/L\b"), "mmol/L", "rn read for m"),
    (re.compile(r"\binEq/L\b"), "mEq/L", "in read for m"),
    (re.compile(r"\bbeats/ruin\b"), "beats/min", "min read as ruin"),
    # Units that do not exist. Electrolytes and their SI twins are per litre;
    # a per-decilitre spelling is always the scan losing the "m" of "mmol".
    (re.compile(r"\bmmol/dL\b"), "mmol/L", "mmol/dL is not a unit"),
    (re.compile(r"\bmEq/dL\b"), "mEq/L", "mEq/dL is not a unit"),
    (re.compile(r"\bmmo/L\b"), "mmol/L", "dropped l"),
    # "q" scanned as "g" — both lines carry an SI twin in mmol/L.
    (re.compile(r"\bmEg/L\b"), "mEq/L", "g read for q"),
    (re.compile(r"\bmOsm/Kg\b"), "mOsm/kg", "casing"),
    # "≤" flattened to "<_" — the underscore is the lower half of the glyph.
    (re.compile(r"<_\s*(?=\d)"), "≤", "<_ read for ≤"),
    (re.compile(r">_\s*(?=\d)"), "≥", ">_ read for ≥"),
    (re.compile(r"\bmMol/L\b"), "mmol/L", "casing"),
    (re.compile(r"\bmg/dl\b"), "mg/dL", "casing"),
    (re.compile(r"\bg/dl\b"), "g/dL", "casing"),
    # µ scanned as p, 4 or lowercase l. Picolitres and "4mol" are never real;
    # the counts they carry are ordinary per-microlitre cell counts.
    (re.compile(r"(?<=\d)/pL\b"), "/µL", "p read for µ"),
    (re.compile(r"\b4mol/L\b"), "µmol/L", "4 read for µ"),
    (re.compile(r"(?<=\d)/µl\b"), "/µL", "casing"),
    # mm Hg is two tokens; a slash between them is a scanning artefact.
    (re.compile(r"\bmm\s*/\s*Hg\b"), "mm Hg", "stray slash"),
    (re.compile(r"\bmmHg\b"), "mm Hg", "missing space"),
    # Superscripts the scan flattened. Cubic units and BMI first, then powers
    # of ten, which need the exponent checked before they can be rewritten.
    (re.compile(r"\bkg/m2\b"), "kg/m²", "lost superscript"),
    (re.compile(r"\b([µμu])m3\b"), r"µm³", "lost superscript"),
    (re.compile(r"\bmm3\b"), "mm³", "lost superscript"),
    # MCV is a volume, not a rate: "90/µm³" means "90 µm³".
    (re.compile(r"(?<=\d)/µm³"), " µm³", "stray slash"),
    (re.compile(
        r"(?P<value>\d[\d,]*(?:\.\d+)?)\s*[x×X]\s*10\s*(?P<exp>\d{1,2})(?P<tail>\s*/)"),
     _repair_power, "lost power-of-ten superscript"),
]


def repair_units(text: str) -> tuple[str, list[str]]:
    """Apply the replacement table, returning the new text and what fired."""
    if not text:
        return text or "", []
    fired = []
    for pattern, replacement, reason in OCR_REPAIRS:
        new_text, count = pattern.subn(replacement, text)
        if new_text != text:
            fired.append(f"{reason} ×{count}")
            text = new_text
    return text, fired


# ---------------------------------------------------------------------------
# Lab panel -> markdown table
# ---------------------------------------------------------------------------

# A value is a number (or a stock qualitative phrase) followed by up to four
# unit-ish tokens and an optional parenthetical. "Unit-ish" is the load-bearing
# part: without it the scanner walks straight out of the panel and swallows the
# sentence that follows the last result.
QUALITATIVE = (r"within normal limits|within the normal range|normal|negative|positive|"
               r"pending|trace|not detected|nondetectable|nonreactive|reactive|absent|present")

# Bare unit words that carry no digit and no slash, so nothing else identifies them.
BARE_UNITS = (r"Hg|fL|fl|mm|cm|kg|mg|g|dL|mL|L|U|IU|sec|s|units?|cells?|mEq|mmol|nmol|pmol|"
              r"[µμu]mol|mOsm|ng|pg|mIU|[µμu]IU|kcal|mcg|[µμu]g|[µμu]L|mL|MoM|mg|dl|per|"
              r"[µμu]m³|mm³|hpf|HPF|LPF|seconds?|minutes?|hours?|days?|weeks?|months?|years?|to")
_UNIT_TOKEN = (rf"(?:[x×X]|[<>≤≥]|\d[\w,./%°³²⁰-⁹\-]*|"
               rf"[A-Za-zµμ³²⁰-⁹]+(?:/[A-Za-zµμ0-9³²⁰-⁹\-]+)+|%|°[CF]|"
               rf"(?:{BARE_UNITS}))\.?(?![\w/])")
_CONTINUATION = re.compile(rf"^{_UNIT_TOKEN}$")

_QUALITATIVE_ONLY = re.compile(rf"^(?:{QUALITATIVE})$", re.I)
_LEAD = rf"(?:[<>≤≥]\s*)?(?:\d[\w,./%°³²⁰-⁹\-]*|{QUALITATIVE})"
# Continuations are separated by a space, never a newline: a line break ends the
# result. Letting \s match it lets the value run into the sentence below the list.
_ITEM = re.compile(
    rf"(?<![A-Za-z])(?P<name>[A-Za-z][A-Za-z0-9 '’\-()/%]{{1,60}}?)\s*[,:]\s*"
    rf"(?P<value>{_LEAD}(?:[ \t]+{_UNIT_TOKEN}){{0,6}})"
    rf"(?P<paren>(?:[ \t]*\([^()]{{1,90}}\)){{0,2}})",
    re.I)

# Words that mean the "name" is really a clause, not an analyte.
_NOT_A_NAME = re.compile(
    r"\b(?:is|are|was|were|has|have|had|shows?|showed|reveals?|revealed|"
    r"includes?|included|and|with|of|the|his|her|he|she|but|which|that|"
    r"who|when|following|about|approximately|over|after|before|during|for)\b", re.I)

# A differential is written the other way round — "29% segmented neutrophils,
# 28% bands" — so reading it as name-then-value pairs each result with the
# NEXT analyte's percentage. Anything sitting directly after a bare percentage
# is part of such a list and is left as prose.
_PERCENT_LABELS_NEXT = re.compile(r"\d\s*%\s*$")

# A test name never carries a measurement; if it does, the split was wrong.
_VALUE_INSIDE_NAME = re.compile(r"\d\s*(?:%|[A-Za-zµμ]+/[A-Za-zµμ])")

# What a panel may start after: a lead-in's colon, a separator, a bullet, or a
# line of its own. Anything else means the "panel" is a sentence — a string of
# vital signs, say — and rewriting a sentence as a table breaks its grammar.
_PANEL_OPENERS = (":", ";", "•", "·", "-", "–", "—", "\n")

def _IS_SI_VALUE(inner: str) -> bool:
    """A parenthetical is the SI twin only if it is one number and its unit.

    "(4% segmented, 83% lymphocytes)" also opens with a digit but is a
    differential, and filing it under SI units would be a lie about the data.
    """
    return (bool(re.match(r"^[<>≤≥]?\s*\d", inner))
            and "," not in inner and len(inner.split()) <= 3)


# A reference range is an annotation on the result above it, not a test.
_RANGE_LABEL = re.compile(r"^(?:reference|normal)\s+(?:range|value)s?$", re.I)

# Enough consecutive results that the block is a panel rather than a passing
# mention of one or two values inside a sentence.
MIN_ITEMS = 4

# How much punctuation may sit between two results and still count as the same
# panel: a separator and a space or two, never a whole clause.
MAX_GAP = 4


def _clean_value(value: str) -> str:
    """Trim the sentence-ending punctuation that belongs to the prose, not the value."""
    return value.strip().rstrip(",;.").strip()


def _accept(match: re.Match) -> tuple[str, str, str] | None:
    """(analyte, result, si) if this match is a real lab result, else None."""
    # A "result" that starts inside a bracket is part of the previous result's
    # reference range ("(normal, 150 to 350)"), not a new analyte.
    if match.start() and match.string[match.start() - 1] in "([":
        return None
    if _PERCENT_LABELS_NEXT.search(match.string[: match.start()]):
        return None

    # The last entry of a written-out list carries the conjunction that joined
    # it: "…; and oxygen saturation, 97%".
    name = re.sub(r"^and\s+", "", match.group("name").strip(" ,:;-"), flags=re.I)
    # A column header the extraction glued on ("Patient Result - Hemoglobin").
    name = name.rpartition(" - ")[2].strip() or name
    if len(name) < 2 or _NOT_A_NAME.search(name):
        return None
    # A name carrying its own value ("Factor II 0.20 U/mL (reference range")
    # means the comma we split on belongs to a parenthetical, not to this test.
    if _VALUE_INSIDE_NAME.search(name):
        return None

    value = match.group("value").strip()
    if not _QUALITATIVE_ONLY.match(value):
        # The first token opens the value; every later one must look like a unit.
        for token in value.split()[1:]:
            if not _CONTINUATION.match(token):
                return None

    si = ""
    for paren in re.findall(r"\([^()]*\)", match.group("paren") or ""):
        inner = paren[1:-1].strip()
        # A parenthetical that opens with a number is the SI twin; anything
        # else ("normal, 150 to 350") is a reference range and stays beside
        # the result, where a reader expects to find it.
        if not si and _IS_SI_VALUE(inner):
            si = inner
        else:
            value = f"{value} {paren}"

    # A bare percentage that labels the words after it ("Differential count, 29%
    # segmented neutrophils") is the same trap from the other side.
    if not si and value.rstrip().endswith("%") and \
            re.match(r"^[ \t]+[a-z]", match.string[match.end():]):
        return None
    return name, _clean_value(value), si


def find_panels(text: str) -> list[tuple[int, int, list[tuple[str, str, str]]]]:
    """Locate runs of consecutive lab results: (start, end, items)."""
    if not text:
        return []
    accepted = []
    for match in _ITEM.finditer(text):
        item = _accept(match)
        if item:
            accepted.append((match.start(), match.end(), item))

    panels, run = [], []
    for entry in accepted:
        if run and entry[0] - run[-1][1] > MAX_GAP:
            panels.append(run)
            run = []
        # Two results for the same analyte mean the run has drifted into prose
        # that repeats a word; end the panel rather than emit a confused table.
        if run and not _RANGE_LABEL.match(entry[2][0]) and \
                entry[2][0].lower() in {r[2][0].lower() for r in run}:
            panels.append(run)
            run = []
        run.append(entry)
    if run:
        panels.append(run)

    out = []
    for panel in panels:
        end = panel[-1][1]
        # A panel whose next word is lowercase was not parsed to its end — the
        # last result trails off into "on arterial blood gas" or a differential.
        # Leave the whole block as prose rather than publish a truncated table.
        if re.match(r"^[ \t\n]*[.;,•·]*[ \t\n]*[a-z]", text[end:]):
            continue
        before = text[: panel[0][0]].rstrip(" \t")
        if before and not before.endswith(_PANEL_OPENERS):
            continue
        items = _fold_reference_ranges([e[2] for e in panel])
        if len(items) >= MIN_ITEMS:
            out.append((panel[0][0], end, items))
    return out


def _fold_reference_ranges(items: list[tuple[str, str, str]]) -> list[tuple[str, str, str]]:
    """"…, 125 U/L; reference range, ≤40 U/L" is one result, not two."""
    folded: list[tuple[str, str, str]] = []
    for name, value, si in items:
        if folded and _RANGE_LABEL.match(name):
            prev_name, prev_value, prev_si = folded[-1]
            folded[-1] = (prev_name, f"{prev_value} ({name.lower()}, {value})", prev_si)
            continue
        folded.append((name, value, si))
    return folded


def render_table(items: list[tuple[str, str, str]]) -> str:
    """A GFM table; the SI column only exists when something fills it."""
    with_si = any(si for _n, _v, si in items)
    header = ["Test", "Result"] + (["SI units"] if with_si else [])
    lines = ["| " + " | ".join(header) + " |",
             "| " + " | ".join(["---"] * len(header)) + " |"]
    for name, value, si in items:
        cells = [name, value] + ([si] if with_si else [])
        lines.append("| " + " | ".join(c.replace("|", "\\|") for c in cells) + " |")
    return "\n".join(lines)


def tabulate_panels(text: str) -> tuple[str, int]:
    """Replace every confidently-parsed lab panel with a markdown table."""
    if not text or "\n|" in text or text.lstrip().startswith("|"):
        return text or "", 0  # already tabulated — keep this idempotent
    panels = find_panels(text)
    if not panels:
        return text, 0
    out, cursor = [], 0
    for start, end, items in panels:
        # The lead-in keeps its colon but loses the bullet that was about to
        # introduce the first result, and the prose after a table loses the
        # separator that used to join it to the last result.
        before = text[cursor:start].rstrip().rstrip("•·-–—").rstrip()
        out.append(_leading_prose(before) if cursor else before)
        out.append("\n\n" + render_table(items) + "\n\n")
        cursor = end
    out.append(_leading_prose(text[cursor:]))
    return "".join(out).strip(), len(panels)


def _leading_prose(text: str) -> str:
    """Prose that followed a result: drop the punctuation that joined it on."""
    return text.lstrip(" \t\n.;,•·").lstrip()


# ---------------------------------------------------------------------------


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--apply", action="store_true", help="write the changes")
    parser.add_argument("--explanations", action="store_true",
                        help="repair units in explanations too")
    parser.add_argument("--tables", action="store_true",
                        help="also convert lab panels to markdown tables")
    parser.add_argument("--limit", type=int, default=0, help="only scan N questions")
    parser.add_argument("--show", type=int, default=6, help="how many examples to print")
    args = parser.parse_args()

    db = SessionLocal()
    try:
        query = db.query(Question).order_by(Question.id)
        if args.limit:
            query = query.limit(args.limit)
        questions = query.all()

        stem_fixes, expl_fixes, tabled = [], [], []
        for question in questions:
            new_stem, stem_reasons = repair_units(question.question_text or "")
            panels = 0
            if args.tables:
                new_stem, panels = tabulate_panels(new_stem)
            new_expl, expl_reasons = (question.explanation or ""), []
            if args.explanations:
                new_expl, expl_reasons = repair_units(question.explanation or "")

            stem_changed = new_stem != (question.question_text or "")
            expl_changed = args.explanations and new_expl != (question.explanation or "")
            if not stem_changed and not expl_changed:
                continue

            if stem_reasons:
                stem_fixes.append((question.id, stem_reasons))
            if expl_reasons:
                expl_fixes.append((question.id, expl_reasons))
            if panels:
                tabled.append((question.id, panels))

            if args.apply:
                # Snapshot the question as it stands, through the same helper
                # the edit endpoint uses, so the version history and the
                # restore button work on these edits exactly as on a human one.
                _snapshot_question(db, question, SCRIPT_EDITOR_ID)
                if stem_changed:
                    question.question_text = new_stem
                if expl_changed:
                    question.explanation = new_expl

        if args.apply:
            db.commit()

        print("APPLIED" if args.apply else "DRY RUN")
        print(f"  questions scanned        : {len(questions)}")
        print(f"  stems with unit repairs  : {len(stem_fixes)}")
        if args.explanations:
            print(f"  explanations repaired    : {len(expl_fixes)}")
        else:
            print("  explanations             : skipped (pass --explanations)")
        if args.tables:
            print(f"  stems tabulated          : {len(tabled)} "
                  f"({sum(n for _i, n in tabled)} panels)")
        else:
            print("  lab tables               : skipped (pass --tables)")

        reasons: dict[str, int] = {}
        for _qid, fired in stem_fixes + expl_fixes:
            for entry in fired:
                reason, _, count = entry.rpartition(" ×")
                reasons[reason] = reasons.get(reason, 0) + int(count)
        print("\n  repairs by cause:")
        for reason, count in sorted(reasons.items(), key=lambda kv: -kv[1]):
            print(f"    {count:6d}  {reason}")

        for qid, fired in stem_fixes[: args.show]:
            print(f"    question {qid}: {', '.join(fired)}")
        for qid, panels in tabled[: args.show]:
            print(f"    question {qid}: {panels} lab table(s)")

        if not args.apply:
            print("\n  Re-run with --apply to write these changes.")
        return 0
    finally:
        db.close()


if __name__ == "__main__":
    sys.exit(main())