clarabel 0.11.1

use super::*;

#[cfg(feature = "sdp")]
use crate::algebra::triangular_number;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};

// ---------------------------------------------------
// We define some machinery here for enumerating the
// different cone types that can live in the composite cone
// ---------------------------------------------------

/// API type describing the type of a conic constraint.
///
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
#[derive(Debug, Clone, PartialEq)]
pub enum SupportedConeT<T> {
    /// The zero cone (used for equality constraints).
    ///
    /// The parameter indicates the cones dimension.
    ZeroConeT(usize),
    /// The nonnegative orthant.  
    ///
    /// The parameter indicates the cones dimension.
    NonnegativeConeT(usize),
    /// The second order cone / Lorenz cone / ice-cream cone.
    ///  
    /// The parameter indicates the cones dimension.
    SecondOrderConeT(usize),
    /// The exponential cone in R^3.
    ///
    /// This cone takes no parameters
    ExponentialConeT(),
    /// The power cone in R^3.
    ///
    /// The parameter indicates the power.
    PowerConeT(T),
    /// The generalized power cone.
    ///
    /// The first vector of parameters supplies the nonnegative powers "alpha" of
    /// the left-hand side of the constraint.  The second scalar parameter provides
    /// the dimension of the 2-norm bounded vector in the right-hand side of the
    /// constraint.   The "alpha" terms must sum to 1.
    GenPowerConeT(Vec<T>, usize),

    /// The positive semidefinite cone in triangular form.
    ///
    /// The parameter indicates the matrix dimension, i.e. size = n
    /// means that the variable is the upper triangle of an nxn matrix.
    #[cfg(feature = "sdp")]
    PSDTriangleConeT(usize),
}

impl<T> SupportedConeT<T> {
    // this reports the number of slack variables that will be generated by
    // this cone.  Equivalent to `numels` for the internal cone representation.
    // Required for user data validation prior to building a problem.

    pub(crate) fn nvars(&self) -> usize {
        match self {
            SupportedConeT::ZeroConeT(dim) => *dim,
            SupportedConeT::NonnegativeConeT(dim) => *dim,
            SupportedConeT::SecondOrderConeT(dim) => *dim,
            SupportedConeT::ExponentialConeT() => 3,
            SupportedConeT::PowerConeT(_) => 3,
            #[cfg(feature = "sdp")]
            SupportedConeT::PSDTriangleConeT(dim) => triangular_number(*dim),
            SupportedConeT::GenPowerConeT(α, dim2) => α.len() + *dim2,
        }
    }
}

impl<T> std::fmt::Display for SupportedConeT<T>
where
    T: FloatT,
{
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", &self.as_tag().as_str())
    }
}

// we will use the SupportedConeT as a user facing marker
// for the constraint types, and then map them through
// make_cone to get the internal cone representations.

pub fn make_cone<T: FloatT>(cone: &SupportedConeT<T>) -> SupportedCone<T> {
    match cone {
        SupportedConeT::NonnegativeConeT(dim) => NonnegativeCone::<T>::new(*dim).into(),
        SupportedConeT::ZeroConeT(dim) => ZeroCone::<T>::new(*dim).into(),
        SupportedConeT::SecondOrderConeT(dim) => SecondOrderCone::<T>::new(*dim).into(),
        SupportedConeT::ExponentialConeT() => ExponentialCone::<T>::new().into(),
        SupportedConeT::PowerConeT(α) => PowerCone::<T>::new(*α).into(),
        SupportedConeT::GenPowerConeT(α, dim2) => {
            GenPowerCone::<T>::new((*α).clone(), *dim2).into()
        }
        #[cfg(feature = "sdp")]
        SupportedConeT::PSDTriangleConeT(dim) => PSDTriangleCone::<T>::new(*dim).into(),
    }
}

