ommx 3.0.0-beta.2

#[cfg(test)]
use crate::PolynomialBase;
use crate::{
    Coefficient, Degree, Linear, MonomialDyn, Polynomial, Quadratic, VariableID, VariableIDPair,
};
use derive_more::From;
use num::traits::Inv;
use std::borrow::Cow;

mod add;
mod approx;
mod arbitrary;
mod div;
mod evaluate;
mod evaluate_bound;
mod logical_memory;
mod mul;
mod parse;
mod reduce_binary_power;
mod serialize;
mod sub;
mod substitute;

/// A real-valued function of decision variables used for objective and constraint functions.
///
/// This can be up to polynomial currently, but it will be extended to exponential and logarithm in the future.
#[derive(Clone, PartialEq, From, Default)]
pub enum Function {
    #[default]
    Zero,
    /// Constant term stored as a [`Coefficient`].
    Constant(Coefficient),
    Linear(Linear),
    Quadratic(Quadratic),
    Polynomial(Polynomial),
}

impl serde::Serialize for Function {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        match self {
            Function::Zero => {
                use serde::ser::SerializeMap;
                let map = serializer.serialize_map(Some(0))?;
                map.end()
            }
            Function::Constant(c) => {
                use serde::ser::SerializeMap;
                let mut map = serializer.serialize_map(Some(1))?;
                map.serialize_entry(&(), &c.into_inner())?;
                map.end()
            }
            Function::Linear(l) => l.serialize(serializer),
            Function::Quadratic(q) => q.serialize(serializer),
            Function::Polynomial(p) => p.serialize(serializer),
        }
    }
}

impl<'de> serde::Deserialize<'de> for Function {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        // Always deserialize as Polynomial
        let polynomial = Polynomial::deserialize(deserializer)?;
        Ok(Function::Polynomial(polynomial))
    }
}

impl TryFrom<f64> for Function {
    type Error = crate::CoefficientError;
    fn try_from(value: f64) -> Result<Self, Self::Error> {
        match Coefficient::try_from(value) {
            Ok(c) => Ok(Function::Constant(c)),
            Err(crate::CoefficientError::Zero) => Ok(Function::Zero),
            Err(e) => Err(e),
        }
    }
}

#[cfg(test)]
impl<M> From<Result<PolynomialBase<M>, crate::CoefficientError>> for Function
where
    M: crate::Monomial,
    PolynomialBase<M>: Into<Function>,
{
    fn from(value: Result<PolynomialBase<M>, crate::CoefficientError>) -> Self {
        value.unwrap().into()
    }
}

#[doc(hidden)]
pub trait IntoFunctionForMacro {
    fn into_function_for_macro(self) -> Function;
}

impl<T> IntoFunctionForMacro for Result<T, crate::CoefficientError>
where
    Function: From<T>,
{
    fn into_function_for_macro(self) -> Function {
        Function::from(self.unwrap())
    }
}

macro_rules! impl_into_function_for_macro {
    ($($ty:ty),* $(,)?) => {
        $(
            impl IntoFunctionForMacro for $ty {
                fn into_function_for_macro(self) -> Function {
                    Function::from(self)
                }
            }
        )*
    };
}

impl_into_function_for_macro!(
    Function,
    Coefficient,
    Linear,
    Quadratic,
    Polynomial,
    crate::LinearMonomial,
    crate::QuadraticMonomial,
    MonomialDyn,
);

impl Function {
    pub fn constant_term(&self) -> f64 {
        match self {
            Function::Zero => 0.0,
            Function::Constant(c) => c.into_inner(),
            Function::Linear(l) => l.constant_term(),
            Function::Quadratic(q) => q.constant_term(),
            Function::Polynomial(p) => p.constant_term(),
        }
    }

