herculesabqp 0.1.2

use std::sync::Arc;

use anyhow::{Result, bail};

use crate::matrix::QuadraticMatrix;

use super::operator::QuadraticOperator;
use super::types::Diagnostics;
use super::utils::{
    DenseVec, add_assign, clip_in_place, compute_diagnostics_from_qx, objective_from_qx,
};

pub(crate) trait FirstOrderProblem {
    fn n(&self) -> usize;
    fn c(&self) -> &DenseVec;
    fn clip_in_place(&self, x: &mut DenseVec);
    fn matvec_into(&self, x: &DenseVec, out: &mut DenseVec);
    fn objective_from_qx(&self, x: &DenseVec, qx: &DenseVec) -> f64;
    fn diagnostics_from_qx(
        &self,
        qx: &DenseVec,
        x: &DenseVec,
        bound_tol: f64,
        dual_certification: bool,
    ) -> Diagnostics;
}

/// Build default `[0, 1]` box bounds for `n` variables.
pub fn default_bounds(n: usize) -> (Vec<f64>, Vec<f64>) {
    (vec![0.0; n], vec![1.0; n])
}

/// Canonical in-memory representation of a box-constrained quadratic program.
///
/// This type gathers the data for problems of the form
/// `min 0.5 x^T Q x + c^T x` subject to `lb <= x <= ub`.
///
/// A few invariants are enforced at construction time:
/// - `Q` is square
/// - `c`, `lb`, and `ub` all have matching dimension
/// - every bound pair satisfies `lb[i] <= ub[i]`
/// - `Q` is symmetrized unless the caller has already opted into trusting
///   symmetry upstream
///
/// The solver treats this as the central "problem interface" internally:
/// matrix-vector products, projections, objective evaluation, diagnostics, and
/// reduced active-set extraction are all routed through this struct.
#[derive(Clone, Debug)]
pub(crate) struct BoxQPProblem {
    // Symmetrized quadratic matrix for the objective.
    pub q: Arc<QuadraticMatrix>,
    // Linear term in the objective.
    pub c: Arc<DenseVec>,
    // Lower box bounds.
    pub lb: DenseVec,
    // Upper box bounds.
    pub ub: DenseVec,
}

impl BoxQPProblem {
    /// Validate and store a box-constrained quadratic program in solver form.
    ///
    /// The returned instance owns dense copies of the linear term and bounds,
    /// and stores the symmetric part of `Q`.
    ///
    /// This is the right constructor to use whenever a caller has raw matrix
    /// and vector data and wants the solver-ready representation with
    /// dimensions and bounds checked up front.
    #[cfg(test)]
    pub fn from_parts(q: QuadraticMatrix, c: &[f64], lb: &[f64], ub: &[f64]) -> Result<Self> {
        Self::from_parts_with_symmetry(q, c, lb, ub, false)
    }

    /// Validate and store a box-constrained quadratic program with an explicit
    /// symmetry contract for `Q`.
    ///
    /// When `assume_symmetric` is `true`, the constructor trusts the caller
    /// and stores `Q` as-is. When `false`, it defensively replaces `Q` by
    /// `0.5 * (Q + Q^T)`.
    #[cfg(test)]
    pub fn from_parts_with_symmetry(
        q: QuadraticMatrix,
        c: &[f64],
        lb: &[f64],
        ub: &[f64],
        assume_symmetric: bool,
    ) -> Result<Self> {
        let (rows, cols) = q.shape();
        if rows != cols {
            bail!("Q must be square, got shape ({rows}, {cols})");
        }
        if c.len() != rows || lb.len() != rows || ub.len() != rows {
            bail!(
                "dimension mismatch: Q is {rows}x{cols}, c has len {}, lb has len {}, ub has len {}",
                c.len(),
                lb.len(),
                ub.len()
            );
        }
        for i in 0..rows {
            if lb[i] > ub[i] {
                bail!("invalid bounds at index {i}: lb={} > ub={}", lb[i], ub[i]);
            }
        }
        let c = DenseVec::from_vec(c.to_vec());
        let lb = DenseVec::from_vec(lb.to_vec());
        let ub = DenseVec::from_vec(ub.to_vec());
        let q = if assume_symmetric { q } else { q.symmetrized() };
        Self::from_shared_parts(Arc::new(q), Arc::new(c), lb, ub)
    }

