staplebridge/training/stack.py

"""Construction of the StapleBridge training stack.

Builds, in dependency order: the hydrocarbon staple catalog, the geometry
oracle, the transition graph, the reference priors (frozen ESM2 peptide prior,
anchor prior, block prior), the reference energy and kernel, the plan-aware
reference sampler, the terminal energy, the policy and value networks, the
controlled kernel, and the optimizer.

The construction is numerically identical to the run that produced the released
checkpoint.
"""
from __future__ import annotations

import random
from functools import wraps
from pathlib import Path
from typing import Any, Callable

import numpy as np
import torch
import yaml

from staplebridge.chemistry.edit_distance import weighted_edit_distance
from staplebridge.chemistry.state import StapleState
from staplebridge.data.dataset import load_leads
from staplebridge.data.schemas import LeadExample
from staplebridge.hydrocarbon.actions import HydrocarbonTransitionGraph
from staplebridge.hydrocarbon.catalog import hydrocarbon_catalog_from_config
from staplebridge.hydrocarbon.endpoint_prior import (
    EmpiricalHydrocarbonEndpointPrior,
    EndpointPriorConfig,
)
from staplebridge.hydrocarbon.geometry import HydrocarbonGeometryOracle
from staplebridge.hydrocarbon.plan_reference import build_plan_aware_sampler
from staplebridge.hydrocarbon.property_energy import (
    HydrocarbonPropertyEnergyConfig,
    HydrocarbonPropertyScorer,
)
from staplebridge.hydrocarbon.terminal_energy import (
    HydrocarbonTerminalEnergy,
    HydrocarbonTerminalEnergyConfig,
)
from staplebridge.integrations.peptiverse import PeptiVerseWrapper
from staplebridge.models.control_kernel import ControlKernelConfig, ControlledKernel
from staplebridge.models.policy_net import PolicyNet
from staplebridge.models.value_net import ValueNet
from staplebridge.oracles.prior_factory import build_reference_priors
from staplebridge.reference.energy import ReferenceEnergy, ReferenceEnergyConfig
from staplebridge.reference.kernel import ReferenceKernel
from staplebridge.utils.paths import PACKAGE_ROOT


def install_embedding_cache(wrapper: PeptiVerseWrapper) -> dict[str, int]:
    predictor = wrapper.predictor
    if predictor is None:
        raise RuntimeError("PeptiVerse predictor unavailable")
    counts = {"hits": 0, "misses": 0}
    for name in ("wt_embedder", "smiles_embedder", "chemberta_embedder"):
        embedder = getattr(predictor, name)
        for method_name in ("pooled", "unpooled"):
            original: Callable[[str], Any] = getattr(embedder, method_name)
            cache: dict[str, Any] = {}

            @wraps(original)
            def cached(value: str, _original=original, _cache=cache):
                if value in _cache:
                    counts["hits"] += 1
                    return _cache[value]
                counts["misses"] += 1
                result = _original(value)
                _cache[value] = result
                return result

            setattr(embedder, method_name, cached)
    return counts


def build_stack(config: dict[str, Any], seed: int) -> dict[str, Any]:
    hydro = dict(config.get("hydrocarbon") or {})
    catalog = hydrocarbon_catalog_from_config(hydro)
    catalog_index = {block.block_id: block for block in catalog}
    geometry = HydrocarbonGeometryOracle(
        catalog,
        sentinel_cgeom=float(
            (hydro.get("geometry") or {}).get("sentinel_cgeom", 10.0)
        ),
    )
    graph = HydrocarbonTransitionGraph(
        catalog, max_neighbors=int(config.get("max_neighbors", 128))
    )
    priors = build_reference_priors(
        config.get("reference_priors"),
        catalog=catalog_index,
        geometry_oracle=geometry,
        device=str((config.get("training") or {}).get("device", "cpu")),
    )
    if (config.get("reference_priors") or {}).get("strict_no_mock"):
        ensure_available = getattr(priors.peptide, "ensure_available", None)
        if callable(ensure_available):
            ensure_available()
    reference_energy = ReferenceEnergy(
        peptide_prior=priors.peptide,
        anchor_prior=priors.anchor,
        block_prior=priors.block,
        geometry_oracle=geometry,
        catalog_index=catalog_index,
        config=ReferenceEnergyConfig(**dict(config.get("reference") or {})),
    )
    base_kernel = ReferenceKernel(
        reference_energy, group_normalize=True, substitution_downweight=0.25
    )
    sampler, plan_cfg = build_plan_aware_sampler(
        graph, base_kernel, config, seed=seed, root=PACKAGE_ROOT
    )

    prior_cfg = EndpointPriorConfig.from_dict(hydro.get("endpoint_prior"))
    if (
        float(prior_cfg.weight_pair) != 0.0
        or prior_cfg.use_length
        or prior_cfg.use_relative_position
        or prior_cfg.use_local_context
    ):
        raise RuntimeError(
            "the main objective forbids empirical endpoint terms in terminal energy; "
            "pair x spacing belongs only to plan-aware reference"
        )
    # The empirical table is consumed by the plan-aware sampler above. Disable
    # the terminal object completely so missing generated analysis artefacts do
    # not trigger a reload and the same evidence cannot be counted twice.
    prior_cfg.enabled = False
    endpoint_prior = EmpiricalHydrocarbonEndpointPrior(prior_cfg, root=PACKAGE_ROOT)
    return {
        "hydro": hydro,
        "catalog": catalog,
        "catalog_index": catalog_index,
        "geometry": geometry,
        "graph": graph,
        "sampler": sampler,
        "plan_cfg": plan_cfg,
        "endpoint_prior": endpoint_prior,
        "reference_priors": priors,
        "reference_priors_manifest": priors.to_manifest(),
    }


