staplebridge/reference/energy.py

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from staplebridge.chemistry.state import StapleState
from staplebridge.data.schemas import BuildingBlock
from staplebridge.oracles.base import AnchorPriorBase, BlockPriorBase, GeometryOracleBase, PeptidePriorBase
from staplebridge.utils.profiling import STAGE_TIMER


ACTION_NOOP = "noop"
ACTION_RESIDUE_SUBSTITUTION_MOTIF = "residue_substitution_motif"
ACTION_RESIDUE_SUBSTITUTION_OTHER = "residue_substitution_other"
ACTION_ASSIGN_ANCHOR = "assign_anchor"
ACTION_REASSIGN_ANCHOR = "reassign_anchor"
ACTION_ASSIGN_BLOCK = "assign_block"
ACTION_ACTIVATE_TOPOLOGY = "activate_topology"
ACTION_UNKNOWN = "unknown"


@dataclass
class ComponentWeights:
    """Per-component multipliers applied on top of the eta_ scalars.

    Defaults are all 1.0 → identical behavior to pre-weight code. Turning any
    weight away from 1.0 lets the config section
    ``reference.weights.{seq,anchor,block,geom}`` re-scale the reference
    energy without changing the base eta_cost / eta_spps / eta_type / eta_geom
    knobs.
    """

    seq: float = 1.0
    anchor: float = 1.0
    block: float = 1.0
    geom: float = 1.0
    progress: float = 1.0


@dataclass
class ActionProgressWeights:
    """Reward/penalty on the action-type level.

    All values are *positive contributions to the reference energy*, so
    negative numbers reward the action and positive numbers penalize it.
    Defaults were picked so that motif-creating actions clearly beat unrelated
    substitutions but do not overpower the peptide/anchor/block priors.
    """

    noop: float = 0.5
    residue_substitution_motif: float = -1.5
    residue_substitution_other: float = 1.0
    assign_anchor: float = -3.0
    reassign_anchor: float = -0.5
    assign_block: float = -2.0
    activate_topology: float = -6.0
    unknown: float = 0.0

    def as_dict(self) -> dict[str, float]:
        return {
            ACTION_NOOP: self.noop,
            ACTION_RESIDUE_SUBSTITUTION_MOTIF: self.residue_substitution_motif,
            ACTION_RESIDUE_SUBSTITUTION_OTHER: self.residue_substitution_other,
            ACTION_ASSIGN_ANCHOR: self.assign_anchor,
            ACTION_REASSIGN_ANCHOR: self.reassign_anchor,
            ACTION_ASSIGN_BLOCK: self.assign_block,
            ACTION_ACTIVATE_TOPOLOGY: self.activate_topology,
            ACTION_UNKNOWN: self.unknown,
        }


@dataclass
class ReferenceEnergyConfig:
    eta_cost: float = 0.5
    eta_spps: float = 0.3
    eta_type: float = 5.0
    eta_geom: float = 1.0
    eta_progress: float = 1.0
    weights: dict[str, float] | None = None
    action_weights: dict[str, float] | None = None
    normalize: dict[str, Any] | None = None  # reserved; disabled by default
    # Stage-aware geometry:
    #   True (default): if candidate lacks anchor OR block, drop full E_geom /
    #   E_type from the reference energy and record geom_status="not_applicable".
    #   False: legacy behavior (any missing anchor/block is treated as full
    #   geometry failure via the sentinel cgeom=10 in the geometry oracle).
    stage_aware_geometry: bool = True

    def resolved_weights(self) -> ComponentWeights:
        w = ComponentWeights()
        if self.weights:
            for k in ("seq", "anchor", "block", "geom", "progress"):
                if k in self.weights:
                    setattr(w, k, float(self.weights[k]))
        return w

    def resolved_action_weights(self) -> ActionProgressWeights:
        aw = ActionProgressWeights()
        if self.action_weights:
            for k, v in self.action_weights.items():
                if hasattr(aw, k):
                    setattr(aw, k, float(v))
        return aw


# ---------------------------------------------------------------------------
# Action-type classification
# ---------------------------------------------------------------------------

def _spans_lactam_motif(seq: list[str], i: int, j: int) -> bool:
    """True if (i, j) is a valid K-(D|E) motif at spacing 3 or 4."""
    if i > j:
        i, j = j, i
    sp = j - i
    if sp not in (3, 4):
        return False
    if i < 0 or j >= len(seq):
        return False
    if sp == 3:
        return seq[i] == "K" and seq[j] == "D"
    return seq[i] == "K" and seq[j] == "E"


