basin 1.13.1

//! Iteration driver. The high-level entry point is [`Executor`];
//! [`Stepper`] exposes one-iteration-at-a-time control, and [`run_loop_with_control`]
//! is the borrowed-problem variant used by composed solvers.
//!
//! # Canonical iteration ordering
//!
//! [`Executor::run`] (and the equivalent [`Stepper`]/[`run_loop_with_control`]
//! paths) drive the solver through this exact sequence, and every
//! contract elsewhere in the framework cross-links here:
//!
//! 1. Convergence history resets, then [`Solver::init`] is called **once**. The
//!    returned state is what iter-0 sees.
//! 2. Then, repeatedly, before each [`Solver::next_iter`] call
//!    (including the first):
//!    1. An executor-attached [`CancellationToken`] is checked, when
//!       configured. Cancellation stops the run with
//!       [`TerminationReason::Cancelled`]. The borrowed [`run_loop_with_control`] path
//!       has no attached token and skips this step.
//!    2. Execution controls check iteration, legacy cost and gradient budgets,
//!       raw evaluation budgets, elapsed time, target cost, improvement stall,
//!       and acceptance stall, followed by application hooks. See [`RunControl`]. Deprecated
//!       criteria retain insertion order within the hook list.
//!    3. [`Solver::check_convergence`] evaluates solver-owned convergence.
//!       Its default delegates to the legacy [`Solver::terminate`] hook.
//!       The first stop ends the run.
//! 3. If nothing fired, [`Solver::next_iter`] is called. It may itself
//!    report a mid-iter termination via its return tuple; in that case
//!    the iteration counter is **not** incremented, so the final
//!    [`State::iter`] reflects the last *fully completed* iteration.
//! 4. Otherwise the iteration counter is incremented and we go back to
//!    step 2.
//!
//! Because checks happen *before* iter 0, an already-optimal initial
//! point exits immediately with the corresponding reason rather than
//! taking one redundant step.
//!
//! With [`Executor::require_evaluated_state`], each successful `init` or
//! `next_iter` return is checked for a complete record before bookkeeping or
//! observation, including clean mid-step stops. Restored checkpoints are also
//! validated before observation. Missing records panic as solver contract
//! violations; hard problem errors bypass publication and validation.
//!
//! [`Executor::resume`] restores state-carried evolution data and evaluation
//! counters, while [`Executor::resume_from_checkpoint`] restores the solver,
//! state, and counters from an [`ExactCheckpoint`] and skips `init`.
//! Consume a paused stepper with [`Stepper::into_checkpoint`], or retain the
//! final solver and raw counters with [`Executor::run_with_solver`].

// Keep the Basin 1.x compatibility bridge and shared check implementations local.
#![allow(deprecated)]

use crate::core::checkpoint::{CheckpointSink, ExactCheckpoint};
use crate::core::observer::{Observe, ObserverMode};
use crate::core::problem::{EvalCounts, Problem};
use crate::core::run_control::RunControl;
use crate::core::solver::Solver;
use crate::core::state::{CountsMirror, ExactResumeState, State};
use crate::core::termination::{TerminationCriterion, TerminationReason};
use std::sync::{
    Arc,
    atomic::{AtomicBool, Ordering},
};

/// Shared, one-shot signal for stopping an [`Executor`] between iterations.
///
/// Clones refer to the same lock-free flag, so a UI or worker thread can keep
/// one handle while the executor owns another. Calling [`cancel`](Self::cancel)
/// is idempotent; a token cannot be reset.
///
/// Cancellation is cooperative: the executor checks after solver
/// initialization and before each new top-level iteration. It does not
/// interrupt an active [`Solver::next_iter`] or problem evaluation. For
/// finer-grained cancellation, return a typed error from the problem method.
///
/// # Example
///
/// ```
/// use basin::CancellationToken;
///
/// let token = CancellationToken::new();
/// let cancel_handle = token.clone();
/// assert!(!token.is_cancelled());
///
/// cancel_handle.cancel();
/// assert!(token.is_cancelled());
/// ```
#[derive(Clone, Debug, Default)]
pub struct CancellationToken {
    cancelled: Arc<AtomicBool>,
}

impl CancellationToken {
    /// Create a token in the active (not cancelled) state.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Request cancellation. Calling this more than once has no further
    /// effect.
    pub fn cancel(&self) {
        self.cancelled.store(true, Ordering::Relaxed);
    }

    /// Whether cancellation has been requested through this token or any of
    /// its clones.
    #[must_use]
    pub fn is_cancelled(&self) -> bool {
        self.cancelled.load(Ordering::Relaxed)
    }
}

/// Outcome of an optimization run.
///
/// Owns the final solver state plus the reason the executor stopped.
/// Delegates `param()`/`cost()`/`iter()` to the underlying state so
/// callers don't need to import `State` for the common reads.
/// Use [`Executor::run_with_solver`] to retain the final solver and raw counts
/// in an [`OptimizationResultWithSolver`].
pub struct OptimizationResult<S> {
    /// Final solver state at termination.
    pub state: S,
    /// Why the executor stopped.
    pub reason: TerminationReason,
}

impl<S: State> OptimizationResult<S> {
    /// Final iterate.
    pub fn param(&self) -> &S::Param {
        self.state.param()
    }

    /// Cost at the final iterate.
    pub fn cost(&self) -> S::Float {
        self.state.cost()
    }

    /// Number of fully completed iterations.
    pub fn iter(&self) -> u64 {
        self.state.iter()
    }

