src/endopath/mapping.py

"""Compile the target data dictionary into a canonical concept list (#91, the A2 pass).

The A2 pass reads the target dictionary and emits one canonical concept per variable:
its category, whether it is asserted or derived, its cardinality, and, for a derived
concept, the input primitives and the named rule that produce it. Every concept is cited
to the registry, and a concept that resolves to no citation fails the pass.

This list is the column spec that field-spec generation (#96) and export (#101) consume,
and that the registry browser (#94) renders.

The registry is the source of truth for the concept list, not `src/endopath/fields.py`:

- `data/taxonomy/dependency_graph.json` supplies the structural part of every concept:
  its id, its category, whether it is asserted or derived, and, for a derived concept, the
  input concepts it is computed from (the `derived-from` and `feeds-stage` edges). That file
  is itself a validated transcription of `docs/taxonomy.md` sections 1 and 2.
- `_ENRICHMENT` below supplies the three facts the graph does not carry: the cardinality
  (from each dimension's closure tag in `docs/taxonomy.md` section 1), the named rule (from
  the derived-values table in section 2 and the stage algorithm in section 4), and the
  registry citation. It is hand-transcribed the same way the dependency graph is, and the
  compiler validates that it covers exactly the graph's nodes, so a node added to the graph
  fails the pass until its cardinality, rule, and citation are supplied here.

The two invariants this pass serves:

1. Numbers stay numbers. The raw millimetre quantities (`invasion_depth`,
   `myometrial_thickness`, and their cervical counterparts) are asserted concepts, and the
   percentage and the category are separate derived concepts computed from them by the
   `ratio_percent` formula and a threshold. A record that stored only the category could not
   answer the other convention's question.
2. No value without its evidence. `cited_to` resolves against the registry: either an
   `extraction_id` in `data/taxonomy/extractions.csv` (a checksummed, page-anchored, verbatim
   span, #29), or a section of `docs/taxonomy.md` for a concept the registry currently
   evidences in prose rather than in a value-level extraction. A `cited_to` that resolves to
   neither raises `CompilationError`.

`rule_kind` is closed at the same set #35 closes its derivations on: `formula`, `threshold`,
and `membership`, plus `algorithm` for FIGO stage, which `docs/taxonomy.md` section 2 and #35
both treat as an edition-specific algorithm rather than one of the three simple kinds. `rule`
carries a named formula where the registry names one (`ratio_percent`, bound once and used for
both the myometrial and the cervical ratio; `derive_stage` for the stage algorithm) and
otherwise repeats the `rule_kind`. No cell holds a free-text formula.

`category` uses the registry's own category ids (`histotype`, `myo`, `lvsi`, `nodes`,
`molecular`, `stage`), the ones `data/taxonomy/dependency_graph.json` and the registry browser
already key on, so the compiled list joins to them without a translation table.
"""

from __future__ import annotations

import csv
import json
import re
from enum import Enum
from pathlib import Path
from typing import Optional

from pydantic import BaseModel, ConfigDict, model_validator

REPO_ROOT = Path(__file__).resolve().parents[2]
TAXONOMY_DIR = REPO_ROOT / "data" / "taxonomy"
DEPENDENCY_GRAPH_PATH = TAXONOMY_DIR / "dependency_graph.json"
EXTRACTIONS_PATH = TAXONOMY_DIR / "extractions.csv"
TAXONOMY_DOC_PATH = REPO_ROOT / "docs" / "taxonomy.md"

# A citation of this form points at a heading in docs/taxonomy.md, e.g.
# "docs/taxonomy.md#1.3" is section 1.3. Used for concepts the registry currently
# evidences in prose rather than in a value-level extraction row.
_DOC_CITATION_RE = re.compile(r"^docs/taxonomy\.md#(?P<section>\d+(?:\.\d+)*)$")
# Heading numbers in docs/taxonomy.md: "## 2. Derived values" -> "2",
# "### 1.3 Myometrial invasion" -> "1.3".
_DOC_HEADING_RE = re.compile(r"^#{2,4}\s+(?P<section>\d+(?:\.\d+)*)\.?\s")


class CompilationError(ValueError):
    """The dictionary did not compile: an uncited concept, an unknown input, a
    rule outside the closed vocabulary, or a drift between the enrichment table
    and the dependency graph."""


class AssertionType(str, Enum):
    """A value a reviewer reads from the report (`asserted`), or one the pass
    computes from asserted values by a named rule (`derived`)."""

    ASSERTED = "asserted"
    DERIVED = "derived"