def _has_lactam_partner(seq: list[str], pos: int, tok: str) -> bool:
    """True if placing `tok` at `pos` completes a K-(D|E) i,i+3 / i,i+4 motif
    given the *current* residues at the partner positions."""
    n = len(seq)
    if tok == "K":
        for sp, partner in ((3, "D"), (4, "E")):
            j = pos + sp
            if j < n and seq[j] == partner:
                return True
            k = pos - sp
            if k >= 0 and seq[k] == partner:
                return True
        return False
    if tok in ("D", "E"):
        sp = 3 if tok == "D" else 4
        i = pos - sp
        if i >= 0 and seq[i] == "K":
            return True
        # K on the other side (K would be j-side, unusual, but be permissive)
        j = pos + sp
        if j < n and seq[j] == "K":
            return True
        return False
    return False


def classify_action(z: StapleState, z_next: StapleState) -> str:
    """Assign an action-type label to the transition z -> z_next.

    Robust to composite actions in :mod:`staplebridge.graph.neighbors`
    (anchor-assign additionally sets block_id in the same transition).
    """
    if z.topology != z_next.topology:
        if z.topology == "linear" and z_next.topology == "stapled":
            return ACTION_ACTIVATE_TOPOLOGY
        return ACTION_UNKNOWN

    if z.sequence_tokens != z_next.sequence_tokens:
        # Which position(s) changed?
        n = min(len(z.sequence_tokens), len(z_next.sequence_tokens))
        diffs = [i for i in range(n) if z.sequence_tokens[i] != z_next.sequence_tokens[i]]
        # Residue substitution action always changes exactly one position.
        for pos in diffs:
            tok = z_next.sequence_tokens[pos]
            if _has_lactam_partner(z.sequence_tokens, pos, tok):
                return ACTION_RESIDUE_SUBSTITUTION_MOTIF
        return ACTION_RESIDUE_SUBSTITUTION_OTHER

    if z.anchor_pair != z_next.anchor_pair:
        if z.anchor_pair is None and z_next.anchor_pair is not None:
            return ACTION_ASSIGN_ANCHOR
        return ACTION_REASSIGN_ANCHOR

    if z.block_id != z_next.block_id:
        return ACTION_ASSIGN_BLOCK

    return ACTION_NOOP


