basin 0.2.0

use crate::core::constraint::BoxConstrained;
use crate::core::math::{ClampInPlace, NegInPlace, ScaledAdd};
use crate::core::problem::{CostFunction, Gradient};
use crate::core::solver::Solver;
use crate::core::state::BasicState;
use crate::core::termination::TerminationReason;
use crate::line_search::{Constant, LineSearch};

/// Projected gradient descent for box-constrained problems.
///
/// Steepest-descent step along `−∇f` followed by an element-wise
/// projection back into `[lower, upper]`. The first n-D constrained
/// solver in basin and the smallest vehicle for the
/// [`BoxConstrained`] trait — handing this solver an unconstrained
/// problem is a compile error per tenet 4.
///
/// # Algorithm
///
/// At [`init`](Solver::init) the iterate is projected onto the feasible
/// box once, so an infeasible starting point is silently corrected
/// (and downstream termination criteria see a feasible iterate at iter
/// 0). Each [`next_iter`](Solver::next_iter) then computes
///
/// ```text
/// d   ← −∇f(x)
/// α   ← line_search.next(...)        # on the unconstrained step
/// x   ← π_C(x + α d)                  # project after the step
/// ```
///
/// where `π_C` clamps each component into `[lower, upper]`.
///
/// # Contract
///
/// - **Caller must:** implement [`BoxConstrained`] on the problem with
///   `lower[i] ≤ upper[i]` for every component. Equal bounds are
///   allowed (the corresponding component is pinned).
/// - **Caller must:** pair with a feasible **or** infeasible initial
///   param; an infeasible start is projected at `init`.
/// - **Implementor (this solver) must:** maintain feasibility across
///   iterations — once the loop has run, every iterate the executor
///   sees is in the box.
///
/// The line search runs against the *unconstrained* trial step
/// `f(x + α d)`. If the projection moves the post-step iterate
/// substantially, Armijo guarantees on the unconstrained step do not
/// transfer to `f(π_C(x + α d))`. For tighter guarantees use a small
/// fixed step ([`Constant`]) or wait on a constraint-aware (SPG-style)
/// line search.
///
/// # Termination
///
/// No solver-internal optimality test; the canonical first-order
/// metric is provided as a framework-level criterion,
/// [`ProjectedGradientTolerance`](crate::core::termination::ProjectedGradientTolerance),
/// which captures the bounds at construction so it does not need
/// problem access in `check`. The framework's
/// [`MaxIter`](crate::core::termination::MaxIter),
/// [`CostTolerance`](crate::core::termination::CostTolerance),
/// [`ParamTolerance`](crate::core::termination::ParamTolerance), and
/// [`MaxTime`](crate::core::termination::MaxTime) work on
/// [`BasicState`] for free.
///
/// # Backends
///
/// Backend-generic — works with any `V` implementing
/// [`ScaledAdd<f64>`](crate::core::math::ScaledAdd) +
/// [`NegInPlace`] + [`ClampInPlace`] + `Clone`. That covers
/// `Vec<f64>`, `nalgebra::DVector<f64>` (feature `nalgebra`),
/// `ndarray::Array1<f64>` (feature `ndarray`), and `faer::Col<f64>`
/// (feature `faer`). The problem must implement [`BoxConstrained`].
pub struct ProjectedGradientDescent<S> {
    line_search: S,
}

impl ProjectedGradientDescent<Constant> {
    /// Projected gradient descent with a fixed step size `alpha`.
    /// Equivalent to `with_line_search(Constant(alpha))`. Recommended
    /// default — the line search variant has the caveat documented on
    /// the type.
    pub fn new(alpha: f64) -> Self {
        Self {
            line_search: Constant(alpha),
        }
    }
}

impl<S> ProjectedGradientDescent<S> {
    /// Projected gradient descent with an explicit line-search strategy.
    ///
    /// Note: the line search runs against the *unconstrained* trial
    /// step (see the type-level rustdoc). Backtracking is honest only
    /// while the projection isn't active; consider a fixed step in
    /// regimes where many components hit their bounds.
    pub fn with_line_search(line_search: S) -> Self {
        Self { line_search }
    }
}

impl<P, V, S> Solver<P, BasicState<V>> for ProjectedGradientDescent<S>
where
    P: CostFunction<Param = V, Output = f64> + Gradient<Param = V, Gradient = V> + BoxConstrained,
    V: ScaledAdd<f64> + NegInPlace + ClampInPlace + Clone,
    S: LineSearch<P, V>,
{
    fn init(&mut self, problem: &P, mut state: BasicState<V>) -> BasicState<V> {
        // Project an infeasible start once so iter-0 termination checks
        // see a feasible iterate. Subsequent iterations preserve
        // feasibility by construction.
        state.param.clamp_in_place(problem.lower(), problem.upper());
        state.cost = Some(problem.cost(&state.param));
        state.gradient = Some(problem.gradient(&state.param));
        state.cost_evals += 1;
        state.gradient_evals += 1;
        state
    }

    fn next_iter(
        &mut self,
        problem: &P,
        mut state: BasicState<V>,
    ) -> (BasicState<V>, Option<TerminationReason>) {
        let grad = state
            .gradient
            .take()
            .expect("gradient not set: Solver::init must run before next_iter");
        let prev_cost = state
            .cost
            .expect("cost not set: Solver::init must run before next_iter");
        let mut direction = grad.clone();
        direction.neg_in_place();
        let step = self
            .line_search
            .next(problem, &state.param, prev_cost, &grad, &direction);
        state.cost_evals += step.cost_evals;
        state.gradient_evals += step.gradient_evals;
        state.param.scaled_add(step.alpha, &direction);
        state.param.clamp_in_place(problem.lower(), problem.upper());
        state.cost = Some(problem.cost(&state.param));
        state.gradient = Some(problem.gradient(&state.param));
        state.cost_evals += 1;
        state.gradient_evals += 1;
        (state, None)
    }
}