src/tools/query_variant_status.py

"""Tool: query_variant_status(genes, source) -> per-sample status + cohort frequency.

Answers "which samples carry a mutation/CNV in these genes, and at what frequency?"
Reads a StatusMatrix (variant_status.build_status_matrix) and summarizes it. Honors the
per-study modality gate: if the cohort has no CNV, `modalities_available['cnv']` is False
and NO cnv figures are emitted — never a null CNV column dressed up as "no alterations".
"""
from __future__ import annotations

import sys
from pathlib import Path

from fastmcp import FastMCP

from ..workflows import lineage
from ..workflows.variant_status import build_status_matrix, load_panel

# repo root on sys.path so the root-level `licenses` module imports under both entry points
# (gradio_ui.py at the root, and `python -m src.server`).
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from licenses import citation_block  # noqa: E402

variant_status_mcp = FastMCP(name="variant_status")

# statuses that count as "altered" for the cohort-frequency numerator.
_MUT_ALTERED = {"missense", "truncating", "hotspot"}
_CNV_ALTERED = {"deep_del", "loss", "gain", "amp"}
# Every non-baseline SV status counts as altered. `rearrangement` is included deliberately: it
# means "an SV touches this gene and the source did not characterize it further", which is still
# an alteration — excluding it would under-report the very cohorts whose annotation is thinnest.
_SV_ALTERED = {"fusion_in_frame", "fusion_out_of_frame", "intragenic", "rearrangement"}

_SV_UNCHARACTERIZED_CAVEAT = (
    "This cohort's structural-variant calls carry NO frame annotation — the source labels events "
    "only as free text (e.g. 'GENEA-GENEB fusion') with `variantClass: NA`. Every event is "
    "therefore reported as `rearrangement`, and this cohort can never return `fusion_in_frame`. "
    "Read an absence of in-frame fusions here as a fact about the annotation, NOT about the "
    "biology. Partner names may also be non-coding transcripts or clone identifiers rather than "
    "genes, so a listed event is not by itself evidence of an actionable fusion."
)


