osmac_optimizer.py

from __future__ import annotations

"""
osmac_optimizer.py
==================

Streaming fairness-constrained OSMAC-style optimizer for DFE-AutoML.

Purpose
-------
Implements a practical streaming constrained Bayesian optimization layer for
DFE-AutoML.

This optimizer:
- keeps separate local surrogates for objective and fairness constraint behavior
- updates these surrogates each window / period
- ranks candidate configurations with constrained acquisition
- returns the next candidate set and current incumbent

Designed to align with:
- config_space.py
- candidate_evaluator.py
- dynamic fairness envelope / DFE logic
- streaming search requirements inspired by OSMAC
- constrained BO logic inspired by FairBO / CBO

Core design
-----------
For each candidate x, we model:

1) Objective surrogate
   f_t(x): predictive utility on current / recent windows

2) Constraint surrogate
   c_t(x): unfairness proxy and/or feasibility probability

Then rank candidates with one of these modes:

A) feasibility_first
   Used when we do not yet have a feasible recent incumbent.
   Select x maximizing P(feasible)

B) cei_static
   Constrained expected improvement with fixed fairness threshold:
       cEI(x) = EI(x) * P(c(x) <= epsilon)

C) cei_envelope
   Constrained expected improvement with dynamic fairness envelope:
       cEI_t(x) = EI_t(x) * P(c_t(x) in E(t))

In practice, candidate_evaluator.py already returns:
- predictive metrics
- fairness metrics
- sfm payload
- feasibility probability
- fairness risk
- constrained utility proxy

This optimizer consumes those outputs and proposes next configurations.

Notes
-----
- This module does not directly evaluate candidates.
- It consumes outputs from candidate_evaluator.py via register_observation().
- Surrogates are lightweight distance-weighted local regressors over encoded configs.
- This keeps the implementation deterministic, stable, and dependency-light.
"""

import copy
import math
import random
from dataclasses import dataclass, field
from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, Union

import numpy as np

from config_space import PipelineConfig, StreamingConfigSpace
from candidate_evaluator import CandidateEvaluationResult


# ============================================================
# Helpers
# ============================================================

def _safe_float(x: Any, default: float = 0.0) -> float:
    try:
        v = float(x)
        if not math.isfinite(v):
            return float(default)
        return float(v)
    except Exception:
        return float(default)


def _safe_int(x: Any, default: int = 0) -> int:
    try:
        if isinstance(x, bool):
            return int(x)
        return int(float(x))
    except Exception:
        return int(default)


def _clip(x: float, lo: float, hi: float) -> float:
    return max(float(lo), min(float(hi), float(x)))


def _clip01(x: Any) -> float:
    return _clip(_safe_float(x, 0.0), 0.0, 1.0)


def _normal_pdf(z: float) -> float:
    return math.exp(-0.5 * z * z) / math.sqrt(2.0 * math.pi)


def _normal_cdf(z: float) -> float:
    return 0.5 * (1.0 + math.erf(z / math.sqrt(2.0)))


def _euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.linalg.norm(a - b))


def _stable_sigmoid(x: float) -> float:
    x = float(_clip(x, -60.0, 60.0))
    return 1.0 / (1.0 + math.exp(-x))


# ============================================================
# Dataclasses
# ============================================================

@dataclass(frozen=True)
class OSMACOptimizerConfig:
    seed: int = 42
    streaming_method: str = "online_fairbo_ensemble"

    # Candidate generation
    pool_size: int = 64
    n_propose: int = 8
    initial_design_size: int = 12
    random_fraction: float = 0.20
    mutation_fraction: float = 0.55
    crossover_fraction: float = 0.25

    # Local surrogates
    neighborhood_size: int = 15
    min_points_for_surrogate: int = 6
    uncertainty_floor: float = 1e-6
    distance_temperature: float = 2.0

    # Acquisition setup
    acquisition_mode: str = "auto"   # auto | feasibility_first | cei_static | cei_envelope
    use_expected_improvement: bool = True
    objective_mode: str = "maximize_accuracy"   # maximize_accuracy | maximize_constrained_utility | minimize_log_loss
    incumbent_metric: str = "constrained_utility"

    # Constraint interpretation
    static_fairness_threshold: float = 0.05
    hard_feasibility_floor: float = 0.05
    use_recent_feasibility_horizon: bool = True
    feasible_horizon_windows: int = 5

    # Acquisition weights
    utility_weight: float = 1.0
    feasibility_weight: float = 1.5
    risk_penalty_weight: float = 0.85
    novelty_weight: float = 0.10
    incumbent_bias: float = 0.05
    reliability_penalty_weight: float = 0.20
    low_support_penalty_weight: float = 0.15

    # Streaming adaptation
    recency_decay: float = 0.98
    drift_reactivity: float = 0.50
    fairness_pressure_scale: float = 1.00
    max_history: int = 2000

    # Diversity
    enforce_unique_configs: bool = True
    diversity_min_distance: float = 0.05

    # Envelope preference / policy
    envelope_violation_penalty: float = 1.00
    recent_feasible_bonus: float = 0.25
    feasible_priority_bonus: float = 2.0
    infeasible_priority_penalty: float = 1.0
    no_feasible_fallback_to_feasibility_first: bool = True

    version: str = "osmac_optimizer.v2"


