tests/test_aggregate.py

"""Tests for the cross-user aggregate prior (`dee.core.aggregate`).

Pins the privacy + correctness invariants we actually promise:
  * aggregates by substitution TYPE, standardized within each library;
  * k-anonymity: a substitution measured by < MIN_USERS distinct users is
    dropped and never appears in the output;
  * the execution gate refuses before the policy effective date;
  * the resulting prior recovers the right sign and folds into the SA pool.
Pure numpy — no DB, no torch.
"""
import datetime as dt

import numpy as np
import pytest

from dee.core import aggregate as agg
from dee.core.aggregate import (
    EFFECTIVE_DATE,
    AggregationGateError,
    GlobalPrior,
    apply_global_prior,
    build_priors,
    parse_sub,
)
from dee.optimizer.search import Mutation


AFTER = EFFECTIVE_DATE                       # gate passes on/after the date
BEFORE = EFFECTIVE_DATE - dt.timedelta(days=1)


# --------------------------------------------------------------------------- #
def test_parse_sub():
    assert parse_sub("W58L") == ("W", "L")
    assert parse_sub(" k204r ") == ("K", "R")
    assert parse_sub("L10L") is None          # synonymous
    assert parse_sub("junk") is None


# --------------------------------------------------------------------------- #
# the gate
# --------------------------------------------------------------------------- #
def test_gate_blocks_before_effective_date():
    with pytest.raises(AggregationGateError):
        build_priors([("u1", [("A2V", 1.0), ("A2V,L3I", 2.0)])], now=BEFORE)


def test_gate_opens_on_effective_date():
    # should not raise (may return empty if no signal, but must not be gated)
    out = build_priors([("u1", [("A2V", 1.0), ("L3I", 2.0)])], now=AFTER)
    assert isinstance(out, GlobalPrior)


def test_enforce_gate_false_for_tests():
    out = build_priors([("u1", [("A2V", 1.0), ("L3I", 2.0)])], now=BEFORE, enforce_gate=False)
    assert isinstance(out, GlobalPrior)


# --------------------------------------------------------------------------- #
# k-anonymity floor
# --------------------------------------------------------------------------- #
def _lib_for_sub(values):
    """A 2-row library exercising substitutions W>L (good) and D>A (bad)."""
    # variant 'W10L' high, 'D20A' low -> standardized gives +/- effects
    return [("W10L", values[0]), ("D20A", values[1])]


def test_kanon_drops_substitutions_below_floor():
    # W>L and D>A measured by only 2 distinct users -> must be dropped (floor=3)
    obls = [
        ("u1", _lib_for_sub((2.0, 0.0))),
        ("u2", _lib_for_sub((2.0, 0.0))),
    ]
    out = build_priors(obls, min_users=3, now=AFTER)
    assert out.effects == {}                  # nothing meets the >=3-user floor

    # add a third distinct user -> now kept
    obls.append(("u3", _lib_for_sub((2.0, 0.0))))
    out3 = build_priors(obls, min_users=3, now=AFTER)
    assert ("W", "L") in out3.effects and ("D", "A") in out3.effects
    assert out3.n_users[("W", "L")] == 3


def test_repeated_user_does_not_inflate_kanon():
    # same user across two libraries is still ONE distinct user -> below floor
    obls = [
        ("u1", _lib_for_sub((2.0, 0.0))),
        ("u1", _lib_for_sub((2.0, 0.0))),
        ("u2", _lib_for_sub((2.0, 0.0))),
    ]
    out = build_priors(obls, min_users=3, now=AFTER)
    assert out.effects == {}                  # only 2 distinct users


# --------------------------------------------------------------------------- #
# correctness: sign recovery + standardization
# --------------------------------------------------------------------------- #
def test_recovers_substitution_sign_across_users():
    # W>L consistently beneficial, D>A consistently deleterious, 4 users,
    # each on a DIFFERENT arbitrary assay scale (tests within-library z-scoring)
    scales = [(1.0, 0.0), (100.0, 50.0), (5.0, -5.0), (0.2, 0.0)]   # (gain, offset)
    obls = []
    for i, (gain, off) in enumerate(scales):
        meas = [
            ("W10L", off + gain * 1.0),      # good
            ("W30L", off + gain * 0.8),      # good (same sub type, other position)
            ("D20A", off + gain * -1.0),     # bad
            ("W10L,D20A", off + gain * 0.0), # mixed
        ]
        obls.append((f"u{i}", meas))
    out = build_priors(obls, min_users=3, now=AFTER)
    assert out.effects[("W", "L")] > 0       # beneficial recovered
    assert out.effects[("D", "A")] < 0       # deleterious recovered
    assert out.n_users[("W", "L")] == 4
    assert out.n_obs[("W", "L")] >= 4


def test_zero_variance_or_tiny_library_contributes_nothing():
    obls = [
        ("u1", [("W10L", 5.0)]),                  # single row -> no effect
        ("u2", [("W10L", 3.0), ("D20A", 3.0)]),   # zero variance -> dropped
    ]
    out = build_priors(obls, min_users=1, now=AFTER)
    assert out.effects == {}