    /// Cumulative cost-function evaluations across the run.
    pub fn cost_evals(&self) -> u64 {
        self.state.cost_evals()
    }

    /// Best iterate observed during the run: the lowest-cost point
    /// the executor ever saw. For sorted-simplex/sorted-population
    /// states this coincides with [`param`](Self::param); for non-
    /// monotone single-iterate runs (Brent's probes, future SA) the
    /// two diverge.
    pub fn best_param(&self) -> &S::Param {
        self.state.best_param()
    }

    /// Cost at [`best_param`](Self::best_param).
    pub fn best_cost(&self) -> S::Float {
        self.state.best_cost()
    }

    /// Iteration at which [`best_param`](Self::best_param) was found.
    pub fn best_iter(&self) -> u64 {
        self.state.best_iter()
    }

    /// Cumulative cost evaluations at the moment
    /// [`best_param`](Self::best_param) was found; answers "how many
    /// evals until the solver hit its best?".
    pub fn best_cost_evals(&self) -> u64 {
        self.state.best_cost_evals()
    }

    /// Consume the result and return the final state.
    pub fn into_state(self) -> S {
        self.state
    }
}

/// Outcome of an optimization run that retains the final solver by ownership.
///
/// Returned by [`Executor::run_with_solver`] and
/// [`Stepper::run_to_end_with_solver`]. The solver retains its final model,
/// history, and other evolving machinery. Neither it nor the state needs to
/// implement `Clone` or serialization.
///
/// Convenience readers use the same state semantics as [`OptimizationResult`].
/// In particular, [`cost_evals`](Self::cost_evals) can fold evaluation
/// categories according to the state's [`CountsMirror`] implementation. Use
/// [`counts`](Self::counts) for the authoritative per-category counters.
pub struct OptimizationResultWithSolver<S, So> {
    /// Final solver state at termination.
    pub state: S,
    /// Final solver, including its evolving machinery and convergence history.
    pub solver: So,
    /// Authoritative evaluation counters, including counts restored on resume.
    pub counts: EvalCounts,
    /// Why the executor stopped.
    pub reason: TerminationReason,
}

impl<S, So> OptimizationResultWithSolver<S, So> {
    /// Consume the result and return the final state, dropping the solver.
    pub fn into_state(self) -> S {
        self.state
    }

    /// Consume the result and retain ordinary progress and its stop reason.
    ///
    /// Drops the solver and the raw counters. State-mirrored counts remain
    /// available through the returned [`OptimizationResult`].
    pub fn into_result(self) -> OptimizationResult<S> {
        OptimizationResult {
            state: self.state,
            reason: self.reason,
        }
    }

    /// Consume the result into a checkpoint for exact continuation.
    ///
    /// Moves the solver and state without cloning or serialization and
    /// discards the termination reason. Pass the checkpoint to
    /// [`Executor::resume_from_checkpoint`] with the same problem and newly
    /// configured execution policy to continue from this boundary.
    pub fn into_checkpoint(self) -> ExactCheckpoint<So, S> {
        ExactCheckpoint::from_parts(self.solver, self.state, self.counts)
    }
}

impl<S: State, So> OptimizationResultWithSolver<S, So> {
    /// Final iterate.
    pub fn param(&self) -> &S::Param {
        self.state.param()
    }

    /// Cost at the final iterate.
    pub fn cost(&self) -> S::Float {
        self.state.cost()
    }

    /// Number of fully completed iterations.
    pub fn iter(&self) -> u64 {
        self.state.iter()
    }

    /// Cumulative evaluation work under the state's [`CountsMirror`] mapping.
    ///
    /// Use `counts.cost_evals` for the raw number of cost-function calls.
    pub fn cost_evals(&self) -> u64 {
        self.state.cost_evals()
    }

    /// Best iterate under the state's incumbent-selection semantics.
    pub fn best_param(&self) -> &S::Param {
        self.state.best_param()
    }

    /// Cost at [`best_param`](Self::best_param).
    pub fn best_cost(&self) -> S::Float {
        self.state.best_cost()
    }

    /// Iteration at which [`best_param`](Self::best_param) was selected.
    pub fn best_iter(&self) -> u64 {
        self.state.best_iter()
    }

    /// State-mirrored evaluation work when the best iterate was selected.
    pub fn best_cost_evals(&self) -> u64 {
        self.state.best_cost_evals()
    }
}

/// Outcome of a single [`Stepper::step`] call.
///
/// `Stopped` carries the same [`TerminationReason`] the executor would
/// have returned. After `Stopped` is returned once, subsequent calls to
/// `step` keep returning the same `Stopped(reason)` so callers don't
/// have to track whether they're done.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StepOutcome {
    /// The step completed without triggering termination.
    Continue,
    /// Termination fired with the given reason. Subsequent
    /// [`Stepper::step`] calls keep returning this same outcome.
    Stopped(TerminationReason),
}

