#[cfg(feature = "nalgebra")]
use crate::core::inner::WarmStart;
use crate::core::math::{
    Dot, GeneralRankOneUpdate, MatVec, MatrixIdentity, NegInPlace, NormSquared, ScaleInPlace,
    ScaledAdd, VectorLen,
};
use crate::core::problem::{CostFunction, Gradient};
use crate::core::solver::Solver;
use crate::core::state::QuasiNewtonState;
use crate::core::termination::TerminationReason;
use crate::line_search::{LineSearch, Wolfe};

/// BFGS quasi-Newton solver.
///
/// Maintains a dense inverse-Hessian approximation `H` updated by the
/// rank-2 BFGS formula. The search direction is `d = −H·∇f(x)`; the step
/// length is set by a configurable line search (default: strong Wolfe,
/// which is what guarantees `yᵀs > 0` so each update preserves positive
/// definiteness).
///
/// On the first accepted step we rescale `H ← (sᵀy / yᵀy)·I` (Nocedal &
/// Wright (6.20)) — cheap, large convergence improvement on poorly scaled
/// problems.
///
/// **Curvature failure (`yᵀs ≤ ε · |y| · |s|`):** the H update is skipped
/// for that iteration. Strong Wolfe with `c2 < 1` guarantees `yᵀs > 0` in
/// exact arithmetic, so this branch is a numerical safeguard, not the
/// primary path. (Damped BFGS / Powell's modification is overkill when
/// strong Wolfe is in place — see plan.)
///
/// # Backends
///
/// Runs on `Vec<f64>` (via the hand-rolled
/// [`DenseMatrix`](crate::core::math::DenseMatrix)), nalgebra
/// (`DVector<f64>` / `DMatrix<f64>`), and faer (`Col<f64>` / `Mat<f64>`).
/// The dense inverse-Hessian needs only matvec, an identity constructor,
/// scaling, and the rank-one update [`GeneralRankOneUpdate`] — no
/// factorization — so it stays backend-generic. `ndarray` is a
/// compile-time error per tenet 5: its `Array2<f64>` implements neither
/// [`GeneralRankOneUpdate`] nor [`MatrixIdentity`].
///
/// # Examples
///
/// BFGS on the 2-D Rosenbrock function over the dependency-free
/// `Vec<f64>` backend. Quasi-Newton solvers iterate a
/// [`QuasiNewtonState`], parameterised by the
/// param vector and the dense matrix type — here `Vec<f64>` and
/// [`DenseMatrix`](crate::DenseMatrix):
///
/// ```
/// use basin::{BasicState, CostFunction, DenseMatrix, Executor, Gradient, QuasiNewtonState, BFGS};
///
/// struct Rosenbrock;
/// impl CostFunction for Rosenbrock {
///     type Param = Vec<f64>;
///     type Output = f64;
///     fn cost(&self, x: &Vec<f64>) -> f64 {
///         (1.0 - x[0]).powi(2) + 100.0 * (x[1] - x[0].powi(2)).powi(2)
///     }
/// }
/// impl Gradient for Rosenbrock {
///     type Param = Vec<f64>;
///     type Gradient = Vec<f64>;
///     fn gradient(&self, x: &Vec<f64>) -> Vec<f64> {
///         vec![
///             -2.0 * (1.0 - x[0]) - 400.0 * x[0] * (x[1] - x[0].powi(2)),
///             200.0 * (x[1] - x[0].powi(2)),
///         ]
///     }
/// }
///
/// let result = Executor::new(
///     Rosenbrock,
///     BFGS::new(),
///     QuasiNewtonState::<Vec<f64>, DenseMatrix>::new(vec![-1.2, 1.0]),
/// )
/// .max_iter(100)
/// .run();
/// assert!(result.cost() < 1e-8);
/// ```
pub struct BFGS<S = Wolfe> {
    line_search: S,
    epsilon: f64,
}

impl Default for BFGS<Wolfe> {
    fn default() -> Self {
        Self::new()
    }
}

impl BFGS<Wolfe> {
    /// BFGS with the strong-Wolfe line search (Nocedal & Wright defaults)
    /// and `ε = 1e-10` for the curvature-condition guard.
    pub fn new() -> Self {
        Self {
            line_search: Wolfe::new(),
            epsilon: 1e-10,
        }
    }
}

impl<S> BFGS<S> {
    /// BFGS with an explicit line-search strategy.
    pub fn with_line_search(line_search: S) -> Self {
        Self {
            line_search,
            epsilon: 1e-10,
        }
    }

    /// Relative threshold for the curvature condition `yᵀs > ε · |y| · |s|`.
    /// Iterations where this fails skip the H update (rare with strong
    /// Wolfe). Default `1e-10`.
    pub fn epsilon(mut self, epsilon: f64) -> Self {
        assert!(epsilon >= 0.0, "epsilon must be ≥ 0");
        self.epsilon = epsilon;
        self
    }
}

