ommx 3.0.0-beta.2

use super::Instance;
use crate::{
    constraint::Equality, ATol, Bound, Bounds, Coefficient, ConstraintID, Evaluate,
    InfeasibleDetected, Kind, Linear, LinearMonomial, VariableID,
};
use anyhow::{bail, Context, Result};
use num::traits::Inv;

/// Signal returned when exact integer-slack conversion is structurally valid
/// but cannot construct an exact finite encoding.
///
/// Callers may recover by selecting an explicitly approximate transformation.
/// Missing constraints, unsupported variable kinds, and other contract or
/// materialization failures are not classified as this signal.
#[derive(Debug, thiserror::Error)]
#[error(transparent)]
pub struct ExactIntegerSlackUnavailable(ExactIntegerSlackFailure);

#[derive(Debug, thiserror::Error)]
enum ExactIntegerSlackFailure {
    #[error("Cannot normalize the coefficients to integers: constraint={id:?}")]
    CoefficientsNotNormalizable { id: ConstraintID },

    #[error(
        "The range of the slack variable exceeds the limit: evaluated({evaluated}) > limit({limit})"
    )]
    RangeTooLarge {
        id: ConstraintID,
        evaluated: f64,
        limit: u64,
    },
}

impl ExactIntegerSlackUnavailable {
    fn coefficients_not_normalizable(id: ConstraintID) -> Self {
        Self(ExactIntegerSlackFailure::CoefficientsNotNormalizable { id })
    }

    fn range_too_large(id: ConstraintID, evaluated: f64, limit: u64) -> Self {
        Self(ExactIntegerSlackFailure::RangeTooLarge {
            id,
            evaluated,
            limit,
        })
    }
}

impl Instance {
    /// Convert an inequality $f(x) \leq 0$ to an equality $a f(x) + s = 0$ with a
    /// newly introduced integer slack variable $s$, where $a$ is the minimal positive
    /// factor that makes every coefficient of $f(x)$ integer.
    ///
    /// # Errors
    ///
    /// Returns [`ExactIntegerSlackUnavailable`] when exact coefficient
    /// normalization or the configured finite slack range is unavailable, and
    /// [`InfeasibleDetected`] when variable bounds prove the inequality
    /// infeasible. Missing constraints, unsupported variable kinds, allocation,
    /// and coefficient-arithmetic failures retain their ordinary error types.
    pub fn convert_inequality_to_equality_with_integer_slack(
        &mut self,
        constraint_id: u64,
        max_integer_range: u64,
        atol: ATol,
    ) -> Result<()> {
        let constraint_id = ConstraintID::from(constraint_id);
        let bounds = self.bounds();
        let kinds = self.kinds();

        let (function, equality) = {
            let constraint = self
                .constraint_collection
                .active()
                .get(&constraint_id)
                .with_context(|| format!("Constraint ID {constraint_id:?} not found"))?;
            (constraint.function().clone(), constraint.equality)
        };

        if equality != Equality::LessThanOrEqualToZero {
            bail!("The constraint is not inequality: ID={constraint_id:?}");
        }

        for id in function.required_ids() {
            let kind = kinds
                .get(&id)
                .with_context(|| format!("Decision variable ID {id:?} not found"))?;
            if !matches!(kind, Kind::Binary | Kind::Integer) {
                bail!("The constraint contains continuous decision variables: ID={id:?}");
            }
        }

        let a = function.content_factor().map_err(|_| {
            ExactIntegerSlackUnavailable::coefficients_not_normalizable(constraint_id)
        })?;
        let af = (function.clone() * a)?;

        let af_bound = af.evaluate_bound(&bounds);
        let af_bound = af_bound.as_integer_bound(atol).ok_or(
            InfeasibleDetected::InequalityConstraintBound {
                id: constraint_id,
                bound: af_bound,
            },
        )?;
        if af_bound.lower() > 0.0 {
            bail!(InfeasibleDetected::InequalityConstraintBound {
                id: constraint_id,
                bound: af_bound,
            });
        }
        if af_bound.upper() <= 0.0 {
            self.relax_constraint(
                constraint_id,
                "ommx.Instance.convert_inequality_to_equality_with_integer_slack".to_string(),
                [],
            )?;
            return Ok(());
        }

        let slack_bound = Bound::new(0.0, -af_bound.lower()).unwrap();
        if slack_bound.width() > max_integer_range as f64 {
            return Err(ExactIntegerSlackUnavailable::range_too_large(
                constraint_id,
                slack_bound.width(),
                max_integer_range,
            )
            .into());
        }

        let slack_id = self.new_decision_variable_with_label(
            Kind::Integer,
            slack_bound,
            crate::ModelingLabel {
                name: Some("ommx.slack".to_string()),
                subscripts: vec![constraint_id.into_inner() as i64],
                ..Default::default()
            },
            None,
            atol,
        )?;

        let slack_term = Linear::single_term(LinearMonomial::Variable(slack_id), a.inv()?);
        let new_function = (function + slack_term)?;

        let mut constraint = self
            .constraint_collection
            .active()
            .get(&constraint_id)
            .cloned()
            .expect("constraint presence was verified above");
        *constraint.function_mut() = new_function;
        constraint.equality = Equality::EqualToZero;
        self.constraint_collection
            .replace_active_row(constraint_id, constraint)?;

        Ok(())
    }

