app/viewer.py

"""OncoDSL sanity viewer.

Run with:
    streamlit run app/viewer.py

Two top-level tabs (Glossary expander sits above both, available everywhere):

- **Dataset** — the five sanity panels: cohort overview, MSI counts, TMB-vs-MSI
  histogram (the load-bearing sanity check), stage/age/missing fields, and a
  small expression-matrix peek.
- **Hypothesis 1** — verification of the H1 program on the named matrix: a
  short intro, the DSL reference (nouns + verbs), the composition diagram,
  and the H1 outputs (conclusion → score distributions → Effect → Fit).

Presentation (palette, fonts, chart theme) lives in app/theme.py and
.streamlit/config.toml. Importing `theme` registers and enables the Altair theme
used by every chart in this file.
"""

from __future__ import annotations

import os
import sys
from pathlib import Path

import altair as alt
import numpy as np
import pandas as pd
import requests
import streamlit as st

# Allow `streamlit run app/viewer.py` from the repo root without installing the
# package: prepend the repo root to sys.path so `import data_pipeline` works.
REPO_ROOT = Path(__file__).resolve().parent.parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from data_pipeline import schema  # noqa: E402
from app import theme  # noqa: E402  — registers the "oncodsl" Altair theme on import.
from airgap import anonymise, reveal  # noqa: E402
from validate.h1 import IMMUNE_GENES, MMR_GENES, run_h1, usable_msi_cohort  # noqa: E402
from dsl import Apply, Load  # noqa: E402

CHART_HEIGHT = 300

st.set_page_config(page_title="OncoDSL — TCGA CRC sanity viewer", layout="wide")


# --- data loading -----------------------------------------------------------

@st.cache_data(show_spinner=False)
def load_clinical() -> pd.DataFrame | None:
    p = schema.PROCESSED_DIR / "clinical.parquet"
    if not p.exists():
        return None
    return pd.read_parquet(p)


@st.cache_data(show_spinner=False)
def load_expression_shape_and_sample() -> tuple[tuple[int, int], pd.DataFrame] | None:
    p = schema.PROCESSED_DIR / "expression.parquet"
    if not p.exists():
        return None
    df = pd.read_parquet(p)
    sample = df.iloc[:20, :10]
    return df.shape, sample


@st.cache_resource(show_spinner=False)
def _h1_result():
    """Run the full H1 program against the processed cohort."""
    return run_h1()


@st.cache_data(show_spinner=False)
def _airgap_demo() -> tuple[pd.DataFrame, list[str], list[str]] | None:
    """Show the anonymised matrix head + reveal() round-trip."""
    p = schema.PROCESSED_DIR / "expression.parquet"
    if not p.exists():
        return None
    cohort = Load("processed")
    anonymised = anonymise(cohort.expression)
    head = anonymised.iloc[:5, :8]
    revealed = reveal(list(head.columns))
    return head, list(head.columns), revealed


def usable_mask(clinical: pd.DataFrame) -> pd.Series:
    return (
        clinical["msi_status"].isin(["MSI-H", "MSS"])
        & clinical["has_expression"]
        & (clinical["stage"] != "NA")
        & clinical["age"].notna()
    )


# --- glossary ---------------------------------------------------------------

def _glossary_md() -> str:
    hi = schema.MSI_SENSOR_HIGH
    lo = schema.MSI_SENSOR_LOW
    return (
        "- **TCGA** — The Cancer Genome Atlas: a large public dataset of tumour genetics.\n"
        "- **Adenocarcinoma** — cancer of gland-forming cells (the usual colorectal type).\n"
        "- **COAD** — Colon Adenocarcinoma (cancer of the colon).\n"
        "- **READ** — Rectum Adenocarcinoma (cancer of the rectum).\n"
        "- **MACR** — Mucinous Adenocarcinoma of the Colon and Rectum: a mucinous tissue "
        "SUBTYPE that can be colon or rectum (so this field is tumour TYPE, not pure "
        "anatomical site).\n"
        "- **MSI** — Microsatellite Instability: instability in short repeating DNA "
        "stretches, caused by a failed DNA \"spell-checker\" (mismatch repair).\n"
        "- **MSI-H (MSI-High)** — many unstable microsatellites; broken repair; high "
        "mutation load.\n"
        "- **MSS (Microsatellite Stable)** — repair working; a clean, stable genome.\n"
        "- **MSI-Indeterminate** — between the thresholds; not clearly classifiable.\n"
        "- **NA** — no MSI value available.\n"
        f"- **MSIsensor score** — a computed measure of how many microsatellites differ "
        f"between the tumour and the patient's own normal DNA; higher = more unstable. "
        f"This viewer's labels use: **≥ {hi} → MSI-H**, **< {lo} → MSS**, "
        f"**between {lo} and {hi} → Indeterminate** (thresholds read from the data layer).\n"
        "- **TMB** — Tumour Mutational Burden: roughly, the number of mutations in the "
        "tumour. MSI-H tumours have high TMB (the key sanity check).\n"
        "- **Stage** — how far the cancer has spread (I–IV).\n"
        "- **RSEM** — the method used to turn raw RNA-sequencing reads into the "
        "per-gene expression numbers shown here. The values are relative (good "
        "for comparing across patients), not absolute counts.\n"
        "\n"
        "**Biology you'll meet in later steps**\n"
        "- **MMR (mismatch repair)** — the DNA spell-checker; genes MLH1, MSH2, MSH6, "
        "PMS2. dMMR = broken, pMMR = working.\n"
        "- **CD8A, GZMA, PRF1** — markers that \"killer\" T-cells are attacking the "
        "tumour.\n"
        "- **PD-L1 / CD274** — a \"shield\" protein tumours use to switch off the "
        "immune attack.\n"
        "- **Checkpoint inhibitor (anti-PD-1)** — a drug that removes the shield so "
        "the immune system can attack.\n"
    )


# --- main -------------------------------------------------------------------

_EXPANDER_CSS = """
<style>
/* Streamlit 1.58 expander testids verified against compiled bundle:
   stExpander (container) / stExpanderDetails (body) / stExpanderIcon (chevron).
   We do NOT style stExpanderIcon — keep the native chevron visible. */
[data-testid="stExpander"] {
    margin-top: 0.4rem;
}
[data-testid="stExpander"] details {
    border: 1px solid #ECEAE4;
    border-radius: 8px;
    background-color: #FCFBF8;
}
[data-testid="stExpander"] summary {
    color: #5A6670;
    font-size: 0.9rem;
    font-weight: 500;
}
</style>
"""


def main() -> None:
    st.markdown(_EXPANDER_CSS, unsafe_allow_html=True)
    st.title("OncoDSL — TCGA CRC sanity viewer")
    st.caption(
        f"Source: cBioPortal study `{schema.STUDY}` "
        f"(TCGA Colorectal Adenocarcinoma, PanCancer Atlas)."
    )

    # Glossary sits ABOVE the tabs so it's available everywhere.
    with st.expander("Glossary — plain-language key", expanded=False):
        st.markdown(_glossary_md())

    clinical = load_clinical()
    expr_loaded = load_expression_shape_and_sample()
    if clinical is None or expr_loaded is None:
        st.error(
            "Processed data not found. Run from the repo root:\n\n"
            "```\n"
            "python -m data_pipeline.download\n"
            "python -m data_pipeline.build\n"
            "```"
        )
        return

    tab_dataset, tab_h1, tab_h2 = st.tabs(
        ["Dataset", "Hypothesis 1", "Hypothesis 2"]
    )
    with tab_dataset:
        _section_dataset(clinical, expr_loaded)
    with tab_h1:
        _section_h1_intro()
        st.divider()
        _section_dsl()
        st.divider()
        _section_h1_walkthrough()
        st.divider()
        _section_h1()
        st.divider()
        _section_airgap_demo()
    with tab_h2:
        _section_h2()


