dee/core/aggregate.py

"""Cross-user aggregate prior for Directed Evolution — the data moat.

Pools wet-lab outcomes ACROSS users into a de-identified, amino-acid
*substitution-type* effect table (e.g. ``W>L`` tends to be tolerated). This is
the field's collective signal; nothing in it is attributable to a person. It
folds into round-1 scoring the same additive way the per-user surrogate does
(``dee.core.active_learning``), so every user's first design benefits from
everyone's measured results.

Privacy invariants (these are the whole point — keep them):

  * **Aggregate by substitution TYPE, never by position or user.** Positions are
    protein-specific (``W58L`` in one protein ≠ another), so only the
    substitution type generalizes and is non-identifying.
  * **Standardize within each (user, library) before pooling.** Each assay's
    scale is arbitrary; z-scoring per library means no raw measured value and no
    user's scale survives into the aggregate.
  * **k-anonymity floor.** A substitution type is kept only if it was measured by
    at least ``MIN_USERS`` *distinct* users. Rows below the floor are dropped
    entirely — they never reach the stored table.
  * **Gated execution.** :func:`build_priors` refuses to run before
    :data:`EFFECTIVE_DATE` (the 30-day-notice effective date of Privacy/Terms
    v2.0). Built now, but cannot produce an aggregate before the policy is live.

numpy only (no sklearn/scipy); runs in milliseconds on CPU.
"""
from __future__ import annotations

import datetime as _dt
import re
from dataclasses import dataclass, replace
from typing import Dict, List, Optional, Sequence, Tuple

import numpy as np

# The Privacy/Terms v2.0 effective date (30-day notice from 2026-06-08).
EFFECTIVE_DATE = _dt.date(2026, 7, 8)
# k-anonymity: keep a substitution only if >= this many DISTINCT users measured it.
MIN_USERS = 3
_RIDGE_LAMBDA = 1.0

_LABEL_RE = re.compile(r"^([A-Za-z])(\d+)([A-Za-z*])$")   # 'W58L'


class AggregationGateError(RuntimeError):
    """Raised when aggregation is attempted before the policy effective date."""


def parse_sub(label: str) -> Optional[Tuple[str, str]]:
    """'W58L' -> ('W', 'L') substitution type. None if malformed or synonymous."""
    m = _LABEL_RE.match((label or "").strip())
    if not m:
        return None
    wt, mut = m.group(1).upper(), m.group(3).upper()
    if wt == mut:
        return None
    return (wt, mut)


def _split_labels(labels) -> List[str]:
    parts = re.split(r"[,\s;]+", labels.strip()) if isinstance(labels, str) else list(labels)
    return [p.strip() for p in parts if p and _LABEL_RE.match(p.strip())]


def _library_single_site_effects(
    measurements: Sequence[Tuple[object, object]],
    ridge_lambda: float = _RIDGE_LAMBDA,
) -> Dict[str, float]:
    """De-noised single-site effects for ONE library, on within-library
    standardized outcomes. ``measurements`` is [(mutation_labels, value), ...].
    Returns {mutation_label: effect}; {} if too little signal."""
    rows: List[Tuple[List[str], float]] = []
    for labels, value in (measurements or []):
        try:
            y = float(value)
        except (TypeError, ValueError):
            continue
        muts = _split_labels(labels)
        if muts:
            rows.append((muts, y))
    if len(rows) < 2:
        return {}
    y = np.array([r[1] for r in rows], dtype=float)
    if y.std() < 1e-9:
        return {}
    y = (y - y.mean()) / (y.std() + 1e-9)              # standardize within library

    keys = sorted({m for muts, _ in rows for m in muts})
    idx = {k: i for i, k in enumerate(keys)}
    A = np.zeros((len(rows), len(keys)))
    for i, (muts, _) in enumerate(rows):
        for m in muts:
            A[i, idx[m]] = 1.0
    # ridge over single-site incidence (y is centered → no intercept needed)
    beta = np.linalg.solve(A.T @ A + ridge_lambda * np.eye(len(keys)), A.T @ y)
    return {k: float(beta[idx[k]]) for k in keys}