/// Drive a solver one iteration at a time.
///
/// Owns the problem, state, solver and termination criteria, runs
/// `solver.init` exactly once on construction, and exposes
/// [`step`](Self::step)/[`run_to_end`](Self::run_to_end) so callers can
/// interleave their own work between iterations: recording trajectories,
/// animating from a UI, pausing on a button press, evaluating a custom
/// budget, etc.
///
/// [`Executor::run`] is `self.into_stepper().run_to_end()`; the stepper
/// is the building block, the executor is the convenience wrapper.
///
/// # Example
///
/// ```ignore
/// let solver = solver.with_absolute_gradient_tolerance(1e-6);
/// let mut stepper = Executor::new(problem, solver, state)
///     .max_iter(100)
///     .into_stepper()?;
///
/// let reason = loop {
///     match stepper.step()? {
///         StepOutcome::Continue => { /* observe `stepper.state()` */ }
///         StepOutcome::Stopped(reason) => break reason,
///     }
/// };
/// ```
pub struct Stepper<P, S, So> {
    problem: Problem<P>,
    // `Option<S>` because `Solver::next_iter` consumes the state by
    // value. A hard error cannot restore that owned value, so the slot
    // remains empty after a failed step. Successful steps restore it.
    state: Option<S>,
    solver: So,
    control: RunControl<S>,
    observers: Vec<(Box<dyn Observe<S>>, ObserverMode)>,
    checkpoints: Vec<(Box<dyn CheckpointSink<So, S>>, ObserverMode)>,
    cancellation_token: Option<CancellationToken>,
    finished: Option<TerminationReason>,
}

impl<P, S, So> Stepper<P, S, So>
where
    S: State + CountsMirror,
    So: Solver<P, S>,
{
    /// Read-only access to the current state, between steps.
    ///
    /// # Panics
    ///
    /// Panics after [`step`](Self::step) returns `Err`, because the
    /// failing solver call consumed the state.
    pub fn state(&self) -> &S {
        self.state
            .as_ref()
            .expect("state slot is Some between steps")
    }

    /// Wrapper-side evaluation counters. These are authoritative:
    /// solvers can only call into the user's problem through the
    /// wrapper, so every cost/gradient/residual/Jacobian /
    /// Hessian call is reflected here. The state mirror under
    /// [`state`](Self::state) is refreshed after every successful
    /// [`Solver::init`] /
    /// [`Solver::next_iter`];
    /// on the typed-`Err` path the state slot is dropped (see
    /// [`step`](Self::step)) but `counts` is still readable here for
    /// diagnostics.
    pub fn counts(&self) -> &EvalCounts {
        self.problem.counts()
    }

    /// Termination reason if the stepper has stopped, else `None`.
    pub fn finished(&self) -> Option<&TerminationReason> {
        self.finished.as_ref()
    }

    /// Total iterations that have completed so far. Convenience read
    /// equivalent to `self.state().iter()`.
    ///
    /// # Panics
    ///
    /// Panics after [`step`](Self::step) returns `Err`, as
    /// [`state`](Self::state) does.
    pub fn iter(&self) -> u64 {
        self.state().iter()
    }

    /// Advance one iteration. Once a `Stopped` outcome has been returned
    /// the stepper is sticky: subsequent calls keep returning the same
    /// `Stopped(reason)` without touching the state or solver.
    ///
    /// Registered observers fire here:
    /// [`observe_iter`](Observe::observe_iter) on
    /// [`StepOutcome::Continue`], gated by each observer's
    /// [`ObserverMode`]; [`observe_final`](Observe::observe_final) once
    /// when this call first returns [`StepOutcome::Stopped`]. See the
    /// [`observer`](crate::core::observer) module for the lifecycle.
    ///
    /// Returns `Err` when the underlying problem returns `Err` from any
    /// cost/gradient/residual/Jacobian/Hessian call during the
    /// step. A hard error consumes the state and may leave solver machinery
    /// partially updated. It does not set [`finished`](Self::finished).
    /// Callers can inspect [`counts`](Self::counts), then drop the stepper or
    /// call [`into_checkpoint`](Self::into_checkpoint), which returns `None`.
    /// State access and further stepping are not supported after the error.
    /// Observers and checkpoint sinks do not fire on the failed transition.
    pub fn step(&mut self) -> Result<StepOutcome, So::Error> {
        if let Some(reason) = self.finished {
            return Ok(StepOutcome::Stopped(reason));
        }
        let outcome = if self
            .cancellation_token
            .as_ref()
            .is_some_and(CancellationToken::is_cancelled)
        {
            StepOutcome::Stopped(TerminationReason::Cancelled)
        } else {
            step_once(
                &mut self.problem,
                &EvalCounts::default(),
                &mut self.state,
                &mut self.solver,
                &mut self.control,
                &mut [],
            )?
        };
        match outcome {
            StepOutcome::Continue => {
                let state = self
                    .state
                    .as_ref()
                    .expect("state slot is Some after Continue");
                let iter = state.iter();
                // `update_best` set `best_iter == iter` iff this iteration
                // strictly improved the incumbent; that is exactly the
                // `NewBest` firing condition.
                let is_new_best = state.best_iter() == iter;
                for (checkpoint, mode) in self.checkpoints.iter_mut() {
                    if mode.fires_on(iter, is_new_best) {
                        checkpoint.save(
                            &self.solver,
                            state,
                            self.problem.counts(),
                        );
                    }
                }
                for (observer, mode) in self.observers.iter_mut() {
                    if mode.fires_on(iter, is_new_best) {
                        observer.observe_iter(state);
                    }
                }
            }
            StepOutcome::Stopped(reason) => {
                self.finished = Some(reason);
                let state =
                    self.state.as_ref().expect("state slot is Some on Stopped");
                for (checkpoint, _mode) in self.checkpoints.iter_mut() {
                    checkpoint.save(&self.solver, state, self.problem.counts());
                }
                for (observer, _mode) in self.observers.iter_mut() {
                    observer.observe_final(state, &reason);
                }
            }
        }
        Ok(outcome)
    }

    /// Drive [`step`](Self::step) to completion and return an
    /// [`OptimizationResult`].
    /// Use [`run_to_end_with_solver`](Self::run_to_end_with_solver) to retain
    /// the final solver and raw evaluation counters as well.
    pub fn run_to_end(self) -> Result<OptimizationResult<S>, So::Error> {
        self.run_to_end_with_solver()
            .map(OptimizationResultWithSolver::into_result)
    }

    /// Drive [`step`](Self::step) to completion, retaining the final solver,
    /// state, raw evaluation counters, and termination reason by ownership.
    ///
    /// Uses the same lifecycle as [`run_to_end`](Self::run_to_end), including
    /// observer and checkpoint callbacks. An already-stopped stepper returns
    /// its recorded reason without repeating final callbacks. No `Clone` or
    /// serialization bounds are required.
    ///
    /// Returns the solver's error on a failed transition, without a partial
    /// result or recoverable checkpoint. Calling this after a previous
    /// [`step`](Self::step) error is unsupported, just as with `run_to_end`.
    pub fn run_to_end_with_solver(
        mut self,
    ) -> Result<OptimizationResultWithSolver<S, So>, So::Error> {
        loop {
            if let StepOutcome::Stopped(reason) = self.step()? {
                return Ok(OptimizationResultWithSolver {
                    state: self
                        .state
                        .take()
                        .expect("state slot is Some on stop"),
                    solver: self.solver,
                    counts: *self.problem.counts(),
                    reason,
                });
            }
        }
    }

    /// Consume the stepper into a checkpoint at its current boundary.
    ///
    /// Returns `Some` after successful initialization, between completed
    /// steps, or after a clean stop (including cancellation and a mid-step
    /// stop). Returns `None` after a hard [`step`](Self::step) error consumed
    /// the state; partial solver machinery cannot form an exact checkpoint.
    ///
    /// Moves the solver and state and copies the authoritative counters
    /// without requiring `Clone` or serialization. Extraction performs no
    /// evaluations, convergence checks, or observer/checkpoint callbacks.
    /// The problem, execution policy, and any recorded termination reason
    /// are dropped. Resume through [`Executor::resume_from_checkpoint`],
    /// which describes the requirements for exact continuation.
    ///
    /// # Example
    ///
    /// ```
    /// use basin::{CostFunction, Executor, NelderMead, State};
    /// # fn main() -> Result<(), std::convert::Infallible> {
    /// struct Sphere;
    /// impl CostFunction for Sphere {
    ///     type Param = Vec<f64>;
    ///     type Output = f64;
    ///     type Error = std::convert::Infallible;
    ///     fn cost(&self, x: &Vec<f64>) -> Result<f64, Self::Error> {
    ///         Ok(x.iter().map(|v| v * v).sum())
    ///     }
    /// }
    /// let mut stepper = Executor::from_start(
    ///     Sphere, NelderMead::new(), vec![2.0, 1.0],
    /// ).into_stepper()?;
    /// stepper.step()?;
    /// let checkpoint = stepper.into_checkpoint().unwrap();
    /// assert_eq!(checkpoint.state().iter(), 1);
    /// let result = Executor::resume_from_checkpoint(Sphere, checkpoint)
    ///     .max_iter(10)
    ///     .run()?;
    /// assert_eq!(result.iter(), 10);
    /// # Ok(())
    /// # }
    /// ```
    pub fn into_checkpoint(self) -> Option<ExactCheckpoint<So, S>> {
        Some(ExactCheckpoint::from_parts(
            self.solver,
            self.state?,
            *self.problem.counts(),
        ))
    }

    /// Consume the stepper and return the final state.
    ///
    /// # Panics
    ///
    /// Panics after [`step`](Self::step) returns `Err`, because the
    /// failing solver call consumed the state.
    pub fn into_state(self) -> S {
        self.state.expect("state slot is Some at drop")
    }
}