def _section_dataset(clinical: pd.DataFrame, expr_loaded) -> None:
    (n_genes, n_expr_samples), expr_sample = expr_loaded

    # ----- Panel 1: dataset overview -----
    st.header("1. Dataset overview")
    st.caption(
        "How big the cohort is, and how many samples are usable for downstream analysis."
    )
    n_total = len(clinical)
    n_coad = int((clinical["site"] == "COAD").sum())
    n_read = int((clinical["site"] == "READ").sum())
    n_macr = int((clinical["site"] == "MACR").sum())
    usable = usable_mask(clinical)
    n_usable = int(usable.sum())

    c1, c2, c3, c4, c5 = st.columns(5)
    c1.metric("Total samples", n_total)
    c2.metric(
        "COAD", n_coad,
        help="Colon Adenocarcinoma — cancer of the colon.",
    )
    c3.metric(
        "READ", n_read,
        help="Rectum Adenocarcinoma — cancer of the rectum.",
    )
    c4.metric(
        "MACR", n_macr,
        help="Mucinous Adenocarcinoma of the Colon and Rectum — a mucinous tissue "
             "subtype that can be colon or rectum.",
    )
    c5.metric(
        "Usable cohort", n_usable,
        help="Samples with an MSI label in {MSI-H, MSS} AND expression AND non-missing "
             "stage AND non-missing age — the rows downstream analysis can use.",
    )
    st.caption(
        f"Genes in expression matrix: **{n_genes}** · expression samples: "
        f"**{n_expr_samples}** · MSI label derived from MSIsensor "
        f"(≥{schema.MSI_SENSOR_HIGH} → MSI-H, <{schema.MSI_SENSOR_LOW} → MSS, "
        "between → Indeterminate — cBioPortal does not ship a clean MSI-H/MSS column)."
    )

    with st.expander("How to read this", expanded=False):
        st.markdown(
            "- This is the whole colorectal cohort (colon + rectum + a mucinous "
            "subtype), pooled on purpose — more data helps when samples are scarce.\n"
            "- READ (155) is the rectal slice; we hold rectum back for validation later.\n"
            "- MACR are mucinous tumours that can be colon or rectum, so they're "
            "counted separately."
        )

    st.divider()

    # ----- Panel 2: MSI counts -----
    st.header("2. MSI status counts")
    st.caption(
        "How many have a broken DNA spell-checker (MSI-H) vs a working one (MSS)."
    )

    overall = (
        clinical["msi_status"]
        .value_counts()
        .reindex(theme.MSI_ORDER, fill_value=0)
        .rename_axis("msi_status")
        .reset_index(name="count")
    )
    by_site = (
        clinical.groupby(["site", "msi_status"]).size().reset_index(name="count")
    )

    left, right = st.columns(2)
    with left:
        st.subheader("Overall")
        st.altair_chart(
            alt.Chart(overall)
            .mark_bar()
            .encode(
                x=alt.X("msi_status:N", sort=theme.MSI_ORDER, title="MSI status"),
                y=alt.Y("count:Q", title="samples"),
                color=alt.Color(
                    "msi_status:N",
                    sort=theme.MSI_ORDER,
                    scale=theme.msi_color_scale(),
                    legend=None,
                ),
                tooltip=["msi_status", "count"],
            )
            .properties(height=CHART_HEIGHT),
            width="stretch",
        )
    with right:
        st.subheader("By tumour type")
        st.caption(
            "COAD = colon, READ = rectum, MACR = mucinous subtype (colon or rectum)."
        )
        st.altair_chart(
            alt.Chart(by_site)
            .mark_bar()
            .encode(
                x=alt.X("site:N", title="tumour type"),
                y=alt.Y("count:Q", stack="zero", title="samples"),
                color=alt.Color(
                    "msi_status:N",
                    sort=theme.MSI_ORDER,
                    scale=theme.msi_color_scale(),
                    legend=alt.Legend(title="MSI status", orient="right"),
                ),
                order=alt.Order(
                    "msi_status_order:Q", sort="ascending",
                ),
                tooltip=["site", "msi_status", "count"],
            )
            .transform_calculate(
                msi_status_order=(
                    "indexof(['" + "','".join(theme.MSI_ORDER) + "'], datum.msi_status)"
                )
            )
            .properties(height=CHART_HEIGHT),
            width="stretch",
        )

    with st.expander("How to read this", expanded=False):
        thresholds = (
            f"≥ {schema.MSI_SENSOR_HIGH} → MSI-H, < {schema.MSI_SENSOR_LOW} → MSS, "
            f"between → Indeterminate"
        )
        st.markdown(
            "- All these are colorectal cancers. MSI-H vs MSS is a subtype, not "
            "cancer-vs-no-cancer.\n"
            "- MSS (the majority, ~85%) = the MMR \"spell-checker\" is working: "
            "stable repeats, relatively few mutations — still cancer, driven by "
            "other routes.\n"
            "- MSI-H = spell-checker broken: unstable repeats, many mutations.\n"
            "- These labels come from the MSIsensor score (instability measured "
            "directly in the DNA repeats), NOT from the MMR genes' expression — "
            "that independence is what makes a later rediscovery meaningful.\n"
            f"- The labels are our interpretation of that score: **{thresholds}** "
            "(thresholds read from the data layer)."
        )

    st.divider()

    # ----- Panel 2½: the biology in one picture -----
    _section_biology_picture()

    st.divider()

    # ----- Panel 3: TMB vs MSI — the load-bearing sanity check -----
    st.header("3. TMB vs MSI (sanity check)")
    st.caption(
        "Broken-repair (MSI-H) tumours should carry far more mutations — they should "
        "sit to the right."
    )
    tmb_df = clinical.loc[
        clinical["msi_status"].isin(["MSI-H", "MSS"]) & clinical["tmb"].notna(),
        ["msi_status", "tmb"],
    ].copy()
    # TMB is heavy-tailed; log1p the x axis so the MSI-H tail is readable.
    tmb_df["tmb_log1p"] = np.log1p(tmb_df["tmb"])

    hist = (
        alt.Chart(tmb_df)
        .mark_bar(opacity=0.85)
        .encode(
            x=alt.X(
                "tmb_log1p:Q",
                bin=alt.Bin(maxbins=40),
                title="log1p(TMB nonsynonymous)",
            ),
            y=alt.Y("count():Q", stack=None, title="samples"),
            color=alt.Color(
                "msi_status:N",
                sort=theme.MSI_ORDER,
                scale=theme.msi_color_scale(include_na=False),
                legend=alt.Legend(title="MSI status", orient="right"),
            ),
            tooltip=["msi_status", "count()"],
        )
        .properties(height=CHART_HEIGHT)
    )
    st.altair_chart(hist, width="stretch")

    medians = tmb_df.groupby("msi_status")["tmb"].median().to_dict()
    st.caption(
        f"Median TMB (nonsynonymous) — MSI-H: **{medians.get('MSI-H', float('nan')):.2f}**, "
        f"MSS: **{medians.get('MSS', float('nan')):.2f}**. "
        "If MSI-H median isn't markedly higher than MSS, something is wrong upstream."
    )

    with st.expander("How to read this", expanded=False):
        st.markdown(
            "- The x-axis is a log mutation count; to translate back, "
            "mutations ≈ e^x − 1.\n"
            "- MSI-H tumours sit to the right (hypermutated); MSS to the left "
            "(quiet) — that clean separation is the sanity check that data, "
            "labels and counts are wired up correctly.\n"
            "- Minor exception: a few MSS tumours with a broken POLE polymerase "
            "are also hypermutated."
        )

    st.divider()

    # ----- Panel 4: stage / age / missing -----
    st.header("4. Stage, age, and missing-field counts")
    st.caption(
        "Stage = how far the cancer has spread (I–IV); age at diagnosis; and where "
        "fields are missing across the cohort. Bars are split by MSI status so we "
        "can eyeball confounding."
    )

    # Restrict to MSI-H vs MSS only for the confounder views.
    h_vs_s = clinical[clinical["msi_status"].isin(["MSI-H", "MSS"])].copy()

    left, mid, right = st.columns([1, 1, 1])
    with left:
        st.subheader("Stage — proportion MSI-H")
        stage_order = ["I", "II", "III", "IV", "NA"]
        grouped = h_vs_s.groupby("stage")["msi_status"]
        stage_df = pd.DataFrame({
            "stage": list(grouped.groups.keys()),
            "n":     grouped.size().values,
            "n_msi_h": grouped.apply(lambda s: int((s == "MSI-H").sum())).values,
        })
        stage_df["pct_msi_h"] = stage_df["n_msi_h"] / stage_df["n"]
        stage_df = (
            stage_df.set_index("stage")
            .reindex(stage_order)
            .reset_index()
            .fillna({"n": 0, "n_msi_h": 0, "pct_msi_h": 0})
        )
        st.altair_chart(
            alt.Chart(stage_df)
            .mark_bar(color=theme.MSI_COLORS["MSI-H"])
            .encode(
                x=alt.X("stage:N", sort=stage_order, title="stage"),
                y=alt.Y(
                    "pct_msi_h:Q",
                    title="% MSI-H",
                    axis=alt.Axis(format="%"),
                    scale=alt.Scale(domain=[0, 1]),
                ),
                tooltip=[
                    alt.Tooltip("stage:N"),
                    alt.Tooltip("n:Q", title="n (MSI-H + MSS)"),
                    alt.Tooltip("n_msi_h:Q", title="n MSI-H"),
                    alt.Tooltip("pct_msi_h:Q", title="% MSI-H", format=".1%"),
                ],
            )
            .properties(height=CHART_HEIGHT),
            width="stretch",
        )
    with mid:
        st.subheader("Age — MSI-H vs MSS (density)")
        age_df = h_vs_s.loc[h_vs_s["age"].notna(), ["age", "msi_status"]]
        age_min = float(age_df["age"].min())
        age_max = float(age_df["age"].max())
        st.altair_chart(
            alt.Chart(age_df)
            .transform_density(
                "age",
                as_=["age", "density"],
                groupby=["msi_status"],
                extent=[age_min, age_max],
                steps=80,
            )
            .mark_area(opacity=0.8)
            .encode(
                x=alt.X("age:Q", title="age at diagnosis"),
                y=alt.Y("density:Q", title="density", stack=None),
                color=alt.Color(
                    "msi_status:N",
                    sort=theme.MSI_ORDER,
                    scale=theme.msi_color_scale(include_na=False),
                    legend=alt.Legend(title="MSI status", orient="right"),
                ),
                tooltip=[
                    alt.Tooltip("msi_status:N"),
                    alt.Tooltip("age:Q", format=".1f"),
                    alt.Tooltip("density:Q", format=".3f"),
                ],
            )
            .properties(height=CHART_HEIGHT),
            width="stretch",
        )
    with right:
        st.subheader("Missing per field")
        st.caption(
            "For each field, how many samples lack a value — this is what trims the "
            "full cohort down to the usable set. Small numbers here are normal."
        )
        missing = pd.DataFrame(
            {
                "field": [
                    "msi_status (NA)",
                    "msi_status (Indeterm.)",
                    "tmb",
                    "age",
                    "stage (NA)",
                    "sex (blank)",
                    "os_event",
                    "os_months",
                    "no expression",
                ],
                "missing": [
                    int((clinical["msi_status"] == "NA").sum()),
                    int((clinical["msi_status"] == "MSI-Indeterminate").sum()),
                    int(clinical["tmb"].isna().sum()),
                    int(clinical["age"].isna().sum()),
                    int((clinical["stage"] == "NA").sum()),
                    int(clinical["sex"].isna().sum()
                        + (clinical["sex"].astype(str) == "").sum()),
                    int(clinical["os_event"].isna().sum()),
                    int(clinical["os_months"].isna().sum()),
                    int((~clinical["has_expression"]).sum()),
                ],
            }
        )
        st.dataframe(missing, hide_index=True, width="stretch")

    with st.expander("How to read this", expanded=False):
        st.markdown(
            "**Stage and age**\n"
            "- Stage and age are the background factors (\"confounders\") the "
            "Effect operator will hold constant.\n"
            "- If MSI-H differs from MSS here (e.g. older or earlier-stage), "
            "that's real confounding — which is exactly why we adjust for it "
            "rather than trust raw correlations.\n"
            "\n"
            "**Missing per field**\n"
            "- This is a data-quality check, not a problem — these are small "
            "counts against the full cohort.\n"
            "- **msi_status (NA)** = no MSIsensor score, so no MSI label; "
            "**(Indeterm.)** = score fell between the thresholds, so it can't "
            "be called MSI-H or MSS.\n"
            "- **tmb** = missing mutation-burden value. This only affects the "
            "TMB chart; the MSI label does not depend on TMB.\n"
            "- **age / stage (NA) / sex (blank)** = missing clinical fields.\n"
            "- **os_event / os_months** = survival fields (os = overall "
            "survival; event = whether the patient died; months = length of "
            "follow-up). Used in later steps, not now.\n"
            "- **no expression** = the sample has no RNA-expression data.\n"
            "- The \"usable cohort\" keeps only samples that have everything "
            "the core task needs: a clear MSI label (MSI-H or MSS), expression "
            "data, and the confounders (stage and age). The missing values "
            "above are mostly what got dropped to reach that usable count."
        )

    st.divider()

    # ----- Panel 5: expression sanity (small table, NOT a heatmap) -----
    st.header("5. Expression matrix sanity")
    st.caption(
        "A small slice of the gene-by-sample expression matrix — enough to confirm "
        "the numbers look reasonable, not a full heatmap."
    )
    st.write(f"Shape: **{n_genes} genes × {n_expr_samples} samples** (RSEM).")
    st.write(
        "First 20 genes × 10 samples (raw values), and per-gene summary stats "
        "across those 10 samples:"
    )
    st.dataframe(expr_sample, width="stretch")
    st.dataframe(
        expr_sample.T.describe().T[["count", "mean", "std", "min", "50%", "max"]],
        width="stretch",
    )

    with st.expander("How to read this", expanded=False):
        st.markdown(
            "**Raw values table**\n"
            "- This is a tiny corner of the full gene-expression matrix the engine "
            "will later search — shown only to confirm it loaded.\n"
            "- Rows = genes (Hugo_Symbol). Columns = patients (TCGA-… barcodes, one "
            "tumour each). Each cell = how active that gene is in that tumour (a "
            "relative RNA-expression value): higher = more active, 0 = effectively off.\n"
            "- Many zeros are normal — in any tumour a large fraction of genes are "
            "simply switched off (pseudogenes and tissue-specific genes especially). "
            "The well-known genes (MLH1, CD8A…) sit elsewhere in the matrix.\n"
            "- The number isn't meaningful on its own; what matters is how it "
            "differs across patients.\n"
            "\n"
            "**Per-gene summary stats table**\n"
            "- Summarises each gene across the sample columns shown.\n"
            "- count = samples with a value; mean = average; std = how much it "
            "varies; min / max = range; 50% = median.\n"
            "- A gene that barely varies (std ≈ 0) can't help tell patients apart; "
            "the signal lives in genes that differ across patients.\n"
            "- These stats cover only the small peek shown here — illustrative, not "
            "the full-cohort numbers."
        )


