staplebridge/training/main_loop.py

"""StapleBridge training epoch and per-epoch validation.

The Full Exact-SB formulation:

* finite lead-specific feasible plan support H(x) via
  ``property_free_hard_plan_support`` (chemistry, geometry, edit budget,
  sequence identity, protected positions and exact committed-plan completion);
* ``q_ref(p|x)`` preserved on the legal plan space then conditioned on the hard
  mask;
* exact finite-support teacher ``q*(p|x) ~ q_ref(p|x) exp(-beta E_T)``;
* amortized plan controller ``q_theta(p|x)`` masked and renormalised on the
  same support;
* ``L_plan = KL(q* || q_theta)``;
* conditional execution policy trained on canonical-completing trajectories
  with lead-local positive path weights.

``validate_enabled`` computes the per-epoch validation metrics, including
``q_star_vs_q_theta_kl``, the checkpoint-selection metric. Post-hoc evaluation
reporting is not part of this training release.
"""
from __future__ import annotations

import json
import random
import time
from collections import Counter
from typing import Any

import numpy as np
import torch

from staplebridge.chemistry.state import StapleState
from staplebridge.hydrocarbon.curriculum import build_hydrocarbon_demonstration_path
from staplebridge.hydrocarbon.exact_sb_cache import (
    energy_only_from_config, resolve_exact_sb_target,
)
from staplebridge.hydrocarbon.plan_control import (
    HydrocarbonPlanControlConfig,
    best_committed_plan_trajectory,
    HydrocarbonPlanHead,
    completes_committed_plan,
    controlled_plan_log_probabilities,
    configured_plan_level_objective,
    describe_plan,
    empirical_log_probabilities,
    hierarchical_plan_ranking_enabled,
    mask_and_renormalize_plan_log_probabilities,
    plan_entropy,
    property_free_hard_plan_support,
    sample_committed_plan_trajectory,
    sample_distinct_plans,
    select_plan_and_trajectory,
)
from staplebridge.hydrocarbon.plan_reference import enumerate_legal_plans
from staplebridge.hydrocarbon.plan_validation import (
    aggregate_q_star_diagnostics,
    exact_sb_validation_enabled,
    lead_q_star_diagnostics,
)
from staplebridge.hydrocarbon.property_energy import (
    HydrocarbonPropertyEnergyConfig,
    HydrocarbonPropertyScorer,
    required_original_lead_properties,
)
from staplebridge.hydrocarbon.tokenizer import tokenize_sequence
from staplebridge.models.control_kernel import state_to_features
from staplebridge.training.losses import (
    assign_lead_local_positive_weights,
    compute_weighted_path_loss,
)
from staplebridge.training.records import (
    beam_decode, summarize, terminal_record,
)
from staplebridge.training.trajectory import TransitionStep, WeightedTrajectory
from staplebridge.utils.profiling import STAGE_TIMER


def plan_mode(plan: Any) -> str:
    return f"{plan.ordered_pair}/i,i+{plan.spacing}"


def legal_plans(lead: Any, stack: dict[str, Any]) -> list[Any]:
    return enumerate_legal_plans(
        tokenize_sequence(lead.linear_sequence),
        stack["catalog"],
        protected_positions=lead.protected_positions,
        filters=stack["sampler"].filters,
    )


def property_free_supported_plan_view(
    lead: Any,
    initial_state: StapleState,
    stack: dict[str, Any],
    config: dict[str, Any],
) -> dict[str, Any]:
    """Build the one hard-constrained coarse-plan support used everywhere."""
    all_plans = legal_plans(lead, stack)
    context = {
        "protected_positions": lead.protected_positions,
        "peptide_ca": (lead.target_context or {}).get("peptide_ca"),
    }
    all_reference_weights = stack["sampler"].plan_selection_weights(
        initial_state, all_plans, context
    )
    support_indices, verdicts = property_free_hard_plan_support(
        initial_state=initial_state,
        lead=lead,
        plans=all_plans,
        catalog=stack["catalog"],
        catalog_index=stack["catalog_index"],
        geometry=stack["geometry"],
        config=config,
    )
    return {
        "all_plans": all_plans,
        "all_reference_weights": all_reference_weights,
        "support_indices": support_indices,
        "plans": [all_plans[index] for index in support_indices],
        # Preserve q_ref on the original legal space, then condition it on the
        # hard mask.  Do not recompute empirical within-mode counts afterward.
        "reference_weights": [
            all_reference_weights[index] for index in support_indices
        ],
        "verdicts": verdicts,
        "context": context,
    }


def plan_key(plan: Any) -> tuple[Any, ...]:
    return (plan.ordered_pair, plan.spacing, tuple(plan.anchor_pair), plan.block_id)


