ommx 3.0.0-beta.2

use crate::{
    ATol, Bound, DecisionVariable, DecisionVariableError, DecisionVariableLabel, Instance, Kind,
    VariableID,
};

impl Instance {
    /// Get all unique decision variable names in this instance
    ///
    /// Returns a set of all unique variable names that have at least one named variable.
    /// Variables without names are not included.
    pub fn decision_variable_names(&self) -> std::collections::BTreeSet<String> {
        self.decision_variables
            .keys()
            .filter_map(|id| self.variable_labels().name(*id).map(|s| s.to_owned()))
            .collect()
    }

    /// Get a decision variable by name and subscripts
    ///
    /// # Arguments
    /// * `name` - The name of the decision variable to find
    /// * `subscripts` - The subscripts of the decision variable (can be empty)
    ///
    /// # Returns
    /// * `Some((VariableID, &DecisionVariable))` if a variable with the given name and subscripts is found
    /// * `None` if no matching variable is found
    ///
    /// # Example
    /// ```
    /// use ommx::Instance;
    ///
    /// let instance = Instance::default();
    /// // Find variable with name "x" and no subscripts
    /// let var = instance.get_decision_variable_by_name("x", vec![]);
    /// // Find variable with name "y" and subscripts [1, 2]
    /// let var_indexed = instance.get_decision_variable_by_name("y", vec![1, 2]);
    /// ```
    pub fn get_decision_variable_by_name(
        &self,
        name: &str,
        subscripts: Vec<i64>,
    ) -> Option<(VariableID, &DecisionVariable)> {
        let store = self.variable_labels();
        self.decision_variables.iter().find_map(|(id, var)| {
            (store.name(*id) == Some(name) && store.subscripts(*id) == subscripts.as_slice())
                .then_some((*id, var))
        })
    }

    /// Returns the next available [`VariableID`].
    ///
    /// Finds the maximum ID from decision variables, then adds 1.
    /// If there are no variables, returns `Ok(VariableID(0))`.
    ///
    /// Note: This method does not track which IDs have been allocated.
    /// Consecutive calls will return the same ID until a variable is actually added.
    ///
    /// # Errors
    ///
    /// Returns [`DecisionVariableError::NoAvailableID`] when the existing
    /// maximum ID is `u64::MAX` and no fresh ID can be allocated.
    pub fn next_variable_id(&self) -> Result<VariableID, DecisionVariableError> {
        self.decision_variables
            .last_key_value()
            .map(|(id, _)| {
                id.into_inner()
                    .checked_add(1)
                    .map(VariableID::from)
                    .ok_or(DecisionVariableError::NoAvailableID)
            })
            .unwrap_or(Ok(VariableID::from(0)))
    }

    pub(super) fn ensure_new_decision_variable_capacity(
        &self,
        count: usize,
    ) -> Result<(), DecisionVariableError> {
        if count == 0 {
            return Ok(());
        }
        let count = u64::try_from(count).map_err(|_| DecisionVariableError::NoAvailableID)?;
        if let Some((id, _)) = self.decision_variables.last_key_value() {
            id.into_inner()
                .checked_add(count)
                .ok_or(DecisionVariableError::NoAvailableID)?;
        }
        Ok(())
    }

    pub fn new_decision_variable(
        &mut self,
        kind: Kind,
        bound: Bound,
        fixed_value: Option<f64>,
        atol: ATol,
    ) -> Result<VariableID, DecisionVariableError> {
        self.new_decision_variable_with_label(
            kind,
            bound,
            DecisionVariableLabel::default(),
            fixed_value,
            atol,
        )
    }

    pub(super) fn new_decision_variable_with_label(
        &mut self,
        kind: Kind,
        bound: Bound,
        label: DecisionVariableLabel,
        fixed_value: Option<f64>,
        atol: ATol,
    ) -> Result<VariableID, DecisionVariableError> {
        let id = self.next_variable_id()?;
        let dv = DecisionVariable::new(kind, bound, atol)?;
        self.decision_variables
            .insert(id, dv, label, fixed_value, atol)?;
        Ok(id)
    }

    pub fn new_binary(&mut self) -> VariableID {
        self.new_decision_variable(Kind::Binary, Bound::of_binary(), None, ATol::default())
            .unwrap()
    }

    pub fn new_integer(&mut self) -> VariableID {
        self.new_decision_variable(Kind::Integer, Bound::default(), None, ATol::default())
            .unwrap()
    }

    pub fn new_continuous(&mut self) -> VariableID {
        self.new_decision_variable(Kind::Continuous, Bound::default(), None, ATol::default())
            .unwrap()
    }

    pub fn new_semi_integer(&mut self) -> VariableID {
        self.new_decision_variable(Kind::SemiInteger, Bound::default(), None, ATol::default())
            .unwrap()
    }

    pub fn new_semi_continuous(&mut self) -> VariableID {
        self.new_decision_variable(
            Kind::SemiContinuous,
            Bound::default(),
            None,
            ATol::default(),
        )
        .unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{coeff, linear, Function, Sense};
    use maplit::btreemap;
    use std::collections::BTreeMap;

    #[test]
    fn test_next_variable_id() {
        // Empty instance should return 0
        let decision_variables = BTreeMap::new();
        let objective = coeff!(1.0).into();
        let instance = Instance::new(
            Sense::Minimize,
            objective,
            decision_variables,
            BTreeMap::new(),
        )
        .unwrap();
        assert_eq!(instance.next_variable_id().unwrap(), VariableID::from(0));

        // Instance with variables should return max_id + 1
        let decision_variables = btreemap! {
            VariableID::from(5) => DecisionVariable::binary(),
            VariableID::from(8) => DecisionVariable::binary(),
            VariableID::from(100) => DecisionVariable::binary(),
        };
        let objective = (linear!(5) + coeff!(1.0)).into();
        let instance = Instance::new(
            Sense::Minimize,
            objective,
            decision_variables,
            BTreeMap::new(),
        )
        .unwrap();

        assert_eq!(instance.next_variable_id().unwrap(), VariableID::from(101));
    }

    #[test]
    fn test_next_variable_id_errors_when_id_space_is_exhausted() {
        let decision_variables = btreemap! {
            VariableID::from(u64::MAX) => DecisionVariable::binary(),
        };
        let instance = Instance::new(
            Sense::Minimize,
            Function::Zero,
            decision_variables,
            BTreeMap::new(),
        )
        .unwrap();

        assert!(matches!(
            instance.next_variable_id(),
            Err(DecisionVariableError::NoAvailableID)
        ));
        assert!(matches!(
            instance.ensure_new_decision_variable_capacity(1),
            Err(DecisionVariableError::NoAvailableID)
        ));
    }

    #[test]
    fn test_next_variable_id_with_new_binary() {
        // Test integration with new_binary
        let decision_variables = BTreeMap::new();
        let objective = coeff!(1.0).into();
        let mut instance = Instance::new(
            Sense::Minimize,
            objective,
            decision_variables,
            BTreeMap::new(),
        )
        .unwrap();

        let var1 = instance.new_binary();
        assert_eq!(var1, VariableID::from(0));

        let var2 = instance.new_binary();
        assert_eq!(var2, VariableID::from(1));

        assert_eq!(instance.next_variable_id().unwrap(), VariableID::from(2));
    }
}