"""H1 verification — known-answer check, NOT blind discovery.

Composes the DSL operators over the NAMED expression matrix (real gene symbols)
with hand-picked MMR and immune gene panels, and reports the four numbers the
H1 program should produce on the usable MSI-H-vs-MSS CRC cohort:

1. medians of `mmr_score` and `immune_score` split by MSI status — MSI-H should
   sit LOW on the MMR score and HIGH on the immune score;
2. `Effect(mmr_score, immune_score, adjust={stage, age})` — the adjusted partial
   correlation vs the unadjusted pearson on the same rows, and the n used;
3. `Fit(state=(mmr_score, immune_score)) → MSI-H` — held-out AUROC and
   balanced accuracy.

This module is biology-aware on purpose. The blind-discovery path in
[`engine`](../engine) will instead consume the airgapped view and is forbidden
from importing these gene lists.
"""

from __future__ import annotations

from dataclasses import dataclass

import pandas as pd

from dsl import Cohort, Effect, EffectResult, Fit, FitResult, Load, Reduce, Select

MMR_GENES: list[str] = ["MLH1", "MSH2", "MSH6", "PMS2"]
IMMUNE_GENES: list[str] = ["CD8A", "GZMA", "PRF1"]

LABEL_COL = "msi_status"
POSITIVE_LABEL = "MSI-H"
NEGATIVE_LABEL = "MSS"


def usable_msi_cohort(cohort: Cohort) -> Cohort:
    """Restrict to MSI-H vs MSS samples with non-missing stage and age."""
    lab = cohort.labels
    clin = cohort.clinical
    mask = (
        lab[LABEL_COL].isin([POSITIVE_LABEL, NEGATIVE_LABEL])
        & (clin["stage"] != "NA")
        & clin["stage"].notna()
        & clin["age"].notna()
    )
    return cohort.restrict(mask)


@dataclass
class H1Result:
    n_used: int
    n_msi_h: int
    n_mss: int

    mmr_score: pd.Series
    immune_score: pd.Series
    msi_status: pd.Series

    mmr_median_msi_h: float
    mmr_median_mss: float
    immune_median_msi_h: float
    immune_median_mss: float

    effect: EffectResult
    fit: FitResult

    @property
    def mmr_separates_correct_direction(self) -> bool:
        return self.mmr_median_msi_h < self.mmr_median_mss

    @property
    def immune_separates_correct_direction(self) -> bool:
        return self.immune_median_msi_h > self.immune_median_mss


def run_h1(cohort: Cohort | None = None) -> H1Result:
    """Run the full H1 program and return the verification numbers."""
    if cohort is None:
        cohort = Load("processed")
    cohort = usable_msi_cohort(cohort)

    M = cohort.expression
    mmr_score = Reduce(Select(M, MMR_GENES), agg="mean")
    immune_score = Reduce(Select(M, IMMUNE_GENES), agg="mean")

    status = cohort.labels[LABEL_COL]
    is_msi_h = status == POSITIVE_LABEL

    mmr_medians = mmr_score.groupby(status).median()
    immune_medians = immune_score.groupby(status).median()

    adjust = pd.DataFrame({
        "stage": cohort.clinical["stage"].astype("category"),
        "age": cohort.clinical["age"].astype(float),
    })
    effect = Effect(mmr_score, immune_score, adjust)

    state = pd.DataFrame({"mmr_score": mmr_score, "immune_score": immune_score})
    fit = Fit(state, is_msi_h.astype(int).values)

    return H1Result(
        n_used=len(cohort.sample_ids),
        n_msi_h=int(is_msi_h.sum()),
        n_mss=int((status == NEGATIVE_LABEL).sum()),
        mmr_score=mmr_score,
        immune_score=immune_score,
        msi_status=status,
        mmr_median_msi_h=float(mmr_medians.get(POSITIVE_LABEL, float("nan"))),
        mmr_median_mss=float(mmr_medians.get(NEGATIVE_LABEL, float("nan"))),
        immune_median_msi_h=float(immune_medians.get(POSITIVE_LABEL, float("nan"))),
        immune_median_mss=float(immune_medians.get(NEGATIVE_LABEL, float("nan"))),
        effect=effect,
        fit=fit,
    )