# --------------------------------------------------------------------------- #
# serialization round-trip + pool blend
# --------------------------------------------------------------------------- #
def test_to_rows_is_deidentified_and_roundtrips():
    obls = [("u%d" % i, _lib_for_sub((2.0, 0.0))) for i in range(3)]
    out = build_priors(obls, min_users=3, now=AFTER)
    rows = out.to_rows()
    assert rows and all(set(r) == {"substitution", "n_users", "n_obs", "mean_effect"} for r in rows)
    # no user ids, no positions, no raw values anywhere in the serialized form
    blob = str(rows)
    assert "u0" not in blob and "u1" not in blob and "10" not in "".join(r["substitution"] for r in rows)
    back = GlobalPrior.from_rows(rows)
    assert set(back.effects) == set(out.effects)


def test_apply_global_prior_nudges_matching_substitutions():
    prior = GlobalPrior(effects={("W", "L"): 2.0}, n_users={("W", "L"): 5}, n_obs={("W", "L"): 9})
    pool = [
        Mutation(position=57, wt_aa="W", mut_aa="L", delta_ll=1.0),   # matches -> nudged
        Mutation(position=99, wt_aa="D", mut_aa="A", delta_ll=1.0),   # no match -> unchanged
    ]
    out = apply_global_prior(pool, prior, weight=0.3)
    assert out[0].delta_ll == pytest.approx(1.0 + 0.3 * 2.0)
    assert out[1].delta_ll == 1.0
    # identity fields preserved
    assert (out[0].position, out[0].wt_aa, out[0].mut_aa) == (57, "W", "L")


def test_apply_global_prior_noop_without_prior():
    pool = [Mutation(position=0, wt_aa="A", mut_aa="V", delta_ll=1.0)]
    assert apply_global_prior(pool, None)[0].delta_ll == 1.0
    empty = GlobalPrior(effects={}, n_users={}, n_obs={})
    assert apply_global_prior(pool, empty)[0].delta_ll == 1.0


# --------------------------------------------------------------------------- #
# Richer aggregate: substitution × ESM ΔLL bin
# --------------------------------------------------------------------------- #
def test_bin_ll_thresholds():
    assert agg.bin_ll(-5) == "lo" and agg.bin_ll(0) == "mid" and agg.bin_ll(5) == "hi"
    assert agg.bin_ll(None) == "mid" and agg.bin_ll("x") == "mid"


def test_binned_build_splits_same_substitution_by_ll():
    # W>L at a high-ΔLL position is beneficial; at a low-ΔLL position deleterious
    units = [(f"u{i}", [("W10L", 2.0), ("W30L", 0.0)]) for i in range(3)]
    ll = [{"W10L": 5.0, "W30L": -5.0}] * 3          # hi vs lo bin
    p = agg.build_priors(units, now=AFTER, ll_maps=ll)
    assert ("W", "L", "hi") in p.effects and ("W", "L", "lo") in p.effects
    assert p.effects[("W", "L", "hi")] > 0 and p.effects[("W", "L", "lo")] < 0
    assert p.n_users[("W", "L", "hi")] == 3
    # plain (unbinned) call is unchanged -> 2-tuple keys
    plain = agg.build_priors(units, now=AFTER)
    assert ("W", "L") in plain.effects and all(len(k) == 2 for k in plain.effects)


def test_binned_serialize_roundtrip():
    p = agg.GlobalPrior(effects={("W", "L", "hi"): 1.5}, n_users={("W", "L", "hi"): 4},
                      n_obs={("W", "L", "hi"): 9})
    rows = p.to_rows()
    assert rows[0]["substitution"] == "W>L@hi"
    back = agg.GlobalPrior.from_rows(rows)
    assert ("W", "L", "hi") in back.effects and back.effects[("W", "L", "hi")] == 1.5


def test_binned_blend_matches_mutation_ll_bin():
    from dee.optimizer.search import Mutation
    prior = agg.GlobalPrior(
        effects={("W", "L", "hi"): 2.0, ("W", "L", "lo"): -2.0},
        n_users={("W", "L", "hi"): 5, ("W", "L", "lo"): 5},
        n_obs={("W", "L", "hi"): 9, ("W", "L", "lo"): 9})
    pool = [
        Mutation(position=9,  wt_aa="W", mut_aa="L", delta_ll=5.0),    # hi bin -> +
        Mutation(position=29, wt_aa="W", mut_aa="L", delta_ll=-5.0),   # lo bin -> -
        Mutation(position=40, wt_aa="W", mut_aa="L", delta_ll=0.0),    # mid -> no matching key -> unchanged
    ]
    out = agg.apply_global_prior(pool, prior, weight=0.3)
    assert out[0].delta_ll == pytest.approx(5.0 + 0.3 * 2.0)
    assert out[1].delta_ll == pytest.approx(-5.0 + 0.3 * -2.0)
    assert out[2].delta_ll == 0.0                                       # mid bin, no fallback key