/// Single-iteration core, shared by [`Stepper::step`] (owned) and
/// [`run_loop`] (borrowed). Reads the current state via `state_slot`,
/// checks termination, and either returns `Stopped` (slot left
/// untouched) or hands the state to `solver.next_iter`, mirrors the
/// wrapper's counter delta (relative to `baseline`) onto the state,
/// increments the iteration counter, and puts the returned state back.
///
/// The `baseline` captures the wrapper count at the start of the
/// containing run so the state mirror always reflects *per-run* work:
/// for [`Stepper::step`]/[`Executor::run`] it is
/// [`EvalCounts::default`] (fresh wrapper), for nested
/// [`run_loop`] calls it is the wrapper count at run-loop entry.
///
/// The state slot must be `Some` on entry and is `Some` after an `Ok`
/// return. If [`Solver::next_iter`] returns `Err`, it has consumed the
/// previous state, so the slot remains empty and execution cannot resume
/// from this pair. The wrapper's charged counts remain available.
fn step_once<P, S, So>(
    problem: &mut Problem<P>,
    baseline: &EvalCounts,
    state_slot: &mut Option<S>,
    solver: &mut So,
    control: &mut RunControl<S>,
    criteria: &mut [Box<dyn TerminationCriterion<S>>],
) -> Result<StepOutcome, So::Error>
where
    S: State + CountsMirror,
    So: Solver<P, S>,
{
    {
        let state = state_slot
            .as_ref()
            .expect("step_once called with empty state slot");
        if let Some(reason) =
            control.check(state, &problem.counts().delta_since(baseline))
        {
            return Ok(StepOutcome::Stopped(reason));
        }
        for criterion in criteria.iter_mut() {
            if let Some(reason) = criterion.check(state) {
                return Ok(StepOutcome::Stopped(reason));
            }
        }
        if let Some(reason) = solver.check_convergence(problem, state) {
            return Ok(StepOutcome::Stopped(reason));
        }
    }
    let prev = state_slot.take().unwrap();
    let next_iter_result = solver.next_iter(problem, prev);
    let (mut next, mid_iter_reason) = match next_iter_result {
        Ok(t) => t,
        Err(e) => {
            // The solver consumed `prev`, so restoring the state would
            // require a separate snapshot. Preserve the hard error and
            // leave the wrapper's charged counts available for diagnosis.
            return Err(e);
        }
    };
    control.validate(&next);
    next.mirror(&problem.counts().delta_since(baseline));
    if let Some(reason) = mid_iter_reason {
        // Refresh best-so-far from the mid-iter state too: the solver
        // may have produced its best iterate on the same step that
        // bailed.
        next.update_best();
        *state_slot = Some(next);
        return Ok(StepOutcome::Stopped(reason));
    }
    next.increment_iter();
    next.update_best();
    *state_slot = Some(next);
    Ok(StepOutcome::Continue)
}