impl<T> SupportedConeT<T>
where
    T: FloatT,
{
    pub(crate) fn new_collapsed(cones: &[SupportedConeT<T>]) -> Vec<SupportedConeT<T>> {
        let mut newcones = Vec::with_capacity(cones.len());
        let mut iter = cones.iter().peekable();

        // rollup a subsequence of nonnegative cones or SOC/PSD singleton
        fn collapse<T>(
            iter: &mut std::iter::Peekable<std::slice::Iter<'_, SupportedConeT<T>>>,
            newcones: &mut Vec<SupportedConeT<T>>,
            init_dim: usize,
        ) where
            T: FloatT,
        {
            let mut total_dim = init_dim;
            while let Some(next_cone) = iter.peek() {
                // skip empty cones
                if next_cone.nvars() != 0 {
                    match next_cone {
                        // collapsible cones.
                        SupportedConeT::NonnegativeConeT(dim) => total_dim += dim,
                        SupportedConeT::SecondOrderConeT(1) => total_dim += 1,
                        #[cfg(feature = "sdp")]
                        SupportedConeT::PSDTriangleConeT(1) => total_dim += 1,

                        // stop when we hit a non-collapsible cone
                        _ => break,
                    }
                }
                iter.next();
            }
            newcones.push(SupportedConeT::NonnegativeConeT(total_dim));
        }

        while let Some(cone) = iter.next() {
            if cone.nvars() != 0 {
                match cone {
                    // collapsible cones.   These are cones that can serve
                    // as the first term in a sequence of cones to be collapsed
                    // into a single nonnegative cone.
                    SupportedConeT::NonnegativeConeT(dim) => {
                        collapse(&mut iter, &mut newcones, *dim)
                    }
                    SupportedConeT::SecondOrderConeT(dim) if *dim == 1 => {
                        collapse(&mut iter, &mut newcones, *dim)
                    }
                    #[cfg(feature = "sdp")]
                    SupportedConeT::PSDTriangleConeT(dim) if *dim == 1 => {
                        collapse(&mut iter, &mut newcones, *dim)
                    }

                    // everything else
                    _ => newcones.push(cone.clone()),
                }
            }
        }
        newcones.shrink_to_fit();
        newcones
    }
}

// -------------------------------------
// Here we make a corresponding internal SupportedCone type that
// uses enum_dispatch to allow for static dispatching against
// all of our internal cone types
// -------------------------------------

#[allow(clippy::enum_variant_names)]
#[enum_dispatch(Cone<T>)]
pub enum SupportedCone<T>
where
    T: FloatT,
{
    ZeroCone(ZeroCone<T>),
    NonnegativeCone(NonnegativeCone<T>),
    SecondOrderCone(SecondOrderCone<T>),
    ExponentialCone(ExponentialCone<T>),
    PowerCone(PowerCone<T>),
    GenPowerCone(GenPowerCone<T>),
    #[cfg(feature = "sdp")]
    PSDTriangleCone(PSDTriangleCone<T>),
}

// -------------------------------------
// we need a tagging enum with no data fields to act
// as a bridge between the SupportedConeT API types and the
// internal SupportedCone enum_dispatch wrapper.   This enum
// has no data attached at all, so we can just convert to a u8.
// This would not be necessary if I could assign matching
// discriminants to both types, but that feature is not yet
// stable.  See:
// https://rust-lang.github.io/rfcs/2363-arbitrary-enum-discriminant.html
// -------------------------------------

#[allow(clippy::enum_variant_names)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[repr(u8)]
pub(crate) enum SupportedConeTag {
    ZeroCone = 0,
    NonnegativeCone,
    SecondOrderCone,
    ExponentialCone,
    PowerCone,
    GenPowerCone,
    #[cfg(feature = "sdp")]
    PSDTriangleCone,
}

pub(crate) trait SupportedConeAsTag {
    fn as_tag(&self) -> SupportedConeTag;
}