    /// Convert an inequality $f(x) \leq 0$ to $f(x) + b s \leq 0$ with an integer
    /// slack variable $s \in [0, \text{slack\_upper\_bound}]$.
    ///
    /// Returns the coefficient $b = -\mathrm{lower}(f(x)) / \text{slack\_upper\_bound}$.
    /// Returns `None` if the constraint was trivially satisfied and was moved to
    /// removed_constraints.
    pub fn add_integer_slack_to_inequality(
        &mut self,
        constraint_id: u64,
        slack_upper_bound: u64,
    ) -> Result<Option<f64>> {
        let constraint_id = ConstraintID::from(constraint_id);
        let bounds = self.bounds();
        let kinds = self.kinds();

        let (function, equality) = {
            let constraint = self
                .constraint_collection
                .active()
                .get(&constraint_id)
                .with_context(|| format!("Constraint ID {constraint_id:?} not found"))?;
            (constraint.function().clone(), constraint.equality)
        };

        if equality != Equality::LessThanOrEqualToZero {
            bail!("The constraint is not inequality: ID={constraint_id:?}");
        }

        for id in function.required_ids() {
            let kind = kinds
                .get(&id)
                .with_context(|| format!("Decision variable ID {id:?} not found"))?;
            if !matches!(kind, Kind::Binary | Kind::Integer) {
                bail!("The constraint contains continuous decision variables: ID={id:?}");
            }
        }

        let f_bound = function.evaluate_bound(&bounds);
        if f_bound.lower() > 0.0 {
            bail!(InfeasibleDetected::InequalityConstraintBound {
                id: constraint_id,
                bound: f_bound,
            });
        }
        if f_bound.upper() <= 0.0 {
            self.relax_constraint(
                constraint_id,
                "add_integer_slack_to_inequality".to_string(),
                [],
            )?;
            return Ok(None);
        }

        let b = -f_bound.lower() / slack_upper_bound as f64;
        let slack_bound = Bound::new(0.0, slack_upper_bound as f64).unwrap();

        // Validate the slack coefficient before mutating the instance so failures
        // (e.g. `slack_upper_bound == 0` giving `b = inf`) do not leave an orphan
        // slack decision variable behind. `b` is non-negative in this branch (we
        // bailed on lower > 0 and relaxed when upper <= 0). `b == 0` only when
        // `f_bound.lower() == 0`, which is a boundary case; adding a zero-coefficient
        // slack term is mathematically a no-op so we skip the term in that case,
        // matching v1 observable behavior (where a zero coefficient is dropped on
        // insertion).
        let b_coeff = match Coefficient::try_from(b) {
            Ok(c) => Some(c),
            Err(crate::CoefficientError::Zero) => None,
            Err(e) => return Err(e).context("Slack coefficient must be finite"),
        };

        let slack_id = self.new_decision_variable_with_label(
            Kind::Integer,
            slack_bound,
            crate::ModelingLabel {
                name: Some("ommx.slack".to_string()),
                subscripts: vec![constraint_id.into_inner() as i64],
                ..Default::default()
            },
            None,
            ATol::default(),
        )?;

        let new_function = match b_coeff {
            Some(c) => {
                let slack_term = Linear::single_term(LinearMonomial::Variable(slack_id), c);
                (function + slack_term)?
            }
            None => function,
        };

        let mut constraint = self
            .constraint_collection
            .active()
            .get(&constraint_id)
            .cloned()
            .expect("constraint presence was verified above");
        *constraint.function_mut() = new_function;
        self.constraint_collection
            .replace_active_row(constraint_id, constraint)?;

        Ok(Some(b))
    }