def encode_lead(policy: Any, kernel: Any, state: StapleState) -> torch.Tensor:
    features = state_to_features(state, kernel.block_to_idx)
    device = kernel.device
    features = {key: value.to(device) for key, value in features.items()}
    with torch.no_grad():
        return policy.encoder(**features).detach()


def selected_top1_rates(selected: list[dict[str, Any]]) -> dict[str, float]:
    n = max(len(selected), 1)
    return {
        "top1_chemistry_valid_rate": sum(bool(row["chemistry_valid"]) for row in selected) / n,
        "top1_stapled_rate": sum(bool(row["stapled"]) for row in selected) / n,
    }


def train_enabled_epoch(
    leads: list[Any],
    config: dict[str, Any],
    stack: dict[str, Any],
    energy_fn: Any,
    policy: Any,
    kernel: Any,
    optimizer: Any,
    parameters: list[torch.nn.Parameter],
    head: HydrocarbonPlanHead,
    plan_cfg: HydrocarbonPlanControlConfig,
    plan_rng: random.Random,
    epoch: int,
    exact_sb_cache: Any = None,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    train_cfg = config["training"]
    horizon = int(train_cfg["horizon"])
    chunk_size = int(train_cfg["chunk_size"])
    rows_all: list[dict[str, Any]] = []
    losses: list[float] = []
    path_losses: list[float] = []
    plan_losses: list[float] = []
    entropies: list[float] = []
    coverage: list[float] = []
    disagreements: list[bool] = []
    exact_reverse_kls: list[float] = []
    exact_forward_objectives: list[float] = []
    exact_q_star_entropies: list[float] = []
    exact_plans_scored = 0
    joint_support_enabled = bool(
        getattr(
            getattr(energy_fn, "property_cfg", None),
            "enable_joint_perm_halflife_support",
            False,
        )
    )
    joint_support_rows: list[dict[str, Any]] = []
    hard_mask_totals: Counter[str] = Counter()
    exact_cache_sources: Counter[str] = Counter()
    controlled_top_modes: Counter[str] = Counter()
    weighting_totals: Counter[str] = Counter()
    # --- profiling (opt-in via config profiling.chunk_timing) ------------
    # Wall-clock accounting only: every perf_counter below is read outside the
    # computation it brackets and never feeds a tensor, and no RNG is drawn, so
    # the trained model, loss and sampling stream are bit-for-bit identical
    # whether profiling is on or off.
    #
    # CUDA attribution: the per-lead GPU stages (plan/qtheta, terminal/PeptiVerse,
    # fine rollout's ESM2 forward) each END in a blocking device->host transfer
    # (`.item()` / `.cpu()` / `.tolist()`) that already synchronises at the stage
    # boundary, so their wall-clock is accurate without an added sync. The three
    # chunk-level GPU stages (policy log-prob, backward, optimizer) have no such
    # trailing transfer, so they get an explicit ``torch.cuda.synchronize()`` --
    # once per 32-lead chunk, never inside the per-lead or per-step loops.
    chunk_timing = bool((config.get("profiling") or {}).get("chunk_timing", False))
    _cuda_sync = chunk_timing and torch.cuda.is_available()

    def _sync() -> None:
        if _cuda_sync:
            torch.cuda.synchronize()

    STAGE_KEYS = (
        "plan_qtheta", "qstar_target", "fine_rollout", "neighbor_gen",
        "terminal_pv", "policy_logprob", "backward", "optimizer",
    )
    stage_seconds: dict[str, float] = {key: 0.0 for key in STAGE_KEYS}
    _profiling_scorer = getattr(energy_fn, "property_scorer", None)
    _pv_wrapper = getattr(_profiling_scorer, "predictor", None)
    _esm2_prev = STAGE_TIMER.snapshot()
    epoch_started = time.perf_counter()
    total_leads = len(leads)

    for chunk_start in range(0, len(leads), chunk_size):
        chunk_leads = leads[chunk_start : chunk_start + chunk_size]
        weighted_by_lead: list[list[WeightedTrajectory]] = []
        chunk_rows: list[dict[str, Any]] = []
        chunk_plan_losses: list[torch.Tensor] = []
        chunk_t: dict[str, float] = {key: 0.0 for key in STAGE_KEYS}
        chunk_counts = {
            "plans": 0,
            "legal_plans_before_hard_mask": 0,
            "hard_masked_plans": 0,
            "leads_without_hard_support": 0,
            "trajectories": 0,
            "transitions": 0,
            "neighbors": 0,
        }
        _cache_hits0 = exact_sb_cache.stats.hits if exact_sb_cache is not None else 0
        _cache_misses0 = exact_sb_cache.stats.misses if exact_sb_cache is not None else 0
        _pv_misses0 = _pv_wrapper.cache_misses if _pv_wrapper is not None else 0
        _pv_hits0 = _pv_wrapper.cache_hits if _pv_wrapper is not None else 0
        chunk_started = time.perf_counter()
        for lead_index, lead in enumerate(chunk_leads, start=chunk_start):
            lead_weighted: list[WeightedTrajectory] = []
            z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence))
            support_view = property_free_supported_plan_view(
                lead, z0, stack, config
            )
            all_plans = support_view["all_plans"]
            support_indices = support_view["support_indices"]
            plans = support_view["plans"]
            ref_weights = support_view["reference_weights"]
            context = support_view["context"]
            chunk_counts["legal_plans_before_hard_mask"] += len(all_plans)
            chunk_counts["plans"] += len(plans)
            chunk_counts["hard_masked_plans"] += len(all_plans) - len(plans)
            hard_mask_totals["legal_plans_before_hard_mask"] += len(all_plans)
            hard_mask_totals["hard_supported_plans"] += len(plans)
            hard_mask_totals["hard_masked_plans"] += len(all_plans) - len(plans)
            if not plans:
                chunk_counts["leads_without_hard_support"] += 1
                hard_mask_totals["leads_without_hard_support"] += 1
                if joint_support_enabled:
                    joint_support_rows.append(
                        {
                            "example_id": str(lead.example_id),
                            "n_hard_supported_plans": 0,
                            "joint_plan_count": 0,
                            "joint_nonempty": False,
                            "joint_fallback": False,
                            "q_star_support_size": 0,
                            "status": "no_hard_support",
                        }
                    )
                weighted_by_lead.append(lead_weighted)
                continue
            sampled = sample_distinct_plans(
                plans, ref_weights, plan_cfg.plans_per_lead, plan_rng
            )
            if not sampled:
                raise RuntimeError(f"no sampled legal plan for {lead.example_id}")
            key_to_index = {plan_key(plan): index for index, plan in enumerate(plans)}
            _t = time.perf_counter()
            lead_embedding = encode_lead(policy, kernel, z0)
            all_controlled_logp = controlled_plan_log_probabilities(
                head,
                lead_embedding,
                all_plans,
                support_view["all_reference_weights"],
                len(z0.sequence_tokens),
                sequence_tokens=z0.sequence_tokens,
                peptide_ca=context["peptide_ca"],
            )
            masked_controlled_logp = mask_and_renormalize_plan_log_probabilities(
                all_controlled_logp, support_indices
            )
            controlled_logp = masked_controlled_logp[
                torch.tensor(
                    support_indices,
                    dtype=torch.long,
                    device=masked_controlled_logp.device,
                )
            ]
            chunk_t["plan_qtheta"] += time.perf_counter() - _t
            joint_train_row: dict[str, Any] | None = None
            if plan_cfg.exact_sb_objective:
                reference_logp = empirical_log_probabilities(
                    ref_weights, controlled_logp.device
                )
                # Deterministic half of the target (plan order, q_ref, terminal
                # energies, log q*) may come from the persistent cache. q_theta
                # is always the live `controlled_logp` computed above.
                _t = time.perf_counter()
                exact_energy_tensor, _, cache_info = resolve_exact_sb_target(
                    lead=lead,
                    plans=plans,
                    reference_log_probabilities=reference_logp,
                    beta=plan_cfg.exact_sb_beta,
                    energy_fn=energy_fn,
                    initial_state=z0,
                    build_terminal=lambda state, plan: build_hydrocarbon_demonstration_path(
                        state, plan, stack["catalog"]
                    )[-1],
                    cache=exact_sb_cache,
                    energy_only=energy_only_from_config(config),
                )
                exact_cache_sources[str(cache_info["source"])] += 1
                target_support_mask = (
                    torch.tensor(
                        cache_info["target_support_mask"],
                        dtype=torch.bool,
                        device=controlled_logp.device,
                    )
                    if cache_info.get("target_support_mask") is not None
                    else None
                )
                if joint_support_enabled:
                    joint_train_row = {
                        "example_id": str(lead.example_id),
                        "n_hard_supported_plans": len(plans),
                        "joint_plan_count": int(cache_info["joint_plan_count"]),
                        "joint_nonempty": bool(cache_info["joint_nonempty"]),
                        "joint_fallback": bool(cache_info["joint_fallback"]),
                        "q_star_support_size": int(cache_info["q_star_support_size"]),
                        "selection_semantics": "q_theta_top1_hard_plan",
                        "status": "scored",
                    }
                    joint_support_rows.append(joint_train_row)
                chunk_t["qstar_target"] += time.perf_counter() - _t
                _t = time.perf_counter()
                exact_loss, forward_objective, log_q_star = configured_plan_level_objective(
                    controlled_logp,
                    [],
                    [],
                    plan_cfg.target_temperature,
                    exact_sb_objective=True,
                    reference_log_probabilities=reference_logp,
                    terminal_energies=exact_energy_tensor,
                    exact_sb_beta=plan_cfg.exact_sb_beta,
                    target_support_mask=target_support_mask,
                )
                chunk_plan_losses.append(exact_loss)
                exact_reverse_kls.append(float(exact_loss.detach().cpu().item()))
                exact_forward_objectives.append(
                    float(forward_objective.detach().cpu().item())
                )
                q_star = log_q_star.exp()
                if target_support_mask is None:
                    # Preserve historical false/fallback arithmetic exactly.
                    q_star_entropy = -(q_star * log_q_star).sum()
                else:
                    q_star_entropy = -torch.where(
                        q_star > 0.0,
                        q_star * log_q_star,
                        torch.zeros_like(q_star),
                    ).sum()
                exact_q_star_entropies.append(
                    float(q_star_entropy.detach().cpu().item())
                )
                exact_plans_scored += len(plans)
                chunk_t["plan_qtheta"] += time.perf_counter() - _t
            reference_top = int(np.argmax(ref_weights))
            controlled_top = int(torch.argmax(controlled_logp).item())
            if joint_train_row is not None:
                joint_train_row.update(
                    {
                        "selected_plan_index": controlled_top,
                        "selected_plan": describe_plan(plans[controlled_top]),
                        "selected_satisfies_joint_condition": bool(
                            target_support_mask is not None
                            and target_support_mask[controlled_top].item()
                        ),
                    }
                )
            disagreements.append(reference_top != controlled_top)
            entropies.append(float(plan_entropy(controlled_logp).detach().cpu().item()))
            coverage.append(len(sampled) / len(plans))
            controlled_top_modes[
                f"{plans[controlled_top].ordered_pair}/i,i+{plans[controlled_top].spacing}"
            ] += 1

            sampled_indices: list[int] = []
            penetrance_values: list[float] = []
            for trajectory_index, plan in enumerate(sampled):
                _t = time.perf_counter()
                trajectory = sample_committed_plan_trajectory(
                    stack["sampler"],
                    z0,
                    plan,
                    lead.protected_positions,
                    context,
                    horizon,
                )
                chunk_t["fine_rollout"] += time.perf_counter() - _t
                chunk_counts["trajectories"] += 1
                trajectory_context = {**context, "hydrocarbon_plan": plan}
                steps: list[TransitionStep] = []
                for t, (state, next_state) in enumerate(
                    zip(trajectory.states[:-1], trajectory.states[1:])
                ):
                    _t = time.perf_counter()
                    candidates = stack["graph"].neighbors(
                        state, protected_positions=lead.protected_positions
                    )
                    chunk_t["neighbor_gen"] += time.perf_counter() - _t
                    chunk_counts["neighbors"] += len(candidates)
                    chunk_counts["transitions"] += 1
                    chosen = next(
                        (index for index, candidate in enumerate(candidates) if candidate == next_state),
                        None,
                    )
                    if chosen is None:
                        raise RuntimeError("committed transition absent from graph")
                    steps.append(TransitionStep(state, next_state, candidates, chosen, t))
                terminal = trajectory.states[-1]
                _t = time.perf_counter()
                row = terminal_record(
                    z0,
                    terminal,
                    lead,
                    energy_fn,
                    stack,
                    epoch=epoch,
                    lead_index=lead_index,
                    trajectory_index=trajectory_index,
                    committed_plan=describe_plan(plan),
                    committed_plan_completed=trajectory.progress.plan_completed,
                    path_length=len(steps),
                )
                chunk_t["terminal_pv"] += time.perf_counter() - _t
                penetrance = row.get(
                    "hydrocarbon_permeability_penetrance_product_mean"
                )
                penetrance_values.append(float(penetrance) if penetrance is not None else -1.0)
                sampled_indices.append(key_to_index[plan_key(plan)])
                exact_completion = completes_committed_plan(terminal, plan)
                failure_reason = (
                    None
                    if exact_completion
                    else "off_plan_completion"
                    if terminal.topology == "stapled"
                    else "unfinished"
                )
                row["trajectory_training_status"] = (
                    "positive" if exact_completion else failure_reason
                )
                lead_weighted.append(
                    WeightedTrajectory(
                        steps=steps,
                        terminal_state=terminal,
                        terminal_energy=float(row["terminal_energy"]),
                        context=trajectory_context,
                        is_positive=exact_completion,
                        failure_reason=failure_reason,
                    )
                )
                chunk_rows.append(row)
            weighted_by_lead.append(lead_weighted)
            if not plan_cfg.exact_sb_objective:
                _t = time.perf_counter()
                legacy_loss, _, _ = configured_plan_level_objective(
                    controlled_logp,
                    sampled_indices,
                    penetrance_values,
                    plan_cfg.target_temperature,
                    exact_sb_objective=False,
                )
                chunk_plan_losses.append(legacy_loss)
                chunk_t["plan_qtheta"] += time.perf_counter() - _t

        weight_diagnostics = assign_lead_local_positive_weights(weighted_by_lead)
        for key in ("n_positive", "n_failure", "n_leads_without_positive"):
            weighting_totals[key] += int(weight_diagnostics[key])
        weighted = [
            trajectory for group in weighted_by_lead for trajectory in group
        ]
        if not chunk_plan_losses:
            rows_all.extend(chunk_rows)
            continue
        _sync()
        _t = time.perf_counter()
        log_probs = [
            kernel.log_prob_of(
                step.state,
                step.candidates,
                step.chosen_idx,
                step.t,
                trajectory.context,
            )
            for trajectory in weighted
            for step in trajectory.steps
        ]
        path_loss = compute_weighted_path_loss(weighted, log_probs)
        plan_loss = torch.stack(chunk_plan_losses).mean()
        total_loss = path_loss + float(plan_cfg.loss_weight) * plan_loss
        _sync()
        chunk_t["policy_logprob"] += time.perf_counter() - _t
        optimizer.zero_grad()
        _t = time.perf_counter()
        total_loss.backward()
        _sync()
        chunk_t["backward"] += time.perf_counter() - _t
        _t = time.perf_counter()
        torch.nn.utils.clip_grad_norm_(
            parameters, float(train_cfg.get("grad_clip_norm", 1.0))
        )
        optimizer.step()
        _sync()
        chunk_t["optimizer"] += time.perf_counter() - _t
        losses.append(float(total_loss.detach().cpu().item()))
        path_losses.append(float(path_loss.detach().cpu().item()))
        plan_losses.append(float(plan_loss.detach().cpu().item()))
        rows_all.extend(chunk_rows)

        for key in STAGE_KEYS:
            stage_seconds[key] += chunk_t[key]
        if chunk_timing:
            _esm2_now = STAGE_TIMER.snapshot()
            esm2_diff = {
                key: _esm2_now.get(key, 0) - _esm2_prev.get(key, 0)
                for key in (
                    "esm2_forward_time", "esm2_cache_hit", "esm2_cache_miss",
                    "esm2_forward_batches", "esm2_forward_calls",
                )
            }
            _esm2_prev = _esm2_now
            leads_done = chunk_start + len(chunk_leads)
            elapsed = time.perf_counter() - epoch_started
            leads_per_min = leads_done / (elapsed / 60.0) if elapsed > 0 else 0.0
            eta_seconds = (
                (elapsed / leads_done) * (total_leads - leads_done) if leads_done else 0.0
            )
            record = {
                "epoch": epoch + 1,
                "chunk": chunk_start // chunk_size + 1,
                "leads_done": leads_done,
                "total_leads": total_leads,
                "chunk_seconds": round(time.perf_counter() - chunk_started, 3),
                "stage_seconds": {key: round(chunk_t[key], 4) for key in STAGE_KEYS},
                "esm2": {
                    key: (round(value, 4) if "time" in key else int(value))
                    for key, value in esm2_diff.items()
                },
                "counts": dict(chunk_counts),
                "exact_sb_cache": {
                    "hits": (exact_sb_cache.stats.hits - _cache_hits0)
                    if exact_sb_cache is not None else 0,
                    "misses": (exact_sb_cache.stats.misses - _cache_misses0)
                    if exact_sb_cache is not None else 0,
                },
                "peptiverse": {
                    "predictions": (_pv_wrapper.cache_misses - _pv_misses0)
                    if _pv_wrapper is not None else 0,
                    "cache_hits": (_pv_wrapper.cache_hits - _pv_hits0)
                    if _pv_wrapper is not None else 0,
                },
                "leads_per_min": round(leads_per_min, 1),
                "epoch_eta_min": round(eta_seconds / 60.0, 1),
            }
            print(f"[chunk timing] {json.dumps(record, ensure_ascii=False)}", flush=True)


    metrics = {
        "loss": float(np.mean(losses)),
        "path_loss": float(np.mean(path_losses)),
        "plan_loss": float(np.mean(plan_losses)),
        "plan_entropy": float(np.mean(entropies)),
        "plan_coverage": float(np.mean(coverage)),
        "reference_vs_controlled_plan_disagreement": float(np.mean(disagreements)),
        "controlled_plan_mode_mix": dict(controlled_top_modes),
        "exact_sb_objective": bool(plan_cfg.exact_sb_objective),
        "exact_sb_beta": float(plan_cfg.exact_sb_beta),
        "plan_encoder_v2": bool(plan_cfg.plan_encoder_v2),
        "legacy_target_temperature": float(plan_cfg.target_temperature),
        "legacy_property_only_equivalent_beta": float(
            1.0
            / (
                plan_cfg.target_temperature
                * float(
                    config["hydrocarbon"]["terminal_energy"]["property"][
                        "penetrance_weight"
                    ]
                )
            )
        ),
        "q_star_vs_q_theta_kl": (
            float(np.mean(exact_reverse_kls)) if exact_reverse_kls else None
        ),
        "exact_sb_forward_objective": (
            float(np.mean(exact_forward_objectives))
            if exact_forward_objectives
            else None
        ),
        "q_star_entropy": (
            float(np.mean(exact_q_star_entropies))
            if exact_q_star_entropies
            else None
        ),
        "exact_sb_plans_scored": int(exact_plans_scored),
        "exact_sb_target_sources": dict(exact_cache_sources),
        "property_free_hard_plan_mask": True,
        "hard_plan_mask_uses_peptiverse": False,
        "hard_plan_mask_totals": dict(hard_mask_totals),
        "trajectory_weighting": {
            "normalization": "lead_local_positive_only",
            **{key: int(value) for key, value in weighting_totals.items()},
        },
        "stage_seconds": {key: float(value) for key, value in stage_seconds.items()},
    }
    if joint_support_enabled:
        supported_joint_rows = [
            row for row in joint_support_rows if row["status"] == "scored"
        ]
        solution_leads = sum(
            bool(row["joint_nonempty"]) for row in supported_joint_rows
        )
        selected_joint = sum(
            bool(row["selected_satisfies_joint_condition"])
            for row in supported_joint_rows
        )
        metrics["joint_perm_halflife_support"] = {
            "enabled": True,
            "leads_total": len(joint_support_rows),
            "leads_with_hard_support": len(supported_joint_rows),
            "leads_without_hard_support": sum(
                row["status"] == "no_hard_support" for row in joint_support_rows
            ),
            "leads_with_joint_solution": solution_leads,
            "fallback_leads": sum(
                bool(row["joint_fallback"]) for row in supported_joint_rows
            ),
            "joint_plan_count": sum(
                int(row["joint_plan_count"]) for row in supported_joint_rows
            ),
            "q_star_support_size": sum(
                int(row["q_star_support_size"]) for row in supported_joint_rows
            ),
            "selected_joint_solutions": selected_joint,
            "joint_solution_recovery": (
                selected_joint / solution_leads if solution_leads else None
            ),
            "per_lead": joint_support_rows,
        }
    return rows_all, metrics