/// Drive a solver to completion against a shared [`Problem`] wrapper.
///
/// `Executor` is a thin owning wrapper over this. Composed solvers
/// (e.g. CG inside CMA, NM inside DE) call `run_loop` directly so the
/// inner solver shares the outer's wrapper: inner cost and gradient
/// calls bump the same [`EvalCounts`] as outer calls, so the eval
/// aggregation contract (`CONTRIBUTING.md` "Solver composition" rule 1) is
/// satisfied automatically for same-problem inners. For composed
/// solvers driving an inner against an **adapter problem** (e.g.
/// [`LogBarrier`](crate::core::barrier::LogBarrier)), construct a
/// fresh `Problem::new(adapter)`, pass `&mut` into `run_loop`, then
/// fold the inner wrapper's [`EvalCounts`] back into the outer's via
/// [`EvalCounts::add`] on [`Problem::counts_mut`].
///
/// The inner state's [`State::cost_evals`] (mirrored via
/// [`CountsMirror`]) reflects only *per-run* work: `run_loop` takes
/// a baseline snapshot of [`Problem::counts`] at entry, and the state
/// mirror computes the delta against that. Nested `run_loop` calls
/// against the same wrapper therefore see clean per-call counters.
///
/// Apart from executor-attached cancellation (which is top-level only),
/// semantics match `Executor::run`: each criterion is
/// [`reset`](crate::core::termination::TerminationCriterion::reset) at
/// entry, so a criteria vector reused across calls (as an
/// [`InnerExecutor`](crate::core::inner::InnerExecutor) does) sees fresh
/// per-run state. Then `init` is called once, then on each iteration
/// framework `criteria` are checked in insertion order before
/// the solver's own `terminate` hook, before stepping. `max_iter` is
/// checked against `state.iter()` and exits with `TerminationReason::MaxIter`.
/// `next_iter` may also report a mid-iter termination via its return tuple;
/// in that case the iteration counter is left untouched so the final
/// `state.iter()` still reflects the last fully completed iteration.
#[deprecated(
    note = "use `run_loop_with_control`; removal scheduled for Basin 2.0"
)]
pub fn run_loop<P, S, So>(
    problem: &mut Problem<P>,
    state: S,
    solver: &mut So,
    criteria: &mut [Box<dyn TerminationCriterion<S>>],
    max_iter: u64,
) -> Result<OptimizationResult<S>, So::Error>
where
    S: State + CountsMirror,
    So: Solver<P, S>,
{
    let mut control = RunControl::new().max_iter(max_iter);
    run_loop_impl(problem, state, solver, &mut control, criteria)
}

/// Drive a borrowed solver using execution budgets and application stops.
///
/// Convergence is configured on `solver`. Controls and convergence history
/// reset before initialization; evaluation counts are relative to run entry.
/// See [`RunControl`] for check ordering and clock semantics.
pub fn run_loop_with_control<P, S, So>(
    problem: &mut Problem<P>,
    state: S,
    solver: &mut So,
    control: &mut RunControl<S>,
) -> Result<OptimizationResult<S>, So::Error>
where
    S: State + CountsMirror,
    So: Solver<P, S>,
{
    run_loop_impl(problem, state, solver, control, &mut [])
}

fn run_loop_impl<P, S, So>(
    problem: &mut Problem<P>,
    mut state: S,
    solver: &mut So,
    control: &mut RunControl<S>,
    criteria: &mut [Box<dyn TerminationCriterion<S>>],
) -> Result<OptimizationResult<S>, So::Error>
where
    S: State + CountsMirror,
    So: Solver<P, S>,
{
    control.reset();
    solver.reset_convergence();
    let baseline = *problem.counts();
    // Reset each criterion's internal per-run state before the run, so a
    // criteria vector reused across `run_loop` calls (e.g. an
    // `InnerExecutor` driven once per outer iter) sees fresh state each
    // call. Stateful criteria (`MaxTime`, `RelativeGradientTolerance`,
    // `NoImprovement`) would otherwise carry state across runs and
    // misbehave; the default `reset` is a no-op for stateless ones.
    for criterion in criteria.iter_mut() {
        criterion.reset();
    }
    // Reset best-so-far so the state always reflects per-run work,
    // matching the snapshot discipline `state.mirror` uses for eval
    // counters. This makes the same state safe to drive across
    // multiple `run_loop` calls (e.g. an outer solver re-driving an
    // inner) without best-so-far bleeding from one run into the next.
    state.reset_best();
    let mut state = solver.init(problem, state)?;
    control.validate(&state);
    // Mirror init's work onto the state before any termination check.
    state.mirror(&problem.counts().delta_since(&baseline));
    state.update_best();
    let mut slot = Some(state);
    let reason = loop {
        match step_once(
            problem, &baseline, &mut slot, solver, control, criteria,
        )? {
            StepOutcome::Continue => continue,
            StepOutcome::Stopped(reason) => break reason,
        }
    };
    Ok(OptimizationResult {
        state: slot.take().expect("state slot is Some on stop"),
        reason,
    })
}

