use crate::{
    Coefficient, Degree, Linear, MonomialDyn, Polynomial, Quadratic, VariableID, VariableIDPair,
};
use derive_more::From;
use num::{traits::Inv, One, Zero};
use std::borrow::Cow;

mod add;
mod approx;
mod arbitrary;
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,
    /// Non-zero constant
    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),
        }
    }
}

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(),
        }
    }
}
