tests/test_active_learning.py

"""Tests for the Directed-Evolution active-learning surrogate
(`dee.core.active_learning`) — the Design→Build→Test→Learn round-2 core.

These pin the behaviour we actually claim to users:
  * mutation-label parsing (1-indexed → 0-indexed),
  * honest fallback to the ESM-2 prior below the measurement floor / with no
    signal (no overclaiming),
  * the headline property: when the zero-shot prior is *wrong*, real bench
    measurements override it and the round-2 ranking recovers the truth,
  * `adjust_pool` plugs the learned scores back into the SA `Mutation` pool
    without disturbing the other fields.
No torch/ESM needed — the surrogate is pure numpy over a supplied pool.
"""
import numpy as np
import pytest

from dee.core.active_learning import (
    MIN_MEASUREMENTS,
    Surrogate,
    fit_surrogate,
    parse_label,
    parse_mutations,
)
from dee.optimizer.search import Mutation


# --------------------------------------------------------------------------- #
# label parsing
# --------------------------------------------------------------------------- #
def test_parse_label_basic():
    assert parse_label("W58L") == (57, "L")          # 1-indexed -> 0-indexed
    assert parse_label("  k204r ") == (203, "R")      # whitespace + case
    assert parse_label("A1*") == (0, "*")             # stop codon, first residue


@pytest.mark.parametrize("bad", ["", "58L", "WL", "W0L", "W-3L", "WfooL", None])
def test_parse_label_rejects_malformed(bad):
    assert parse_label(bad) is None


def test_parse_mutations_string_and_list():
    assert parse_mutations("W58L,K204R") == [(57, "L"), (203, "R")]
    assert parse_mutations("W58L K204R; T2A") == [(57, "L"), (203, "R"), (1, "A")]
    assert parse_mutations(["W58L", "K204R"]) == [(57, "L"), (203, "R")]
    # malformed tokens are silently dropped, valid ones kept
    assert parse_mutations("W58L, junk, K204R") == [(57, "L"), (203, "R")]


# --------------------------------------------------------------------------- #
# helpers
# --------------------------------------------------------------------------- #
def _pool(prior_by_key):
    """Build a frozen-Mutation pool from {(pos, aa): delta_ll}."""
    return [
        Mutation(position=pos, wt_aa="A", mut_aa=aa, delta_ll=float(dll))
        for (pos, aa), dll in prior_by_key.items()
    ]


def _label(pos, aa):
    return f"A{pos + 1}{aa}"  # inverse of parse_label, wt 'A' is irrelevant


def _spearman(a, b):
    a, b = np.asarray(a, float), np.asarray(b, float)
    ra = np.argsort(np.argsort(a))
    rb = np.argsort(np.argsort(b))
    return float(np.corrcoef(ra, rb)[0, 1])


# --------------------------------------------------------------------------- #
# honest fallback
# --------------------------------------------------------------------------- #
def test_fallback_below_measurement_floor():
    pool = _pool({(0, "L"): 1.0, (1, "R"): 0.5})
    meas = [([_label(0, "L")], 3.0), ([_label(1, "R")], 1.0)]  # only 2 < floor
    s = fit_surrogate(pool, meas)
    assert s.learned is False
    assert s.n_effects == 0
    assert s.w_prior == 1.0
    # adjusted == the pure prior, so round 2 == a fresh search on ESM-2
    assert s.adjusted == {(0, "L"): 1.0, (1, "R"): 0.5}
    assert str(MIN_MEASUREMENTS) in s.note


def test_fallback_when_no_spread_in_values():
    pool = _pool({(i, "L"): 0.1 * i for i in range(5)})
    meas = [([_label(i, "L")], 2.0) for i in range(5)]  # ≥ floor but zero variance
    s = fit_surrogate(pool, meas)
    assert s.learned is False
    assert s.adjusted == {(i, "L"): 0.1 * i for i in range(5)}


def test_empty_inputs_dont_crash():
    s = fit_surrogate([], [])
    assert s.learned is False
    assert s.adjusted == {}