    /// Store a validated problem from shared matrix / linear-term state and
    /// owned bounds.
    ///
    /// This is primarily used by the prepared-solver path so repeated solves
    /// can reuse expensive matrix preprocessing without rebuilding the shared
    /// `Q` and `c` data each time.
    pub fn from_shared_parts(
        q: Arc<QuadraticMatrix>,
        c: Arc<DenseVec>,
        lb: DenseVec,
        ub: DenseVec,
    ) -> Result<Self> {
        let (rows, cols) = q.shape();
        if rows != cols {
            bail!("Q must be square, got shape ({rows}, {cols})");
        }
        if c.len() != rows || lb.len() != rows || ub.len() != rows {
            bail!(
                "dimension mismatch: Q is {rows}x{cols}, c has len {}, lb has len {}, ub has len {}",
                c.len(),
                lb.len(),
                ub.len()
            );
        }
        for i in 0..rows {
            if lb[i] > ub[i] {
                bail!("invalid bounds at index {i}: lb={} > ub={}", lb[i], ub[i]);
            }
        }
        Ok(Self { q, c, lb, ub })
    }

    /// Apply diagonal variable scaling and return the transformed problem.
    ///
    /// If the original variables are `x = D z`, this builds the equivalent
    /// problem in `z` coordinates:
    /// - `Q_scaled = D Q D`
    /// - `c_scaled = D c`
    /// - `lb_scaled = D^{-1} lb`
    /// - `ub_scaled = D^{-1} ub`
    ///
    /// The solve path uses this to improve conditioning on badly scaled
    /// problems while still reporting the final answer in the original `x`
    /// coordinates.
    #[cfg(test)]
    pub fn scaled_from(&self, scale: &DenseVec) -> Result<Self> {
        let scale_slice = scale
            .as_slice()
            .expect("dense scaling vector should be contiguous");
        let mut c_scaled = DenseVec::zeros(scale.len());
        let mut lb_scaled = DenseVec::zeros(scale.len());
        let mut ub_scaled = DenseVec::zeros(scale.len());
        for i in 0..scale.len() {
            c_scaled[i] = self.c[i] * scale[i];
            lb_scaled[i] = self.lb[i] / scale[i];
            ub_scaled[i] = self.ub[i] / scale[i];
        }
        Self::from_parts_with_symmetry(
            self.q.as_ref().scaled(scale_slice),
            c_scaled
                .as_slice()
                .expect("dense vector should be contiguous"),
            lb_scaled
                .as_slice()
                .expect("dense vector should be contiguous"),
            ub_scaled
                .as_slice()
                .expect("dense vector should be contiguous"),
            true,
        )
    }

    /// Project a dense vector onto the stored box bounds.
    ///
    /// This is the projection operator used by the first-order method and by
    /// the polish line search. It mutates `x` in place to avoid extra
    /// allocations in the hot path.
    pub fn clip_in_place(&self, x: &mut DenseVec) {
        clip_in_place(x, &self.lb, &self.ub);
    }

    /// Multiply the quadratic matrix by a dense vector into a caller-owned buffer.
    ///
    /// This is the preferred form inside iterative algorithms because it avoids
    /// allocating a fresh dense vector for every matrix-vector product.
    pub fn matvec_into(&self, x: &DenseVec, out: &mut DenseVec) {
        self.q.as_ref().matvec_into(x, out);
    }

    /// Compute the gradient `Qx + c` into a caller-owned buffer.
    ///
    /// Like [`Self::matvec_into`], this exists mainly so the accelerated
    /// projected-gradient loop and
    /// polish step can reuse memory rather than allocate a new gradient vector
    /// every time they need one.
    pub fn gradient_into(&self, x: &DenseVec, out: &mut DenseVec) {
        self.matvec_into(x, out);
        add_assign(out, &self.c);
    }

    /// Evaluate the objective given both `x` and a precomputed `Qx`.
    ///
    /// This is useful whenever the caller already has `Qx` from a recent
    /// matrix-vector product and wants to avoid repeating that work just to get
    /// the scalar objective value.
    pub fn objective_from_qx(&self, x: &DenseVec, qx: &DenseVec) -> f64 {
        objective_from_qx(&self.c, x, qx)
    }

    /// Form the solver's stopping diagnostics from cached `Qx`.
    ///
    /// The returned internal diagnostics bundle contains the relative gap and
    /// KKT-style box residuals used by the solve loop and the polish gate.
    ///
    /// This path is written around cached `Qx` so the solver can evaluate
    /// objective quality without paying for another matrix-vector product.
    pub fn diagnostics_from_qx(
        &self,
        qx: &DenseVec,
        x: &DenseVec,
        bound_tol: f64,
        dual_certification: bool,
    ) -> Diagnostics {
        compute_diagnostics_from_qx(
            qx,
            &self.c,
            x,
            &self.lb,
            &self.ub,
            bound_tol,
            dual_certification,
        )
    }

    /// Extract a reduced matrix view for active-set polishing.
    ///
    /// This is used by the polish step to build free-variable reduced systems
    /// such as `Q_ff` and `Q_fa` without exposing matrix-format-specific logic
    /// to the rest of the solver.
    pub fn principal(&self, rows: &[usize], cols: &[usize]) -> QuadraticMatrix {
        self.q.as_ref().principal(rows, cols)
    }
}

