src/endopath/longitudinal.py

"""Project a participant's joined GDC longitudinal records into the cohort export's outcome columns.

The cohort export projects the confirmed pathology values onto the concept list (`export.py`). Those
concepts answer "what the pathology report states". The outcome of the disease, whether the tumour
recurred, how long the patient was followed, what adjuvant treatment they received, arrives from a
different source: the GDC clinical record joined to the case by participant barcode, not the report. This
module produces that second block of columns, kept separate from the concept list so the report-concept
schema and its checksum keep meaning "what the report answers", and so the mapping-review crosswalk never
shows an outcome as if a report finding mapped to it. The join is the same one `docs/prd.md` section 4.8
names for age at diagnosis: a value the report omits, supplied by join from the clinical record.

The two invariants hold here as they do for the pathology values (`docs/prd.md` section 4.1-4.3):

1. Numbers stay numbers. The day-counts (`days_to_last_follow_up`, `days_to_recurrence`,
   `days_to_first_treatment`) and the GDC-asserted primitives (`progression_or_recurrence`,
   `disease_response_latest`, the treatment-received flags) are stored as they arrive. `recurrence_observed`
   is the one derived column: it unifies the explicit recurrence flag with the disease-response signal by
   the named rule `recurrence_from_followup`, and it sits beside the primitives it is computed from rather
   than replacing them.

2. No value without its evidence. A join-sourced value's evidence is the identity of the GDC record it came
   from, not a page span. `export.py` records the contributing record ids per case in the bundle's
   `join_sources` section. A present value carries no absent reason; the column descriptor marks the whole
   block join-sourced.

3. Every absence names its reason, from the same `FieldState` vocabulary the concept columns use
   (`docs/prd.md` section 4.3). A participant the join reached but who has no follow-up reads `not_stated`.
   A field gated off by a prior value (a recurrence type when no recurrence occurred, days to death for a
   living patient) reads `not_applicable`. A participant the join did not reach at all reads
   `determined_by_join`: the value would arrive by join, and the join found no record.

The GDC recurrence signal lives inside the follow-up record's `structured_data`, not the top-level column.
`progression_or_recurrence == "Yes"` marks a documented recurrence or progression, with
`progression_or_recurrence_type` (Distant / Locoregional / Unknown) and, on some, `days_to_recurrence` or
`days_to_progression`. `disease_response` (TF-Tumor Free / WT-With Tumor) is a corroborating second signal.
Survival arrives from the `demographic` record (`gdc_longitudinal.py` pulls it): `vital_status` and, for a
death, `days_to_death`. `overall_survival_days` derives the survival-analysis time from them: the day of
death for a deceased patient, the last-follow-up day for a living one.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Optional

from endopath.validation import FieldState

# The category the outcome columns carry in the column spec, distinct from the report-concept categories
# (histotype, myo, lvsi, nodes, molecular, stage). The frontend tags this block as join-sourced.
OUTCOME_CATEGORY = "outcome"
JOIN_SOURCE = "GDC clinical/longitudinal record, joined by participant barcode"

# The named derivation rules for the derived outcome columns, recorded in each column descriptor and the
# export bundle so a second centre re-derives them identically.
RECURRENCE_RULE = "recurrence_from_followup"
OVERALL_SURVIVAL_RULE = "overall_survival_from_vital_status"

_YES = "yes"


@dataclass(frozen=True)
class OutcomeCell:
    """One outcome column's value on one case: a concrete value with no reason, or a coded absence with no
    value. `origin` records how it was reached ('join' for a GDC-asserted primitive, 'derived' for the
    recurrence rule, 'absent' for a coded absence)."""

    value: Optional[Any]
    absent_reason: Optional[str]
    origin: str


def _present(value: Any, origin: str = "join") -> OutcomeCell:
    return OutcomeCell(value, None, origin)


def _absent(reason: FieldState) -> OutcomeCell:
    return OutcomeCell(None, reason.value, "absent")


@dataclass(frozen=True)
class Record:
    """One longitudinal record reduced to what the projection reads: its type, its id (for provenance), and
    its parsed `structured_data`. `export.py` builds these from the `longitudinal_records` rows so this
    module stays free of the database and is tested on plain dicts."""

    record_type: str
    record_id: str
    data: dict


def _followups(records: list[Record]) -> list[Record]:
    return [r for r in records if r.record_type == "follow_up"]


def _treatments(records: list[Record]) -> list[Record]:
    return [r for r in records if r.record_type.startswith("treatment_")]


def _demographic(records: list[Record]) -> Optional[Record]:
    for r in records:
        if r.record_type == "demographic":
            return r
    return None


def _int(value: Any) -> Optional[int]:
    """A day-count coerced to int, or None. GDC writes some as floats ('1125.0') and some as strings."""
    if value is None or value == "":
        return None
    try:
        return int(float(value))
    except (TypeError, ValueError):
        return None


def _followup_day(record: Record) -> Optional[int]:
    return _int(record.data.get("days_to_follow_up"))


def _last_follow_up_day(records: list[Record]) -> Optional[int]:
    """The censoring/observation time: the explicit `last_followup` record's day when present, else the
    latest day any follow-up or lost-to-follow-up record carries."""
    candidates: list[int] = []
    for record in records:
        if record.record_type == "last_followup":
            day = _int(record.data.get("days_to_last_follow_up"))
            if day is not None:
                candidates.append(day)
        elif record.record_type == "follow_up":
            day = _followup_day(record)
            if day is not None:
                candidates.append(day)
        elif record.record_type == "lost_to_followup":
            day = _int(record.data.get("days_to_lost_to_followup"))
            if day is not None:
                candidates.append(day)
    return max(candidates) if candidates else None


def _latest_disease_response(followups: list[Record]) -> Optional[str]:
    """The `disease_response` on the follow-up with the largest day that carries one, so the latest known
    status wins. Falls back to any follow-up with a response when none carry a day."""
    dated = [
        (_followup_day(f), f.data.get("disease_response"))
        for f in followups
        if f.data.get("disease_response") not in (None, "")
    ]
    if not dated:
        return None
    with_day = [(day, resp) for day, resp in dated if day is not None]
    if with_day:
        return max(with_day, key=lambda pair: pair[0])[1]
    return dated[-1][1]


def _recurrence_flag(followups: list[Record]) -> bool:
    return any(
        str(f.data.get("progression_or_recurrence", "")).strip().lower() == _YES for f in followups
    )


def _recurrence_type(followups: list[Record]) -> Optional[str]:
    for f in followups:
        if str(f.data.get("progression_or_recurrence", "")).strip().lower() == _YES:
            kind = f.data.get("progression_or_recurrence_type")
            if kind not in (None, ""):
                return kind
    return None


def _days_to_recurrence(followups: list[Record]) -> Optional[int]:
    """The recorded time to recurrence or progression on a follow-up that flagged one. GDC records
    `days_to_recurrence` on some and `days_to_progression` on others; the earliest recorded wins."""
    days: list[int] = []
    for f in followups:
        if str(f.data.get("progression_or_recurrence", "")).strip().lower() != _YES:
            continue
        for key in ("days_to_recurrence", "days_to_progression"):
            day = _int(f.data.get(key))
            if day is not None:
                days.append(day)
    return min(days) if days else None


def _administered(treatments: list[Record]) -> list[Record]:
    """The treatments actually given: `treatment_or_therapy == "yes"`. GDC emits a template treatment
    record on nearly every case for a therapy that was considered but not administered, carrying
    `treatment_or_therapy = "no"`, so the raw record list is not the set of delivered treatments."""
    return [
        t for t in treatments if str(t.data.get("treatment_or_therapy", "")).strip().lower() == _YES
    ]


def _received(records: list[Record], *prefixes: str) -> OutcomeCell:
    """Whether an adjuvant treatment of one of `prefixes` was administered. Because GDC emits a
    `treatment_<type>` template record marked `treatment_or_therapy = "no"` on cases where the therapy was
    not given, record presence overcounts; administration is read from the flag. Any matching record marked
    "yes" is True; an explicit "no" with no "yes" is False; matching records that carry only "unknown" or no
    status leave administration unknown, a coded `not_stated` rather than an asserted withholding (invariant
    3). No matching record at all stays False: the modality is simply absent for this case."""
    matched = [
        r for r in records if any(r.record_type.startswith(f"treatment_{prefix}") for prefix in prefixes)
    ]
    if not matched:
        return _present(False)
    statuses = [str(r.data.get("treatment_or_therapy", "")).strip().lower() for r in matched]
    if _YES in statuses:
        return _present(True)
    if "no" in statuses:
        return _present(False)
    return _absent(FieldState.NOT_STATED)


def _days_to_first_treatment(treatments: list[Record]) -> Optional[int]:
    days = [_int(t.data.get("days_to_treatment_start")) for t in treatments]
    days = [d for d in days if d is not None]
    return min(days) if days else None


def _therapeutic_agents(treatments: list[Record]) -> list[str]:
    agents: set[str] = set()
    for t in treatments:
        agent = t.data.get("therapeutic_agents")
        if agent not in (None, ""):
            agents.add(str(agent))
    return sorted(agents)


def project_outcomes(records: list[Record]) -> dict[str, OutcomeCell]:
    """Project one participant's longitudinal records onto the outcome columns.

    `records` is the participant's GDC records (empty when the join reached no record for them). Returns one
    `OutcomeCell` per outcome column key, each either a value or a coded absence, never a bare blank.
    """
    # The join reached no record for this participant: every outcome is `determined_by_join`, the value that
    # would arrive from the clinical record but did not for this case.
    if not records:
        return {column["key"]: _absent(FieldState.DETERMINED_BY_JOIN) for column in OUTCOME_COLUMNS}

    followups = _followups(records)
    treatments = _treatments(records)
    has_followup = bool(followups)

    cells: dict[str, OutcomeCell] = {}

    # --- follow-up / observation time ---
    last_day = _last_follow_up_day(records)
    cells["days_to_last_follow_up"] = (
        _present(last_day) if last_day is not None else _absent(FieldState.NOT_STATED)
    )
    cells["follow_up_count"] = _present(len(followups))

    latest_response = _latest_disease_response(followups)
    cells["disease_response_latest"] = (
        _present(latest_response) if latest_response is not None else _absent(FieldState.NOT_STATED)
    )

    # --- recurrence (GDC-asserted primitives) ---
    recurred = _recurrence_flag(followups)
    if has_followup:
        # Follow-up was recorded, so a "Yes" flag means recurrence and its absence means none observed.
        cells["progression_or_recurrence"] = _present(recurred)
    else:
        cells["progression_or_recurrence"] = _absent(FieldState.NOT_STATED)

    if not has_followup:
        cells["recurrence_type"] = _absent(FieldState.NOT_STATED)
        cells["days_to_recurrence"] = _absent(FieldState.NOT_STATED)
    elif not recurred:
        # A recurrence type and its timing have no meaning when no recurrence occurred: a gate, not a miss.
        cells["recurrence_type"] = _absent(FieldState.NOT_APPLICABLE)
        cells["days_to_recurrence"] = _absent(FieldState.NOT_APPLICABLE)
    else:
        kind = _recurrence_type(followups)
        cells["recurrence_type"] = (
            _present(kind) if kind is not None else _absent(FieldState.NOT_STATED)
        )
        day = _days_to_recurrence(followups)
        cells["days_to_recurrence"] = (
            _present(day) if day is not None else _absent(FieldState.NOT_STATED)
        )

    # --- recurrence (derived): unify the explicit flag with the disease-response signal (invariant 1) ---
    cells["recurrence_observed"] = _derive_recurrence_observed(
        cells["progression_or_recurrence"], cells["disease_response_latest"]
    )

    # --- adjuvant treatment received ---
    cells["received_chemotherapy"] = _received(records, "chemotherapy")
    cells["received_radiation"] = _received(records, "radiation", "brachytherapy")
    cells["received_hormone_therapy"] = _received(records, "hormone")

    administered = _administered(treatments)
    if administered:
        first_day = _days_to_first_treatment(administered)
        cells["days_to_first_treatment"] = (
            _present(first_day) if first_day is not None else _absent(FieldState.NOT_STATED)
        )
        agents = _therapeutic_agents(administered)
        cells["therapeutic_agents"] = (
            _present(", ".join(agents)) if agents else _absent(FieldState.NOT_STATED)
        )
    else:
        cells["days_to_first_treatment"] = _absent(FieldState.NOT_APPLICABLE)
        cells["therapeutic_agents"] = _absent(FieldState.NOT_APPLICABLE)

    # --- survival: from the demographic record (GDC-asserted primitives) ---
    demographic = _demographic(records)
    vital = demographic.data.get("vital_status") if demographic else None
    death_day = _int(demographic.data.get("days_to_death")) if demographic else None
    is_dead = isinstance(vital, str) and vital.strip().lower() == "dead"

    cells["vital_status"] = _present(vital) if vital is not None else _absent(FieldState.NOT_STATED)
    if death_day is not None:
        cells["days_to_death"] = _present(death_day)
    elif is_dead:
        # Recorded dead, but no day: searched the join source and it is silent.
        cells["days_to_death"] = _absent(FieldState.NOT_STATED)
    elif vital is not None:
        # Alive (or a non-death status): days to death has no meaning, a gate not a miss.
        cells["days_to_death"] = _absent(FieldState.NOT_APPLICABLE)
    else:
        cells["days_to_death"] = _absent(FieldState.NOT_STATED)

    # --- survival (derived): the event/censoring time for a survival analysis (invariant 1) ---
    cells["overall_survival_days"] = _derive_overall_survival(
        is_dead, death_day, cells["days_to_last_follow_up"]
    )

    return cells


def _derive_overall_survival(
    is_dead: bool, death_day: Optional[int], last_follow_up: OutcomeCell
) -> OutcomeCell:
    """Rule `overall_survival_from_vital_status`: the survival-analysis time is the day of death for a
    deceased patient and the last-follow-up day for a living one. Derived from the survival primitives and
    the observation time, which are stored beside it (invariant 1). Absent when neither is known."""
    if is_dead and death_day is not None:
        return OutcomeCell(death_day, None, "derived")
    if last_follow_up.value is not None:
        return OutcomeCell(last_follow_up.value, None, "derived")
    if is_dead:
        return _absent(FieldState.NOT_STATED)
    return _absent(FieldState.NOT_STATED)


def _median(values: list[int]) -> Optional[float]:
    if not values:
        return None
    ordered = sorted(values)
    mid = len(ordered) // 2
    if len(ordered) % 2:
        return float(ordered[mid])
    return (ordered[mid - 1] + ordered[mid]) / 2


def aggregate_cohort(outcomes_per_case: list[dict[str, OutcomeCell]]) -> dict:
    """Tally the projected outcomes across a cohort for the dashboard's longitudinal module.

    Takes the per-case `project_outcomes` results (so the cohort figures use the same recurrence and
    survival definitions the export does, never a second parallel definition) and returns the recurrence,
    survival, and adjuvant-treatment counts, plus the median follow-up and overall-survival times and the
    follow-up-day values for a histogram. Only cases the join reached carry these, so the caller passes the
    cases that have longitudinal records.
    """
    recurrence = {"observed": 0, "none": 0, "unknown": 0}
    survival = {"dead": 0, "alive": 0, "unknown": 0}
    treatment = {"chemotherapy": 0, "radiation": 0, "hormone": 0}
    follow_up_days: list[int] = []
    overall_survival_days: list[int] = []

    for outcomes in outcomes_per_case:
        recurred = outcomes["recurrence_observed"].value
        if recurred is True:
            recurrence["observed"] += 1
        elif recurred is False:
            recurrence["none"] += 1
        else:
            recurrence["unknown"] += 1

        vital = outcomes["vital_status"].value
        if vital == "Dead":
            survival["dead"] += 1
        elif vital == "Alive":
            survival["alive"] += 1
        else:
            survival["unknown"] += 1

        if outcomes["received_chemotherapy"].value is True:
            treatment["chemotherapy"] += 1
        if outcomes["received_radiation"].value is True:
            treatment["radiation"] += 1
        if outcomes["received_hormone_therapy"].value is True:
            treatment["hormone"] += 1

        follow_up = outcomes["days_to_last_follow_up"].value
        if follow_up is not None:
            follow_up_days.append(follow_up)
        survival_days = outcomes["overall_survival_days"].value
        if survival_days is not None:
            overall_survival_days.append(survival_days)

    return {
        "cases": len(outcomes_per_case),
        "recurrence": recurrence,
        "survival": survival,
        "treatment": treatment,
        "median_follow_up_days": _median(follow_up_days),
        "median_overall_survival_days": _median(overall_survival_days),
        "follow_up_days": follow_up_days,
    }


def _derive_recurrence_observed(
    flag: OutcomeCell, disease_response: OutcomeCell
) -> OutcomeCell:
    """Rule `recurrence_from_followup`: the tumour recurred if GDC flagged a progression or recurrence, or
    if the latest disease response reads as with-tumour. Derived from the two primitive columns, so the
    stored primitives still answer the question this column pre-computes (invariant 1). When neither
    primitive is present, the derived column inherits their absence."""
    if flag.value is True:
        return OutcomeCell(True, None, "derived")
    response = disease_response.value
    if isinstance(response, str) and response.strip().lower().startswith("wt"):
        return OutcomeCell(True, None, "derived")
    if flag.value is False:
        return OutcomeCell(False, None, "derived")
    # Neither primitive present: follow-up was not recorded, so recurrence was not assessed.
    return _absent(FieldState.NOT_STATED)


def contributing_record_ids(records: list[Record]) -> list[str]:
    """The ids of the records that fed this case's outcome columns, for the export bundle's join provenance
    (invariant 2). Follow-up, last-follow-up, lost-to-follow-up, and treatment records are read; other
    types are not, so they do not claim provenance they did not supply."""
    read_types = {"follow_up", "last_followup", "lost_to_followup", "demographic"}
    return sorted(
        r.record_id
        for r in records
        if r.record_type in read_types or r.record_type.startswith("treatment_")
    )


# The outcome column spec, appended after the concept columns in the export (`export.py`). Each entry mirrors
# the concept-column descriptor shape (key, label, kind, category, assertion_type) so the same renderer and
# absent-reason pairing apply, plus `sourced_via`/`join_source` marking the block join-sourced and, for the
# one derived column, its rule and inputs.
OUTCOME_COLUMNS: list[dict] = [
    {
        "key": "days_to_last_follow_up",
        "label": "Days to last follow-up",
        "assertion_type": "asserted",
        "note": "Observation/censoring time from the GDC follow-up record.",
    },
    {
        "key": "follow_up_count",
        "label": "Follow-up events",
        "assertion_type": "asserted",
        "note": "Number of GDC follow-up records for the participant.",
    },
    {
        "key": "disease_response_latest",
        "label": "Latest disease response",
        "assertion_type": "asserted",
        "note": "disease_response on the latest dated follow-up (TF-Tumor Free / WT-With Tumor / Unknown).",
    },
    {
        "key": "progression_or_recurrence",
        "label": "Progression or recurrence",
        "assertion_type": "asserted",
        "note": "GDC-asserted flag; true when any follow-up records a progression or recurrence.",
    },
    {
        "key": "recurrence_type",
        "label": "Recurrence type",
        "assertion_type": "asserted",
        "note": "Distant / Locoregional / Unknown, from the follow-up that flagged recurrence.",
    },
    {
        "key": "days_to_recurrence",
        "label": "Days to recurrence",
        "assertion_type": "asserted",
        "note": "Earliest recorded days_to_recurrence or days_to_progression on a recurrence follow-up.",
    },
    {
        "key": "recurrence_observed",
        "label": "Recurrence observed (derived)",
        "assertion_type": "derived",
        "derived": {
            "rule": RECURRENCE_RULE,
            "inputs": ["progression_or_recurrence", "disease_response_latest"],
        },
        "note": "Unifies the explicit recurrence flag with a with-tumour disease response.",
    },
    {
        "key": "received_chemotherapy",
        "label": "Received chemotherapy",
        "assertion_type": "asserted",
        "note": "A GDC chemotherapy treatment record is present.",
    },
    {
        "key": "received_radiation",
        "label": "Received radiation",
        "assertion_type": "asserted",
        "note": "A GDC radiation or brachytherapy treatment record is present.",
    },
    {
        "key": "received_hormone_therapy",
        "label": "Received hormone therapy",
        "assertion_type": "asserted",
        "note": "A GDC hormone-therapy treatment record is present.",
    },
    {
        "key": "days_to_first_treatment",
        "label": "Days to first treatment",
        "assertion_type": "asserted",
        "note": "Earliest days_to_treatment_start across the GDC treatment records.",
    },
    {
        "key": "therapeutic_agents",
        "label": "Therapeutic agents",
        "assertion_type": "asserted",
        "note": "Distinct therapeutic agents named across the GDC treatment records.",
    },
    {
        "key": "vital_status",
        "label": "Vital status",
        "assertion_type": "asserted",
        "note": "From the GDC demographic record: Alive or Dead.",
    },
    {
        "key": "days_to_death",
        "label": "Days to death",
        "assertion_type": "asserted",
        "note": "From the GDC demographic record; not_applicable for a living patient.",
    },
    {
        "key": "overall_survival_days",
        "label": "Overall survival days (derived)",
        "assertion_type": "derived",
        "derived": {
            "rule": OVERALL_SURVIVAL_RULE,
            "inputs": ["vital_status", "days_to_death", "days_to_last_follow_up"],
        },
        "note": "Day of death for a deceased patient, else the last-follow-up day (the survival time axis).",
    },
]

OUTCOME_COLUMN_KEYS: list[str] = [column["key"] for column in OUTCOME_COLUMNS]