def test_non_numeric_and_out_of_pool_measurements_ignored():
    pool = _pool({(0, "L"): 1.0})
    meas = [
        ([_label(0, "L")], "not-a-number"),  # bad value -> dropped
        (["Z9Q"], 5.0),                       # mutation not in pool -> dropped
    ]
    s = fit_surrogate(pool, meas)
    assert s.n_train == 0
    assert s.learned is False


# --------------------------------------------------------------------------- #
# the headline property: data overrides a misleading prior
# --------------------------------------------------------------------------- #
def test_data_overrides_misleading_prior():
    """Ground-truth additive effects exist; the ESM-2 prior is deliberately
    *anti-correlated* with them. After logging enough bench measurements, the
    round-2 acquisition ranking should recover the truth far better than the
    prior did — that's the whole value proposition."""
    rng = np.random.default_rng(7)
    keys = [(i, "L") for i in range(8)]
    true_effect = {k: v for k, v in zip(keys, [3.0, -2.0, 2.5, -1.0, 1.5, -2.5, 0.5, -0.5])}
    # Prior is the NEGATION of the truth (worst case): high-prior = actually bad.
    prior = {k: -true_effect[k] for k in keys}
    pool = _pool(prior)

    # Simulate the user measuring random multi-site variants; value = additive
    # truth + small noise (scale is arbitrary; surrogate standardizes it).
    measurements = []
    for _ in range(40):
        n_sites = int(rng.integers(1, 4))
        chosen = [keys[i] for i in rng.choice(len(keys), size=n_sites, replace=False)]
        val = sum(true_effect[k] for k in chosen) + rng.normal(0, 0.25)
        measurements.append(([_label(p, a) for (p, a) in chosen], 10.0 + val))

    s = fit_surrogate(pool, measurements)
    assert s.learned is True
    assert s.n_train == 40
    assert s.n_effects > 0

    true_vec = [true_effect[k] for k in keys]
    prior_vec = [prior[k] for k in keys]
    learned_vec = [s.adjusted[k] for k in keys]

    rho_prior = _spearman(prior_vec, true_vec)
    rho_learned = _spearman(learned_vec, true_vec)

    assert rho_prior < 0                  # prior is anti-correlated (by construction)
    assert rho_learned > 0.6              # learned ranking recovers the truth
    assert rho_learned > rho_prior + 1.0  # and is a large improvement over the prior


def test_adjust_pool_preserves_fields_and_swaps_score():
    pool = _pool({(0, "L"): 1.0, (3, "R"): -0.5, (5, "K"): 0.2})
    # enough varied data to learn
    meas = [
        ([_label(0, "L")], 9.0),
        ([_label(3, "R")], 2.0),
        ([_label(5, "K")], 5.0),
        ([_label(0, "L"), _label(5, "K")], 8.0),
        ([_label(3, "R"), _label(5, "K")], 4.0),
    ]
    s = fit_surrogate(pool, meas)
    new_pool = s.adjust_pool(pool)

    assert len(new_pool) == len(pool)
    for old, new in zip(pool, new_pool):
        # identity fields untouched...
        assert (new.position, new.wt_aa, new.mut_aa) == (old.position, old.wt_aa, old.mut_aa)
        # ...delta_ll replaced by the learned acquisition score
        assert new.delta_ll == pytest.approx(s.adjusted[(old.position, old.mut_aa)])
    # frozen dataclass: originals are unchanged
    assert pool[0].delta_ll == 1.0


def test_unmeasured_mutation_falls_back_toward_prior():
    """A mutation never measured gets β≈0, so its adjusted score stays driven by
    the (re-weighted) prior plus only the exploration bonus — it isn't invented."""
    pool = _pool({(0, "L"): 1.0, (1, "R"): 0.8, (2, "K"): 0.6, (3, "D"): 0.4, (9, "Q"): 5.0})
    # measure everything EXCEPT (9, 'Q')
    meas = [
        ([_label(0, "L")], 3.0),
        ([_label(1, "R")], 2.0),
        ([_label(2, "K")], 1.0),
        ([_label(3, "D")], 0.5),
        ([_label(0, "L"), _label(1, "R")], 4.5),
    ]
    s = fit_surrogate(pool, meas)
    assert s.learned is True
    # the unmeasured high-prior mutation is still present and finite
    assert (9, "Q") in s.adjusted
    assert np.isfinite(s.adjusted[(9, "Q")])