basin 0.1.0

use crate::core::executor::OptimizationResult;
use crate::core::inner::InnerExecutor;
use crate::core::math::{
    ComponentMulAssign, MatTransposeVec, MatVec, MatrixIdentity, NormSquared, RankOneUpdate,
    SampleStandardNormal, ScaleInPlace, ScaledAdd, SymmetricEigen, VectorLen,
};
use crate::core::problem::CostFunction;
use crate::core::solver::Solver;
use crate::core::state::{
    BasicPopulationState, BasicSimplexState, BasicState, GradientState, LbfgsState, State,
};
use crate::core::termination::{TerminationCriterion, TerminationReason};
use crate::solver::cma_es::{sort_population_ascending, CmaEs};
use crate::solver::lbfgsb::LBFGSB;
use crate::solver::levenberg_marquardt::LevenbergMarquardt;
use crate::solver::nelder_mead::NelderMead;

/// An inner solver eligible to plug into a CMA-ES injection wrapper
/// ([`CmaInject`] / [`BoundedCmaInject`](crate::solver::BoundedCmaInject)).
///
/// The trait is the contract between the outer memetic glue and the
/// inner local solver: it says "given a CMA-ES candidate `x` and the
/// current step-size `σ`, build a fresh inner state; given a final
/// inner state, report the total work units accumulated."
///
/// # Implementations
///
/// Shipped impls for [`NelderMead`], [`LevenbergMarquardt`], and
/// [`LBFGSB`]. To plug in something else, either impl this trait on
/// your solver, or wrap a `Solver<P, S>` in [`ClosureInner`] with
/// inline seeder/work closures (escape hatch for one-off experiments
/// and the `AlwaysFails`-style failure-bubbling tests).
///
/// # Why an associated state type
///
/// Each inner has a natural state shape: NM wants a simplex (`n + 1`
/// vertices), LM wants a single iterate with cached residual / Jacobian,
/// L-BFGS-B wants the limited-memory history. Tying `State` to the
/// trait lets the memetic factory write
/// `BoundedCmaInject::with_inner_solver(cma, LBFGSB::new())` without
/// the caller having to spell out `LbfgsState<V>` in turbofish — `I`
/// determines it.
///
/// # Eval aggregation
///
/// `work_units(&self, state)` is what the outer rolls into its
/// `cost_evals` counter (AGENTS.md "Solver composition" rule 1). For
/// gradient/Jacobian inners it should sum `state.cost_evals() +
/// state.gradient_evals()` — CMA-ES outer state has no separate
/// derivative-eval counter, so derivative work collapses into
/// `cost_evals` honestly.
pub trait MemeticInner<V> {
    /// State shape this inner iterates against.
    type State: State<Param = V>;

    /// Build a fresh inner state seeded at CMA-ES candidate `x` with
    /// the current step-size `sigma`. Called once per refined
    /// candidate per outer generation.
    fn seed(&self, x: &V, sigma: f64) -> Self::State;

    /// Total inner work units to roll into the outer's `cost_evals`.
    /// Typically `state.cost_evals() + state.gradient_evals()` for
    /// derivative-based inners, `state.cost_evals()` alone for
    /// derivative-free inners. Takes `&self` so closure-based inners
    /// ([`ClosureInner`]) can dispatch through captured state.
    fn work_units(&self, state: &Self::State) -> u64;
}

/// Closure type for `ClosureInner`'s state seeder.
type ClosureSeedFn<V, S> = Box<dyn Fn(&V, f64) -> S>;
/// Closure type for `ClosureInner`'s work-unit aggregator.
type ClosureWorkFn<S> = Box<dyn Fn(&S) -> u64>;

/// Closure-based [`MemeticInner`] wrapper for custom inners that don't
/// have a native impl. Holds an inner solver plus the two closures
/// `MemeticInner` would otherwise express directly.
///
/// Intended use is one-off experiments and contract tests (e.g. the
/// `AlwaysFails` harness verifying `SolverFailed` bubbling). For
/// shipping configurations, prefer impl-ing `MemeticInner` on your
/// solver type — it's a five-line trait.
pub struct ClosureInner<I, S, V> {
    inner: I,
    seed_fn: ClosureSeedFn<V, S>,
    work_fn: ClosureWorkFn<S>,
}