@dataclass
class ObservationRecord:
    config: Dict[str, Any]
    config_id: str
    x: List[float]

    # Objective and acquisition-facing pieces
    utility_score: float
    cost_score: float
    fairness_score: float
    feasibility_probability: float
    fairness_risk: float
    constrained_utility: float

    # Explicit CBO fields
    objective_value: float = 0.0
    constraint_value: float = 0.0
    constraint_feasible: bool = False
    feasibility_mode: str = "unknown"  # static_threshold | envelope | unknown

    # Predictive / fairness details
    accuracy: float = 0.0
    log_loss: float = 0.0
    fm_smooth: float = 0.0
    fm_raw: float = 0.0
    fm_drift: float = 0.0
    fm_uncertainty: float = 1.0
    envelope_lower: float = 0.0
    envelope_upper: float = 0.1
    envelope_violation: bool = False
    envelope_severity: float = 0.0
    is_reliable: bool = False
    insufficient_support: bool = True
    effective_support: int = 0

    timestep: int = 0
    window_index: int = 0
    recency_weight: float = 1.0

    raw_result: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "config": copy.deepcopy(self.config),
            "config_id": str(self.config_id),
            "x": list(self.x),
            "utility_score": float(self.utility_score),
            "cost_score": float(self.cost_score),
            "fairness_score": float(self.fairness_score),
            "feasibility_probability": float(self.feasibility_probability),
            "fairness_risk": float(self.fairness_risk),
            "constrained_utility": float(self.constrained_utility),
            "objective_value": float(self.objective_value),
            "constraint_value": float(self.constraint_value),
            "constraint_feasible": bool(self.constraint_feasible),
            "feasibility_mode": str(self.feasibility_mode),
            "accuracy": float(self.accuracy),
            "log_loss": float(self.log_loss),
            "fm_smooth": float(self.fm_smooth),
            "fm_raw": float(self.fm_raw),
            "fm_drift": float(self.fm_drift),
            "fm_uncertainty": float(self.fm_uncertainty),
            "envelope_lower": float(self.envelope_lower),
            "envelope_upper": float(self.envelope_upper),
            "envelope_violation": bool(self.envelope_violation),
            "envelope_severity": float(self.envelope_severity),
            "is_reliable": bool(self.is_reliable),
            "insufficient_support": bool(self.insufficient_support),
            "effective_support": int(self.effective_support),
            "timestep": int(self.timestep),
            "window_index": int(self.window_index),
            "recency_weight": float(self.recency_weight),
            "raw_result": copy.deepcopy(self.raw_result),
        }


@dataclass
class AcquisitionRecord:
    config: Dict[str, Any]
    config_id: str
    x: List[float]

    # Objective surrogate
    pred_objective_mean: float = 0.0
    pred_objective_std: float = 0.0

    # Constraint surrogate
    pred_constraint_mean: float = 0.0
    pred_constraint_std: float = 0.0
    pred_feasibility_mean: float = 0.0
    pred_feasibility_std: float = 0.0

    # Auxiliary surrogates
    pred_risk_mean: float = 0.0
    pred_risk_std: float = 0.0
    pred_fairness_mean: float = 0.0
    pred_fairness_std: float = 0.0

    # Acquisition
    acquisition_mode: str = "unknown"
    best_reference: float = 0.0
    improvement_mean: float = 0.0
    ei_like: float = 0.0
    p_feasible: float = 0.0
    novelty: float = 0.0

    constrained_acquisition: float = 0.0
    reliability_penalty: float = 0.0
    low_support_penalty: float = 0.0
    final_score: float = 0.0

    rationale: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "config": copy.deepcopy(self.config),
            "config_id": str(self.config_id),
            "x": list(self.x),
            "pred_objective_mean": float(self.pred_objective_mean),
            "pred_objective_std": float(self.pred_objective_std),
            "pred_constraint_mean": float(self.pred_constraint_mean),
            "pred_constraint_std": float(self.pred_constraint_std),
            "pred_feasibility_mean": float(self.pred_feasibility_mean),
            "pred_feasibility_std": float(self.pred_feasibility_std),
            "pred_risk_mean": float(self.pred_risk_mean),
            "pred_risk_std": float(self.pred_risk_std),
            "pred_fairness_mean": float(self.pred_fairness_mean),
            "pred_fairness_std": float(self.pred_fairness_std),
            "acquisition_mode": str(self.acquisition_mode),
            "best_reference": float(self.best_reference),
            "improvement_mean": float(self.improvement_mean),
            "ei_like": float(self.ei_like),
            "p_feasible": float(self.p_feasible),
            "novelty": float(self.novelty),
            "constrained_acquisition": float(self.constrained_acquisition),
            "reliability_penalty": float(self.reliability_penalty),
            "low_support_penalty": float(self.low_support_penalty),
            "final_score": float(self.final_score),
            "rationale": copy.deepcopy(self.rationale),
        }


@dataclass
class ProposalBatch:
    proposed_configs: List[Dict[str, Any]]
    acquisition_records: List[Dict[str, Any]]
    incumbent: Optional[Dict[str, Any]]
    optimizer_state: Dict[str, Any]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "proposed_configs": copy.deepcopy(self.proposed_configs),
            "acquisition_records": copy.deepcopy(self.acquisition_records),
            "incumbent": copy.deepcopy(self.incumbent),
            "optimizer_state": copy.deepcopy(self.optimizer_state),
        }


# ============================================================
# OSMAC optimizer
# ============================================================