class Cardinality(str, Enum):
    """How many values a single case may take for this concept, from the
    closure tag in docs/taxonomy.md section 1. A quantity that is either stated
    once or absent is `zero_or_one`; a closed value set with a residual is
    `exactly_one`."""

    EXACTLY_ONE = "exactly_one"
    ZERO_OR_ONE = "zero_or_one"
    ZERO_OR_MORE = "zero_or_more"


class RuleKind(str, Enum):
    """The closed set of derivation kinds (#35). `algorithm` is the FIGO-stage
    exception docs/taxonomy.md section 2 names: an edition-specific algorithm,
    not one of the three simple kinds."""

    FORMULA = "formula"
    THRESHOLD = "threshold"
    MEMBERSHIP = "membership"
    ALGORITHM = "algorithm"


# The named rules the registry gives a stable identifier beyond the bare rule_kind.
# `ratio_percent` is registered once and bound twice (docs/taxonomy.md section 2, #35);
# `derive_stage` is the stage algorithm (docs/taxonomy.md section 4). Every other derived
# value repeats its rule_kind as its rule, so no cell holds a free-text formula.
_NAMED_FORMULAS: frozenset[str] = frozenset({"ratio_percent"})
_NAMED_ALGORITHMS: frozenset[str] = frozenset({"derive_stage"})


class DerivedRule(BaseModel):
    """The derivation of a derived concept: the input concepts and the named rule."""

    model_config = ConfigDict(frozen=True)

    inputs: tuple[str, ...]
    rule: str
    rule_kind: RuleKind

    @model_validator(mode="after")
    def _rule_is_in_the_closed_vocabulary(self) -> "DerivedRule":
        if not self.inputs:
            raise CompilationError("a derived concept must name at least one input")
        kind = self.rule_kind
        if kind is RuleKind.FORMULA and self.rule not in _NAMED_FORMULAS:
            raise CompilationError(
                f"formula rule {self.rule!r} is not a named formula in {sorted(_NAMED_FORMULAS)}"
            )
        if kind is RuleKind.ALGORITHM and self.rule not in _NAMED_ALGORITHMS:
            raise CompilationError(
                f"algorithm rule {self.rule!r} is not a named algorithm in {sorted(_NAMED_ALGORITHMS)}"
            )
        if kind in (RuleKind.THRESHOLD, RuleKind.MEMBERSHIP) and self.rule != kind.value:
            raise CompilationError(
                f"{kind.value} rule must repeat its rule_kind (the registry names no formula "
                f"for it), got rule={self.rule!r}"
            )
        return self


class CompiledConcept(BaseModel):
    """One canonical concept: the #91 contract row."""

    model_config = ConfigDict(frozen=True)

    concept_id: str
    category: str
    assertion_type: AssertionType
    cardinality: Cardinality
    derived: Optional[DerivedRule] = None
    cited_to: str

    @model_validator(mode="after")
    def _derived_block_matches_assertion_type(self) -> "CompiledConcept":
        if self.assertion_type is AssertionType.DERIVED and self.derived is None:
            raise CompilationError(f"{self.concept_id} is derived but names no rule or inputs")
        if self.assertion_type is AssertionType.ASSERTED and self.derived is not None:
            raise CompilationError(f"{self.concept_id} is asserted but carries a derived block")
        return self