impl<I, S, V> ClosureInner<I, S, V> {
    /// Wrap `inner` with explicit seeder and work-unit closures.
    pub fn new(
        inner: I,
        seed_fn: impl Fn(&V, f64) -> S + 'static,
        work_fn: impl Fn(&S) -> u64 + 'static,
    ) -> Self {
        Self {
            inner,
            seed_fn: Box::new(seed_fn),
            work_fn: Box::new(work_fn),
        }
    }
}

impl<P, I, S, V> Solver<P, S> for ClosureInner<I, S, V>
where
    I: Solver<P, S>,
    S: State,
{
    fn init(&mut self, problem: &P, state: S) -> S {
        self.inner.init(problem, state)
    }
    fn next_iter(&mut self, problem: &P, state: S) -> (S, Option<TerminationReason>) {
        self.inner.next_iter(problem, state)
    }
    fn terminate(&self, state: &S) -> Option<TerminationReason> {
        self.inner.terminate(state)
    }
}

impl<I, S, V> MemeticInner<V> for ClosureInner<I, S, V>
where
    S: State<Param = V>,
{
    type State = S;
    fn seed(&self, x: &V, sigma: f64) -> S {
        (self.seed_fn)(x, sigma)
    }
    fn work_units(&self, state: &S) -> u64 {
        (self.work_fn)(state)
    }
}

// -----------------------------------------------------------------------
// MemeticInner impls for the three shipped inners.
// -----------------------------------------------------------------------

impl<V> MemeticInner<V> for NelderMead
where
    V: VectorLen + Clone + std::ops::IndexMut<usize, Output = f64>,
{
    type State = BasicSimplexState<V>;
    fn seed(&self, x: &V, sigma: f64) -> BasicSimplexState<V> {
        // σ-scaled axis-aligned simplex: edge = current CMA step-size,
        // so the inner's exploration tracks the outer distribution's
        // spread and shrinks with σ. Hansen 2011 doesn't prescribe a
        // specific simplex; this matches the S11 default that the
        // existing tests validate against.
        let n = x.vec_len();
        let mut vertices = Vec::with_capacity(n + 1);
        vertices.push(x.clone());
        for j in 0..n {
            let mut v = x.clone();
            v[j] += sigma;
            vertices.push(v);
        }
        BasicSimplexState::from_simplex(vertices)
    }
    fn work_units(&self, state: &BasicSimplexState<V>) -> u64 {
        state.cost_evals()
    }
}

impl<V> MemeticInner<V> for LevenbergMarquardt
where
    V: Clone,
{
    type State = BasicState<V>;
    fn seed(&self, x: &V, _sigma: f64) -> BasicState<V> {
        BasicState::new(x.clone())
    }
    fn work_units(&self, state: &BasicState<V>) -> u64 {
        state.cost_evals() + state.gradient_evals()
    }
}

impl<V, LS> MemeticInner<V> for LBFGSB<LS>
where
    V: Clone,
{
    type State = LbfgsState<V>;
    fn seed(&self, x: &V, _sigma: f64) -> LbfgsState<V> {
        LbfgsState::new(x.clone(), self.m_capacity)
    }
    fn work_units(&self, state: &LbfgsState<V>) -> u64 {
        state.cost_evals() + state.gradient_evals()
    }
}

// -----------------------------------------------------------------------
// CmaInject — memetic CMA-ES with Hansen-2011 injection.
// -----------------------------------------------------------------------

