staplebridge/training/records.py

"""Terminal-state records, validation summaries and constrained beam decode.

``terminal_record`` scores a finished trajectory and flattens the result into
one auditable row; ``summarize`` aggregates those rows into the validation
metrics, including the selection metric; ``beam_decode`` runs plan-conditioned
constrained decoding.
"""
from __future__ import annotations

import json
import math
from collections import Counter
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import torch

from staplebridge.chemistry.edit_distance import weighted_edit_distance
from staplebridge.chemistry.state import StapleState
from staplebridge.data.schemas import LeadExample
from staplebridge.hydrocarbon.actions import (
    FailureReason,
    validate_hydrocarbon_staple,
)
from staplebridge.hydrocarbon.terminal_energy import HydrocarbonTerminalEnergy
from staplebridge.hydrocarbon.tokenizer import tokenize_sequence
from staplebridge.models.control_kernel import ControlledKernel

def state_key(state: StapleState) -> tuple[Any, ...]:
    return (tuple(state.sequence_tokens), state.anchor_pair, state.block_id, state.topology)


def terminal_record(
    z0: StapleState,
    terminal: StapleState,
    lead: LeadExample,
    energy_fn: HydrocarbonTerminalEnergy,
    stack: dict[str, Any],
    **extra: Any,
) -> dict[str, Any]:
    block = stack["catalog_index"].get(terminal.block_id) if terminal.block_id else None
    verdict = validate_hydrocarbon_staple(
        terminal.sequence_tokens,
        None if terminal.anchor_pair is None else tuple(terminal.anchor_pair),
        block,
        stack["catalog"],
    )
    energy, info = energy_fn(z0, terminal, lead)
    topology_label = None
    if verdict is FailureReason.OK and terminal.anchor_pair is not None:
        i, j = terminal.anchor_pair
        topology_label = f"{terminal.sequence_tokens[i]}-{terminal.sequence_tokens[j]}/i,i+{j-i}"
    record: dict[str, Any] = {
        "example_id": lead.example_id,
        "terminal_sequence": terminal.to_sequence(),
        "anchor_pair": terminal.anchor_pair,
        "block_id": terminal.block_id,
        "topology_label": topology_label,
        "stapled": terminal.topology == "stapled",
        "chemistry_valid": verdict is FailureReason.OK,
        "topology_status": verdict.value,
        "weighted_edit_distance": float(weighted_edit_distance(terminal, z0)),
        "terminal_energy": float(energy),
        "base_terminal_energy": float(info.get("base_terminal_energy", 0.0)),
        "endpoint_prior_energy": float(info.get("E_endpoint_total", 0.0)),
        "property_status": info.get("hydrocarbon_property_status"),
    }
    for key, value in info.items():
        if key.startswith("hydrocarbon_"):
            record[key] = value
    scorer = getattr(energy_fn, "property_scorer", None)
    if scorer is not None:
        original = scorer.score_original_linear(
            z0.sequence_tokens, lead_key=str(lead.example_id)
        )
        record.update(original)
        product = record.get(
            "hydrocarbon_permeability_penetrance_product_mean"
        )
        baseline = record.get(
            "hydrocarbon_permeability_penetrance_original_linear"
        )
        if product is not None and baseline is not None:
            record[
                "hydrocarbon_permeability_penetrance_delta_vs_original_lead"
            ] = float(product) - float(baseline)
        for prop in ("solubility", "toxicity", "halflife"):
            product = record.get(f"hydrocarbon_{prop}_product_mean")
            baseline = record.get(f"hydrocarbon_{prop}_original_linear")
            if product is not None and baseline is not None:
                record[f"hydrocarbon_{prop}_delta_vs_original_lead"] = (
                    float(product) - float(baseline)
                )
    record.update(extra)
    return record


def finite_values(rows: Iterable[dict[str, Any]], key: str) -> list[float]:
    values: list[float] = []
    for row in rows:
        value = row.get(key)
        if isinstance(value, (int, float)) and math.isfinite(float(value)):
            values.append(float(value))
    return values


def average(rows: Iterable[dict[str, Any]], key: str) -> float | None:
    values = finite_values(rows, key)
    return float(np.mean(values)) if values else None