class ReferenceEnergy:
    def __init__(
        self,
        peptide_prior: PeptidePriorBase,
        anchor_prior: AnchorPriorBase,
        block_prior: BlockPriorBase,
        geometry_oracle: GeometryOracleBase,
        catalog_index: dict[str, BuildingBlock],
        config: ReferenceEnergyConfig,
    ) -> None:
        self.peptide_prior = peptide_prior
        self.anchor_prior = anchor_prior
        self.block_prior = block_prior
        self.geometry_oracle = geometry_oracle
        self.catalog_index = catalog_index
        self.cfg = config
        self._weights = config.resolved_weights()
        self._action_weights = config.resolved_action_weights()
        # Optional debug ring buffer, off by default.
        self._debug_last: dict[str, float] | None = None

    # ------------------------------------------------------------------
    # Decomposition (used by diagnostics; structure mirrors the paper's
    # reference process: peptide prior + anchor prior + block prior +
    # synthesis cost/SPPS + geometry term + action-type progress prior).
    # ------------------------------------------------------------------
    def decompose(
        self,
        z: StapleState,
        z_next: StapleState,
        context: dict[str, Any] | None = None,
        precomputed_seq_score: float | None = None,
    ) -> dict[str, float]:
        context = context or {}
        block = self.catalog_index.get(z_next.block_id) if z_next.block_id else None
        peptide_ca = context.get("peptide_ca")
        w = self._weights

        if precomputed_seq_score is not None:
            e_seq_score = float(precomputed_seq_score)
        else:
            with STAGE_TIMER.section("ESM2_prior_time"):
                e_seq_score = self.peptide_prior.score_transition(z, z_next, context)
        e_seq = -w.seq * e_seq_score
        with STAGE_TIMER.section("anchor_prior_time"):
            anchor_score = self.anchor_prior.score_anchor(
                z_next.sequence_tokens, z_next.anchor_pair, context
            )
        e_anchor = -w.anchor * anchor_score
        with STAGE_TIMER.section("block_prior_time"):
            block_score = self.block_prior.score_block(
                z_next.sequence_tokens, z_next.anchor_pair, block, context
            )
        cost = block.cost_score if block else 1.0
        spps = block.spps_score if block else 1.0
        e_block = -w.block * block_score
        e_cost = self.cfg.eta_cost * cost
        e_spps = self.cfg.eta_spps * spps

        # Stage-aware geometry: only apply full ctype / cgeom when the
        # candidate is a real "stapling geometry" state — anchor + block set.
        # Otherwise the sentinel cgeom=10 for missing anchor/block would
        # dominate the reference energy and trap intermediate transitions.
        stage_aware = bool(getattr(self.cfg, "stage_aware_geometry", True))
        geom_applicable = (
            z_next.anchor_pair is not None and z_next.block_id is not None
        )
        if stage_aware and not geom_applicable:
            ctype_bool = False
            cgeom_val = 0.0
            e_type = 0.0
            e_geom = 0.0
            geom_status = "not_applicable"
        else:
            with STAGE_TIMER.section("geometry_time"):
                ctype_bool = bool(
                    self.geometry_oracle.ctype(
                        z_next.sequence_tokens,
                        z_next.anchor_pair,
                        block,
                        peptide_ca=peptide_ca,
                    )
                )
                cgeom_val = float(
                    self.geometry_oracle.cgeom(
                        z_next.sequence_tokens,
                        z_next.anchor_pair,
                        block,
                        peptide_ca=peptide_ca,
                    )
                )
            e_type = self.cfg.eta_type * (1.0 - float(ctype_bool))
            e_geom = w.geom * self.cfg.eta_geom * cgeom_val
            geom_status = "applied"

        # Action-type progress prior — encourages the reference process to
        # actually build a stapled Khard state instead of drifting through
        # unrelated residue substitutions.
        action_type = classify_action(z, z_next)
        action_bias = self._action_weights.as_dict().get(action_type, 0.0)
        # For activate_topology, only reward if geometry is truly applicable
        # AND ctype passes; otherwise fall back to a small positive penalty.
        if action_type == ACTION_ACTIVATE_TOPOLOGY and not (geom_applicable and ctype_bool):
            action_bias = abs(self._action_weights.as_dict().get(action_type, 0.0)) * 0.5
        e_progress = w.progress * self.cfg.eta_progress * action_bias

        # Decomposition convention: every term is a *positive contribution*
        # to the reference energy (so that "dominant" is just argmax).
        decomp = {
            "E_seq": float(e_seq),
            "E_anchor": float(e_anchor),
            "E_block": float(e_block),
            "E_cost": float(e_cost),
            "E_spps": float(e_spps),
            "E_type": float(e_type),
            "E_geom": float(e_geom),
            "E_progress": float(e_progress),
            "raw_cost": float(cost),
            "raw_spps": float(spps),
            "raw_cgeom": float(cgeom_val),
            "ctype_ok": bool(ctype_bool),
            "geom_status": geom_status,
            "action_type": action_type,
        }
        decomp["E_total"] = float(
            decomp["E_seq"]
            + decomp["E_anchor"]
            + decomp["E_block"]
            + decomp["E_cost"]
            + decomp["E_spps"]
            + decomp["E_type"]
            + decomp["E_geom"]
            + decomp["E_progress"]
        )
        self._debug_last = decomp
        return decomp

    def compute_reference_energy(
        self,
        z: StapleState,
        z_next: StapleState,
        context: dict[str, Any] | None = None,
    ) -> float:
        return float(self.decompose(z, z_next, context)["E_total"])

    def decompose_batch(
        self,
        z: StapleState,
        candidates: list[StapleState],
        context: dict[str, Any] | None = None,
    ) -> list[dict[str, float]]:
        """Decompose all candidates, sharing one batched peptide-prior call.

        Motivation: with the ESM2 delta prior each candidate would otherwise
        trigger its own model forward, even though most candidates in a
        neighborhood share the same source sequence. ``batch_score_transitions``
        lets the prior amortize over the whole neighborhood.
        """
        if not candidates:
            return []
        with STAGE_TIMER.section("ESM2_prior_time"):
            seq_scores = self.peptide_prior.batch_score_transitions(
                z, candidates, context
            )
        STAGE_TIMER.bump("candidates_scored", len(candidates))
        return [
            self.decompose(z, c, context, precomputed_seq_score=seq_scores[i])
            for i, c in enumerate(candidates)
        ]

    @property
    def component_weights(self) -> ComponentWeights:
        return self._weights

    @property
    def action_weights(self) -> ActionProgressWeights:
        return self._action_weights

    @property
    def last_decomposition(self) -> dict[str, float] | None:
        return dict(self._debug_last) if self._debug_last else None