/// Memetic CMA-ES with Hansen (2011) injection: outer CMA-ES proposes
/// `λ` candidates per generation, an inner local solver
/// ([`MemeticInner`]) refines the best `k`, and the refined points are
/// Mahalanobis-clipped and injected back into the population for the
/// next CMA update.
///
/// The only departure from the standard
/// [`CmaEs`](crate::solver::cma_es::CmaEs) update is clipping each
/// injected point's normalised step in Mahalanobis distance:
///
/// ```text
///   y_i ← min(1, c_y / ‖C^{-1/2} y_i‖) · y_i        (Hansen 2011 eq. 4)
///   c_y = √n + 2n/(n+2)                              (Table 1 default)
/// ```
///
/// with `y_i = (x_i − m)/σ` and `C^{-1/2} = B D^{-1} Bᵀ` from the
/// post-update eigendecomposition CMA-ES already maintains. After
/// clipping, replaced candidates re-enter the population on equal
/// footing with regular samples — all subsequent CMA updates
/// (m, p_σ, p_c, C, σ) run the standard equations unchanged. Lamarckian
/// by construction; no Baldwinian mode in the paper.
///
/// # Inner solver
///
/// Generic over any `I: MemeticInner<V>`. The associated `I::State`
/// determines the inner state shape. Shipped impls cover
/// [`NelderMead`], [`LevenbergMarquardt`], and [`LBFGSB`]. For
/// L-BFGS-B inner with consistent bound flow, use the bounded sibling
/// [`BoundedCmaInject`](crate::solver::BoundedCmaInject) over
/// [`BoundedCmaEs`](crate::solver::BoundedCmaEs).
///
/// # Eval aggregation
///
/// The outer's `state.cost_evals` aggregates total inner work units
/// per `I::work_units(state)`. For derivative-based inners (LM,
/// L-BFGS-B) the impl sums `cost_evals + gradient_evals`; CMA-ES
/// outer state has no `gradient_evals` field
/// (`BasicPopulationState` extends `State`, not `GradientState`), so
/// derivative-eval counts collapse into `cost_evals` honestly per
/// AGENTS.md "Solver composition" rule 1.
///
/// # Backends
///
/// Same coverage as [`CmaEs`]: nalgebra (`DVector` / `DMatrix`) and
/// faer (`Col` / `Mat`). `Vec<f64>` and `ndarray` produce a
/// compile-time error per tenet 5.
pub struct CmaInject<I, V, M>
where
    I: MemeticInner<V>,
{
    cma: CmaEs<V, M>,
    inner: InnerExecutor<I::State, I>,
    k: usize,
    c_y_override: Option<f64>,
}

impl<I, V, M> CmaInject<I, V, M>
where
    I: MemeticInner<V>,
{
    /// Wrap a configured [`CmaEs`] with `inner` as the local
    /// refinement step. Defaults: `k = 1` refinement per generation,
    /// inner `max_iter = 50`, `c_y` = Hansen-2011 Table 1 default.
    pub fn with_inner_solver(cma: CmaEs<V, M>, inner: I) -> Self {
        Self {
            cma,
            inner: InnerExecutor::new(inner).max_iter(50),
            k: 1,
            c_y_override: None,
        }
    }

    /// Number of best-ranked candidates to refine and inject each
    /// generation. Default `1`.
    ///
    /// # Panics
    ///
    /// Panics if `k == 0`. `k > λ` is silently clamped at runtime.
    pub fn with_k(mut self, k: usize) -> Self {
        assert!(k >= 1, "CmaInject requires k >= 1, got {}", k);
        self.k = k;
        self
    }

    /// Override the Hansen-2011 clipping threshold `c_y` (default
    /// `√n + 2n/(n+2)`).
    ///
    /// # Panics
    ///
    /// Panics if `c_y <= 0`.
    pub fn with_c_y(mut self, c_y: f64) -> Self {
        assert!(c_y > 0.0, "CmaInject requires c_y > 0, got {}", c_y);
        self.c_y_override = Some(c_y);
        self
    }

    /// Inner solver iteration budget per outer generation (default `50`).
    pub fn with_inner_max_iter(self, n: u64) -> Self {
        let Self {
            cma,
            inner,
            k,
            c_y_override,
        } = self;
        Self {
            cma,
            inner: inner.max_iter(n),
            k,
            c_y_override,
        }
    }

    /// Register a stateless termination criterion on the inner loop.
    /// Criteria are reused across every outer iteration's inner run,
    /// so they MUST be stateless across calls — `MaxIter`, the
    /// `*Tolerance` family, and `MaxCostEvals` are safe;
    /// [`MaxTime`](crate::core::termination::MaxTime) is **not**.
    /// See AGENTS.md "Solver composition" rule 2.
    pub fn inner_terminate_on<C>(self, criterion: C) -> Self
    where
        C: TerminationCriterion<I::State> + 'static,
    {
        let Self {
            cma,
            inner,
            k,
            c_y_override,
        } = self;
        Self {
            cma,
            inner: inner.terminate_on(criterion),
            k,
            c_y_override,
        }
    }
}

/// Hansen 2011 Table 1: `c_y = √n + 2n/(n+2)`, chosen so <10% of
/// regular `y_i` would be clipped at typical `n` and <1% for `n > 10`.
///
/// `pub(crate)` so the sibling
/// [`BoundedCmaInject`](crate::solver::BoundedCmaInject) can share
/// this default without re-deriving it.
pub(crate) fn default_c_y(n: usize) -> f64 {
    let n = n as f64;
    n.sqrt() + 2.0 * n / (n + 2.0)
}