def summarize(rows: list[dict[str, Any]], selected: list[dict[str, Any]], n_leads: int) -> dict[str, Any]:
    selected_scored = [row for row in selected if row.get("property_status") == "scored"]
    failures = Counter()
    for row in selected:
        if not row["stapled"]:
            failures["not_stapled"] += 1
        elif not row["chemistry_valid"]:
            failures[f"invalid:{row['topology_status']}"] += 1
        elif row.get("property_status") != "scored":
            failures[str(row.get("property_status", "property_not_scored"))] += 1
    mode_counts = Counter(row.get("topology_label") for row in selected if row.get("topology_label"))
    tox_threshold = 0.49
    tox_values = finite_values(selected_scored, "hydrocarbon_toxicity_product_mean")
    return {
        "n_candidates": len(rows),
        "n_selected": len(selected),
        "chemistry_valid_rate": sum(bool(row["chemistry_valid"]) for row in rows) / len(rows) if rows else 0.0,
        "stapled_rate": sum(bool(row["stapled"]) for row in rows) / len(rows) if rows else 0.0,
        "accepted_at_lead_count": sum(bool(row["chemistry_valid"] and row["stapled"] and row.get("property_status") == "scored") for row in selected),
        "accepted_at_lead_rate": sum(bool(row["chemistry_valid"] and row["stapled"] and row.get("property_status") == "scored") for row in selected) / n_leads,
        "mean_original_linear_penetrance": average(
            selected_scored,
            "hydrocarbon_permeability_penetrance_original_linear",
        ),
        "mean_product_penetrance": average(
            selected_scored,
            "hydrocarbon_permeability_penetrance_product_mean",
        ),
        "mean_delta_penetrance_vs_original_lead": average(
            selected_scored,
            "hydrocarbon_permeability_penetrance_delta_vs_original_lead",
        ),
        "mean_delta_penetrance_vs_plan_precursor": average(
            selected_scored,
            "hydrocarbon_permeability_penetrance_delta_vs_plan_precursor",
        ),
        "mean_delta_penetrance": average(
            selected_scored,
            "hydrocarbon_permeability_penetrance_delta_vs_original_lead",
        ),
        "mean_delta_penetrance_definition": "vs_original_lead",
        "mean_delta_toxicity_log_only": average(selected_scored, "hydrocarbon_toxicity_delta_product_linear"),
        "mean_product_toxicity_log_only": average(selected_scored, "hydrocarbon_toxicity_product_mean"),
        "toxicity_violation_rate_log_only": sum(value > tox_threshold for value in tox_values) / len(tox_values) if tox_values else None,
        "mean_hemolysis_monitor": average(selected_scored, "hydrocarbon_hemolysis_product_mean"),
        "mean_halflife_monitor": average(selected_scored, "hydrocarbon_halflife_product_mean"),
        "mean_weighted_edit_distance": average(selected, "weighted_edit_distance"),
        "mean_penetrance_ez_sensitivity": average(selected_scored, "hydrocarbon_permeability_penetrance_EZ_abs_diff"),
        "max_penetrance_ez_sensitivity": max(finite_values(selected_scored, "hydrocarbon_permeability_penetrance_EZ_abs_diff"), default=None),
        "s5_s5_i4_count": mode_counts.get("S5-S5/i,i+4", 0),
        "s5_s5_i4_rate": mode_counts.get("S5-S5/i,i+4", 0) / n_leads,
        "r8_s5_i7_count": mode_counts.get("R8-S5/i,i+7", 0),
        "r8_s5_i7_rate": mode_counts.get("R8-S5/i,i+7", 0) / n_leads,
        "topology_distribution": dict(mode_counts),
        "failure_reasons": dict(failures),
        "backend_distribution": dict(Counter(row.get("hydrocarbon_property_backend", "not_scored") for row in selected)),
        "mode_distribution": dict(Counter(row.get("hydrocarbon_property_mode", "not_scored") for row in selected)),
        "endpoint_prior_energy_max_abs": max((abs(float(row["endpoint_prior_energy"])) for row in rows), default=0.0),
    }


def select_lowest_energy(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
    grouped: dict[str, list[dict[str, Any]]] = {}
    for row in rows:
        grouped.setdefault(str(row["example_id"]), []).append(row)
    return [min(group, key=lambda row: (row["terminal_energy"], row.get("trajectory_index", 0))) for _, group in sorted(grouped.items())]


def beam_decode(
    lead: LeadExample,
    stack: dict[str, Any],
    kernel: ControlledKernel,
    horizon: int,
    beam_size: int,
    committed_plan: Any | None = None,
) -> list[tuple[StapleState, float, int]]:
    z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence))
    peptide_ca = (lead.target_context or {}).get("peptide_ca")
    if committed_plan is None:
        plan_probe = stack["sampler"].sample_trajectory(
            z0,
            protected_positions=lead.protected_positions,
            context={
                "protected_positions": lead.protected_positions,
                "peptide_ca": peptide_ca,
            },
            horizon=0,
        )
        committed_plan = plan_probe.plan
    context = {
        "protected_positions": lead.protected_positions,
        "peptide_ca": peptide_ca,
        "hydrocarbon_plan": committed_plan,
    }
    beam: list[tuple[StapleState, float, int]] = [(z0, 0.0, 0)]
    for t in range(horizon):
        expanded: list[tuple[StapleState, float, int]] = []
        for state, log_probability, path_length in beam:
            candidates = stack["graph"].neighbors(state, protected_positions=lead.protected_positions)
            if not candidates:
                expanded.append((state, log_probability, path_length))
                continue
            log_probs = torch.log_softmax(kernel.controlled_logits(state, candidates, t=t, context=context), dim=0)
            keep = min(beam_size, len(candidates))
            values, indices = torch.topk(log_probs, k=keep)
            expanded.extend((candidates[int(index)], log_probability + float(value), path_length + 1) for value, index in zip(values.cpu(), indices.cpu()))
        best: dict[tuple[Any, ...], tuple[StapleState, float, int]] = {}
        for item in expanded:
            key = state_key(item[0])
            if key not in best or item[1] > best[key][1]:
                best[key] = item
        beam = sorted(best.values(), key=lambda item: item[1], reverse=True)[:beam_size]
    return beam

def write_json(path: Path, value: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(value, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")


def write_jsonl(path: Path, rows: Iterable[dict[str, Any]], mode: str = "w") -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open(mode, encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")