# --- Dataset tab: biology in one picture -----------------------------------

# Two parallel stories side-by-side. MSI-H (warm) sits on top: broken
# spell-checker -> MMR genes off, mutations pile up, immune system attracted.
# MSS (cool) below: working spell-checker, few mutations, immune-cold; with
# a "note"-shaped reminder that MSS is still cancer, just driven by other
# mechanisms.
_BIOLOGY_DOT = """\
digraph biology {
    rankdir=LR;
    bgcolor="transparent";
    nodesep=0.45;
    ranksep=0.9;
    splines=spline;
    fontname="Helvetica";

    node [shape=box, style="rounded,filled", penwidth=1.1,
          fontname="Helvetica", fontcolor="#23303A", fontsize=12,
          margin="0.22,0.14", color="#3A6B7E", fillcolor="#F4F2EE"];
    edge [color="#C2CACF", penwidth=1.3, arrowsize=0.7,
          fontname="Helvetica", fontsize=11, fontcolor="#5A6670"];

    msih      [label="MSI-H\\nbroken spell-checker (dMMR)",
               fillcolor="#FBEFE2", color="#BC6B2E"];
    msih_mmr  [label="MMR genes LOW\\nMLH1 · MSH2 · MSH6 · PMS2"];
    msih_mut  [label="mutations pile up\\n→ neoantigens"];
    msih_imm  [label="immune markers HIGH\\nCD8A · GZMA · PRF1"];
    msih -> msih_mmr [label="switched off"];
    msih -> msih_mut;
    msih_mut -> msih_imm [label="immune-hot"];

    mss       [label="MSS\\nworking spell-checker",
               fillcolor="#EAF0F2", color="#3A6B7E"];
    mss_mmr   [label="MMR genes normal / HIGH"];
    mss_mut   [label="few mutations\\n(stable genome)"];
    mss_imm   [label="immune markers LOW"];
    mss_other [label="cancer driven by OTHER\\nmechanisms (not broken MMR)",
               shape=note, fillcolor="#EFEFEA", color="#9AA0A6"];
    mss -> mss_mmr [label="active"];
    mss -> mss_mut;
    mss_mut -> mss_imm [label="immune-cold"];
    mss -> mss_other;
}
"""


