use super::*;
use crate::algebra::FloatT;

#[cfg(feature = "sdp")]
use crate::algebra::triangular_number;

// ---------------------------------------------------
// 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.
///  
#[derive(Debug, Clone)]
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(),
    }
}

// -------------------------------------
// 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",
        }
    }
}