class OSMACOptimizer:
    """
    Streaming fairness-constrained optimizer for DFE-AutoML.

    Workflow
    --------
    1. Create optimizer with a shared StreamingConfigSpace
    2. Ask for candidates with propose_next()
    3. Evaluate them externally with candidate_evaluator.py
    4. Feed results back via register_observation()
    5. Repeat each window
    """

    def __init__(
        self,
        config_space: StreamingConfigSpace,
        config: Optional[OSMACOptimizerConfig] = None,
    ) -> None:
        self.config_space = config_space
        self.cfg = config or OSMACOptimizerConfig()
        self.rng = random.Random(int(self.cfg.seed))

        self.window_index: int = 0
        self.observations: List[ObservationRecord] = []
        self.observation_by_id: Dict[str, List[ObservationRecord]] = {}
        self.incumbent: Optional[ObservationRecord] = None
        self.last_optimizer_state: Dict[str, Any] = {}
        self._vector_schema = self.config_space.get_vector_schema()

    # --------------------------------------------------------
    # Public API
    # --------------------------------------------------------
    def reset(self) -> None:
        self.window_index = 0
        self.observations.clear()
        self.observation_by_id.clear()
        self.incumbent = None
        self.last_optimizer_state = {}

    def step_window(self) -> None:
        self.window_index += 1
        self._refresh_recency_weights()

    def register_observation(
        self,
        evaluation_result: Union[CandidateEvaluationResult, Mapping[str, Any]],
        *,
        timestep: Optional[int] = None,
        window_index: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        Register one evaluated candidate.
        """
        payload = evaluation_result.to_dict() if isinstance(evaluation_result, CandidateEvaluationResult) else dict(evaluation_result or {})

        config = dict(payload.get("config", {}) or {})
        config_id = str(payload.get("config_id", "")) or PipelineConfig.from_dict(config).stable_id()
        x = [float(v) for v in payload.get("config_vector", [])]

        predictive = dict(payload.get("predictive_metrics", {}) or {})
        fairness = dict(payload.get("fairness_metrics", {}) or {})
        feasibility = dict(payload.get("feasibility", {}) or {})
        score_payload = dict(payload.get("score_payload", {}) or {})
        sfm_payload = dict(payload.get("sfm_payload", {}) or {})
        envelope = dict(payload.get("envelope", {}) or {})
        raw_artifacts = dict(payload.get("raw_artifacts", {}) or {})

        fm_smooth = _safe_float(fairness.get("fm_smooth", sfm_payload.get("fm_smooth", fairness.get("primary_metric", 0.0))), 0.0)
        fm_raw = _safe_float(fairness.get("fm_raw", sfm_payload.get("fm_raw", fm_smooth)), fm_smooth)
        feasible_prob = _clip01(score_payload.get("feasibility_probability", feasibility.get("feasibility_probability", 0.0)))
        envelope_lower = _safe_float(
            envelope.get("lower", envelope.get("envelope_lower", 0.0)),
            0.0,
        )
        envelope_upper = _safe_float(
            envelope.get("upper", envelope.get("envelope_upper", self.cfg.static_fairness_threshold)),
            self.cfg.static_fairness_threshold,
        )
        envelope_violation = bool(fairness.get("envelope_violation", envelope.get("envelope_violation", False)))
        constraint_feasible = bool(feasibility.get("feasible_hard", not envelope_violation))
        feasibility_mode = "envelope" if ("upper" in envelope or "envelope_upper" in envelope) else "static_threshold"

        rec = ObservationRecord(
            config=copy.deepcopy(config),
            config_id=config_id,
            x=list(x),

            utility_score=_safe_float(score_payload.get("utility_score", predictive.get("accuracy", 0.0)), 0.0),
            cost_score=_safe_float(score_payload.get("cost_score", predictive.get("log_loss", 0.0)), 0.0),
            fairness_score=_safe_float(score_payload.get("fairness_score", 1.0 - fm_smooth), 0.0),
            feasibility_probability=feasible_prob,
            fairness_risk=_clip01(score_payload.get("fairness_risk", feasibility.get("fairness_risk", 1.0))),
            constrained_utility=_safe_float(score_payload.get("constrained_utility", 0.0), 0.0),

            objective_value=self._extract_objective_from_payload(
                predictive=predictive,
                score_payload=score_payload,
            ),
            constraint_value=fm_smooth,
            constraint_feasible=constraint_feasible,
            feasibility_mode=feasibility_mode,

            accuracy=_safe_float(predictive.get("accuracy", 0.0), 0.0),
            log_loss=_safe_float(predictive.get("log_loss", 0.0), 0.0),
            fm_smooth=fm_smooth,
            fm_raw=fm_raw,
            fm_drift=_safe_float(fairness.get("fm_drift", sfm_payload.get("fm_drift", 0.0)), 0.0),
            fm_uncertainty=_clip01(fairness.get("fm_uncertainty", sfm_payload.get("fm_uncertainty", 1.0))),
            envelope_lower=envelope_lower,
            envelope_upper=envelope_upper,
            envelope_violation=envelope_violation,
            envelope_severity=_safe_float(
                fairness.get("envelope_severity", envelope.get("envelope_severity", envelope.get("severity", 0.0))),
                0.0,
            ),
            is_reliable=bool(fairness.get("is_reliable", sfm_payload.get("is_reliable", False))),
            insufficient_support=bool(fairness.get("insufficient_support", sfm_payload.get("insufficient_support", True))),
            effective_support=_safe_int(fairness.get("effective_support", sfm_payload.get("effective_support", 0)), 0),

            timestep=int(self.window_index if timestep is None else timestep),
            window_index=int(self.window_index if window_index is None else window_index),
            raw_result=copy.deepcopy(payload),
        )

        # Optional standardized payload from candidate_evaluator, if present.
        objective_block = dict(raw_artifacts.get("optimizer_targets", {}) or {})
        if objective_block:
            rec.objective_value = _safe_float(objective_block.get("objective_value", rec.objective_value), rec.objective_value)
            rec.constraint_value = _safe_float(objective_block.get("constraint_value", rec.constraint_value), rec.constraint_value)
            rec.constraint_feasible = bool(objective_block.get("constraint_feasible", rec.constraint_feasible))
            rec.feasibility_mode = str(objective_block.get("constraint_probability_target", rec.feasibility_mode))

        self.observations.append(rec)
        self.observation_by_id.setdefault(rec.config_id, []).append(rec)

        if len(self.observations) > int(self.cfg.max_history):
            excess = len(self.observations) - int(self.cfg.max_history)
            self.observations = self.observations[excess:]
            self._rebuild_observation_index()

        self._refresh_recency_weights()
        self._update_incumbent()

        return {
            "n_observations": int(len(self.observations)),
            "incumbent": self.get_incumbent_summary(),
            "has_recent_feasible": bool(self._has_recent_feasible()),
        }

    def register_batch(
        self,
        evaluation_results: Sequence[Union[CandidateEvaluationResult, Mapping[str, Any]]],
        *,
        timestep: Optional[int] = None,
        window_index: Optional[int] = None,
    ) -> Dict[str, Any]:
        for res in evaluation_results:
            self.register_observation(res, timestep=timestep, window_index=window_index)
        return {
            "n_observations": int(len(self.observations)),
            "incumbent": self.get_incumbent_summary(),
            "has_recent_feasible": bool(self._has_recent_feasible()),
        }

    def propose_next(
        self,
        *,
        n_propose: Optional[int] = None,
        include_incumbent: bool = True,
        fairness_context: Optional[Mapping[str, Any]] = None,
        force_exploration: bool = False,
    ) -> ProposalBatch:
        """
        Propose the next candidate set.
        """
        n_propose_eff = max(1, int(n_propose if n_propose is not None else self.cfg.n_propose))
        fairness_context = dict(fairness_context or {})

        if len(self.observations) < int(self.cfg.initial_design_size):
            proposals = self._initial_design(n=n_propose_eff, include_incumbent=include_incumbent)
            self.last_optimizer_state = {
                "mode": "initial_design",
                "window_index": int(self.window_index),
                "n_observations": int(len(self.observations)),
                "has_recent_feasible": bool(self._has_recent_feasible()),
                "version": str(self.cfg.version),
            }
            return ProposalBatch(
                proposed_configs=[p["config"] for p in proposals],
                acquisition_records=proposals,
                incumbent=self.get_incumbent_summary(),
                optimizer_state=copy.deepcopy(self.last_optimizer_state),
            )

        candidate_pool = self._generate_candidate_pool(
            pool_size=max(n_propose_eff, int(self.cfg.pool_size)),
            fairness_context=fairness_context,
            force_exploration=force_exploration,
        )

        ranked = self.rank_candidates(
            candidate_pool=candidate_pool,
            fairness_context=fairness_context,
        )

        selected = self._select_diverse_topk(
            ranked=ranked,
            k=n_propose_eff,
            include_incumbent=include_incumbent,
        )

        acquisition_mode = self._resolve_acquisition_mode(fairness_context=fairness_context)

        optimizer_state = self._build_optimizer_state(
            fairness_context=fairness_context,
            candidate_pool_size=len(candidate_pool),
            mode=acquisition_mode,
        )
        self.last_optimizer_state = optimizer_state

        return ProposalBatch(
            proposed_configs=[r["config"] for r in selected],
            acquisition_records=selected,
            incumbent=self.get_incumbent_summary(),
            optimizer_state=optimizer_state,
        )

    def rank_candidates(
        self,
        *,
        candidate_pool: Sequence[Union[PipelineConfig, Mapping[str, Any]]],
        fairness_context: Optional[Mapping[str, Any]] = None,
    ) -> List[Dict[str, Any]]:
        fairness_context = dict(fairness_context or {})
        acquisition_mode = self._resolve_acquisition_mode(fairness_context=fairness_context)
        records: List[AcquisitionRecord] = []

        best_reference = self._get_best_reference()

        for cand in candidate_pool:
            cfg = self.config_space.validate_config(cand)
            x = np.asarray(self.config_space.encode_config(cfg, as_numpy=True), dtype=float)

            obj_mean, obj_std = self._predict_target(x=x, target="objective")
            cons_mean, cons_std = self._predict_target(x=x, target="constraint")
            feas_mean, feas_std = self._predict_target(x=x, target="feasibility")
            risk_mean, risk_std = self._predict_target(x=x, target="risk")
            fair_mean, fair_std = self._predict_target(x=x, target="fairness")

            feas_mean = max(float(feas_mean), float(self.cfg.hard_feasibility_floor))
            risk_mean = _clip01(risk_mean)
            fair_mean = max(0.0, float(fair_mean))

            fairness_pressure = self._compute_fairness_pressure(fairness_context)
            envelope_upper = _safe_float(
                fairness_context.get("envelope_upper", fairness_context.get("upper", self.cfg.static_fairness_threshold)),
                self.cfg.static_fairness_threshold,
            )
            static_eps = _safe_float(
                fairness_context.get("static_fairness_threshold", self.cfg.static_fairness_threshold),
                self.cfg.static_fairness_threshold,
            )

            # Derive p_feasible according to mode
            if acquisition_mode == "cei_static":
                p_feasible = self._probability_feasible_static(
                    constraint_mean=cons_mean,
                    constraint_std=cons_std,
                    epsilon=static_eps,
                )
            elif acquisition_mode == "cei_envelope":
                lower = _safe_float(fairness_context.get("envelope_lower", fairness_context.get("lower", 0.0)), 0.0)
                upper = _safe_float(fairness_context.get("envelope_upper", fairness_context.get("upper", envelope_upper)), envelope_upper)
                p_feasible = self._probability_feasible_envelope(
                    constraint_mean=cons_mean,
                    constraint_std=cons_std,
                    lower=lower,
                    upper=upper,
                )
            else:
                p_feasible = feas_mean

            p_feasible = float(max(self.cfg.hard_feasibility_floor, _clip01(p_feasible)))

            improvement_mean = float(obj_mean - best_reference)
            ei_like = self._compute_ei_like(
                mean=float(obj_mean),
                std=float(obj_std),
                best=float(best_reference),
            )

            novelty = self._compute_novelty(x)
            reliability_penalty = float(
                self.cfg.reliability_penalty_weight *
                max(0.0, fair_mean - p_feasible)
            )
            low_support_penalty = float(
                self.cfg.low_support_penalty_weight *
                max(0.0, fairness_pressure * risk_mean)
            )

            if acquisition_mode == "feasibility_first":
                constrained_acquisition = float(
                    p_feasible
                    + 0.10 * novelty
                    - 0.25 * risk_mean
                )
            else:
                constrained_acquisition = float(
                    self.cfg.utility_weight * ei_like *
                    (p_feasible ** self.cfg.feasibility_weight)
                )

            final_score = float(
                constrained_acquisition
                - self.cfg.risk_penalty_weight * (1.0 + fairness_pressure) * risk_mean
                - reliability_penalty
                - low_support_penalty
                + self.cfg.novelty_weight * novelty
            )

            if self.incumbent is not None and cfg.stable_id() == self.incumbent.config_id:
                final_score += float(self.cfg.incumbent_bias)

            rec = AcquisitionRecord(
                config=cfg.to_dict(),
                config_id=cfg.stable_id(),
                x=[float(v) for v in x.tolist()],
                pred_objective_mean=float(obj_mean),
                pred_objective_std=float(obj_std),
                pred_constraint_mean=float(cons_mean),
                pred_constraint_std=float(cons_std),
                pred_feasibility_mean=float(feas_mean),
                pred_feasibility_std=float(feas_std),
                pred_risk_mean=float(risk_mean),
                pred_risk_std=float(risk_std),
                pred_fairness_mean=float(fair_mean),
                pred_fairness_std=float(fair_std),
                acquisition_mode=acquisition_mode,
                best_reference=float(best_reference),
                improvement_mean=float(improvement_mean),
                ei_like=float(ei_like),
                p_feasible=float(p_feasible),
                novelty=float(novelty),
                constrained_acquisition=float(constrained_acquisition),
                reliability_penalty=float(reliability_penalty),
                low_support_penalty=float(low_support_penalty),
                final_score=float(final_score),
                rationale={
                    "fairness_pressure": float(fairness_pressure),
                    "objective_mode": str(self.cfg.objective_mode),
                    "incumbent_metric": str(self.cfg.incumbent_metric),
                    "static_fairness_threshold": float(static_eps),
                    "envelope_upper": float(envelope_upper),
                    "has_recent_feasible": bool(self._has_recent_feasible()),
                },
            )
            records.append(rec)

        records.sort(key=lambda r: r.final_score, reverse=True)
        return [r.to_dict() for r in records]

    def get_incumbent_summary(self) -> Optional[Dict[str, Any]]:
        if self.incumbent is None:
            return None
        return self.incumbent.to_dict()

    def get_history(self) -> List[Dict[str, Any]]:
        return [obs.to_dict() for obs in self.observations]

    def get_recent_history(self, horizon_windows: Optional[int] = None) -> List[Dict[str, Any]]:
        records = self._get_recent_observations(horizon_windows=horizon_windows)
        return [obs.to_dict() for obs in records]

    def get_state(self) -> Dict[str, Any]:
        return {
            "window_index": int(self.window_index),
            "n_observations": int(len(self.observations)),
            "has_recent_feasible": bool(self._has_recent_feasible()),
            "incumbent": self.get_incumbent_summary(),
            "last_optimizer_state": copy.deepcopy(self.last_optimizer_state),
            "config": {
                "seed": int(self.cfg.seed),
                "pool_size": int(self.cfg.pool_size),
                "n_propose": int(self.cfg.n_propose),
                "initial_design_size": int(self.cfg.initial_design_size),
                "objective_mode": str(self.cfg.objective_mode),
                "incumbent_metric": str(self.cfg.incumbent_metric),
                "acquisition_mode": str(self.cfg.acquisition_mode),
                "static_fairness_threshold": float(self.cfg.static_fairness_threshold),
                "version": str(self.cfg.version),
            },
        }

    # --------------------------------------------------------
    # Payload parsing helpers
    # --------------------------------------------------------
    def _extract_objective_from_payload(
        self,
        *,
        predictive: Mapping[str, Any],
        score_payload: Mapping[str, Any],
    ) -> float:
        mode = str(self.cfg.objective_mode).strip().lower()
        if mode == "maximize_accuracy":
            return _safe_float(predictive.get("accuracy", score_payload.get("utility_score", 0.0)), 0.0)
        if mode == "maximize_constrained_utility":
            return _safe_float(score_payload.get("constrained_utility", 0.0), 0.0)
        if mode == "minimize_log_loss":
            return -_safe_float(predictive.get("log_loss", score_payload.get("cost_score", 0.0)), 0.0)
        return _safe_float(score_payload.get("constrained_utility", predictive.get("accuracy", 0.0)), 0.0)

    # --------------------------------------------------------
    # Surrogate core
    # --------------------------------------------------------
    def _predict_target(self, x: np.ndarray, target: str) -> Tuple[float, float]:
        """
        Distance-weighted neighborhood regression.

        target:
            objective   -> predictive objective target
            constraint  -> unfairness constraint value
            feasibility -> feasibility probability
            risk        -> fairness risk
            fairness    -> fm_smooth-like fairness burden proxy
        """
        if len(self.observations) < max(2, int(self.cfg.min_points_for_surrogate)):
            return self._cold_start_prediction(target=target)

        neighbors = self._get_neighbors(x=x, k=int(self.cfg.neighborhood_size))
        if not neighbors:
            return self._cold_start_prediction(target=target)

        vals: List[float] = []
        weights: List[float] = []

        for dist, obs in neighbors:
            val = self._extract_target(obs=obs, target=target)
            rec_w = float(obs.recency_weight)
            sim_w = math.exp(-float(self.cfg.distance_temperature) * dist)
            w = max(self.cfg.uncertainty_floor, rec_w * sim_w)

            vals.append(float(val))
            weights.append(float(w))

        if not vals:
            return self._cold_start_prediction(target=target)

        arr_vals = np.asarray(vals, dtype=float)
        arr_w = np.asarray(weights, dtype=float)
        arr_w = np.maximum(arr_w, self.cfg.uncertainty_floor)
        arr_w = arr_w / np.maximum(arr_w.sum(), self.cfg.uncertainty_floor)

        mean = float(np.sum(arr_w * arr_vals))
        var = float(np.sum(arr_w * ((arr_vals - mean) ** 2)))
        std = float(max(math.sqrt(max(var, 0.0)), self.cfg.uncertainty_floor))

        if target in {"feasibility", "risk"}:
            mean = _clip01(mean)

        if target == "constraint":
            mean = max(0.0, mean)

        return mean, std

    def _cold_start_prediction(self, target: str) -> Tuple[float, float]:
        if target == "objective":
            if self.cfg.objective_mode == "maximize_accuracy":
                return 0.50, 0.25
            if self.cfg.objective_mode == "maximize_constrained_utility":
                return 0.25, 0.25
            if self.cfg.objective_mode == "minimize_log_loss":
                return -0.69, 0.25
        if target == "constraint":
            return self.cfg.static_fairness_threshold, 0.10
        if target == "feasibility":
            return 0.50, 0.25
        if target == "risk":
            return 0.50, 0.25
        if target == "fairness":
            return self.cfg.static_fairness_threshold, 0.10
        return 0.0, 1.0

    def _extract_target(self, obs: ObservationRecord, target: str) -> float:
        if target == "objective":
            return float(self._observation_objective(obs))
        if target == "constraint":
            return float(obs.constraint_value)
        if target == "feasibility":
            return float(obs.feasibility_probability)
        if target == "risk":
            return float(obs.fairness_risk)
        if target == "fairness":
            return float(obs.fm_smooth)
        return 0.0

    def _observation_objective(self, obs: ObservationRecord) -> float:
        mode = str(self.cfg.objective_mode).strip().lower()

        if mode == "maximize_accuracy":
            return float(obs.accuracy)
        if mode == "maximize_constrained_utility":
            return float(obs.constrained_utility)
        if mode == "minimize_log_loss":
            return float(-obs.log_loss)

        return float(obs.constrained_utility)

    def _get_best_reference(self) -> float:
        obs_list = self._get_reference_observations_for_best()
        if not obs_list:
            return 0.0

        metric = str(self.cfg.incumbent_metric).strip().lower()
        vals: List[float] = []
        for obs in obs_list:
            if metric == "accuracy":
                vals.append(float(obs.accuracy))
            elif metric == "constrained_utility":
                vals.append(float(obs.constrained_utility))
            elif metric == "utility":
                vals.append(float(obs.utility_score))
            elif metric == "negative_log_loss":
                vals.append(float(-obs.log_loss))
            else:
                vals.append(float(obs.constrained_utility))
        return max(vals) if vals else 0.0

    def _get_reference_observations_for_best(self) -> List[ObservationRecord]:
        if not self.observations:
            return []
        if self._has_recent_feasible():
            recent = self._get_recent_observations()
            feasible_recent = [obs for obs in recent if obs.constraint_feasible]
            if feasible_recent:
                return feasible_recent
        return list(self.observations)

    def _compute_ei_like(self, mean: float, std: float, best: float) -> float:
        if not self.cfg.use_expected_improvement:
            return max(0.0, mean - best)

        sigma = max(float(std), self.cfg.uncertainty_floor)
        delta = float(mean - best)
        z = delta / sigma
        ei = delta * _normal_cdf(z) + sigma * _normal_pdf(z)
        return float(max(ei, 0.0))

    def _get_neighbors(self, x: np.ndarray, k: int) -> List[Tuple[float, ObservationRecord]]:
        scored: List[Tuple[float, ObservationRecord]] = []
        denom = max(1.0, math.sqrt(len(x)))

        for obs in self.observations:
            xo = np.asarray(obs.x, dtype=float)
            dist = _euclidean_distance(x, xo) / denom
            scored.append((float(dist), obs))

        scored.sort(key=lambda t: t[0])
        return scored[:max(1, int(k))]

    def _compute_novelty(self, x: np.ndarray) -> float:
        if not self.observations:
            return 1.0

        denom = max(1.0, math.sqrt(len(x)))
        dists = []
        for obs in self.observations:
            xo = np.asarray(obs.x, dtype=float)
            d = _euclidean_distance(x, xo) / denom
            dists.append(float(d))

        if not dists:
            return 1.0

        min_dist = min(dists)
        return float(_clip(min_dist / max(self.cfg.diversity_min_distance, 1e-6), 0.0, 1.0))

    # --------------------------------------------------------
    # Feasibility / acquisition logic
    # --------------------------------------------------------
    def _probability_feasible_static(
        self,
        *,
        constraint_mean: float,
        constraint_std: float,
        epsilon: float,
    ) -> float:
        sigma = max(float(constraint_std), self.cfg.uncertainty_floor)
        z = (float(epsilon) - float(constraint_mean)) / sigma
        return float(_clip01(_normal_cdf(z)))

    def _probability_feasible_envelope(
        self,
        *,
        constraint_mean: float,
        constraint_std: float,
        lower: float,
        upper: float,
    ) -> float:
        sigma = max(float(constraint_std), self.cfg.uncertainty_floor)
        z_upper = (float(upper) - float(constraint_mean)) / sigma
        z_lower = (float(lower) - float(constraint_mean)) / sigma
        return float(_clip01(_normal_cdf(z_upper) - _normal_cdf(z_lower)))

    def _resolve_acquisition_mode(
        self,
        *,
        fairness_context: Mapping[str, Any],
    ) -> str:
        mode = str(self.cfg.acquisition_mode).strip().lower()

        if mode in {"feasibility_first", "cei_static", "cei_envelope"}:
            return mode

        # auto mode
        has_recent_feasible = self._has_recent_feasible()
        has_envelope = (
            "envelope_upper" in fairness_context or
            "upper" in fairness_context or
            (self.incumbent is not None and self.incumbent.feasibility_mode == "envelope")
        )

        if self.cfg.no_feasible_fallback_to_feasibility_first and not has_recent_feasible:
            return "feasibility_first"

        if has_envelope:
            return "cei_envelope"

        return "cei_static"

    def _has_recent_feasible(self) -> bool:
        recent = self._get_recent_observations()
        return any(obs.constraint_feasible for obs in recent)

    def _get_recent_observations(self, horizon_windows: Optional[int] = None) -> List[ObservationRecord]:
        if not self.observations:
            return []

        if not self.cfg.use_recent_feasibility_horizon:
            return list(self.observations)

        horizon = int(self.cfg.feasible_horizon_windows if horizon_windows is None else horizon_windows)
        current = max(obs.window_index for obs in self.observations)
        lo = current - max(0, horizon) + 1
        return [obs for obs in self.observations if int(obs.window_index) >= lo]

    # --------------------------------------------------------
    # Candidate generation
    # --------------------------------------------------------
    def _initial_design(self, n: int, include_incumbent: bool) -> List[Dict[str, Any]]:
        configs = self.config_space.sample_initial_population(
            budget=max(int(n), int(self.cfg.initial_design_size)),
            include_feature_selectors=True,
            seed=self.cfg.seed + self.window_index,
        )
        selected = configs[:max(1, int(n))]

        if include_incumbent and self.incumbent is not None:
            inc_cfg = PipelineConfig.from_dict(self.incumbent.config)
            if not any(PipelineConfig.from_dict(c.to_dict()).stable_id() == inc_cfg.stable_id() for c in selected):
                if selected:
                    selected[0] = inc_cfg
                else:
                    selected = [inc_cfg]

        out: List[Dict[str, Any]] = []
        for cfg in selected:
            x = self.config_space.encode_config(cfg, as_numpy=True)
            out.append({
                "config": cfg.to_dict(),
                "config_id": cfg.stable_id(),
                "x": [float(v) for v in np.asarray(x, dtype=float).tolist()],
                "pred_objective_mean": None,
                "pred_objective_std": None,
                "pred_constraint_mean": None,
                "pred_constraint_std": None,
                "pred_feasibility_mean": None,
                "pred_feasibility_std": None,
                "pred_risk_mean": None,
                "pred_risk_std": None,
                "pred_fairness_mean": None,
                "pred_fairness_std": None,
                "acquisition_mode": "initial_design",
                "best_reference": None,
                "improvement_mean": None,
                "ei_like": None,
                "p_feasible": None,
                "novelty": None,
                "constrained_acquisition": None,
                "reliability_penalty": None,
                "low_support_penalty": None,
                "final_score": None,
                "rationale": {"mode": "initial_design"},
            })
        return out

    def _generate_candidate_pool(
        self,
        *,
        pool_size: int,
        fairness_context: Mapping[str, Any],
        force_exploration: bool,
    ) -> List[PipelineConfig]:
        pool: List[PipelineConfig] = []
        target_size = max(1, int(pool_size))

        n_random = max(1, int(round(target_size * self.cfg.random_fraction)))
        n_mutation = max(1, int(round(target_size * self.cfg.mutation_fraction)))
        n_crossover = max(0, target_size - n_random - n_mutation)

        if force_exploration:
            n_random = max(n_random, int(0.50 * target_size))
            n_mutation = max(1, target_size - n_random)
            n_crossover = 0

        # 1) random exploration
        for _ in range(n_random):
            pool.append(self.config_space.sample_random_config(rng=self.rng, allow_feature_selector=True))

        # 2) mutate elite configs
        elite = self._get_elite_configs(max(2, n_mutation))
        if not elite:
            elite = [self.config_space.sample_random_config(rng=self.rng, allow_feature_selector=True)]

        fairness_pressure = self._compute_fairness_pressure(fairness_context)
        mutation_boost = float(1.0 + self.cfg.drift_reactivity * fairness_pressure)

        for _ in range(n_mutation):
            parent = self.rng.choice(elite)
            child = self.config_space.mutate_config(
                config=parent,
                rng=self.rng,
                p_change_component=_clip(0.20 * mutation_boost, 0.10, 0.60),
                p_change_hparam=_clip(0.50 * mutation_boost, 0.20, 0.90),
                force_valid=True,
            )
            pool.append(child)

        # 3) crossover among elite
        for _ in range(n_crossover):
            if len(elite) >= 2:
                pa, pb = self.rng.sample(elite, 2)
            else:
                pa = elite[0]
                pb = self.config_space.sample_random_config(rng=self.rng, allow_feature_selector=True)

            child = self.config_space.crossover(pa, pb, rng=self.rng)
            child = self.config_space.mutate_config(
                child,
                rng=self.rng,
                p_change_component=0.10,
                p_change_hparam=0.25,
                force_valid=True,
            )
            pool.append(child)

        if self.incumbent is not None:
            pool.append(PipelineConfig.from_dict(self.incumbent.config))

        if self.cfg.enforce_unique_configs:
            dedup: Dict[str, PipelineConfig] = {}
            for cfg in pool:
                dedup[cfg.stable_id()] = cfg
            pool = list(dedup.values())

        while len(pool) < target_size:
            pool.append(self.config_space.sample_random_config(rng=self.rng, allow_feature_selector=True))
            if self.cfg.enforce_unique_configs:
                dedup = {cfg.stable_id(): cfg for cfg in pool}
                pool = list(dedup.values())

        return pool[:target_size]

    def _get_elite_configs(self, k: int) -> List[PipelineConfig]:
        if not self.observations:
            return []

        ranked = sorted(
            self.observations,
            key=lambda obs: self._elite_sort_key(obs),
            reverse=True,
        )

        out: List[PipelineConfig] = []
        seen: set = set()
        for obs in ranked:
            if obs.config_id in seen:
                continue
            seen.add(obs.config_id)
            out.append(PipelineConfig.from_dict(obs.config))
            if len(out) >= int(k):
                break
        return out

    def _elite_sort_key(self, obs: ObservationRecord) -> float:
        return float(
            obs.constrained_utility
            - self.cfg.risk_penalty_weight * obs.fairness_risk
            - (self.cfg.envelope_violation_penalty if obs.envelope_violation else 0.0)
            + (self.cfg.recent_feasible_bonus if obs.constraint_feasible else 0.0)
        )

    def _select_diverse_topk(
        self,
        *,
        ranked: Sequence[Mapping[str, Any]],
        k: int,
        include_incumbent: bool,
    ) -> List[Dict[str, Any]]:
        selected: List[Dict[str, Any]] = []
        selected_x: List[np.ndarray] = []
        selected_ids: set = set()

        if include_incumbent and self.incumbent is not None:
            inc_cfg = dict(self.incumbent.config)
            inc_x = np.asarray(self.config_space.encode_config(inc_cfg, as_numpy=True), dtype=float)
            selected.append({
                "config": copy.deepcopy(inc_cfg),
                "config_id": str(self.incumbent.config_id),
                "x": [float(v) for v in inc_x.tolist()],
                "pred_objective_mean": None,
                "pred_objective_std": None,
                "pred_constraint_mean": None,
                "pred_constraint_std": None,
                "pred_feasibility_mean": None,
                "pred_feasibility_std": None,
                "pred_risk_mean": None,
                "pred_risk_std": None,
                "pred_fairness_mean": None,
                "pred_fairness_std": None,
                "acquisition_mode": "incumbent_carryover",
                "best_reference": self._get_best_reference(),
                "improvement_mean": None,
                "ei_like": None,
                "p_feasible": None,
                "novelty": None,
                "constrained_acquisition": None,
                "reliability_penalty": None,
                "low_support_penalty": None,
                "final_score": float("inf"),
                "rationale": {"mode": "incumbent_carryover"},
            })
            selected_x.append(inc_x)
            selected_ids.add(str(self.incumbent.config_id))

        for rec in ranked:
            if len(selected) >= int(k):
                break

            cfg_id = str(rec.get("config_id", ""))
            if cfg_id in selected_ids:
                continue

            x = np.asarray(rec.get("x", []), dtype=float)
            if x.size == 0:
                continue

            if selected_x:
                min_dist = min(
                    _euclidean_distance(x, sx) / max(1.0, math.sqrt(len(x)))
                    for sx in selected_x
                )
            else:
                min_dist = 1.0

            if min_dist < float(self.cfg.diversity_min_distance):
                continue

            selected.append(copy.deepcopy(dict(rec)))
            selected_x.append(x)
            selected_ids.add(cfg_id)

        if len(selected) < int(k):
            for rec in ranked:
                if len(selected) >= int(k):
                    break
                cfg_id = str(rec.get("config_id", ""))
                if cfg_id in selected_ids:
                    continue
                selected.append(copy.deepcopy(dict(rec)))
                selected_ids.add(cfg_id)

        return selected[:int(k)]

    # --------------------------------------------------------
    # Incumbent logic
    # --------------------------------------------------------
    def _update_incumbent(self) -> None:
        if not self.observations:
            self.incumbent = None
            return

        ranked = sorted(self.observations, key=self._incumbent_key, reverse=True)
        self.incumbent = ranked[0]

    def _incumbent_key(self, obs: ObservationRecord) -> float:
        metric = str(self.cfg.incumbent_metric).strip().lower()
        is_feasible = bool(obs.constraint_feasible and not obs.envelope_violation)
        feasibility_priority = self.cfg.feasible_priority_bonus if is_feasible else -self.cfg.infeasible_priority_penalty
        reliability_bonus = 0.03 if obs.is_reliable else 0.0

        if metric == "accuracy":
            return float(
                feasibility_priority
                + obs.accuracy
                - 0.25 * obs.fairness_risk
                - (self.cfg.envelope_violation_penalty if obs.envelope_violation else 0.0)
                + reliability_bonus
            )

        if metric == "utility":
            return float(
                feasibility_priority
                + obs.utility_score
                - 0.25 * obs.fairness_risk
                - (self.cfg.envelope_violation_penalty if obs.envelope_violation else 0.0)
                + reliability_bonus
            )

        if metric == "negative_log_loss":
            return float(
                feasibility_priority
                -obs.log_loss
                - 0.25 * obs.fairness_risk
                - (self.cfg.envelope_violation_penalty if obs.envelope_violation else 0.0)
                + reliability_bonus
            )

        # default: constrained utility with fairness-aware adjustments
        return float(
            feasibility_priority
            + obs.constrained_utility
            - self.cfg.risk_penalty_weight * obs.fairness_risk
            - (self.cfg.envelope_violation_penalty if obs.envelope_violation else 0.0)
            + (self.cfg.recent_feasible_bonus if obs.constraint_feasible else 0.0)
            + reliability_bonus
        )

    # --------------------------------------------------------
    # Streaming adaptation
    # --------------------------------------------------------
    def _refresh_recency_weights(self) -> None:
        if not self.observations:
            return

        current = max(obs.window_index for obs in self.observations)
        decay = float(self.cfg.recency_decay)

        for obs in self.observations:
            lag = max(0, current - int(obs.window_index))
            obs.recency_weight = float(decay ** lag)

    def _rebuild_observation_index(self) -> None:
        self.observation_by_id = {}
        for obs in self.observations:
            self.observation_by_id.setdefault(obs.config_id, []).append(obs)

    def _compute_fairness_pressure(self, fairness_context: Mapping[str, Any]) -> float:
        """
        Estimate current fairness pressure from envelope/drift context.

        Expected keys may include:
        - envelope_violation
        - envelope_severity
        - fm_drift
        - fm_uncertainty
        - feasibility_probability
        """
        violation = 1.0 if bool(fairness_context.get("envelope_violation", False)) else 0.0
        severity = _clip01(fairness_context.get("envelope_severity", fairness_context.get("severity", 0.0)))
        fm_drift = max(0.0, _safe_float(fairness_context.get("fm_drift", 0.0), 0.0))
        fm_uncertainty = _clip01(fairness_context.get("fm_uncertainty", 0.0))
        feas = _clip01(fairness_context.get("feasibility_probability", 1.0))

        pressure = (
            0.40 * violation
            + 0.25 * severity
            + 0.20 * _clip01(fm_drift)
            + 0.15 * fm_uncertainty
            + 0.20 * (1.0 - feas)
        )
        pressure = _clip01(self.cfg.fairness_pressure_scale * pressure)
        return float(pressure)

    def _build_optimizer_state(
        self,
        *,
        fairness_context: Mapping[str, Any],
        candidate_pool_size: int,
        mode: str,
    ) -> Dict[str, Any]:
        return {
            "mode": str(mode),
            "window_index": int(self.window_index),
            "n_observations": int(len(self.observations)),
            "candidate_pool_size": int(candidate_pool_size),
            "best_reference": float(self._get_best_reference()),
            "fairness_pressure": float(self._compute_fairness_pressure(fairness_context)),
            "has_recent_feasible": bool(self._has_recent_feasible()),
            "incumbent_id": None if self.incumbent is None else str(self.incumbent.config_id),
            "objective_mode": str(self.cfg.objective_mode),
            "incumbent_metric": str(self.cfg.incumbent_metric),
            "acquisition_mode": str(mode),
            "version": str(self.cfg.version),
        }


# ============================================================
# Convenience factory
# ============================================================

def build_default_osmac_optimizer(
    config_space: StreamingConfigSpace,
    seed: int = 42,
) -> OSMACOptimizer:
    return OSMACOptimizer(
        config_space=config_space,
        config=OSMACOptimizerConfig(seed=int(seed)),
    )


# ============================================================
# Minimal smoke test
# ============================================================

if __name__ == "__main__":
    from config_space import build_default_config_space
    from candidate_evaluator import evaluate_candidate

    space = build_default_config_space(seed=42)
    opt = build_default_osmac_optimizer(space, seed=42)

    batch0 = opt.propose_next(n_propose=4)
    print("Initial proposals:", len(batch0.proposed_configs))

    X = [
        {"f1": 0.1, "f2": 1.0},
        {"f1": 0.4, "f2": 0.0},
        {"f1": 0.9, "f2": 1.0},
        {"f1": 0.2, "f2": 0.0},
        {"f1": 0.8, "f2": 1.0},
        {"f1": 0.3, "f2": 0.0},
    ]
    y = [0, 0, 1, 1, 1, 0]
    s = [
        {"sex": "F", "race": "A"},
        {"sex": "M", "race": "A"},
        {"sex": "F", "race": "B"},
        {"sex": "M", "race": "B"},
        {"sex": "F", "race": "B"},
        {"sex": "M", "race": "A"},
    ]

    for cfg in batch0.proposed_configs:
        res = evaluate_candidate(
            config_space=space,
            config=cfg,
            X_window=X,
            y_window=y,
            sensitive_window=s,
            protected_attributes=["sex", "race"],
            candidate_name="smoke_candidate",
        )
        opt.register_observation(res)

    opt.step_window()

    batch1 = opt.propose_next(
        n_propose=4,
        fairness_context={
            "fm_drift": 0.08,
            "fm_uncertainty": 0.20,
            "feasibility_probability": 0.70,
            "envelope_violation": False,
            "envelope_severity": 0.02,
            "envelope_lower": 0.0,
            "envelope_upper": 0.10,
        },
    )

    print("Next proposals:", len(batch1.proposed_configs))
    print("Incumbent:", batch1.incumbent["config_id"] if batch1.incumbent else None)
    if batch1.acquisition_records:
        print("Top acquisition mode:", batch1.acquisition_records[0]["acquisition_mode"])
        print("Top acquisition score:", batch1.acquisition_records[0]["final_score"])