    /// Snapshot of bounds for every decision variable.
    fn bounds(&self) -> Bounds {
        self.decision_variables
            .iter()
            .map(|(id, dv)| (*id, dv.bound()))
            .collect()
    }

    /// Snapshot of kinds for every decision variable.
    fn kinds(&self) -> fnv::FnvHashMap<VariableID, Kind> {
        self.decision_variables
            .iter()
            .map(|(id, dv)| (*id, dv.kind()))
            .collect()
    }
}

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

    #[test]
    fn converts_integer_inequality_to_equality_with_slack() {
        // min x1 + x2 s.t. x1 + x2 - 4 <= 0, with x1, x2 integer in [0, 3]
        let dv = btreemap! {
            VariableID::from(1) => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
            VariableID::from(2) => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
        };
        let objective = (Function::from(linear!(1)) + Function::from(linear!(2))).unwrap();
        let constraint_fn = ((Function::from(linear!(1)) + Function::from(linear!(2))).unwrap()
            + coeff!(-4.0))
        .unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::less_than_or_equal_to_zero(constraint_fn,
            ),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();

        instance
            .convert_inequality_to_equality_with_integer_slack(0, 32, ATol::default())
            .unwrap();

        let constraint = instance
            .constraints()
            .get(&ConstraintID::from(0))
            .expect("constraint should still be present");
        assert_eq!(constraint.equality, Equality::EqualToZero);
        // Slack var should have been added
        let store = instance.variable_labels();
        assert!(instance
            .decision_variables
            .keys()
            .any(|id| store.name(*id) == Some("ommx.slack")));
    }

    #[test]
    fn exact_integer_slack_range_limit_is_a_recoverable_signal() {
        let id = VariableID::from(1);
        let dv = btreemap! {
            id => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
        };
        let objective = Function::from(linear!(1));
        let constraint_fn = (Function::from(linear!(1)) + coeff!(-2.0)).unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::less_than_or_equal_to_zero(constraint_fn),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();

        let err = instance
            .convert_inequality_to_equality_with_integer_slack(0, 1, ATol::default())
            .unwrap_err();

        assert!(err.is::<ExactIntegerSlackUnavailable>());
        assert!(instance.constraints().contains_key(&ConstraintID::from(0)));
    }

    #[test]
    fn add_integer_slack_updates_function_but_keeps_inequality() {
        // min x1 s.t. x1 - 2 <= 0, x1 integer in [0, 3]
        let dv = btreemap! {
            VariableID::from(1) => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
        };
        let objective = Function::from(linear!(1));
        let constraint_fn = (Function::from(linear!(1)) + coeff!(-2.0)).unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::less_than_or_equal_to_zero(constraint_fn,
            ),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();

        let b = instance
            .add_integer_slack_to_inequality(0, 2)
            .unwrap()
            .expect("constraint should still be active");
        assert!(b > 0.0);

        let constraint = instance
            .constraints()
            .get(&ConstraintID::from(0))
            .expect("constraint should still be present");
        assert_eq!(constraint.equality, Equality::LessThanOrEqualToZero);
        let store = instance.variable_labels();
        assert!(instance
            .decision_variables
            .keys()
            .any(|id| store.name(*id) == Some("ommx.slack")));
    }

    #[test]
    fn always_satisfied_inequality_is_relaxed() {
        // x1 - 10 <= 0 with x1 in [0, 3] is always satisfied
        let dv = btreemap! {
            VariableID::from(1) => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
        };
        let objective = Function::from(linear!(1));
        let constraint_fn = (Function::from(linear!(1)) + coeff!(-10.0)).unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::less_than_or_equal_to_zero(constraint_fn,
            ),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();

        let result = instance.add_integer_slack_to_inequality(0, 2).unwrap();
        assert!(result.is_none());
        assert!(instance.constraints().is_empty());
        assert_eq!(instance.removed_constraints().len(), 1);
    }

    #[test]
    fn rejects_zero_slack_upper_bound_without_mutating_instance() {
        // f(x) = x1 - 2 with x1 in [0, 3] gives a finite non-zero lower, so
        // `slack_upper_bound == 0` drives `b = inf`. The call must fail without
        // leaving behind an orphan slack decision variable.
        let dv = btreemap! {
            VariableID::from(1) => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
        };
        let objective = Function::from(linear!(1));
        let constraint_fn = (Function::from(linear!(1)) + coeff!(-2.0)).unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::less_than_or_equal_to_zero(constraint_fn,
            ),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();
        let before = instance.decision_variables.len();

        let err = instance.add_integer_slack_to_inequality(0, 0).unwrap_err();
        assert!(err.to_string().to_lowercase().contains("finite"));
        // No slack variable should have been added on the failure path.
        assert_eq!(instance.decision_variables.len(), before);
        // The original constraint is still the untouched inequality.
        let constraint = instance.constraints().get(&ConstraintID::from(0)).unwrap();
        assert_eq!(constraint.equality, Equality::LessThanOrEqualToZero);
    }

    #[test]
    fn convert_inequality_rejects_equality_constraint() {
        // An `EqualToZero` constraint must be rejected by
        // `convert_inequality_to_equality_with_integer_slack`, which names the
        // contract in its identifier. Matches the guard in the sibling
        // `add_integer_slack_to_inequality`.
        let dv = btreemap! {
            VariableID::from(1) => DecisionVariable::new(
                Kind::Integer,
                Bound::new(0.0, 3.0).unwrap(),
                ATol::default(),
            ).unwrap(),
        };
        let objective = Function::from(linear!(1));
        let constraint_fn = (Function::from(linear!(1)) + coeff!(-2.0)).unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::equal_to_zero(constraint_fn,
            ),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();

        let err = instance
            .convert_inequality_to_equality_with_integer_slack(0, 32, ATol::default())
            .unwrap_err();
        assert!(!err.is::<ExactIntegerSlackUnavailable>());
        assert!(err.to_string().contains("not inequality"));
    }

    #[test]
    fn rejects_constraint_with_continuous_variable() {
        let dv = btreemap! {
            VariableID::from(1) => DecisionVariable::continuous(),
        };
        let objective = Function::from(linear!(1));
        let constraint_fn = (Function::from(linear!(1)) + coeff!(-2.0)).unwrap();
        let constraints = btreemap! {
            ConstraintID::from(0) => crate::Constraint::less_than_or_equal_to_zero(constraint_fn,
            ),
        };
        let mut instance = Instance::new(Sense::Minimize, objective, dv, constraints).unwrap();

        let err = instance
            .convert_inequality_to_equality_with_integer_slack(0, 32, ATol::default())
            .unwrap_err();
        assert!(!err.is::<ExactIntegerSlackUnavailable>());
        assert!(err.to_string().contains("continuous decision variables"));
    }
}