# id -> the three facts dependency_graph.json does not carry: cardinality (docs/taxonomy.md
# section 1 closure tags), the named rule for a derived concept (section 2 derived-values
# table and section 4 stage algorithm), and the registry citation. An asserted concept omits
# `rule`/`rule_kind`; a derived concept carries both. The compiler validates this table covers
# exactly the graph's nodes.
#
# A citation is an extraction_id (data/taxonomy/extractions.csv) where the registry has a
# value-level or element-level span, and a docs/taxonomy.md section where it currently has only
# prose. The prose-only concepts are the raw myometrial and cervical quantities, their ratios
# and category, the molecular assays, and the ProMisE and TCGA classes: #29 has not yet cut
# element-level extraction rows for them, and this pass picks up the extraction automatically
# once one lands.
_ENRICHMENT: dict[str, dict] = {
    # --- Histotype ---------------------------------------------------------------------
    "histologic_type": {"cardinality": "exactly_one", "cited_to": "ex_cap_histotype_element"},
    "figo_grade": {"cardinality": "exactly_one", "cited_to": "ex_cap_grade_element"},
    "grade_binary": {
        "cardinality": "exactly_one",
        "rule": "membership",
        "rule_kind": "membership",
        "cited_to": "docs/taxonomy.md#2",
    },
    "histotype_aggressiveness": {
        "cardinality": "exactly_one",
        "rule": "membership",
        "rule_kind": "membership",
        "cited_to": "ex_figo2023_aggressive_types",
    },
    # --- Myometrial and cervical invasion ----------------------------------------------
    # Numbers stay numbers: the two millimetre quantities are asserted, the percentage and
    # the category are derived from them.
    "invasion_depth": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.3"},
    "myometrial_thickness": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.3"},
    "myometrial_invasion_percent": {
        "cardinality": "zero_or_one",
        "rule": "ratio_percent",
        "rule_kind": "formula",
        "cited_to": "docs/taxonomy.md#2",
    },
    "myometrial_invasion_category": {
        "cardinality": "exactly_one",
        "rule": "threshold",
        "rule_kind": "threshold",
        "cited_to": "docs/taxonomy.md#2",
    },
    "cervical_stromal_invasion": {"cardinality": "exactly_one", "cited_to": "docs/taxonomy.md#1.4"},
    "cervical_invasion_depth": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.4"},
    "cervical_wall_thickness": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.4"},
    "cervical_wall_percent": {
        "cardinality": "zero_or_one",
        "rule": "ratio_percent",
        "rule_kind": "formula",
        "cited_to": "docs/taxonomy.md#2",
    },
    # --- Lymphovascular space invasion -------------------------------------------------
    "lvsi_status": {"cardinality": "exactly_one", "cited_to": "ex_cap_lvsi_equivalence"},
    "lvsi_foci_count": {"cardinality": "zero_or_one", "cited_to": "ex_cap_lvsi_count_instruction"},
    "lvsi_extent": {
        "cardinality": "exactly_one",
        "rule": "threshold",
        "rule_kind": "threshold",
        "cited_to": "ex_figo2023_lvsi_substantial_rule",
    },
    # --- Regional lymph nodes ----------------------------------------------------------
    "pn_category": {"cardinality": "exactly_one", "cited_to": "docs/taxonomy.md#1.6"},
    # --- Molecular ---------------------------------------------------------------------
    "pole_exonuclease_mutation": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.7"},
    "mmr_ihc": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.7"},
    "p53_ihc": {"cardinality": "zero_or_one", "cited_to": "docs/taxonomy.md#1.7"},
    "promise_class": {
        "cardinality": "exactly_one",
        "rule": "membership",
        "rule_kind": "membership",
        "cited_to": "docs/taxonomy.md#1.7",
    },
    "tcga_class": {"cardinality": "exactly_one", "cited_to": "docs/taxonomy.md#1.7"},
    # --- Stage -------------------------------------------------------------------------
    "figo_stage": {
        "cardinality": "exactly_one",
        "rule": "derive_stage",
        "rule_kind": "algorithm",
        "cited_to": "ex_cap_figo2009_element",
    },
}


# Which built-in CAP checklist variable (src/endopath/fields.py) each compiles to. The
# derived variables map to their derived concept, whose inputs the compiled list carries
# separately. Every checklist field the MVP carries maps to a concept here; a field with no
# concept would be listed in UNMODELED_CAP_FIELDS as an explicit gap rather than fabricated.
CAP_CHECKLIST_COVERAGE: dict[str, str] = {
    "histologic_type": "histologic_type",
    "histologic_grade": "figo_grade",
    "myometrial_invasion_percent": "myometrial_invasion_percent",
    "cervical_stromal_invasion": "cervical_stromal_invasion",
    "lymphovascular_space_invasion": "lvsi_extent",
    "regional_lymph_node_status": "pn_category",
    "pathologic_stage": "figo_stage",
    "molecular_classification": "promise_class",
}

# CAP fields fields.py lists that the taxonomy registry does not model, and so route to the reviewer
# queue as unresolved rather than to a fabricated concept. Empty in the MVP: every checklist field maps
# to a concept in CAP_CHECKLIST_COVERAGE.
UNMODELED_CAP_FIELDS: frozenset[str] = frozenset()


def _load_graph() -> tuple[dict[str, dict], dict[str, list[str]], frozenset[str]]:
    """Return (nodes by id, inputs by derived concept id, category id set) from the
    dependency graph. Inputs are the `derived-from` and `feeds-stage` targets, in graph order."""
    data = json.loads(DEPENDENCY_GRAPH_PATH.read_text(encoding="utf-8"))
    nodes = {node["id"]: node for node in data["nodes"]}
    inputs: dict[str, list[str]] = {}
    for edge in data["edges"]:
        if edge["relation"] in ("derived-from", "feeds-stage"):
            inputs.setdefault(edge["from"], []).append(edge["to"])
    categories = frozenset(node["category"] for node in data["nodes"])
    return nodes, inputs, categories