# Richer aggregate: split each substitution by the ESM-2 ΔLL bin it sat in, so
# the bench data can say "ESM was right (or wrong) when it was confident here."
# Wild-type-marginal ΔLL is ~negative for deleterious, ~positive for tolerated.
LL_LO, LL_HI = -2.0, 2.0


def bin_ll(ll: Optional[float]) -> str:
    """ESM ΔLL → coarse confidence bin ('lo' | 'mid' | 'hi')."""
    try:
        x = float(ll)
    except (TypeError, ValueError):
        return "mid"
    return "lo" if x < LL_LO else ("hi" if x > LL_HI else "mid")


@dataclass
class GlobalPrior:
    """The aggregated, de-identified substitution-effect prior.

    Keys are (wt, mut) 2-tuples or, in the richer ESM-binned aggregate,
    (wt, mut, bin) 3-tuples. apply_global_prior handles both."""
    effects: Dict[Tuple[str, str], float]   # (wt, mut) -> pooled standardized effect
    n_users: Dict[Tuple[str, str], int]     # distinct users behind each (>= MIN_USERS)
    n_obs: Dict[Tuple[str, str], int]       # total single-site observations behind each

    def __bool__(self) -> bool:
        return bool(self.effects)

    def to_rows(self) -> List[dict]:
        """Serialize for storage in public.mutation_priors (de-identified).
        A 3-tuple key (wt, mut, bin) serializes as 'W>L@hi'."""
        out = []
        for sub, eff in sorted(self.effects.items(), key=lambda kv: tuple(map(str, kv[0]))):
            label = f"{sub[0]}>{sub[1]}" + (f"@{sub[2]}" if len(sub) > 2 and sub[2] else "")
            out.append({
                "substitution": label,
                "n_users": self.n_users[sub],
                "n_obs": self.n_obs[sub],
                "mean_effect": round(eff, 6),
            })
        return out

    @classmethod
    def from_rows(cls, rows: Sequence[dict]) -> "GlobalPrior":
        eff, nu, no = {}, {}, {}
        for r in (rows or []):
            s = str(r.get("substitution", ""))
            if ">" not in s:
                continue
            core, _, b = s.partition("@")
            wt, mut = core.split(">", 1)
            key = (wt.strip().upper(), mut.strip().upper())
            if b.strip():
                key = key + (b.strip().lower(),)
            try:
                eff[key] = float(r["mean_effect"])
            except (TypeError, ValueError, KeyError):
                continue
            nu[key] = int(r.get("n_users", 0) or 0)
            no[key] = int(r.get("n_obs", 0) or 0)
        return cls(effects=eff, n_users=nu, n_obs=no)