impl<P, I, V, M> Solver<P, BasicPopulationState<V>> for CmaInject<I, V, M>
where
    P: CostFunction<Param = V, Output = f64>,
    I: MemeticInner<V> + Solver<P, <I as MemeticInner<V>>::State>,
    I::State: State<Param = V, Float = f64>,
    V: VectorLen
        + Clone
        + ScaledAdd<f64>
        + ScaleInPlace
        + ComponentMulAssign
        + NormSquared
        + SampleStandardNormal
        + std::ops::Index<usize, Output = f64>
        + std::ops::IndexMut<usize, Output = f64>,
    M: MatrixIdentity
        + MatVec<V>
        + MatTransposeVec<V>
        + ScaleInPlace
        + RankOneUpdate<V>
        + SymmetricEigen<V>
        + Clone,
{
    fn init(&mut self, problem: &P, state: BasicPopulationState<V>) -> BasicPopulationState<V> {
        // Hansen's preliminary experiments inject from iter 1 onward,
        // so we delegate the initial population to vanilla CMA-ES.
        self.cma.init(problem, state)
    }

    fn next_iter(
        &mut self,
        problem: &P,
        state: BasicPopulationState<V>,
    ) -> (BasicPopulationState<V>, Option<TerminationReason>) {
        // 1. Vanilla CMA-ES iteration: update m, σ, C from the
        //    previous generation, sample λ fresh candidates sorted by
        //    cost ascending.
        let (mut state, reason) = self.cma.next_iter(problem, state);
        if let Some(r) = reason {
            return (state, Some(r));
        }

        // Snapshot internals for clipping.
        let (n, m, sigma) = {
            let w = self
                .cma
                .working()
                .expect("CmaEs::init must run before CmaInject::next_iter");
            (w.n, w.m.clone(), w.sigma)
        };
        let c_y = self.c_y_override.unwrap_or_else(|| default_c_y(n));
        let refine = self.k.min(state.candidates.len());

        for i in 0..refine {
            // 2. Seed the inner state via the trait. The σ argument
            //    lets seeders that scale with the CMA distribution
            //    (NM's σ-scaled simplex) track the current spread.
            let inner_state = self.inner.solver().seed(&state.candidates[i], sigma);

            // 3. Drive the inner.
            let inner_result: OptimizationResult<I::State> = self.inner.run(problem, inner_state);

            // 4. Eval aggregation: roll inner total-work into outer
            //    cost_evals via the trait (AGENTS.md "Solver
            //    composition" rule 1).
            state.cost_evals += self.inner.solver().work_units(&inner_result.state);

            // 5. Failure routing: bubble SolverFailed only (rule 3).
            if inner_result.reason.is_failure() {
                return (state, Some(inner_result.reason));
            }

            // 6. Extract refined point.
            let x_refined = inner_result.state.param().clone();

            // 7. y = (x_refined − m) / σ.
            let mut y = x_refined;
            y.scaled_add(-1.0, &m);
            y.scale_in_place(1.0 / sigma);

            // 8. ‖C^{-1/2} y‖ = ‖D^{-1} ⊙ Bᵀ y‖.
            let inv_sqrt_norm = {
                let w = self.cma.working().expect("working still populated");
                let mut bt_y = w.b.mat_transpose_vec(&y);
                bt_y.component_mul_assign(&w.d_inv);
                bt_y.norm_squared().sqrt()
            };

            // 9. Clipping factor α (Hansen 2011 eq. 4 + eq. 10).
            if inv_sqrt_norm > 0.0 {
                let alpha = (c_y / inv_sqrt_norm).min(1.0);
                if alpha < 1.0 {
                    y.scale_in_place(alpha);
                }
            }

            // 10. x_inj = m + σ · y_clipped.
            let mut x_inj = m.clone();
            x_inj.scaled_add(sigma, &y);

            // 11. Re-evaluate: clipping moves the point in original
            //     space, so the cost field has to match.
            let cost_new = problem.cost(&x_inj);
            state.cost_evals += 1;

            state.candidates[i] = x_inj;
            state.costs[i] = cost_new;
        }

        // 12. Re-sort: rank-µ update depends on the order.
        if refine > 0 {
            sort_population_ascending(&mut state.candidates, &mut state.costs);
        }

        (state, None)
    }

    fn terminate(&self, state: &BasicPopulationState<V>) -> Option<TerminationReason> {
        <CmaEs<V, M> as Solver<P, BasicPopulationState<V>>>::terminate(&self.cma, state)
    }
}