    pub fn linear_terms(&self) -> Box<dyn Iterator<Item = (VariableID, Coefficient)> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(_) => Box::new(std::iter::empty()),
            Function::Linear(l) => Box::new(l.linear_terms()),
            Function::Quadratic(q) => Box::new(q.linear_terms()),
            Function::Polynomial(p) => Box::new(p.linear_terms()),
        }
    }

    pub fn quadratic_terms(&self) -> Box<dyn Iterator<Item = (VariableIDPair, Coefficient)> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(_) => Box::new(std::iter::empty()),
            Function::Linear(l) => Box::new(l.quadratic_terms()),
            Function::Quadratic(q) => Box::new(q.quadratic_terms()),
            Function::Polynomial(p) => Box::new(p.quadratic_terms()),
        }
    }

    pub fn as_linear(&self) -> Option<Cow<'_, Linear>> {
        match self {
            Function::Zero => Some(Cow::Owned(Linear::zero())),
            Function::Constant(c) => Some(Cow::Owned((*c).into())),
            Function::Linear(l) => Some(Cow::Borrowed(l)),
            Function::Quadratic(q) => q.try_into().map(Cow::Owned).ok(),
            Function::Polynomial(p) => p.try_into().map(Cow::Owned).ok(),
        }
    }

    pub fn as_quadratic(&self) -> Option<Cow<'_, Quadratic>> {
        match self {
            Function::Zero => Some(Cow::Owned(Quadratic::zero())),
            Function::Constant(c) => Some(Cow::Owned((*c).into())),
            Function::Linear(l) => Some(Cow::Owned(l.clone().into())),
            Function::Quadratic(q) => Some(Cow::Borrowed(q)),
            Function::Polynomial(p) => p.try_into().map(Cow::Owned).ok(),
        }
    }

    pub fn num_terms(&self) -> usize {
        match self {
            Function::Zero => 0,
            Function::Constant(_) => 1,
            Function::Linear(l) => l.num_terms(),
            Function::Quadratic(q) => q.num_terms(),
            Function::Polynomial(p) => p.num_terms(),
        }
    }

    pub fn degree(&self) -> Degree {
        match self {
            Function::Zero => 0.into(),
            Function::Constant(_) => 0.into(),
            Function::Linear(l) => l.degree(),
            Function::Quadratic(q) => q.degree(),
            Function::Polynomial(p) => p.degree(),
        }
    }

    pub fn iter(&self) -> Box<dyn Iterator<Item = (MonomialDyn, &Coefficient)> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(c) => Box::new(std::iter::once((MonomialDyn::default(), c))),
            Function::Linear(l) => Box::new(l.iter().map(|(k, v)| (MonomialDyn::from(*k), v))),
            Function::Quadratic(q) => Box::new(q.iter().map(|(k, v)| (MonomialDyn::from(*k), v))),
            Function::Polynomial(p) => Box::new(p.iter().map(|(k, v)| (k.clone(), v))),
        }
    }

    pub fn iter_mut(&mut self) -> Box<dyn Iterator<Item = (MonomialDyn, &mut Coefficient)> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(c) => Box::new(std::iter::once((MonomialDyn::default(), c))),
            Function::Linear(l) => Box::new(l.iter_mut().map(|(k, v)| (MonomialDyn::from(*k), v))),
            Function::Quadratic(q) => {
                Box::new(q.iter_mut().map(|(k, v)| (MonomialDyn::from(*k), v)))
            }
            Function::Polynomial(p) => Box::new(p.iter_mut().map(|(k, v)| (k.clone(), v))),
        }
    }

    pub fn values(&self) -> Box<dyn Iterator<Item = &Coefficient> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(c) => Box::new(std::iter::once(c)),
            Function::Linear(l) => Box::new(l.values()),
            Function::Quadratic(q) => Box::new(q.values()),
            Function::Polynomial(p) => Box::new(p.values()),
        }
    }

    pub fn values_mut(&mut self) -> Box<dyn Iterator<Item = &mut Coefficient> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(c) => Box::new(std::iter::once(c)),
            Function::Linear(l) => Box::new(l.values_mut()),
            Function::Quadratic(q) => Box::new(q.values_mut()),
            Function::Polynomial(p) => Box::new(p.values_mut()),
        }
    }

    pub fn keys(&self) -> Box<dyn Iterator<Item = MonomialDyn> + '_> {
        match self {
            Function::Zero => Box::new(std::iter::empty()),
            Function::Constant(_) => Box::new(std::iter::once(MonomialDyn::default())),
            Function::Linear(l) => Box::new(l.keys().map(|k| MonomialDyn::from(*k))),
            Function::Quadratic(q) => Box::new(q.keys().map(|k| MonomialDyn::from(*k))),
            Function::Polynomial(p) => Box::new(p.keys().cloned()),
        }
    }

    /// Get a minimal positive factor `a` which make all coefficients of `a * self` integer.
    ///
    /// This returns `1` for zero function. See also <https://en.wikipedia.org/wiki/Primitive_part_and_content>.
    pub fn content_factor(&self) -> crate::Result<Coefficient> {
        match self {
            Function::Zero => Ok(Coefficient::one()),
            // The factor must be positive so that multiplying it preserves the sign of
            // the constraint, matching `PolynomialBase::content_factor` for Linear/Quadratic/Polynomial.
            Function::Constant(c) => Ok(c.inv()?.abs()),
            Function::Linear(l) => l.content_factor(),
            Function::Quadratic(q) => q.content_factor(),
            Function::Polynomial(p) => p.content_factor(),
        }
    }
}