def _section_biology_picture() -> None:
    st.header("The biology in one picture")
    st.caption(
        "The whole story: the spell-checker's state drives both the MMR-gene "
        "activity we measure and (via mutations) the immune response."
    )
    st.graphviz_chart(_BIOLOGY_DOT, width="stretch")
    st.caption(
        "MMR-gene activity is our readout of the spell-checker — low activity "
        "≈ broken."
    )


# --- Hypothesis 1 tab -------------------------------------------------------

_H1_INTRO_MD = (
    "**Hypothesis 1 — Verification.** Before trusting the system to discover "
    "anything, we check the tooling on a known answer. We hand it the established "
    "MSI genes (we don't make it find them) and confirm the operators reproduce "
    "the textbook biology: MSI-H tumours have low MMR-gene activity and high "
    "immune activity, the MMR→immune link survives adjusting for stage and age, "
    "and the two scores tell MSI-H from MSS apart. A strong result here isn't "
    "a discovery — it's calibration. The blind discovery (Hypothesis 2) comes "
    "next."
)


def _section_h1_intro() -> None:
    st.markdown(_H1_INTRO_MD)


# --- DSL panel --------------------------------------------------------------

_DSL_NOUNS = [
    ("Cohort",     "a set of patients with their data"),
    ("Matrix",     "patients × genes"),
    ("Vector",     "one number per patient (a score)"),
    ("FeatureSet", "a chosen set of genes"),
    ("Scalar",     "a single number"),
    ("Outcome",    "the thing predicted (e.g. MSI status)"),
]

_DSL_VERBS = [
    ("Load",        "read the cohort (expression + clinical + labels) into one shape"),
    ("Select",      "pick a subset of the matrix's columns by their labels"),
    ("Reduce",      "collapse the chosen columns to one score per patient (mean)"),
    ("Split",       "partition the cohort by a per-patient predicate"),
    ("Associate",   "observed correlation between two per-patient series"),
    ("Effect",      "partial correlation after holding the listed confounders constant"),
    ("Search",      "rank features by an objective and return the top-k (placeholder)"),
    ("Fit / Apply", "logistic regression to MSI-H vs MSS; Apply gives probabilities"),
]

_H1_PROGRAM_CODE = """\
mmr_score    = Reduce(Select(M, MMR), "mean")
immune_score = Reduce(Select(M, IMMUNE), "mean")

effect = Effect(mmr_score, immune_score, adjust={stage, age})
fit    = Fit((mmr_score, immune_score), msi_h_label)
"""

# Composition diagram. Operator nodes are stone-filled; the two SCORE nodes
# (MMR score, immune score) are cream "meta-concepts"; the two OUTPUT nodes
# (Effect, Fit) are cool-tinted to mark them as the program's results; the
# two "Given:" note-shaped nodes make explicit that we supply the gene panels
# in H1 (we don't make the system discover them).
_H1_DOT = """\
digraph H1 {
    rankdir=LR;
    bgcolor="transparent";
    nodesep=0.5;
    ranksep=1.0;
    splines=spline;
    fontname="Helvetica";

    node [shape=box, style="rounded,filled", penwidth=1.2,
          fontname="Helvetica", fontcolor="#23303A", fontsize=12,
          margin="0.22,0.14", color="#3A6B7E", fillcolor="#F4F2EE"];
    edge [color="#C2CACF", penwidth=1.3, arrowsize=0.7];

    GivenMMR [label="Given: MLH1, MSH2, MSH6, PMS2",
              shape=note, fillcolor="#EFEFEA", color="#9AA0A6"];
    GivenIMM [label="Given: CD8A, GZMA, PRF1",
              shape=note, fillcolor="#EFEFEA", color="#9AA0A6"];

    M       [label="Expression matrix"];
    SelMMR  [label="Select(MMR genes)"];
    RedMMR  [label="Reduce(mean)"];
    MMR     [label="MMR score", fillcolor="#FBEFE2", color="#BC6B2E"];
    SelI    [label="Select(immune genes)"];
    RedI    [label="Reduce(mean)"];
    IMM     [label="immune score", fillcolor="#FBEFE2", color="#BC6B2E"];
    Eff     [label="Effect(adjust: stage, age)",
             fillcolor="#EAF0F2", color="#3A6B7E"];
    Ft      [label="Fit → MSI-H probability",
             fillcolor="#EAF0F2", color="#3A6B7E"];

    GivenMMR -> SelMMR;
    GivenIMM -> SelI;
    M -> SelMMR -> RedMMR -> MMR;
    M -> SelI   -> RedI   -> IMM;
    MMR -> Eff;
    IMM -> Eff;
    MMR -> Ft;
    IMM -> Ft;
}
"""


def _section_dsl() -> None:
    st.header("The DSL")
    st.caption(
        "Two short labelled lists: the kinds of thing the engine works with, "
        "and the operations it composes."
    )

    left, right = st.columns(2)
    with left:
        st.markdown(
            "**Nouns — the kinds of thing**\n\n"
            + "\n".join(f"- **{n}** — {d}" for n, d in _DSL_NOUNS)
        )
    with right:
        st.markdown(
            "**Verbs — the operations**\n\n"
            + "\n".join(f"- **{v}** — {d}" for v, d in _DSL_VERBS)
        )

    st.markdown("&nbsp;")  # gentle vertical breathing room
    st.markdown(
        "**The H1 program** — we supply the known genes and check the "
        "operators reproduce the biology."
    )
    st.graphviz_chart(_H1_DOT, width="stretch")

    st.markdown("**The H1 program — as a one-line composition**")
    st.code(_H1_PROGRAM_CODE, language="python")