def build_priors(
    outcomes_by_library: Sequence[Tuple[str, Sequence[Tuple[object, object]]]],
    *,
    min_users: int = MIN_USERS,
    now: Optional[_dt.date] = None,
    enforce_gate: bool = True,
    ll_maps: Optional[Sequence[Optional[Dict[str, float]]]] = None,
) -> GlobalPrior:
    """Aggregate per-library outcomes into the de-identified substitution prior.

    ``outcomes_by_library``: one entry per (user, library):
        (user_id, [(mutation_labels, measured_value), ...]).
    Only substitutions measured by >= ``min_users`` distinct users are kept.

    Raises :class:`AggregationGateError` if run before :data:`EFFECTIVE_DATE`
    (unless ``enforce_gate=False``, used only in tests).
    """
    today = now or _dt.date.today()
    if enforce_gate and today < EFFECTIVE_DATE:
        raise AggregationGateError(
            f"Cross-user aggregation is gated until {EFFECTIVE_DATE.isoformat()} "
            f"(Privacy/Terms v2.0 effective date); today is {today.isoformat()}."
        )

    effects: Dict[Tuple[str, str], List[float]] = {}
    users: Dict[Tuple[str, str], set] = {}
    obs: Dict[Tuple[str, str], int] = {}

    for i, (user_id, measurements) in enumerate(outcomes_by_library or []):
        per = _library_single_site_effects(measurements)
        # Optional ESM ΔLL map for this library → richer (substitution × ESM-bin)
        # key. When absent, keys stay plain (wt, mut) — fully backward compatible.
        ll_map = ll_maps[i] if (ll_maps and i < len(ll_maps) and ll_maps[i]) else None
        # collapse to substitution (× bin) within this library (mean of repeats)
        sub_eff: Dict[tuple, List[float]] = {}
        for label, e in per.items():
            sub = parse_sub(label)
            if sub is None:
                continue
            key = sub + ((bin_ll(ll_map.get(label)),) if ll_map is not None else ())
            sub_eff.setdefault(key, []).append(e)
        for sub, es in sub_eff.items():
            effects.setdefault(sub, []).append(float(np.mean(es)))
            users.setdefault(sub, set()).add(user_id)
            obs[sub] = obs.get(sub, 0) + len(es)

    keep_eff, keep_nu, keep_no = {}, {}, {}
    for sub, es in effects.items():
        if len(users[sub]) >= min_users:           # k-anonymity floor
            keep_eff[sub] = float(np.mean(es))
            keep_nu[sub] = len(users[sub])
            keep_no[sub] = obs[sub]
    return GlobalPrior(effects=keep_eff, n_users=keep_nu, n_obs=keep_no)


def apply_global_prior(pool: list, prior: Optional[GlobalPrior], weight: float = 0.3) -> list:
    """Return a copy of a ``search.Mutation`` pool with each delta_ll softly
    nudged by the global substitution effect: ``delta_ll += weight * effect``.
    No-op (returns the pool unchanged) when there is no prior — so this is inert
    until an aggregate exists (which cannot happen before EFFECTIVE_DATE)."""
    if not prior or not prior.effects:
        return list(pool)
    binned = any(len(k) > 2 for k in prior.effects)   # ESM-binned aggregate?
    out = []
    for m in pool:
        wt = str(getattr(m, "wt_aa", "")).upper()
        mut = str(getattr(m, "mut_aa", "")).upper()
        adj = None
        if binned:                                     # match the mutation's own ΔLL bin first
            adj = prior.effects.get((wt, mut, bin_ll(getattr(m, "delta_ll", None))))
        if adj is None:
            adj = prior.effects.get((wt, mut), 0.0)
        out.append(replace(m, delta_ll=float(m.delta_ll) + weight * adj) if adj else m)
    return out


# In-process cache shared by any caller that wants to blend the stored
# aggregate into ESM-2 scores — the REST DE job (dee/server.py /api/run) and
# the chat tool (dee/core/agent_tools.py design_variant_library) both need
# it, and living here means one Supabase read per TTL window, not two
# independent ones.
_PRIOR_CACHE: Dict[str, object] = {"data": None, "ts": 0.0}
_PRIOR_CACHE_TTL_S = 600.0


def load_cached_global_prior() -> GlobalPrior:
    """Cached GlobalPrior from public.mutation_priors. Empty (inert) prior on
    any error or before the first post-effective-date rebuild — callers
    should treat that as "no blend", never raise."""
    import time as _t
    now = _t.time()
    if _PRIOR_CACHE["data"] is not None and (now - _PRIOR_CACHE["ts"]) < _PRIOR_CACHE_TTL_S:
        return _PRIOR_CACHE["data"]
    try:
        from dee import auth as _auth
        prior = GlobalPrior.from_rows(_auth.get_mutation_priors())
    except Exception:  # noqa: BLE001
        prior = GlobalPrior(effects={}, n_users={}, n_obs={})
    _PRIOR_CACHE["data"] = prior
    _PRIOR_CACHE["ts"] = now
    return prior


def bust_cached_global_prior() -> None:
    """Call after a rebuild writes fresh rows so the next read picks them up
    immediately instead of waiting out the TTL."""
    _PRIOR_CACHE["data"] = None