@torch.no_grad()
def validate_enabled(
    leads: list[Any],
    config: dict[str, Any],
    stack: dict[str, Any],
    energy_fn: Any,
    policy: Any,
    kernel: Any,
    head: HydrocarbonPlanHead,
    exact_sb_cache: Any = None,
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
    validation_cfg = config["validation"]
    selected: list[dict[str, Any]] = []
    entropies: list[float] = []
    coverages: list[float] = []
    disagreements: list[bool] = []
    top1_plans: dict[str, str] = {}
    mode_mix: Counter[str] = Counter()
    # Diagnostics-only: opt-in exact-SB q* on the same property-free hard
    # support used by q_theta. It never changes fine decoding or candidate
    # ranking within the committed plan selected below.
    exact_sb_diagnostics = exact_sb_validation_enabled(config)
    exact_sb_beta = HydrocarbonPlanControlConfig.from_config(config).exact_sb_beta
    q_star_rows: list[dict[str, Any]] = []
    q_star_cache_sources: Counter[str] = Counter()
    joint_support_enabled = bool(
        getattr(
            getattr(energy_fn, "property_cfg", None),
            "enable_joint_perm_halflife_support",
            False,
        )
    )
    joint_support_rows: list[dict[str, Any]] = []
    legal_plans_before_hard_mask = 0
    hard_supported_plans = 0
    leads_without_hard_support = 0
    # Wall-clock accounting only; see the note in train_enabled_epoch.
    q_star_diagnostics_seconds = 0.0

    for lead_index, lead in enumerate(leads):
        z0 = StapleState(sequence_tokens=tokenize_sequence(lead.linear_sequence))
        support_view = property_free_supported_plan_view(lead, z0, stack, config)
        all_plans = support_view["all_plans"]
        support_indices = support_view["support_indices"]
        plans = support_view["plans"]
        ref_weights = support_view["reference_weights"]
        context = support_view["context"]
        legal_plans_before_hard_mask += len(all_plans)
        hard_supported_plans += len(plans)
        strict_support_at_lead = bool(plans)
        if not plans:
            leads_without_hard_support += 1
            if joint_support_enabled:
                joint_support_rows.append(
                    {
                        "example_id": str(lead.example_id),
                        "n_hard_supported_plans": 0,
                        "joint_plan_count": 0,
                        "joint_nonempty": False,
                        "joint_fallback": False,
                        "q_star_support_size": 0,
                        "selected_satisfies_joint_condition": False,
                        "status": "no_hard_support",
                    }
                )
            print(
                f"[plan validation] leads={lead_index + 1}/{len(leads)} "
                "status=no_strict_feasible_plan",
                flush=True,
            )
            continue
        lead_embedding = encode_lead(policy, kernel, z0)
        all_controlled_logp = controlled_plan_log_probabilities(
            head,
            lead_embedding,
            all_plans,
            support_view["all_reference_weights"],
            len(z0.sequence_tokens),
            sequence_tokens=z0.sequence_tokens,
            peptide_ca=context["peptide_ca"],
        )
        masked_controlled_logp = mask_and_renormalize_plan_log_probabilities(
            all_controlled_logp, support_indices
        )
        controlled_logp = masked_controlled_logp[
            torch.tensor(
                support_indices,
                dtype=torch.long,
                device=masked_controlled_logp.device,
            )
        ]
        controlled_probabilities = controlled_logp.detach().exp().cpu()
        plan_modes = [plan_mode(plan) for plan in plans]
        q_theta_metrics = {
            "q_theta_entropy": float(plan_entropy(controlled_logp).cpu().item()),
            "q_theta_top1_mode": plan_modes[int(torch.argmax(controlled_logp).item())],
            "q_theta_s5_s5_i4_probability_mass": float(
                sum(
                    controlled_probabilities[index].item()
                    for index, mode in enumerate(plan_modes)
                    if mode == "S5-S5/i,i+4"
                )
            ),
            "q_theta_r8_s5_i7_probability_mass": float(
                sum(
                    controlled_probabilities[index].item()
                    for index, mode in enumerate(plan_modes)
                    if mode == "R8-S5/i,i+7"
                )
            ),
        }
        entropies.append(float(plan_entropy(controlled_logp).cpu().item()))
        reference_top = int(np.argmax(ref_weights))
        controlled_top = int(torch.argmax(controlled_logp).item())
        disagreements.append(reference_top != controlled_top)
        coverages.append(1.0)

        lead_q_star_row: dict[str, Any] | None = None
        if exact_sb_diagnostics:
            _diagnostics_started = time.perf_counter()
            reference_logp = empirical_log_probabilities(
                ref_weights, controlled_logp.device
            )
            # Same deterministic half as training, same cache. q_theta below is
            # always recomputed from the live head.
            plan_energy_tensor, _, cache_info = resolve_exact_sb_target(
                lead=lead,
                plans=plans,
                reference_log_probabilities=reference_logp,
                beta=exact_sb_beta,
                energy_fn=energy_fn,
                initial_state=z0,
                build_terminal=lambda state, plan: build_hydrocarbon_demonstration_path(
                    state, plan, stack["catalog"]
                )[-1],
                cache=exact_sb_cache,
                energy_only=energy_only_from_config(config),
            )
            q_star_cache_sources[str(cache_info["source"])] += 1
            target_support_mask = (
                torch.tensor(
                    cache_info["target_support_mask"],
                    dtype=torch.bool,
                    device=controlled_logp.device,
                )
                if cache_info.get("target_support_mask") is not None
                else None
            )
            diagnostics = lead_q_star_diagnostics(
                controlled_logp,
                reference_logp,
                plan_energy_tensor,
                exact_sb_beta,
                plan_modes=plan_modes,
                target_support_mask=target_support_mask,
            )
            lead_q_star_row = {
                "example_id": str(lead.example_id),
                "q_star_top1_plan": describe_plan(
                    plans[diagnostics["q_star_top1_index"]]
                ),
                "q_theta_top1_plan": describe_plan(
                    plans[diagnostics["q_theta_top1_index"]]
                ),
                "joint_perm_halflife_support_enabled": bool(
                    cache_info["joint_perm_halflife_support_enabled"]
                ),
                "joint_plan_count": cache_info["joint_plan_count"],
                "joint_nonempty": cache_info["joint_nonempty"],
                "joint_fallback": cache_info["joint_fallback"],
                "q_star_support_size": int(cache_info["q_star_support_size"]),
                **diagnostics,
            }
            q_star_rows.append(lead_q_star_row)
            q_star_diagnostics_seconds += time.perf_counter() - _diagnostics_started

        hierarchical = hierarchical_plan_ranking_enabled(config)
        decoded_by_plan: list[tuple[StapleState, float, int] | None] = [None] * len(plans)
        decode_indices = [controlled_top] if hierarchical else range(len(plans))
        for plan_index in decode_indices:
            plan = plans[plan_index]
            decoded = beam_decode(
                lead,
                stack,
                kernel,
                int(validation_cfg["horizon"]),
                int(validation_cfg["beam_size"]),
                committed_plan=plan,
            )
            if decoded:
                decoded_by_plan[plan_index] = (
                    best_committed_plan_trajectory(decoded, plan) if hierarchical else decoded[0]
                )
        try:
            joint, plan_index, terminal, path_logp, path_length = select_plan_and_trajectory(
                controlled_logp.detach().cpu(), decoded_by_plan, hierarchical=hierarchical
            )
        except RuntimeError as exc:
            raise RuntimeError(f"no decoded plans for {lead.example_id}") from exc
        selected_plan = plans[plan_index]
        label = describe_plan(selected_plan)
        top1_plans[lead.example_id] = label
        mode_mix[
            f"{selected_plan.ordered_pair}/i,i+{selected_plan.spacing}"
        ] += 1
        selected_record = terminal_record(
                z0,
                terminal,
                lead,
                energy_fn,
                stack,
                lead_index=lead_index,
                committed_plan=label,
                plan_log_probability=float(controlled_logp[plan_index].cpu().item()),
                path_log_probability=path_logp,
                joint_log_probability=joint,
                path_length=path_length,
            )
        selected.append(selected_record)

        print(f"[plan validation] leads={lead_index + 1}/{len(leads)}", flush=True)

    summary = summarize(selected, selected, len(leads))
    summary.update(selected_top1_rates(selected))
    summary.update(
        {
            "plan_control_enabled": True,
            "hierarchical_plan_ranking": hierarchical_plan_ranking_enabled(config),
            "plan_entropy": float(np.mean(entropies)) if entropies else None,
            "plan_top1": top1_plans,
            "plan_coverage": float(np.mean(coverages)) if coverages else 0.0,
            "mode_mix": dict(mode_mix),
            "reference_vs_controlled_plan_disagreement": (
                float(np.mean(disagreements)) if disagreements else None
            ),
            "property_free_hard_plan_mask": True,
            "hard_plan_mask_uses_peptiverse": False,
            "legal_plans_before_hard_mask": int(legal_plans_before_hard_mask),
            "hard_supported_plans": int(hard_supported_plans),
            "hard_masked_plans": int(
                legal_plans_before_hard_mask - hard_supported_plans
            ),
            "leads_without_hard_support": int(leads_without_hard_support),
        }
    )
    summary["exact_sb_validation_diagnostics"] = exact_sb_diagnostics
    if exact_sb_diagnostics:
        # Additive only: no existing validation metric above is overwritten.
        summary.update(aggregate_q_star_diagnostics(q_star_rows))
        summary["q_star_per_lead"] = q_star_rows
        # Provenance only; not a metric and not used by any selection.
        summary["q_star_target_sources"] = dict(q_star_cache_sources)
        summary["q_star_diagnostics_seconds"] = float(q_star_diagnostics_seconds)
    if joint_support_enabled:
        supported_joint_rows = [
            row for row in joint_support_rows if row["status"] == "scored"
        ]
        solution_leads = sum(
            bool(row["joint_nonempty"]) for row in supported_joint_rows
        )
        selected_joint = sum(
            bool(row["selected_satisfies_joint_condition"])
            for row in supported_joint_rows
        )
        summary["joint_perm_halflife_support"] = {
            "enabled": True,
            "leads_total": len(joint_support_rows),
            "leads_with_hard_support": len(supported_joint_rows),
            "leads_without_hard_support": sum(
                row["status"] == "no_hard_support" for row in joint_support_rows
            ),
            "leads_with_joint_solution": solution_leads,
            "fallback_leads": sum(
                bool(row["joint_fallback"]) for row in supported_joint_rows
            ),
            "joint_plan_count": sum(
                int(row["joint_plan_count"]) for row in supported_joint_rows
            ),
            "q_star_support_size": sum(
                int(row["q_star_support_size"]) for row in supported_joint_rows
            ),
            "selected_joint_solutions": selected_joint,
            "joint_solution_recovery": (
                selected_joint / solution_leads if solution_leads else None
            ),
            "per_lead": joint_support_rows,
        }
    return selected, summary