def query_variant_status(genes: list[str], source: str) -> dict:
    """Return per-sample status + cohort frequency for `genes` in `source`.

    An empty `genes` list means the whole driver-gene panel.

    Result:
      {
        source, grounded, n_samples,
        n_profiled: {mutation, cnv},         # the actual denominators; <= n_samples
        samples_profiled: {mutation: [...], cnv: [...]},  # the denominators as ID SETS, so a
                                                          # caller can reconstruct WT from the
                                                          # sparse per_sample below
        modalities_available: {mutation, cnv, expression},
        genes: {
          GENE: {
            mutation: {frequency, n_altered, n_profiled, per_sample: {...}, provenance: {...}}
              (omitted if mutation modality absent). per_sample is SPARSE: altered samples
              only. Recover WT as samples_profiled["mutation"] minus per_sample's keys,
            cnv:      {frequency, n_altered, n_profiled, per_sample: {sample: status}}
              (omitted if cnv modality absent — the coverage gate),
          }, ...
        },
        unavailable_modalities: [ ... ],   # explicit, so a caller can SAY what's missing
        genes_assayed: [ ... ] | null,     # null = exome/genome-wide; a list = targeted panel
        unassayed_genes: [ ... ],          # requested genes this cohort never sequenced
        uncurated_genes: [ ... ],          # requested genes absent from OUR curated artifact
      }

    A gene the cohort's panel does not cover yields ``{"assayed": False, "note": ...}`` instead
    of a frequency — the per-gene half of the coverage gate. A gene the curated artifact does
    not carry yields ``{"curated": False, "note": ...}``, which is a different thing: the
    cohort may well have it, we just have not curated it. Neither is ever silently dropped.
    """
    # Empty `genes` means the full panel, the same as the machine endpoint's blank `genes`
    # (`_split_genes` in gradio_ui.py). The two used to disagree: the endpoint expanded to the
    # panel while calling this function directly with [] returned an empty `genes` map, which
    # reads as "nothing altered" rather than "you asked about nothing".
    if not genes:
        genes = load_panel()

    restrict = None
    if source.startswith("cbioportal:"):
        study = source.split(":", 1)[1]
        if lineage.needs_lineage_filter(study):
            restrict = lineage.pancreatic_sample_ids(study)

    sm = build_status_matrix(source, sample_ids=restrict)

    # A requested gene the ARTIFACT does not carry is a stated gap, never a silent omission.
    # This used to be `[g for g in genes if g in sm.genes]`, which dropped such a gene from the
    # payload entirely: a caller asking for the panel got back a shorter `genes` map with no
    # indication anything was missing, which reads as "we looked and found nothing". It is the
    # same confident-false-negative the `assayed: False` gate exists to prevent, reached through
    # a different door — and it fires exactly when the panel is widened, because every curated
    # artifact bakes in the `panel_genes` it was curated against.
    #
    # Kept DISTINCT from `assayed: False` on purpose. `assayed: False` means the cohort's assay
    # never interrogated the gene — a fact about the data, unfixable here. `curated: False` means
    # WE have not curated it yet — a fact about our snapshot, fixed by re-running
    # `python -m src.curate <study>`. Collapsing the two would tell a user their cohort cannot
    # answer a question that a re-curation would answer fine.
    requested = [g for g in genes if g in sm.genes]
    not_curated = [g for g in genes if g not in sm.genes]
    n = len(sm.samples)
    # Frequency denominators are PER MODALITY and are the profiled subset, not the cohort — a
    # sample the study never sequenced is not wild-type. Both the count and the sample set are
    # kept: the set is what the numerator is computed over, so the two can never drift apart.
    denom = {m: sm.denominator(m) for m in ("mutation", "cnv", "sv")}
    denom_set = {m: set(ids) for m, ids in denom.items()}
    out_genes: dict = {}
    # Per-GENE coverage gate, alongside the per-study modality gate. A targeted cohort
    # (MSK-IMPACT) does not carry every panel gene, and a gene nobody sequenced would otherwise
    # come back "0% altered" — a confident false negative, which is worse than a refusal.
    assayed = set(sm.genes_assayed) if sm.genes_assayed is not None else None
    unassayed: list[str] = []

    for g in requested:
        entry: dict = {}
        # `genes_assayed` is derived from the cohort's DNA gene panel (IMPACT341/468/505…), so it
        # answers "was this gene sequenced for mutations and copy number?" — and ONLY that. A
        # targeted cohort's fusion calling is a separate assay with a different gene set: MSK
        # reports 6 somatic NRG1 rearrangements while NRG1 is absent from the DNA-panel
        # intersection. Applying a mutation-panel restriction to the SV modality therefore
        # suppressed data we actually hold, which is the mirror image of the false negative the
        # gate exists to prevent — refusing an answer we have, rather than inventing one we don't.
        #
        # So the gate is MODALITY-SCOPED. Off the DNA panel: mutation and CNV are refused, while
        # SV is still reported IF events were observed, because an observed event is itself proof
        # the gene was interrogated. A gene off the DNA panel with no observed events stays fully
        # refused — there we have no evidence of interrogation from either direction.
        off_dna_panel = assayed is not None and g not in assayed
        has_sv_modality = bool(sm.modalities.get("sv")) and sm.sv is not None
        sv_events = has_sv_modality and any(v != "none" for v in sm.sv.loc[g])
        if off_dna_panel and not sv_events:
            unassayed.append(g)
            out_genes[g] = {
                "assayed": False,
                "note": (
                    f"{g} is not on this cohort's sequencing panel — it was never interrogated, "
                    "so no status or frequency can be reported. This is absence of measurement, "
                    "not absence of alteration."
                ),
            }
            continue
        if off_dna_panel:
            unassayed.append(g)
            entry["assayed"] = False
            entry["note"] = (
                f"{g} is NOT on this cohort's DNA sequencing panel, so no mutation or copy-number "
                "status can be reported for it. Structural variants ARE reported below: the "
                "fusion assay covers genes the DNA panel does not, and the events observed here "
                "are themselves evidence the gene was interrogated for rearrangements. Do not "
                "read the SV frequency as a mutation frequency."
            )
        if sm.modalities["mutation"] and sm.mutation is not None and not off_dna_panel:
            keep = denom_set["mutation"]
            row = sm.mutation.loc[g]
            per_sample = {s: v for s, v in row.items() if v != "WT" and s in keep}
            n_prof = len(denom["mutation"])
            n_alt = sum(1 for s, v in row.items() if v in _MUT_ALTERED and s in keep)
            entry["mutation"] = {
                "frequency": (n_alt / n_prof) if n_prof else 0.0,
                "n_altered": n_alt,
                "n_profiled": n_prof,  # the denominator, stated — never left to be inferred
                "per_sample": per_sample,
                "provenance": {
                    s: sm.provenance[(g, s)] for s in per_sample if (g, s) in sm.provenance
                },
            }
        if sm.modalities["cnv"] and sm.cnv is not None and not off_dna_panel:  # the coverage gate
            keep = denom_set["cnv"]
            row = sm.cnv.loc[g]
            per_sample = {s: v for s, v in row.items() if v != "neutral" and s in keep}
            n_prof = len(denom["cnv"])
            n_alt = sum(1 for s, v in row.items() if v in _CNV_ALTERED and s in keep)
            entry["cnv"] = {
                "frequency": (n_alt / n_prof) if n_prof else 0.0,
                "n_altered": n_alt,
                "n_profiled": n_prof,
                "per_sample": per_sample,
            }
        if has_sv_modality:  # third gate: five of the seven cohorts have no SV data at all
            keep = denom_set["sv"]
            row = sm.sv.loc[g]
            per_sample = {s: v for s, v in row.items() if v != "none" and s in keep}
            n_prof = len(denom["sv"])
            n_alt = sum(1 for s, v in row.items() if v in _SV_ALTERED and s in keep)
            entry["sv"] = {
                "frequency": (n_alt / n_prof) if n_prof else 0.0,
                "n_altered": n_alt,
                "n_profiled": n_prof,
                "per_sample": per_sample,
                "provenance": {
                    s: sm.sv_provenance[(g, s)] for s in per_sample if (g, s) in sm.sv_provenance
                },
            }
        out_genes[g] = entry

    study_ref = sm.source.split(":", 1)[1] if ":" in sm.source else sm.source
    for g in not_curated:
        out_genes[g] = {
            "curated": False,
            "note": (
                f"{g} is not present in this cohort's curated artifact, so no status or "
                "frequency can be reported. This is a gap in OUR snapshot, not a finding about "
                f"the cohort — the artifact was curated against an older gene panel. Re-run "
                f"`python -m src.curate {study_ref}` to fill it in."
            ),
        }

    unavailable = [m for m, present in sm.modalities.items() if not present]

    return {
        "source": sm.source,
        "grounded": sm.grounded,
        "n_samples": n,
        # n_samples is the COHORT; these are what each frequency is actually over. They differ
        # whenever a study assayed only part of its cohort, which is common and easy to miss.
        "n_profiled": {m: len(ids) for m, ids in denom.items() if sm.modalities.get(m)},
        # The profiled sample UNIVERSE, not just its size. `per_sample` below is sparse — it
        # carries only altered samples — so without this a machine caller cannot tell a WT
        # sample from one the study never sequenced, and cannot reconstruct the WT column at
        # all. (The orchestrator's cross-modality join hit exactly this: every gene scored a
        # WT count of 0 and was reported as "no variation in alteration status".)
        #
        # The safe rule this enables, and the ONLY one callers should use: for a gene with
        # `assayed != False`, a sample in `samples_profiled[modality]` but absent from that
        # gene's `per_sample` is WT (mutation) / neutral (cnv). A sample NOT in this list was
        # never interrogated and must never be counted as either. Inferring "absent ⇒ WT" from
        # the cohort instead would invent 36 WT calls on `paad_tcga` alone, which profiles 150
        # of its 186 samples for mutation.
        #
        # Per MODALITY, not per gene: within an assayed gene the denominator is the modality's
        # profiled set. Genes a targeted panel never carried are refused separately via
        # `assayed: False` / `unassayed_genes`, so they never reach this rule.
        # Sorted for deterministic payloads. Invariant: len(samples_profiled[m]) == n_profiled[m].
        "samples_profiled": {
            m: sorted(ids) for m, ids in denom.items() if sm.modalities.get(m)
        },
        "modalities_available": sm.modalities,
        "unavailable_modalities": unavailable,
        # None = genome/exome-wide cohort (no per-gene restriction); a list = targeted panel.
        "genes_assayed": sm.genes_assayed,
        "unassayed_genes": unassayed,  # explicit, so a caller can SAY what was not measured
        # Requested genes missing from the curated artifact — our gap, not the cohort's. Kept
        # separate from `unassayed_genes` because the remedy differs: re-curate vs. cannot be
        # answered by this cohort at all.
        "uncurated_genes": not_curated,
        # "characterized" | "uncharacterized" | None (cohort has no SV modality). The SV analogue
        # of `genes_assayed`: it says how much the source's fusion calls can be asked to support.
        "sv_annotation": sm.sv_annotation if sm.modalities.get("sv") else None,
        "caveats": (
            [_SV_UNCHARACTERIZED_CAVEAT]
            if sm.modalities.get("sv") and sm.sv_annotation == "uncharacterized"
            else []
        ),
        "genes": out_genes,
        # C1 (ADR-0005 / cBioPortal terms): attribution rides the response contract itself.
        "citation": citation_block(
            sm.source, getattr(sm, "terms", None), getattr(sm, "curation", None)
        ),
    }


# Name pinned explicitly: the registered MCP tool name is part of the published contract
# (`deploy/orchestrator_registration.yaml`), so it must not inherit the wrapper's Python name.
@variant_status_mcp.tool(name="query_variant_status")
def query_variant_status_tool(genes: list[str], source: str) -> dict:
    """Per-sample mutation/CNV status + cohort frequency for `genes` in a cohort.

    `source` is a cBioPortal study reference, e.g. ``"cbioportal:paad_tcga"``. A cohort with
    no CNV profile yields NO cnv figures (the per-study coverage gate) and names the absence
    in ``unavailable_modalities``.
    """
    return query_variant_status(genes, source)