impl FirstOrderProblem for BoxQPProblem {
    fn n(&self) -> usize {
        self.c.len()
    }

    fn c(&self) -> &DenseVec {
        &self.c
    }

    fn clip_in_place(&self, x: &mut DenseVec) {
        self.clip_in_place(x);
    }

    fn matvec_into(&self, x: &DenseVec, out: &mut DenseVec) {
        self.matvec_into(x, out);
    }

    fn objective_from_qx(&self, x: &DenseVec, qx: &DenseVec) -> f64 {
        self.objective_from_qx(x, qx)
    }

    fn diagnostics_from_qx(
        &self,
        qx: &DenseVec,
        x: &DenseVec,
        bound_tol: f64,
        dual_certification: bool,
    ) -> Diagnostics {
        self.diagnostics_from_qx(qx, x, bound_tol, dual_certification)
    }
}

/// First-order-only box-QP problem built from an implicit quadratic operator.
///
/// This form supports projected-gradient solves without requiring explicit
/// matrix extraction for symmetrization or polishing.
pub(crate) struct ImplicitBoxQPProblem {
    pub q: Arc<dyn QuadraticOperator>,
    pub c: Arc<DenseVec>,
    pub lb: DenseVec,
    pub ub: DenseVec,
}

impl ImplicitBoxQPProblem {
    pub fn from_shared_parts(
        q: Arc<dyn QuadraticOperator>,
        c: Arc<DenseVec>,
        lb: DenseVec,
        ub: DenseVec,
    ) -> Result<Self> {
        let n = q.n();
        if c.len() != n || lb.len() != n || ub.len() != n {
            bail!(
                "dimension mismatch: implicit Q has size {n}, c has len {}, lb has len {}, ub has len {}",
                c.len(),
                lb.len(),
                ub.len()
            );
        }
        for i in 0..n {
            if lb[i] > ub[i] {
                bail!("invalid bounds at index {i}: lb={} > ub={}", lb[i], ub[i]);
            }
        }
        Ok(Self { q, c, lb, ub })
    }
}

impl FirstOrderProblem for ImplicitBoxQPProblem {
    fn n(&self) -> usize {
        self.c.len()
    }

    fn c(&self) -> &DenseVec {
        &self.c
    }

    fn clip_in_place(&self, x: &mut DenseVec) {
        clip_in_place(x, &self.lb, &self.ub);
    }

    fn matvec_into(&self, x: &DenseVec, out: &mut DenseVec) {
        self.q.matvec_into(x, out);
    }

    fn objective_from_qx(&self, x: &DenseVec, qx: &DenseVec) -> f64 {
        objective_from_qx(&self.c, x, qx)
    }

    fn diagnostics_from_qx(
        &self,
        qx: &DenseVec,
        x: &DenseVec,
        bound_tol: f64,
        dual_certification: bool,
    ) -> Diagnostics {
        compute_diagnostics_from_qx(
            qx,
            &self.c,
            x,
            &self.lb,
            &self.ub,
            bound_tol,
            dual_certification,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ndarray::{Array1, Array2};

    #[test]
    fn scaling_transforms_bounds_and_linear_term_consistently() {
        let mut q = Array2::<f64>::zeros((2, 2));
        q[[0, 0]] = 4.0;
        q[[1, 1]] = 9.0;
        let q = QuadraticMatrix::dense(q);
        let c = vec![2.0, -3.0];
        let lb = vec![-2.0, 0.5];
        let ub = vec![6.0, 2.5];
        let scale = Array1::from_vec(vec![0.5, 2.0]);

        let original = BoxQPProblem::from_parts(q, &c, &lb, &ub).unwrap();
        let scaled = original.scaled_from(&scale).unwrap();

        assert_eq!(scaled.c.as_slice(), Some(&[1.0, -6.0][..]));
        assert_eq!(scaled.lb.as_slice(), Some(&[-4.0, 0.25][..]));
        assert_eq!(scaled.ub.as_slice(), Some(&[12.0, 1.25][..]));
        assert_eq!(scaled.q.diagonal().as_slice(), Some(&[1.0, 36.0][..]));
    }

    #[test]
    fn rejects_invalid_bounds() {
        let mut q = Array2::<f64>::zeros((2, 2));
        q[[0, 0]] = 1.0;
        q[[1, 1]] = 1.0;
        let q = QuadraticMatrix::dense(q);
        let err = BoxQPProblem::from_parts(q, &[0.0, 0.0], &[1.0, 0.0], &[0.0, 1.0]).unwrap_err();
        assert!(err.to_string().contains("invalid bounds"));
    }
}