def _load_extraction_ids() -> frozenset[str]:
    with EXTRACTIONS_PATH.open(newline="", encoding="utf-8") as handle:
        return frozenset(row["extraction_id"] for row in csv.DictReader(handle))


def _load_doc_sections() -> frozenset[str]:
    sections: set[str] = set()
    for line in TAXONOMY_DOC_PATH.read_text(encoding="utf-8").splitlines():
        match = _DOC_HEADING_RE.match(line)
        if match:
            sections.add(match.group("section"))
    return frozenset(sections)


def citation_resolves(cited_to: str, extraction_ids: frozenset[str], doc_sections: frozenset[str]) -> bool:
    """A citation resolves when it names a known extraction_id or a real docs/taxonomy.md section."""
    if cited_to in extraction_ids:
        return True
    doc_match = _DOC_CITATION_RE.match(cited_to)
    return bool(doc_match) and doc_match.group("section") in doc_sections


def compile_concepts() -> list[CompiledConcept]:
    """The A2 pass: compile the registry's target model into the canonical concept list.

    Raises CompilationError if the enrichment table and the dependency graph have drifted, if
    a derived concept names an input that is not itself a concept, if a rule falls outside the
    closed vocabulary, or if any concept resolves to no registry citation.
    """
    nodes, inputs, categories = _load_graph()
    extraction_ids = _load_extraction_ids()
    doc_sections = _load_doc_sections()

    graph_ids = set(nodes)
    enriched_ids = set(_ENRICHMENT)
    if graph_ids != enriched_ids:
        missing = sorted(graph_ids - enriched_ids)
        extra = sorted(enriched_ids - graph_ids)
        raise CompilationError(
            "enrichment table and dependency graph disagree on the concept set: "
            f"missing cardinality/rule/citation for {missing}; enrichment names unknown concepts {extra}"
        )

    concepts: list[CompiledConcept] = []
    for concept_id, node in nodes.items():
        enrichment = _ENRICHMENT[concept_id]
        is_derived = node["kind"] == "derived"
        has_rule = "rule" in enrichment or "rule_kind" in enrichment
        if is_derived != has_rule:
            raise CompilationError(
                f"{concept_id}: dependency graph says kind={node['kind']!r} but enrichment "
                f"{'names' if has_rule else 'omits'} a rule"
            )

        derived: Optional[DerivedRule] = None
        if is_derived:
            concept_inputs = inputs.get(concept_id)
            if not concept_inputs:
                raise CompilationError(f"{concept_id} is derived but the graph gives it no inputs")
            derived = DerivedRule(
                inputs=tuple(concept_inputs),
                rule=enrichment["rule"],
                rule_kind=RuleKind(enrichment["rule_kind"]),
            )

        concept = CompiledConcept(
            concept_id=concept_id,
            category=node["category"],
            assertion_type=AssertionType(node["kind"]),
            cardinality=Cardinality(enrichment["cardinality"]),
            derived=derived,
            cited_to=enrichment["cited_to"],
        )

        if concept.category not in categories:
            raise CompilationError(f"{concept_id} has category {concept.category!r} not in the registry")
        if not citation_resolves(concept.cited_to, extraction_ids, doc_sections):
            raise CompilationError(
                f"{concept_id} cites {concept.cited_to!r}, which is neither an extraction_id in "
                "extractions.csv nor a section of docs/taxonomy.md"
            )
        concepts.append(concept)

    # Input closure: every input of every derived concept is itself a compiled concept.
    known = {concept.concept_id for concept in concepts}
    for concept in concepts:
        if concept.derived is None:
            continue
        for source in concept.derived.inputs:
            if source not in known:
                raise CompilationError(
                    f"{concept.concept_id} derives from {source!r}, which is not a concept in the list"
                )
    return concepts


def concept_list() -> list[dict]:
    """The canonical concept list as plain JSON-ready dicts, one per concept."""
    return [concept.model_dump(mode="json") for concept in compile_concepts()]


def concept_list_json() -> str:
    """The canonical concept list as strict JSON, the #91 artifact."""
    return json.dumps(concept_list(), indent=2) + "\n"


def cap_checklist_coverage() -> dict[str, str]:
    """Which compiled concept each built-in CAP checklist variable maps to, validated against
    the compiled list so a renamed concept fails here rather than silently."""
    known = {concept.concept_id for concept in compile_concepts()}
    for field_name, concept_id in CAP_CHECKLIST_COVERAGE.items():
        if concept_id not in known:
            raise CompilationError(
                f"CAP field {field_name!r} maps to {concept_id!r}, which is not in the compiled list"
            )
    return dict(CAP_CHECKLIST_COVERAGE)


def main() -> None:
    print(concept_list_json(), end="")


if __name__ == "__main__":
    main()