# --- H1 walkthrough: the operators on 6 real, held-out patients ------------

_WALKTHROUGH_N_PER_GROUP = 3


@st.cache_resource(show_spinner=False)
def _usable_cohort_for_walkthrough():
    """The same usable MSI-H-vs-MSS cohort the H1 program runs on (cached)."""
    return usable_msi_cohort(Load("processed"))


@st.cache_data(show_spinner=False)
def _walkthrough_tables():
    """Build the three step tables for 3 MSI-H + 3 MSS held-out patients.

    The patients are pulled from `res.fit.test_index` so their predicted
    probabilities are genuinely out-of-sample (they were never seen during
    fitting). Picking deterministically (the first 3 of each class in the
    test set) keeps the panel stable across reloads.
    """
    res = _h1_result()
    cohort = _usable_cohort_for_walkthrough()
    M = cohort.expression
    status = res.msi_status

    test_ids = list(res.fit.test_index)
    msi_h = [s for s in test_ids if status.loc[s] == "MSI-H"][:_WALKTHROUGH_N_PER_GROUP]
    mss = [s for s in test_ids if status.loc[s] == "MSS"][:_WALKTHROUGH_N_PER_GROUP]
    chosen = msi_h + mss

    mmr_vals = M.loc[chosen, MMR_GENES].round(1)
    immune_vals = M.loc[chosen, IMMUNE_GENES].round(1)
    mmr_score = res.mmr_score.loc[chosen].round(1)
    immune_score = res.immune_score.loc[chosen].round(1)

    step1 = pd.DataFrame(index=chosen)
    step1.insert(0, "MSI", status.loc[chosen].values)
    for g in MMR_GENES:
        step1[g] = mmr_vals[g].values
    step1["MMR score"] = mmr_score.values
    for g in IMMUNE_GENES:
        step1[g] = immune_vals[g].values
    step1["Immune score"] = immune_score.values
    step1.index.name = "patient"

    step2 = pd.DataFrame(
        {
            "MSI": status.loc[chosen].values,
            "MMR score": mmr_score.values,
            "Immune score": immune_score.values,
        },
        index=chosen,
    )
    step2.index.name = "patient"

    state6 = pd.DataFrame(
        {"mmr_score": res.mmr_score.loc[chosen],
         "immune_score": res.immune_score.loc[chosen]}
    )
    probs = Apply(res.fit, state6)
    step3 = pd.DataFrame(
        {
            "MMR score": mmr_score.values,
            "Immune score": immune_score.values,
            "predicted MSI-H probability": probs.values.round(3),
            "actual MSI status": status.loc[chosen].values,
        },
        index=chosen,
    )
    step3.index.name = "patient"
    return step1, step2, step3


def _section_h1_walkthrough() -> None:
    st.header("Working through the program on real patients")
    st.markdown(
        f"Here are {_WALKTHROUGH_N_PER_GROUP} MSI-H and "
        f"{_WALKTHROUGH_N_PER_GROUP} MSS patients **from the held-out test "
        "set** — the model never saw them during fitting. We follow the same "
        "operators through, one step at a time."
    )

    res = _h1_result()
    step1, step2, step3 = _walkthrough_tables()

    st.subheader("Step 1 — Select + Reduce")
    st.dataframe(step1, width="stretch")
    st.markdown(
        "**Select** pulls out these genes' values; **Reduce** averages them "
        "into one score per patient. Notice the MSI-H patients tend to have a "
        "lower MMR score and a higher immune score."
    )

    st.subheader("Step 2 — Effect")
    st.dataframe(step2, width="stretch")
    st.markdown(
        f"**Effect** asks whether a lower MMR score actually drives a higher "
        f"immune score — measured across ALL patients, holding stage and age "
        f"constant. Here that adjusted correlation is "
        f"**{res.effect.partial_corr:+.3f}**: close to zero, so the direct "
        f"link between the two scores is weak."
    )

    st.subheader("Step 3 — Fit")
    st.dataframe(step3, width="stretch")
    st.markdown(
        f"**Fit** learns to turn the two scores into a probability that a "
        f"patient is MSI-H. Here are its predictions vs the truth for these "
        f"held-out patients. Across all held-out patients it scores "
        f"AUROC **{res.fit.auroc:.3f}**."
    )


# --- H1 verification panel --------------------------------------------------

_AUROC_THRESHOLD = 0.75

_HOW_TO_READ_H1 = (
    "- **Score distributions:** each curve is how a score is spread across "
    "one group. We want MSI-H (amber) shifted LOW on MMR score and HIGH on "
    "immune score vs MSS — broken repair plus an immune-hot tumour. Medians "
    "are noted above each chart.\n"
    "    - The y-axis (density) is a smoothed, normalised histogram — how "
    "common a score is within each group. Each curve's area sums to 1, so "
    "the two groups can be compared fairly despite very different sizes "
    "(MSS has far more patients than MSI-H). The number itself isn't "
    "meaningful; what matters is where each curve peaks and how the shapes "
    "differ.\n"
    "- **Effect:** does a low MMR score drive a high immune score? "
    "\"Unadjusted\" is the raw correlation; \"adjusted\" holds stage and age "
    "constant. Both run −1 to +1; here both are small (near zero), so the "
    "direct linear link between the two scores is weak — the subtype "
    "separation comes through more than this single correlation.\n"
    "    - **Why −1 to +1:** Effect is a correlation, always between −1 and "
    "+1. +1 = the two scores rise and fall together in perfect lockstep; "
    "−1 = they move in perfect opposition (one up while the other goes "
    "down); 0 = no linear relationship. Sign = direction, size = strength. "
    "For this biology we'd expect a negative value (low MMR ↔ high immune); "
    "near-zero here means that direct link is weak.\n"
    "- **Fit:** we train on the two scores to predict MSI-H vs MSS, then test "
    "on held-out patients. AUROC: 0.5 = chance, 1.0 = perfect. Raw accuracy "
    "is hidden because always guessing \"MSS\" would already score ~86% on "
    "this imbalanced cohort.\n"
    "    - **Why two metrics:** AUROC measures how well the scores SEPARATE "
    "the two groups when ranked, at any threshold, robust to the size "
    "imbalance — \"can the scores tell MSI-H from MSS at all?\" Balanced "
    "accuracy is the average of how many MSI-H we correctly catch "
    "(sensitivity) and how many MSS we correctly catch (specificity) at a "
    "single yes/no cutoff — \"if forced to decide, how well on BOTH groups, "
    "not just the majority?\" A high AUROC with a lower balanced accuracy "
    "(as here) means the ranking separates well but the default cutoff "
    "isn't tuned to call both classes evenly."
)