/// User-facing driver. Owns the problem, solver, initial state, and the
/// execution controls; [`run`](Self::run) drives the iteration
/// loop to completion. See the [module docs](self) for the canonical
/// ordering and [`into_stepper`](Self::into_stepper) for one-step-at-a-
/// time control.
///
/// # Examples
///
/// Minimize the 2-D sphere and read the outcome off the
/// [`OptimizationResult`]:
///
/// ```
/// use basin::{
///     BasicState, CostFunction, Executor, Gradient, GradientDescent,
/// };
///
/// struct Sphere;
/// impl CostFunction for Sphere {
///     type Param = Vec<f64>;
///     type Output = f64;
///     type Error = std::convert::Infallible;
///     fn cost(&self, x: &Vec<f64>) -> Result<f64, std::convert::Infallible> {
///         Ok(x.iter().map(|xi| xi * xi).sum())
///     }
/// }
/// impl Gradient for Sphere {
///     type Gradient = Vec<f64>;
///     fn gradient(
///         &self,
///         x: &Vec<f64>,
///     ) -> Result<Vec<f64>, std::convert::Infallible> {
///         Ok(x.iter().map(|xi| 2.0 * xi).collect())
///     }
/// }
///
/// let result = Executor::new(
///     Sphere,
///     (GradientDescent::new(0.1)).with_absolute_gradient_tolerance(1e-9),
///     BasicState::new(vec![3.0, -4.0]),
/// )
/// .max_iter(1_000)
/// .run()
/// .unwrap();
///
/// assert!(result.cost() < 1e-12);
/// ```
pub struct Executor<P, S, So> {
    problem: P,
    state: S,
    solver: So,
    control: RunControl<S>,
    observers: Vec<(Box<dyn Observe<S>>, ObserverMode)>,
    checkpoints: Vec<(Box<dyn CheckpointSink<So, S>>, ObserverMode)>,
    cancellation_token: Option<CancellationToken>,
    resume_counts: Option<EvalCounts>,
    skip_init: bool,
}