// user facing API type.   Just gives dimensions / exponents
impl<T> SupportedConeAsTag for SupportedConeT<T> {
    fn as_tag(&self) -> SupportedConeTag {
        match self {
            SupportedConeT::NonnegativeConeT(_) => SupportedConeTag::NonnegativeCone,
            SupportedConeT::ZeroConeT(_) => SupportedConeTag::ZeroCone,
            SupportedConeT::SecondOrderConeT(_) => SupportedConeTag::SecondOrderCone,
            SupportedConeT::ExponentialConeT() => SupportedConeTag::ExponentialCone,
            SupportedConeT::PowerConeT(_) => SupportedConeTag::PowerCone,
            #[cfg(feature = "sdp")]
            SupportedConeT::PSDTriangleConeT(_) => SupportedConeTag::PSDTriangleCone,
            SupportedConeT::GenPowerConeT(_, _) => SupportedConeTag::GenPowerCone,
        }
    }
}

// internal enum_dispatch container.   Each of the (_) contains the cone data objects
impl<T: FloatT> SupportedConeAsTag for SupportedCone<T> {
    fn as_tag(&self) -> SupportedConeTag {
        match self {
            SupportedCone::NonnegativeCone(_) => SupportedConeTag::NonnegativeCone,
            SupportedCone::ZeroCone(_) => SupportedConeTag::ZeroCone,
            SupportedCone::SecondOrderCone(_) => SupportedConeTag::SecondOrderCone,
            SupportedCone::ExponentialCone(_) => SupportedConeTag::ExponentialCone,
            SupportedCone::PowerCone(_) => SupportedConeTag::PowerCone,
            #[cfg(feature = "sdp")]
            SupportedCone::PSDTriangleCone(_) => SupportedConeTag::PSDTriangleCone,
            SupportedCone::GenPowerCone(_) => SupportedConeTag::GenPowerCone,
        }
    }
}

/// Returns the name of the cone from its tag.  Used for printing progress.
impl SupportedConeTag {
    pub fn as_str(&self) -> &'static str {
        match self {
            SupportedConeTag::ZeroCone => "ZeroCone",
            SupportedConeTag::NonnegativeCone => "NonnegativeCone",
            SupportedConeTag::SecondOrderCone => "SecondOrderCone",
            SupportedConeTag::ExponentialCone => "ExponentialCone",
            SupportedConeTag::PowerCone => "PowerCone",
            #[cfg(feature = "sdp")]
            SupportedConeTag::PSDTriangleCone => "PSDTriangleCone",
            SupportedConeTag::GenPowerCone => "GenPowerCone",
        }
    }
}

// ----------------------------------------------
// Iterator for the range of indices of the cone

//PJG: type names are not satisfactory.   Try to combine
//with the internal cone generators.

#[cfg_attr(not(feature = "sdp"), allow(dead_code))]
pub(crate) struct RangeSupportedConesIterator<'a, T> {
    cones: &'a [SupportedConeT<T>],
    index: usize,
    start: usize,
}

#[cfg_attr(not(feature = "sdp"), allow(dead_code))]
impl<T> Iterator for RangeSupportedConesIterator<'_, T> {
    type Item = std::ops::Range<usize>;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < self.cones.len() {
            let cone = &self.cones[self.index];
            let stop = self.start + cone.nvars();
            let range = self.start..stop;
            self.index += 1;
            self.start = stop;
            Some(range)
        } else {
            None
        }
    }
}
#[cfg_attr(not(feature = "sdp"), allow(dead_code))]
pub(crate) trait ConeRanges<'a, T> {
    fn rng_cones_iter(&'a self) -> RangeSupportedConesIterator<'a, T>;
}

#[cfg_attr(not(feature = "sdp"), allow(dead_code))]
impl<'a, T> ConeRanges<'a, T> for [SupportedConeT<T>] {
    fn rng_cones_iter(&'a self) -> RangeSupportedConesIterator<'a, T> {
        RangeSupportedConesIterator::<'a, T> {
            cones: self,
            index: 0,
            start: 0,
        }
    }
}