def _section_h1() -> None:
    st.header("H1 verification")
    res = _h1_result()

    # --- (a) Conclusion block --------------------------------------------
    a_pass = res.mmr_separates_correct_direction       # MMR median lower in MSI-H
    b_pass = res.immune_separates_correct_direction    # immune median higher in MSI-H
    c_pass = res.fit.auroc >= _AUROC_THRESHOLD

    if a_pass and b_pass and c_pass:
        st.success(
            f"**Conclusion:** Hypothesis 1 is verified. The operators reproduce "
            f"the known MSI biology — MSI-H tumours separate from MSS with "
            f"AUROC **{res.fit.auroc:.3f}**, and the scores point the expected "
            f"way (MMR lower, immune higher in MSI-H). The instrument is "
            f"calibrated, so we can trust it for the blind discovery in "
            f"Hypothesis 2. This is a tooling check, not a discovery — we were "
            f"given the genes. Caveat: the direct MMR-score↔immune-score "
            f"correlation is weak (adjusted **{res.effect.partial_corr:+.3f}**), "
            f"so we rely on the overall subtype separation rather than that "
            f"single link."
        )
    else:
        lines = ["**Conclusion:** Hypothesis 1 is not fully verified —"]
        if not a_pass:
            lines.append(
                f"- Check A failed: MMR median MSI-H "
                f"**{res.mmr_median_msi_h:.1f}** vs MSS "
                f"**{res.mmr_median_mss:.1f}** (expected MSI-H lower)."
            )
        if not b_pass:
            lines.append(
                f"- Check B failed: immune median MSI-H "
                f"**{res.immune_median_msi_h:.1f}** vs MSS "
                f"**{res.immune_median_mss:.1f}** (expected MSI-H higher)."
            )
        if not c_pass:
            lines.append(
                f"- Check C failed: held-out AUROC **{res.fit.auroc:.3f}** "
                f"below the **{_AUROC_THRESHOLD:.2f}** threshold."
            )
        st.warning("\n".join(lines))

    # --- (b) Score distributions -----------------------------------------
    st.subheader("Score distributions, split by MSI status")
    score_df = pd.concat([
        pd.DataFrame({"score_name": "MMR score",
                      "value": res.mmr_score.values,
                      "msi_status": res.msi_status.values}),
        pd.DataFrame({"score_name": "Immune score",
                      "value": res.immune_score.values,
                      "msi_status": res.msi_status.values}),
    ], ignore_index=True)

    left, right = st.columns(2)
    for col, score_name, summary in [
        (left,  "MMR score",
         f"median MSI-H **{res.mmr_median_msi_h:.1f}** vs MSS **{res.mmr_median_mss:.1f}** "
         f"({'lower in MSI-H ✓' if res.mmr_separates_correct_direction else 'wrong direction'})"),
        (right, "Immune score",
         f"median MSI-H **{res.immune_median_msi_h:.1f}** vs MSS **{res.immune_median_mss:.1f}** "
         f"({'higher in MSI-H ✓' if res.immune_separates_correct_direction else 'wrong direction'})"),
    ]:
        with col:
            st.markdown(f"**{score_name}** — {summary}")
            sub = score_df[score_df["score_name"] == score_name]
            v_min = float(sub["value"].min())
            v_max = float(sub["value"].max())
            chart = (
                alt.Chart(sub)
                .transform_density(
                    "value",
                    as_=["value", "density"],
                    groupby=["msi_status"],
                    extent=[v_min, v_max],
                    steps=80,
                )
                .mark_area(opacity=0.8)
                .encode(
                    x=alt.X("value:Q", title=score_name),
                    y=alt.Y("density:Q", title="density", stack=None),
                    color=alt.Color(
                        "msi_status:N",
                        sort=theme.MSI_ORDER,
                        scale=theme.msi_color_scale(include_na=False),
                        legend=alt.Legend(title="MSI status", orient="right"),
                    ),
                )
                .properties(height=CHART_HEIGHT)
            )
            st.altair_chart(chart, width="stretch")

    st.markdown(
        "**Why this matters:** MSI-H tumours have a broken spell-checker, so "
        "the MMR genes are less active (MMR score LOW) and the resulting "
        "pile-up of mutations makes the tumour immune-hot (immune score "
        "HIGH). MSS is the reverse — working repair, few mutations, "
        "immune-cold. Seeing exactly that split here confirms our scores "
        "capture the real biology; if they didn't, the tooling would be wrong."
    )

    # --- (c) Effect numbers ----------------------------------------------
    st.markdown(
        "*Step: does the MMR score causally relate to the immune score? "
        "(a check on the mechanism)*"
    )
    st.subheader("Effect: MMR score → immune score, adjusting for stage and age")
    e = res.effect
    e1, e2, e3 = st.columns(3)
    e1.metric(
        "Unadjusted pearson", f"{e.unadjusted:+.3f}",
        help=(
            "Plain correlation between the two scores across patients, from "
            "−1 to +1. Near 0 = little linear association."
        ),
    )
    e2.metric(
        "Adjusted partial corr.", f"{e.partial_corr:+.3f}",
        help=(
            "The same correlation after statistically holding stage and age "
            "constant — what's left once those background factors are removed."
        ),
    )
    e3.metric(
        "n used", e.n_used,
        help="Patients with all required values (both scores + stage + age).",
    )
    st.caption(e.note)

    # --- (d) Fit numbers -------------------------------------------------
    st.markdown(
        "*Step: can the two scores together predict MSI-H vs MSS? "
        "(a check that the signal is usable)*"
    )
    st.subheader("Fit (MMR + immune score → MSI-H) held-out performance")
    f1, f2, f3, f4 = st.columns(4)
    f1.metric(
        "AUROC", f"{res.fit.auroc:.3f}",
        help=(
            "Area Under the ROC Curve: how well the model ranks an MSI-H "
            "patient above an MSS one, from 0.5 (coin-flip) to 1.0 (perfect). "
            "Robust to the MSI-H/MSS imbalance."
        ),
    )
    f2.metric(
        "Balanced accuracy", f"{res.fit.balanced_acc:.3f}",
        help=(
            "Average of sensitivity (MSI-H correctly flagged) and specificity "
            "(MSS correctly flagged); 0.5 = chance."
        ),
    )
    f3.metric(
        "Train n", res.fit.n_train,
        help="Patients used to fit the model.",
    )
    f4.metric(
        "Test n", res.fit.n_test,
        help="Held-out patients used to score it — never seen in training.",
    )
    _auroc_pct = round(res.fit.auroc * 100)
    st.markdown(
        f"**AUROC 0.5 = guessing, 1.0 = perfect.** "
        f"**{res.fit.auroc:.3f}** means: take a random MSI-H patient and a "
        f"random MSS patient — about **{_auroc_pct}%** of the time the model "
        f"gives the MSI-H one the higher score."
    )

    # --- (e) How to read this -------------------------------------------
    with st.expander("How to read this", expanded=False):
        st.markdown(_HOW_TO_READ_H1)


# --- H2 preview: the airgap. Outside the H1 verification flow; sits at the
# bottom of the Hypothesis 1 tab as a setup for the next hypothesis. -----------

def _section_airgap_demo() -> None:
    st.header("Preview — the airgap (setup for Hypothesis 2)")
    st.markdown(
        "This isn't part of Hypothesis 1 — here we used the known gene names "
        "openly. It's a preview of the machinery that will make Hypothesis 2 "
        "honest: the discovery engine is handed a matrix with gene names "
        "hidden (opaque IDs), and only after it picks genes do we \"reveal\" "
        "what they were — proving the next step rediscovered the biology "
        "rather than recalling it."
    )
    demo = _airgap_demo()
    if demo is None:
        st.info("Processed expression matrix not found — skipping airgap demo.")
        return
    head, opaque_ids, revealed = demo
    st.markdown("**Anonymised matrix head** (samples × opaque feature IDs):")
    st.dataframe(head, width="stretch")
    st.markdown("**`reveal()` round-trip** — opaque ID → real gene symbol:")
    st.dataframe(
        pd.DataFrame({"opaque_id": opaque_ids, "real_symbol": revealed}),
        hide_index=True, width="stretch",
    )


# --- Hypothesis 2 tab (reads from the FastAPI) -----------------------------

H2_API_BASE = os.getenv("H2_API_BASE", "http://localhost:8000")

_H2_AUROC_THRESHOLD = 0.75
_H2_P_THRESHOLD = 0.05

