tests/test_dsl.py

"""Unit tests for the DSL operators.

Each operator gets a tiny hand-checkable case; no network, no parquet I/O.
"""

from __future__ import annotations

import numpy as np
import pandas as pd
import pytest

from dsl import (
    Apply,
    Associate,
    Cohort,
    Effect,
    Fit,
    Reduce,
    Search,
    Select,
    Split,
)


# --- fixtures ---------------------------------------------------------------


def _tiny_cohort() -> Cohort:
    """Six samples, five features, with known structure.

    `geneA` is high in samples 4..6 and low in 1..3; `geneB` is the reverse.
    `msi_status` is MSI-H in 4..6 and MSS in 1..3 so a single-gene rule should
    perfectly separate the two groups.
    """
    idx = pd.Index([f"s{i}" for i in range(1, 7)], name="sample_id")
    expr = pd.DataFrame(
        {
            "geneA": [1.0, 1.0, 1.0, 9.0, 9.0, 9.0],
            "geneB": [9.0, 9.0, 9.0, 1.0, 1.0, 1.0],
            "geneC": [5.0, 5.0, 5.0, 5.0, 5.0, 5.0],
            "geneD": [2.0, 3.0, 4.0, 5.0, 6.0, 7.0],
            "geneE": [7.0, 6.0, 5.0, 4.0, 3.0, 2.0],
        },
        index=idx,
    )
    clinical = pd.DataFrame(
        {
            "stage": ["I", "II", "III", "II", "III", "I"],
            "age": [55.0, 60.0, 65.0, 70.0, 50.0, 75.0],
        },
        index=idx,
    )
    labels = pd.DataFrame(
        {"msi_status": ["MSS"] * 3 + ["MSI-H"] * 3},
        index=idx,
    )
    return Cohort(expression=expr, clinical=clinical, labels=labels)


# --- Select / Reduce -------------------------------------------------------


def test_select_picks_columns_in_order():
    c = _tiny_cohort()
    out = Select(c.expression, ["geneB", "geneA"])
    assert list(out.columns) == ["geneB", "geneA"]
    assert out.shape == (6, 2)


def test_select_missing_column_raises():
    c = _tiny_cohort()
    with pytest.raises(KeyError):
        Select(c.expression, ["geneA", "geneZ"])


def test_reduce_mean_collapses_rows():
    c = _tiny_cohort()
    score = Reduce(Select(c.expression, ["geneA", "geneB"]), agg="mean")
    # (1+9)/2 = 5 for every row.
    assert list(score.values) == pytest.approx([5.0] * 6)


def test_reduce_supports_full_agg_vocab():
    c = _tiny_cohort()
    for agg in ("mean", "median", "max", "min", "var"):
        s = Reduce(c.expression, agg=agg)
        assert list(s.index) == list(c.expression.index)


def test_reduce_rejects_unknown_agg():
    c = _tiny_cohort()
    with pytest.raises(ValueError):
        Reduce(c.expression, agg="prod")


# --- Split -----------------------------------------------------------------


def test_split_partitions_by_predicate():
    c = _tiny_cohort()
    a, b = Split(c, lambda co: co.labels["msi_status"] == "MSI-H")
    assert list(a.sample_ids) == ["s4", "s5", "s6"]
    assert list(b.sample_ids) == ["s1", "s2", "s3"]


# --- Associate -------------------------------------------------------------


def test_associate_pearson_perfect_positive():
    s = pd.Series([1.0, 2.0, 3.0, 4.0])
    assert Associate(s, s * 2 + 7, kind="pearson") == pytest.approx(1.0)


def test_associate_pearson_perfect_negative():
    s = pd.Series([1.0, 2.0, 3.0, 4.0])
    assert Associate(s, -s, kind="pearson") == pytest.approx(-1.0)


def test_associate_spearman_monotone_only():
    # Strictly monotone but non-linear -> spearman 1, pearson < 1.
    s = pd.Series([1.0, 2.0, 3.0, 4.0])
    t = s ** 3
    assert Associate(s, t, kind="spearman") == pytest.approx(1.0)
    assert Associate(s, t, kind="pearson") < 1.0


def test_associate_unknown_kind_raises():
    s = pd.Series([1.0, 2.0])
    with pytest.raises(ValueError):
        Associate(s, s, kind="kendall")


# --- Effect ----------------------------------------------------------------


def test_effect_partial_correlation_removes_confounder():
    """y caused entirely by z; x also caused by z but no direct link to y.

    Unadjusted x-vs-y correlation should be ~1, partial correlation given z ~0.
    """
    rng = np.random.default_rng(0)
    n = 400
    z = rng.standard_normal(n)
    x = z + 0.01 * rng.standard_normal(n)
    y = z + 0.01 * rng.standard_normal(n)
    idx = pd.Index([f"s{i}" for i in range(n)], name="sample_id")
    res = Effect(
        pd.Series(x, index=idx),
        pd.Series(y, index=idx),
        pd.DataFrame({"z": z}, index=idx),
    )
    assert res.unadjusted > 0.95
    assert abs(res.partial_corr) < 0.2
    assert res.n_used == n


def test_effect_drops_nans_and_reports_n_used():
    idx = pd.Index([f"s{i}" for i in range(5)], name="sample_id")
    cause = pd.Series([1.0, 2.0, 3.0, 4.0, np.nan], index=idx)
    effect = pd.Series([2.0, 4.0, 6.0, np.nan, 10.0], index=idx)
    adj = pd.DataFrame({"stage": ["I", "II", "II", "III", "I"]}, index=idx)
    res = Effect(cause, effect, adj)
    assert res.n_used == 3


# --- Search ----------------------------------------------------------------


def test_search_rejects_non_opaque_columns():
    df = pd.DataFrame({"GENE1": [1.0], "GENE2": [2.0]})
    with pytest.raises(ValueError, match="opaque IDs"):
        Search(df, lambda s: float(s.sum()), 1)


def test_search_ranks_by_objective_and_takes_top_k():
    df = pd.DataFrame({
        "g00001": [1.0, 1.0, 1.0],
        "g00002": [1.0, 2.0, 3.0],
        "g00003": [0.0, 0.0, 9.0],
    })
    # objective = variance — g00003 highest, g00002 next, g00001 zero.
    top = Search(df, lambda s: float(s.var()), 2)
    assert top == ["g00003", "g00002"]


# --- Fit / Apply -----------------------------------------------------------


def test_fit_separable_features_score_above_random():
    """Two linearly separable classes -> AUROC = 1.0, balanced acc = 1.0."""
    rng = np.random.default_rng(0)
    n = 60
    x_pos = rng.standard_normal(n) + 3.0
    x_neg = rng.standard_normal(n) - 3.0
    state = pd.Series(np.concatenate([x_pos, x_neg]), name="x")
    y = np.array([1] * n + [0] * n)
    res = Fit(state, y, random_state=0)
    assert res.auroc == pytest.approx(1.0)
    assert res.balanced_acc == pytest.approx(1.0)
    assert res.n_train + res.n_test == 2 * n


def test_apply_returns_probabilities_indexed_like_input():
    rng = np.random.default_rng(0)
    n = 40
    state = pd.Series(
        np.concatenate([rng.standard_normal(n) + 2, rng.standard_normal(n) - 2]),
        index=[f"s{i}" for i in range(2 * n)],
        name="x",
    )
    y = np.array([1] * n + [0] * n)
    res = Fit(state, y, random_state=0)
    proba = Apply(res, state)
    assert list(proba.index) == list(state.index)
    assert ((proba >= 0.0) & (proba <= 1.0)).all()