impl<P, S, So> Executor<P, S, So>
where
    S: State + CountsMirror,
    So: Solver<P, S>,
{
    /// Build an executor from a problem, solver, and initial state. The
    /// default `MaxIter` budget is 1000; override with
    /// [`max_iter`](Self::max_iter).
    pub fn new(problem: P, solver: So, state: S) -> Self {
        Self {
            problem,
            state,
            solver,
            control: RunControl::new(),
            observers: Vec::new(),
            checkpoints: Vec::new(),
            cancellation_token: None,
            resume_counts: None,
            skip_init: false,
        }
    }

    /// Build an executor that continues an exact state snapshot.
    ///
    /// Unlike [`new`](Self::new), this constructor preserves the state's
    /// best-so-far history and restores the problem wrapper's cumulative
    /// evaluation counters. The state must implement [`ExactResumeState`],
    /// which promises that it contains all solver evolution data required for
    /// an unchanged continuation. The solver's `init` implementation must be
    /// resume-idempotent.
    ///
    /// `max_iter` remains an absolute iteration limit: resuming a state at
    /// iteration 40 with `.max_iter(100)` performs at most 60 more iterations.
    /// Exact continuation also requires the same deterministic problem,
    /// solver configuration, scalar type, and code. Termination criteria are
    /// configured anew; state-derived criteria such as
    /// [`NoAcceptance`](crate::NoAcceptance) and zero-tolerance
    /// [`NoImprovement`](crate::NoImprovement) preserve history stored in the
    /// state, while criteria with private clocks or anchors begin a new
    /// criterion run.
    pub fn resume(problem: P, solver: So, state: S) -> Self
    where
        S: ExactResumeState,
    {
        let resume_counts = state.resume_counts();
        let mut executor = Self::new(problem, solver, state);
        executor.resume_counts = Some(resume_counts);
        executor
    }

    /// Build an executor that continues an exact solver/state checkpoint.
    ///
    /// Unlike [`new`](Self::new), this constructor restores the solver and
    /// problem wrapper's cumulative evaluation counters, preserves the
    /// state's best-so-far history, and does not call [`Solver::init`].
    ///
    /// `max_iter` remains an absolute iteration limit: resuming a state at
    /// iteration 40 with `.max_iter(100)` performs at most 60 more iterations.
    /// Exact continuation also requires the same deterministic problem,
    /// scalar type, backend behavior, and code. The checkpoint does not contain
    /// the problem, execution limits, application hooks, observers,
    /// cancellation token, or checkpoint sinks; configure that execution
    /// policy anew. State-derived criteria such as
    /// [`NoAcceptance`](crate::NoAcceptance) and zero-tolerance
    /// [`NoImprovement`](crate::NoImprovement) preserve history stored in the
    /// state, while criteria with private clocks or anchors begin a new
    /// criterion run.
    /// Obtain an owned checkpoint with [`Stepper::into_checkpoint`] or
    /// [`OptimizationResultWithSolver::into_checkpoint`], without cloning
    /// or serializing the solver and state.
    pub fn resume_from_checkpoint(
        problem: P,
        checkpoint: ExactCheckpoint<So, S>,
    ) -> Self {
        let (solver, state, resume_counts) = checkpoint.into_parts();
        let mut executor = Self::new(problem, solver, state);
        executor.resume_counts = Some(resume_counts);
        executor.skip_init = true;
        executor
    }

    /// Build an executor seeding the solver's natural initial state at the
    /// starting point `x0`, instead of constructing the [`State`] by hand.
    ///
    /// `Executor::from_start(problem, solver, x0)` calls
    /// [`InitialState::seed`](crate::core::inner::InitialState::seed), so the
    /// caller never names the concrete state type: the common case reads
    /// `Executor::from_start(problem, TrustRegion::new(), x0).run()`. The
    /// seeded state uses the solver's natural default scale (identity inverse
    /// Hessian, default simplex edge, the solver's default trust radius, …).
    ///
    /// Use [`new`](Self::new) directly to supply a custom initial state (a
    /// pre-built simplex, a warm-started inverse Hessian, an anisotropic
    /// CMA-ES covariance). Solvers whose natural initialization needs more
    /// than a point, namely CMA-ES (step-size σ), the population GA, DE, or
    /// random search (they sample the box), and the bracketing scalar solvers
    /// (Brent, golden-section), deliberately do not implement
    /// [`InitialState`](crate::core::inner::InitialState), so calling
    /// `from_start` with one is a compile error pointing back to
    /// [`new`](Self::new).
    pub fn from_start<V>(problem: P, solver: So, x0: V) -> Self
    where
        So: crate::core::inner::InitialState<V, State = S>,
    {
        let state = solver.seed(&x0);
        Self::new(problem, solver, state)
    }

    /// Convenience setter for the default `MaxIter` criterion. Equivalent
    /// effect to `terminate_on(MaxIter(n))` but mutates a dedicated field
    /// so subsequent calls replace rather than stack.
    pub fn max_iter(mut self, n: u64) -> Self {
        self.control.max_iter = n;
        self
    }

    crate::core::run_control::control_methods!();

    /// Append an application stop evaluated before solver convergence.
    /// The closure sees an initialized state, including iteration zero.
    pub fn stop_when<C>(mut self, check: C) -> Self
    where
        C: FnMut(&S) -> Option<TerminationReason> + 'static,
    {
        self.control = std::mem::take(&mut self.control).stop_when(check);
        self
    }

    /// Add a termination criterion. Criteria are checked in insertion
    /// order before each iteration (and before iter 0); the first to
    /// return `Some(_)` stops the run. See the [module docs](self) for
    /// the full per-iteration ordering.
    #[deprecated(
        note = "configure solver convergence or use executor budgets and `stop_when`; removal scheduled for Basin 2.0"
    )]
    pub fn terminate_on<C>(mut self, criterion: C) -> Self
    where
        C: TerminationCriterion<S> + 'static,
    {
        self.control.push_legacy(Box::new(criterion));
        self
    }

    /// Attach a cooperative cancellation token to this run.
    ///
    /// The executor checks the token after [`Solver::init`] and before every
    /// top-level iteration. A cancellation request returns
    /// `Ok(OptimizationResult)` with [`TerminationReason::Cancelled`]; the
    /// state remains at the last fully completed iteration, including its
    /// best-so-far fields. An in-progress iteration or problem evaluation is
    /// allowed to finish before the token is observed.
    ///
    /// Calling this method again replaces the previously configured token.
    ///
    /// # Example
    ///
    /// ```
    /// use basin::{
    ///     BasicState, CancellationToken, CostFunction, Executor, Gradient,
    ///     GradientDescent, TerminationReason,
    /// };
    ///
    /// struct Sphere;
    /// impl CostFunction for Sphere {
    ///     type Param = Vec<f64>;
    ///     type Output = f64;
    ///     type Error = std::convert::Infallible;
    ///     fn cost(&self, x: &Vec<f64>) -> Result<f64, Self::Error> {
    ///         Ok(x.iter().map(|xi| xi * xi).sum())
    ///     }
    /// }
    /// impl Gradient for Sphere {
    ///     type Gradient = Vec<f64>;
    ///     fn gradient(&self, x: &Vec<f64>) -> Result<Vec<f64>, Self::Error> {
    ///         Ok(x.iter().map(|xi| 2.0 * xi).collect())
    ///     }
    /// }
    ///
    /// let token = CancellationToken::new();
    /// let cancel_handle = token.clone();
    /// cancel_handle.cancel(); // A UI callback or worker may hold this clone.
    ///
    /// let result = Executor::new(
    ///     Sphere,
    ///     GradientDescent::new(0.1),
    ///     BasicState::new(vec![1.0, 1.0]),
    /// )
    /// .with_cancellation_token(token)
    /// .run()
    /// .unwrap();
    ///
    /// assert_eq!(result.reason, TerminationReason::Cancelled);
    /// assert_eq!(result.iter(), 0);
    /// ```
    pub fn with_cancellation_token(mut self, token: CancellationToken) -> Self {
        self.cancellation_token = Some(token);
        self
    }

    /// Register an [`Observe`] hook. Observers fire in registration order;
    /// `mode` gates [`Observe::observe_iter`] only;
    /// [`Observe::observe_init`] and [`Observe::observe_final`] always
    /// fire. See the [`observer`](crate::core::observer) module for the
    /// lifecycle.
    ///
    /// Observers cannot fail the run. Use
    /// [`stop_when`](Self::stop_when) for application stopping, or let
    /// an observer cancel a cloned [`CancellationToken`] for a clean
    /// user-requested stop.
    pub fn observe_with<O>(mut self, observer: O, mode: ObserverMode) -> Self
    where
        O: Observe<S> + 'static,
    {
        self.observers.push((Box::new(observer), mode));
        self
    }

    /// Register a solver-aware checkpoint destination.
    ///
    /// `mode` controls saves after completed iterations. Every sink also saves
    /// once when the run stops cleanly, regardless of the mode. Checkpoint
    /// sinks are separate from state-only [`Observe`] implementations because
    /// exact continuation requires both the solver and state.
    ///
    /// Sink failures cannot change the optimization result; implementations
    /// must record or report their own failures. Calling this method again
    /// adds another sink.
    pub fn checkpoint_with<C>(
        mut self,
        checkpoint: C,
        mode: ObserverMode,
    ) -> Self
    where
        C: CheckpointSink<So, S> + 'static,
    {
        self.checkpoints.push((Box::new(checkpoint), mode));
        self
    }

    /// Convert the executor into a [`Stepper`] for one-iteration-at-a-time
    /// control. On a fresh executor or a legacy state-resume executor,
    /// `solver.init` runs here so the returned stepper sits at iter 0 with a
    /// complete state. A solver-aware checkpoint resume skips `init` and uses
    /// its restored iteration boundary directly. All registered observers'
    /// `observe_init` hooks fire in either case. Cancellation is first checked
    /// by the returned stepper's initial [`step`](Stepper::step).
    ///
    /// Returns `Err` when [`Solver::init`] does (e.g. the problem's
    /// initial cost/gradient evaluation `Err`-ed). Observers do *not* fire
    /// on that error path.
    pub fn into_stepper(self) -> Result<Stepper<P, S, So>, So::Error> {
        let Self {
            problem,
            mut state,
            mut solver,
            control,
            mut observers,
            checkpoints,
            cancellation_token,
            resume_counts,
            skip_init,
        } = self;
        let mut problem = Problem::new(problem);
        if let Some(counts) = resume_counts {
            *problem.counts_mut() = counts;
        } else {
            // Fresh top-level wrapper: reset best-so-far so it tracks
            // this run's iterates only, matching the `state.mirror`
            // per-run snapshot discipline.
            state.reset_best();
        }
        let state = if skip_init {
            control.validate(&state);
            state
        } else {
            solver.reset_convergence();
            let mut state = solver.init(&mut problem, state)?;
            control.validate(&state);
            // Mirror init's work onto the state before any termination
            // check. Baseline is zero: this is a fresh top-level wrapper.
            state.mirror(problem.counts());
            state.update_best();
            state
        };
        for (observer, _mode) in observers.iter_mut() {
            observer.observe_init(&state);
        }
        Ok(Stepper {
            problem,
            state: Some(state),
            solver,
            control,
            observers,
            checkpoints,
            cancellation_token,
            finished: None,
        })
    }

    /// Drive the iteration loop to completion and return the
    /// [`OptimizationResult`].
    /// Use [`run_with_solver`](Self::run_with_solver) to retain the final
    /// solver and raw evaluation counters as well.
    ///
    /// Returns `Err` when the underlying problem returns `Err` from any
    /// cost/gradient/residual/Jacobian/Hessian call (the
    /// `P::Error`-flavored hard-abort path; see the
    /// [`problem`](crate::core::problem) module docs).
    pub fn run(self) -> Result<OptimizationResult<S>, So::Error> {
        self.into_stepper()?.run_to_end()
    }

    /// Drive the iteration loop to completion, retaining the final solver,
    /// state, raw evaluation counters, and termination reason by ownership.
    ///
    /// Follows the same initialization, stopping, and callback lifecycle as
    /// [`run`](Self::run). Neither the solver nor the state needs to implement
    /// `Clone` or serialization. Errors propagate with the same type and
    /// without a partial result.
    ///
    /// # Example
    ///
    /// ```
    /// use basin::{CostFunction, Executor, NelderMead};
    /// # fn main() -> Result<(), std::convert::Infallible> {
    /// struct Sphere;
    /// impl CostFunction for Sphere {
    ///     type Param = Vec<f64>;
    ///     type Output = f64;
    ///     type Error = std::convert::Infallible;
    ///     fn cost(&self, x: &Vec<f64>) -> Result<f64, Self::Error> {
    ///         Ok(x.iter().map(|v| v * v).sum())
    ///     }
    /// }
    /// let result = Executor::from_start(
    ///     Sphere, NelderMead::new(), vec![2.0, 1.0],
    /// ).max_iter(3).run_with_solver()?;
    /// assert_eq!(result.iter(), 3);
    /// let counts = result.counts;
    /// let checkpoint = result.into_checkpoint();
    /// let continued = Executor::resume_from_checkpoint(Sphere, checkpoint)
    ///     .max_iter(10)
    ///     .run_with_solver()?;
    /// assert!(continued.counts.cost_evals > counts.cost_evals);
    /// # Ok(())
    /// # }
    /// ```
    pub fn run_with_solver(
        self,
    ) -> Result<OptimizationResultWithSolver<S, So>, So::Error> {
        self.into_stepper()?.run_to_end_with_solver()
    }
}