_H2_HYPOTHESIS_MD = (
    "**Hypothesis 2 — Blind discovery.** Claim: if we hide the gene names "
    "and show the engine only the numbers, it can rediscover the genes that "
    "define the MSI subtype on its own — recovering the MMR spell-checker "
    "genes (MLH1, MSH2, MSH6, PMS2) out of ~20,000, without being told them. "
    "Confirmed if: (1) the program it finds separates MSI-H from MSS on "
    "patients it never trained on (high AUROC), (2) that beats a chance "
    "baseline (permutation test), and (3) when we reveal the genes, they "
    "overlap the known MMR set. This is the real test — H1 only checked the "
    "tooling on genes we supplied; here the engine finds them itself."
)


@st.cache_data(show_spinner=False, ttl=60)
def _h2_api_get(path: str):
    r = requests.get(f"{H2_API_BASE}{path}", timeout=10)
    r.raise_for_status()
    return r.json()


def _h2_api_post_reveal(gene_ids: list[str]) -> list[str]:
    r = requests.post(
        f"{H2_API_BASE}/reveal",
        json={"gene_ids": gene_ids},
        timeout=10,
    )
    r.raise_for_status()
    return r.json()["symbols"]


def _h2_card(num: int, title: str) -> "st.delta_generator.DeltaGenerator":
    card = st.container(border=True)
    card.caption(f"Step {num} of 5")
    card.subheader(title)
    return card


def _h2_program_dot(feature_sets: list[list[str]]) -> str:
    """Render a program (1 or 2 feature sets) as a graphviz DOT string."""
    lines = [
        "digraph p {",
        '    rankdir=LR;',
        '    bgcolor="transparent";',
        '    nodesep=0.5; ranksep=0.9; splines=spline;',
        '    fontname="Helvetica";',
        '    node [shape=box, style="rounded,filled", penwidth=1.2,',
        '          fontname="Helvetica", fontcolor="#23303A", fontsize=11,',
        '          margin="0.22,0.14", color="#3A6B7E", fillcolor="#F4F2EE"];',
        '    edge [color="#C2CACF", penwidth=1.3, arrowsize=0.7];',
        '    M [label="Expression matrix"];',
    ]
    score_nodes: list[str] = []
    for i, fs in enumerate(feature_sets, start=1):
        head = ", ".join(fs[:3])
        suffix = ", …" if len(fs) > 3 else ""
        sel_label = f"Select({len(fs)} genes:\\n{head}{suffix})"
        score_label = f"score {i}" if len(feature_sets) > 1 else "score"
        lines += [
            f'    Sel{i} [label="{sel_label}"];',
            f'    Red{i} [label="Reduce(mean)"];',
            f'    Score{i} [label="{score_label}", fillcolor="#FBEFE2", color="#BC6B2E"];',
            f"    M -> Sel{i} -> Red{i} -> Score{i};",
        ]
        score_nodes.append(f"Score{i}")
    lines.append(
        '    Fit [label="Fit → MSI-H probability", fillcolor="#EAF0F2", color="#3A6B7E"];'
    )
    for s in score_nodes:
        lines.append(f"    {s} -> Fit;")
    lines.append("}")
    return "\n".join(lines)


def _h2_show_api_error(exc: Exception) -> None:
    st.error(
        f"Cannot reach the H2 API at `{H2_API_BASE}`. To bring it up:\n\n"
        "```bash\n"
        "python -m scripts.run_h2          # run the engine once (~1-2 min)\n"
        "uvicorn api.app:app --reload      # then start the API\n"
        "```\n\n"
        f"Error: `{type(exc).__name__}: {exc}`"
    )


def _section_h2() -> None:
    try:
        run_log = _h2_api_get("/run")
        result = _h2_api_get("/result")
    except Exception as exc:
        _h2_show_api_error(exc)
        return

    _h2_step_hypothesis()
    _h2_step_setup(run_log)
    _h2_step_evolve(run_log)
    _h2_step_result(result)
    _h2_step_reveal_and_conclusion(result)


def _h2_step_hypothesis() -> None:
    card = _h2_card(0, "The hypothesis")
    with card:
        _, mid, _ = st.columns([1, 6, 1])
        with mid:
            st.markdown(_H2_HYPOTHESIS_MD)


def _h2_step_setup(run_log: dict) -> None:
    card = _h2_card(1, "The setup — why blind")
    with card:
        _, mid, _ = st.columns([1, 6, 1])
        with mid:
            run = run_log["run"]
            st.markdown(
                f"The cohort: **{run['n_train']} train + {run['n_test']} held-out "
                f"test** patients across **{run['n_genes']:,}** anonymised "
                f"feature IDs (dropped from {run.get('n_genes_input', run['n_genes']):,} "
                f"to {run['n_genes']:,} after removing features with any missing "
                f"value). The prefilter shortlists "
                f"**{run['prefilter_N']:,}** features by |AUROC − 0.5| computed "
                "on the train split only; the GP composes within that shortlist."
            )
        demo = _airgap_demo()
        if demo is not None:
            head, _ids, _syms = demo
            st.markdown("**Anonymised matrix head** (samples × opaque feature IDs):")
            st.dataframe(head, width="stretch")
        with mid:
            st.markdown(
                "The genes are hidden behind meaningless IDs, and the objective "
                "is computed only from MSI status — never from a gene name. So "
                "whatever the engine finds, it found in the data, not by "
                "recalling a name."
            )


def _h2_step_evolve(run_log: dict) -> None:
    card = _h2_card(2, "Watch it evolve")
    with card:
        _, mid, _ = st.columns([1, 6, 1])
        with mid:
            st.markdown(
                "The engine starts with random small programs, scores each on "
                "how well it separates MSI-H from MSS (on held-out folds), "
                "keeps the best, then breeds and mutates them over generations, "
                "discarding the weak — like selective breeding for programs."
            )

        generations = run_log["generations"]
        max_gen = generations[-1]["generation"]
        elitism = int(generations[0].get("elitism", 5))

        gen = st.slider(
            "Generation",
            min_value=0,
            max_value=max_gen,
            value=max_gen,
            step=1,
            key="h2_evolve_gen",
        )

        # Fitness curve.
        gens_df = pd.DataFrame({
            "generation": [g["generation"] for g in generations],
            "best": [g["best_fitness"] for g in generations],
            "median": [g["median_fitness"] for g in generations],
        })
        long_df = gens_df.melt(
            "generation", var_name="series", value_name="fitness",
        )
        fit_curve = (
            alt.Chart(long_df)
            .mark_line(point=True)
            .encode(
                x=alt.X("generation:Q", title="generation"),
                y=alt.Y("fitness:Q", title="fitness (CV AUROC − λ × n_genes)"),
                color=alt.Color(
                    "series:N",
                    scale=alt.Scale(
                        domain=["best", "median"],
                        range=["#BC6B2E", "#6E7F8C"],
                    ),
                    legend=alt.Legend(title=None, orient="top-right"),
                ),
                tooltip=[
                    alt.Tooltip("generation:Q"),
                    alt.Tooltip("series:N"),
                    alt.Tooltip("fitness:Q", format=".3f"),
                ],
            )
            .properties(height=240)
        )
        rule = (
            alt.Chart(pd.DataFrame({"generation": [gen]}))
            .mark_rule(color="#BC6B2E", strokeDash=[4, 4])
            .encode(x="generation:Q")
        )
        st.altair_chart(fit_curve + rule, width="stretch")

        # Top-K candidates at the selected generation.
        cand = generations[gen]["candidates"]
        pop_df = pd.DataFrame([
            {
                "rank": i + 1,
                "id": c["id"],
                "n_genes": c["n_genes"],
                "fitness": c["fitness"],
                "first_ids": ", ".join(c["gene_ids"][:3])
                + ("…" if len(c["gene_ids"]) > 3 else ""),
            }
            for i, c in enumerate(cand)
        ])
        pop_df["status"] = pop_df["rank"].apply(
            lambda r: "survivor (elite)" if r <= elitism else "discarded"
        )

        pop_chart = (
            alt.Chart(pop_df)
            .mark_bar()
            .encode(
                y=alt.Y(
                    "id:N",
                    sort=alt.SortField("rank", order="ascending"),
                    title="candidate",
                ),
                x=alt.X("fitness:Q", title="fitness"),
                color=alt.Color(
                    "status:N",
                    scale=alt.Scale(
                        domain=["survivor (elite)", "discarded"],
                        range=["#3A6B7E", "#9AA0A6"],
                    ),
                    legend=alt.Legend(title=None, orient="top"),
                ),
                opacity=alt.condition(
                    f"datum.rank <= {elitism}",
                    alt.value(1.0),
                    alt.value(0.35),
                ),
                tooltip=[
                    alt.Tooltip("rank:Q"),
                    alt.Tooltip("id:N"),
                    alt.Tooltip("n_genes:Q"),
                    alt.Tooltip("fitness:Q", format=".4f"),
                    alt.Tooltip("first_ids:N", title="first IDs"),
                ],
            )
            .properties(height=220)
        )
        st.markdown(
            f"**Top {len(cand)} candidates at generation {gen}.** "
            f"Top {elitism} (solid) survive as elites into the next generation; "
            "the rest fade."
        )
        st.altair_chart(pop_chart, width="stretch")

        best = cand[0]
        st.markdown(
            f"**Best of generation {gen}:** `{best['id']}` · "
            f"fitness **{best['fitness']:.4f}** · {best['n_genes']} genes"
        )
        st.graphviz_chart(_h2_program_dot(best["feature_sets"]), width="stretch")