#[test]
fn test_cone_ranges() {
    let cones = [
        SupportedConeT::NonnegativeConeT::<f64>(3),
        SupportedConeT::NonnegativeConeT::<f64>(0),
        SupportedConeT::SecondOrderConeT::<f64>(4),
    ];

    let rngs: Vec<std::ops::Range<usize>> = vec![0..3, 3..3, 3..7];

    for (rng, conerng) in std::iter::zip(rngs.iter(), cones.rng_cones_iter()) {
        assert_eq!(*rng, conerng);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_new_collapsed_no_changes() {
        let cones = vec![
            SupportedConeT::<f64>::NonnegativeConeT(3),
            SupportedConeT::SecondOrderConeT(4),
            SupportedConeT::ExponentialConeT(),
        ];

        let expected = cones.clone();
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }

    #[test]
    fn test_new_collapsed_consolidate_nonnegative() {
        let cones = vec![
            SupportedConeT::<f64>::NonnegativeConeT(3),
            SupportedConeT::NonnegativeConeT(2),
            SupportedConeT::SecondOrderConeT(4),
        ];

        let expected = vec![
            SupportedConeT::<f64>::NonnegativeConeT(5),
            SupportedConeT::SecondOrderConeT(4),
        ];
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }

    #[test]
    fn test_new_collapsed_remove_empty() {
        let cones = vec![
            SupportedConeT::<f64>::NonnegativeConeT(3),
            SupportedConeT::ZeroConeT(0),
            SupportedConeT::SecondOrderConeT(4),
            SupportedConeT::NonnegativeConeT(0),
        ];

        let expected = vec![
            SupportedConeT::NonnegativeConeT(3),
            SupportedConeT::SecondOrderConeT(4),
        ];
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }

    #[test]
    fn test_new_collapsed_soc_to_nonnegative() {
        let cones = vec![
            SupportedConeT::<f64>::SecondOrderConeT(1),
            SupportedConeT::SecondOrderConeT(4),
        ];

        let expected = vec![
            SupportedConeT::NonnegativeConeT(1),
            SupportedConeT::SecondOrderConeT(4),
        ];
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }

    #[cfg(feature = "sdp")]
    #[test]
    fn test_new_collapsed_psd_to_nonnegative() {
        let cones = vec![
            SupportedConeT::<f64>::PSDTriangleConeT(1),
            SupportedConeT::SecondOrderConeT(4),
        ];

        let expected = vec![
            SupportedConeT::NonnegativeConeT(1),
            SupportedConeT::SecondOrderConeT(4),
        ];
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }

    #[test]
    fn test_new_collapsed_mixed() {
        let cones = vec![
            SupportedConeT::<f64>::SecondOrderConeT(1),
            SupportedConeT::NonnegativeConeT(3),
            SupportedConeT::NonnegativeConeT(2),
            SupportedConeT::ExponentialConeT(),
            SupportedConeT::NonnegativeConeT(0),
            SupportedConeT::SecondOrderConeT(1),
        ];

        let expected = vec![
            SupportedConeT::NonnegativeConeT(6),
            SupportedConeT::ExponentialConeT(),
            SupportedConeT::NonnegativeConeT(1),
        ];
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }

    #[cfg(feature = "sdp")]
    #[test]
    fn test_new_collapsed_mixed_sdp() {
        let cones = vec![
            SupportedConeT::<f64>::NonnegativeConeT(3),
            SupportedConeT::NonnegativeConeT(2),
            SupportedConeT::ZeroConeT(0),
            SupportedConeT::SecondOrderConeT(1),
            SupportedConeT::PSDTriangleConeT(1),
            SupportedConeT::SecondOrderConeT(4),
            SupportedConeT::NonnegativeConeT(0),
        ];

        let expected = vec![
            SupportedConeT::NonnegativeConeT(7),
            SupportedConeT::SecondOrderConeT(4),
        ];
        let result = SupportedConeT::new_collapsed(&cones);

        assert_eq!(result, expected);
    }
}