def base_terminal_factory(
    config: dict[str, Any], geometry: HydrocarbonGeometryOracle, catalog_index: dict
):
    section = dict(config.get("terminal_energy") or {})
    lambda_close = float(section.get("lambda_close", 1.0))
    lambda_edit = float(section.get("lambda_edit", 0.2))
    lambda_cost = float(section.get("lambda_cost", 0.2))
    infeasible = float(section.get("infeasible_penalty", 20.0))

    def base(z0: StapleState, zt: StapleState, lead: LeadExample):
        block = catalog_index.get(zt.block_id) if zt.block_id else None
        peptide_ca = (lead.target_context or {}).get("peptide_ca")
        ctype = geometry.ctype(
            zt.sequence_tokens, zt.anchor_pair, block, peptide_ca=peptide_ca
        )
        cgeom = geometry.cgeom(
            zt.sequence_tokens, zt.anchor_pair, block, peptide_ca=peptide_ca
        )
        edit = weighted_edit_distance(zt, z0)
        cost = block.cost_score if block else 1.0
        energy = lambda_close * cgeom + lambda_edit * edit + lambda_cost * cost
        if not ctype or zt.topology != "stapled":
            energy += infeasible
        return float(energy), {
            "cgeom": float(cgeom),
            "edit": float(edit),
            "ctype_ok": bool(ctype),
            "cost": float(cost),
        }

    return base


#: Config keys holding filesystem paths. Relative values are resolved against
#: PACKAGE_ROOT so the release runs correctly from any working directory;
#: absolute values are honoured untouched. This is a packaging concern only --
#: in the official run every one of these was already an absolute path, so
#: resolution is a no-op there and no numerical behaviour depends on it.
_PATH_KEYS: tuple[tuple[str, ...], ...] = (
    ("data", "root"),
    ("reference_priors", "peptide", "model_name_or_path"),
    ("reference_priors", "peptide", "cache_path"),
    ("hydrocarbon", "plan_control", "exact_sb_cache", "path"),
    ("property_predictor", "peptiverse_root"),
    ("property_predictor", "classifier_weight_root"),
    ("property_predictor", "manifest_path"),
    ("property_predictor", "hf_cache_dir"),
    ("property_predictor", "esm_model_name_or_path"),
    ("property_predictor", "peptideclm_model_name_or_path"),
    ("property_predictor", "chemberta_model_name_or_path"),
)


def resolve_config_paths(config: dict[str, Any], root: Path | None = None) -> dict[str, Any]:
    """Make every relative path in ``config`` absolute against ``root``."""
    root = root or PACKAGE_ROOT
    for keys in _PATH_KEYS:
        node = config
        for key in keys[:-1]:
            node = node.get(key) if isinstance(node, dict) else None
            if not isinstance(node, dict):
                break
        else:
            raw = node.get(keys[-1])
            if isinstance(raw, str) and raw and not Path(raw).is_absolute():
                node[keys[-1]] = str((root / raw).resolve())
    return config


def load_config(path: Path) -> dict[str, Any]:
    with path.open("r", encoding="utf-8") as handle:
        config = yaml.safe_load(handle) or {}
    if not isinstance(config, dict):
        raise ValueError("configuration root must be a mapping")
    return resolve_config_paths(config)


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)
    torch.use_deterministic_algorithms(False)


def select_leads(path: Path, n: int, max_length: int) -> list[LeadExample]:
    leads = [lead for lead in load_leads(path) if len(lead.linear_sequence) <= max_length][:n]
    if len(leads) != n:
        raise RuntimeError(f"need exactly {n} eligible leads from {path}, found {len(leads)}")
    return leads