impl<P, S, V, M> Solver<P, QuasiNewtonState<V, M>> for BFGS<S>
where
    P: CostFunction<Param = V, Output = f64> + Gradient<Param = V, Gradient = V>,
    S: LineSearch<P, V>,
    V: Clone + Dot + NormSquared + ScaledAdd<f64> + ScaleInPlace + NegInPlace + VectorLen,
    M: MatVec<V> + MatrixIdentity + ScaleInPlace + GeneralRankOneUpdate<V>,
{
    fn init(&mut self, problem: &P, mut state: QuasiNewtonState<V, M>) -> QuasiNewtonState<V, M> {
        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: QuasiNewtonState<V, M>,
    ) -> (QuasiNewtonState<V, M>, Option<TerminationReason>) {
        let g = state
            .gradient
            .take()
            .expect("gradient not set: Solver::init must run before next_iter");
        let cost_old = state
            .cost
            .expect("cost not set: Solver::init must run before next_iter");

        // Quasi-Newton direction: d = −H g. With H positive definite this
        // is automatically a descent direction (gᵀd = −gᵀHg < 0).
        let mut direction = state.inverse_hessian.matvec(&g);
        direction.neg_in_place();

        let step = self
            .line_search
            .next(problem, &state.param, cost_old, &g, &direction);
        state.cost_evals += step.cost_evals;
        state.gradient_evals += step.gradient_evals;

        // Line search bailed (α = 0): direction wasn't descent, or we're
        // at numerical convergence. Restore gradient/cost so the state
        // stays consistent and report it as a mid-iter termination so the
        // executor halts immediately. NaN routes here too
        // (`NaN > 0.0` is false).
        if !(step.alpha.is_finite() && step.alpha > 0.0) {
            state.gradient = Some(g);
            state.cost = Some(cost_old);
            return (state, Some(TerminationReason::SolverConverged));
        }

        // s = α d, x ← x + s.
        let mut s = direction;
        s.scale_in_place(step.alpha);
        state.param.scaled_add(1.0, &s);

        let g_new = problem.gradient(&state.param);
        state.gradient_evals += 1;

        // y = g_new − g.
        let mut y = g_new.clone();
        y.scaled_add(-1.0, &g);
        let sy = s.dot(&y);
        let s_norm = s.norm_squared().sqrt();
        let y_norm = y.norm_squared().sqrt();

        if sy > self.epsilon * s_norm * y_norm {
            // Initial-Hessian rescaling: align H₀ with the local curvature
            // before applying the first BFGS update. Without this, the
            // identity-initialized H produces a unit step that's far too
            // large or small on poorly scaled problems.
            if !state.initial_scaling_done {
                let yy = y.dot(&y);
                if yy > 0.0 {
                    let scale = sy / yy;
                    let n = state.param.vec_len();
                    let mut h0 = M::identity(n);
                    h0.scale_in_place(scale);
                    state.inverse_hessian = h0;
                }
                state.initial_scaling_done = true;
            }

            let rho = 1.0 / sy;
            let hy = state.inverse_hessian.matvec(&y);
            let yhy = y.dot(&hy);
            let coef = rho * (1.0 + rho * yhy);

            // H ← H + coef · s sᵀ − ρ · (s (Hy)ᵀ + (Hy) sᵀ).
            // Three rank-1 updates, all in place.
            state.inverse_hessian.general_rank_one_update(coef, &s, &s);
            state.inverse_hessian.general_rank_one_update(-rho, &s, &hy);
            state.inverse_hessian.general_rank_one_update(-rho, &hy, &s);
        }
        // else: curvature failure (very rare with strong Wolfe). Skip the
        // H update; the line search still produced a descent step, so we
        // continue. If this persists, max_iter / GradientTolerance halt.

        state.cost = Some(problem.cost(&state.param));
        state.cost_evals += 1;
        state.gradient = Some(g_new);
        (state, None)
    }
}

/// Lets [`BFGS`] serve as the inner of a composed solver
/// (e.g. [`BarrierMethod`](crate::solver::BarrierMethod) /
/// [`AugmentedLagrangianMethod`](crate::solver::AugmentedLagrangianMethod)),
/// seeding a fresh [`QuasiNewtonState`] (identity inverse-Hessian) at the
/// warm-start point.
///
/// nalgebra only for now — the composed (barrier / AL) solvers seed their
/// inner on the nalgebra backend. A `Vec<f64>` / faer `WarmStart` is a
/// follow-up if a `Vec`-backed barrier inner is wanted.
#[cfg(feature = "nalgebra")]
impl<S> WarmStart<nalgebra::DVector<f64>> for BFGS<S> {
    type State = QuasiNewtonState<nalgebra::DVector<f64>, nalgebra::DMatrix<f64>>;
    fn seed(&self, x: &nalgebra::DVector<f64>) -> Self::State {
        QuasiNewtonState::<nalgebra::DVector<f64>, nalgebra::DMatrix<f64>>::new(x.clone())
    }
}