def _h2_step_result(result: dict) -> None:
    card = _h2_card(3, "The result")
    with card:
        _, mid, _ = st.columns([1, 6, 1])
        win = result["winning"]
        base = result["baseline"]
        with mid:
            st.markdown(
                f"The engine's pick after evolving: **{win['program_repr']}** — "
                f"a {len(win['gene_ids'])}-gene program."
            )

        st.graphviz_chart(_h2_program_dot(win["feature_sets"]), width="stretch")

        m1, m2, m3, m4 = st.columns(4)
        m1.metric(
            "Held-out AUROC", f"{win['holdout_auroc']:.3f}",
            help=(
                "Area Under the ROC Curve on the held-out test patients the "
                "engine never trained on. 0.5 = chance, 1.0 = perfect."
            ),
        )
        m2.metric(
            "Permutation p", f"{win['permutation_p']:.4f}",
            help=(
                "Fraction of label-shuffled null runs whose AUROC matched or "
                "beat the real winner. Small = unlikely to be a fluke."
            ),
        )
        m3.metric(
            "Baseline AUROC", f"{base['holdout_auroc']:.3f}",
            help=(
                "Deterministic top-k by univariate prefilter, then "
                "Reduce(mean) + Fit. Internal sanity check the GP should match."
            ),
        )
        m4.metric(
            "n_genes used", len(win["gene_ids"]),
            help="Number of opaque feature IDs the winning program touches.",
        )

        with mid:
            auroc_pct = round(win["holdout_auroc"] * 100)
            st.markdown(
                f"**AUROC 0.5 = guessing, 1.0 = perfect.** "
                f"**{win['holdout_auroc']:.3f}** means it ranks a random MSI-H "
                f"above a random MSS about **{auroc_pct}%** of the time."
            )
            st.markdown(
                f"**Permutation p {win['permutation_p']:.4f}** — how often pure "
                "chance matched this; small = unlikely a fluke."
            )


_H2_REVEAL_KEY = "h2_revealed_symbols"
_H2_REVEAL_FOR_IDS_KEY = "h2_revealed_for_ids"


def _h2_step_reveal_and_conclusion(result: dict) -> None:
    win = result["winning"]
    winning_ids = list(win["gene_ids"])

    card = _h2_card(4, "The reveal")
    with card:
        _, mid, _ = st.columns([1, 6, 1])
        with mid:
            st.markdown(
                "We have not looked at any gene name yet. The button below "
                "calls the only endpoint allowed to translate the opaque IDs "
                "back — pressing it is the moment the airgap opens."
            )
            clicked = st.button("Reveal the genes", type="primary")

        if clicked:
            try:
                symbols = _h2_api_post_reveal(winning_ids)
            except Exception as exc:
                _h2_show_api_error(exc)
                return
            st.session_state[_H2_REVEAL_KEY] = symbols
            st.session_state[_H2_REVEAL_FOR_IDS_KEY] = winning_ids

        revealed = st.session_state.get(_H2_REVEAL_KEY)
        revealed_for = st.session_state.get(_H2_REVEAL_FOR_IDS_KEY)
        if not revealed or revealed_for != winning_ids:
            with mid:
                st.info("Click **Reveal the genes** to open the airgap.")
            return

        mmr_set = set(MMR_GENES)
        reveal_df = pd.DataFrame({
            "opaque ID": winning_ids,
            "real symbol": revealed,
            "MMR": [s in mmr_set for s in revealed],
        })
        styled = reveal_df.style.apply(
            lambda row: [
                "background-color: #FBEFE2; color: #BC6B2E; font-weight: 600"
                if row["MMR"] else "" for _ in row
            ],
            axis=1,
        ).hide(axis="index")
        st.dataframe(styled, width="stretch")
        with mid:
            st.markdown(
                "Only now do we look at the names. These are the genes the "
                "engine chose — blind. Any row highlighted in amber matches "
                "the known MMR set."
            )

    recovered = [s for s in revealed if s in set(MMR_GENES)]
    _h2_step_conclusion(result, recovered)


def _h2_step_conclusion(result: dict, recovered: list[str]) -> None:
    card = _h2_card(5, "Conclusion")
    with card:
        _, mid, _ = st.columns([1, 6, 1])
        with mid:
            win = result["winning"]
            auroc = float(win["holdout_auroc"])
            p_val = float(win["permutation_p"])
            k = len(recovered)

            a_pass = auroc >= _H2_AUROC_THRESHOLD
            b_pass = p_val < _H2_P_THRESHOLD
            c_pass = k >= 1

            if a_pass and b_pass and c_pass:
                recovered_str = ", ".join(recovered) if recovered else "—"
                st.success(
                    f"**Conclusion:** Hypothesis 2 is verified. Blind — with "
                    f"gene names hidden — the engine found a program "
                    f"separating MSI-H from MSS (held-out AUROC "
                    f"**{auroc:.3f}**, permutation p **{p_val:.4f}**), and on "
                    f"reveal it recovered **{k}** of the 4 MMR genes "
                    f"(**{recovered_str}**). It rediscovered the biology "
                    f"rather than recalling it. (Like H1, this is the "
                    f"mechanism, not drug response.)"
                )
            else:
                lines = ["**Conclusion:** Hypothesis 2 is not fully verified —"]
                if not a_pass:
                    lines.append(
                        f"- Check A failed: held-out AUROC **{auroc:.3f}** "
                        f"below the **{_H2_AUROC_THRESHOLD:.2f}** threshold."
                    )
                if not b_pass:
                    lines.append(
                        f"- Check B failed: permutation p **{p_val:.4f}** "
                        f"not below **{_H2_P_THRESHOLD:.2f}**."
                    )
                if not c_pass:
                    lines.append(
                        f"- Check C failed: **{k}** of the 4 MMR genes "
                        f"recovered (need at least 1)."
                    )
                st.warning("\n".join(lines))


if __name__ == "__main__":
    main()