def build_predictor(config: dict[str, Any]) -> tuple[PeptiVerseWrapper, dict[str, int]]:
    cfg = dict(config["property_predictor"])
    hf_cache_dir = Path(
        str(cfg.get("hf_cache_dir") or (PACKAGE_ROOT / "models/hf"))
    )
    wrapper = PeptiVerseWrapper(
        peptiverse_root=Path(cfg["peptiverse_root"]),
        classifier_weight_root=Path(cfg["classifier_weight_root"]),
        manifest_path=Path(cfg["manifest_path"]),
        device=str(cfg["device"]),
        strict=bool(cfg["strict"]),
        uncertainty=bool(cfg.get("uncertainty", False)),
        cache_enabled=bool(cfg.get("cache_enabled", True)),
        hf_cache_dir=hf_cache_dir,
        esm_model_name_or_path=cfg.get("esm_model_name_or_path"),
        peptideclm_model_name_or_path=cfg.get(
            "peptideclm_model_name_or_path"
        ),
        chemberta_model_name_or_path=cfg.get("chemberta_model_name_or_path"),
        offline=bool(cfg.get("offline", True)),
        batch_size=int(cfg.get("batch_size", 32)),
    )
    if not wrapper.available or wrapper.predictor is None:
        raise RuntimeError("strict PeptiVerse backend is unavailable; fallback is forbidden")
    return wrapper, install_embedding_cache(wrapper)


def build_energy(config: dict[str, Any], stack: dict[str, Any], scorer: HydrocarbonPropertyScorer) -> HydrocarbonTerminalEnergy:
    # The base terminal factory reads the generic terminal section. Map the
    # unchanged current hydrocarbon training coefficients into that interface.
    base_config = dict(config)
    train_cfg = dict(config.get("training") or {})
    base_config["terminal_energy"] = {
        "lambda_close": train_cfg.get("lambda_close", 1.0),
        "lambda_edit": train_cfg.get("lambda_edit", 0.2),
        "lambda_cost": train_cfg.get("lambda_cost", 0.2),
        "infeasible_penalty": train_cfg.get("infeasible_penalty", 20.0),
    }
    base = base_terminal_factory(base_config, stack["geometry"], stack["catalog_index"])
    terminal = dict(stack["hydro"].get("terminal_energy") or {})
    prop_cfg = HydrocarbonPropertyEnergyConfig.from_dict(terminal.get("property"))
    return HydrocarbonTerminalEnergy(
        base,
        stack["catalog"],
        stack["endpoint_prior"],
        HydrocarbonTerminalEnergyConfig(
            invalid_topology_penalty=float(terminal.get("invalid_topology_penalty", 10.0)),
            penalize_unstapled=bool(terminal.get("penalize_unstapled", True)),
        ),
        property_scorer=scorer,
        property_config=prop_cfg,
    )


class PlanAwareKernelAdapter:
    """Expose plan-conditioned probabilities through the generic kernel API."""

    def __init__(self, plan_kernel: Any) -> None:
        self.plan_kernel = plan_kernel

    def reference_logits(
        self,
        state: StapleState,
        candidates: list[StapleState],
        context: dict[str, Any] | None = None,
    ) -> torch.Tensor:
        context = dict(context or {})
        probabilities, _ = self.plan_kernel.plan_probs(
            state,
            candidates,
            context.get("hydrocarbon_plan"),
            context=context,
        )
        return torch.log(
            probabilities.clamp_min(torch.finfo(probabilities.dtype).tiny)
        )


def build_models(config: dict[str, Any], stack: dict[str, Any], device: torch.device):
    model_cfg = dict(config.get("model") or {})
    train_cfg = dict(config.get("training") or {})
    emb_dim = int(model_cfg.get("emb_dim", 32))
    encoder_cfg = dict(config.get("sequence_encoder") or {}) or None
    policy = PolicyNet(emb_dim=emb_dim, sequence_encoder_cfg=encoder_cfg).to(device)
    value = ValueNet(emb_dim=emb_dim, sequence_encoder_cfg=encoder_cfg).to(device)
    block_to_idx = {"": 0, **{block.block_id: index + 1 for index, block in enumerate(stack["catalog"])}}
    kernel = ControlledKernel(
        reference_kernel=PlanAwareKernelAdapter(stack["sampler"].kernel),
        policy_net=policy,
        value_net=value,
        block_to_idx=block_to_idx,
        cfg=ControlKernelConfig(mode=str(model_cfg.get("mode", "policy_tilt"))),
        horizon=int(train_cfg["horizon"]),
    )
    optimizer = torch.optim.Adam(
        list(policy.parameters()) + list(value.parameters()),
        lr=float(train_cfg.get("lr", 1e-3)),
    )
    return policy, value, kernel, optimizer