ommx 3.0.0-beta.5

use super::*;
use std::collections::BTreeMap;

impl Instance {
    /// Reduce binary powers in the instance.
    ///
    /// This method replaces binary powers in the instance with their equivalent linear expressions.
    /// For binary variables, x^n = x for any n >= 1, so we can reduce higher powers to linear terms.
    ///
    /// # Postconditions
    ///
    /// Reduction rewrites active functions while preserving the output objective.
    ///
    /// ```
    /// use ommx::{
    ///     quadratic, v1::State, ATol, DecisionVariable, Evaluate, Function, Instance,
    ///     Sense, VariableID,
    /// };
    /// use std::collections::{BTreeMap, HashMap};
    ///
    /// let mut instance = Instance::builder()
    ///     .sense(Sense::Maximize)
    ///     .objective(Function::from(quadratic!(1, 1)))
    ///     .decision_variables(BTreeMap::from([(
    ///         VariableID::from(1),
    ///         DecisionVariable::binary(),
    ///     )]))
    ///     .constraints(BTreeMap::new())
    ///     .build()
    ///     .unwrap();
    /// let original = instance.objective().clone();
    ///
    /// assert!(instance.reduce_binary_power().unwrap());
    /// let output = instance.output_objective().unwrap();
    /// assert_eq!(output.sense(), Sense::Maximize);
    /// assert_eq!(output.function(), &original);
    /// assert!(!instance.reduce_binary_power().unwrap());
    /// let state = State::from(HashMap::from([(1, 1.0)]));
    /// assert_eq!(instance.objective().evaluate(&state, ATol::default()).unwrap(), 1.0);
    /// let solution = instance.evaluate(&state, ATol::default()).unwrap();
    /// assert_eq!(*solution.objective(), 1.0);
    /// ```
    pub fn reduce_binary_power(&mut self) -> Result<bool, crate::CoefficientError> {
        let binary_ids = self.binary_ids();
        if binary_ids.is_empty() {
            return Ok(false);
        }

        // Plan: compute the rewritten objective and active-constraint
        // replacements from borrowed values, without cloning unchanged rows,
        // removed constraints, special constraints, named functions, or
        // metadata. All fallible steps happen here, before `self` is
        // mutated, so a `CoefficientError` never leaves `self` partially
        // rewritten.
        let objective_replacement = self.objective.plan_binary_power_reduction(&binary_ids)?;
        let mut replacements = BTreeMap::new();
        for (&id, constraint) in self.constraint_collection.active() {
            if let Some(replacement) = constraint.plan_binary_power_reduction(&binary_ids)? {
                replacements.insert(id, replacement);
            }
        }
        let changed = objective_replacement.is_some() || !replacements.is_empty();

        // Commit: only the touched fields are written, so this drops just
        // the old objective and the old versions of the touched active
        // constraints, not the whole instance.
        if let Some(replacement) = objective_replacement {
            self.replace_active_objective_preserving_output(replacement);
        }
        if !replacements.is_empty() {
            self.constraint_collection
                .replace_active_rows(replacements)
                .expect("replacement IDs were read from active constraints");
        }
        // Note: We don't need to reduce in removed_constraints since they are not active
        Ok(changed)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        coeff, constraint::CreatedData, linear, quadratic, Bound, Coefficient, Constraint,
        DecisionVariable, Equality, Kind, Sense,
    };
    use ::approx::assert_abs_diff_eq;
    use proptest::prelude::*;

    #[test]
    fn test_instance_reduce_binary_power() {
        // Create instance with binary and continuous variables
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(VariableID::from(1), DecisionVariable::binary());
        decision_variables.insert(
            VariableID::from(2),
            DecisionVariable::new(
                Kind::Continuous,
                Bound::new(0.0, 10.0).unwrap(),
                crate::ATol::default(),
            )
            .unwrap(),
        );

        // Objective: x1^2 + x1*x2 + x2^2
        let mut objective_poly =
            (quadratic!(1, 1) + (coeff!(2.0) * quadratic!(1, 2)).unwrap()).unwrap();
        objective_poly = (objective_poly + (coeff!(3.0) * quadratic!(2, 2)).unwrap()).unwrap();
        let objective = Function::Quadratic(objective_poly);

        // Constraint: x1^2 + x2 <= 5  (i.e., x1^2 + x2 - 5 <= 0)
        let mut constraints = BTreeMap::new();
        let constraint_func = Function::Quadratic(
            ((quadratic!(1, 1) + quadratic!(2)).unwrap() + coeff!(-5.0)).unwrap(),
        );
        constraints.insert(
            ConstraintID::from(1),
            Constraint {
                equality: Equality::LessThanOrEqualToZero,
                stage: CreatedData {
                    function: constraint_func,
                },
            },
        );

        let mut instance = Instance::new(
            Sense::Minimize,
            objective.clone(),
            decision_variables,
            constraints,
        )
        .unwrap();

        // Apply reduction
        let changed = instance.reduce_binary_power().unwrap();
        assert!(changed);
        assert_eq!(instance.output_objective().unwrap().function(), &objective);

        // Check objective: x1^2 -> x1
        let mut expected_objective_poly =
            (quadratic!(1) + (coeff!(2.0) * quadratic!(1, 2)).unwrap()).unwrap();
        expected_objective_poly =
            (expected_objective_poly + (coeff!(3.0) * quadratic!(2, 2)).unwrap()).unwrap();
        let expected_objective = Function::Quadratic(expected_objective_poly);
        assert_abs_diff_eq!(instance.objective(), &expected_objective);

        // Check constraint: x1^2 -> x1
        let expected_constraint_func =
            Function::Quadratic(((quadratic!(1) + quadratic!(2)).unwrap() + coeff!(-5.0)).unwrap());
        assert_eq!(
            instance
                .constraints()
                .get(&ConstraintID::from(1))
                .unwrap()
                .function(),
            &expected_constraint_func
        );

        // Apply reduction again - should not change
        let changed2 = instance.reduce_binary_power().unwrap();
        assert!(!changed2);
    }

    #[test]
    fn test_instance_reduce_binary_power_no_binary() {
        // Create instance with only continuous variables
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(VariableID::from(1), DecisionVariable::continuous());
        decision_variables.insert(VariableID::from(2), DecisionVariable::integer());

        let objective = Function::Quadratic(
            (quadratic!(1, 1) + (coeff!(2.0) * quadratic!(2, 2)).unwrap()).unwrap(),
        );

        let mut instance = Instance::new(
            Sense::Minimize,
            objective.clone(),
            decision_variables,
            BTreeMap::new(),
        )
        .unwrap();

        // Apply reduction - should not change since no binary variables
        let changed = instance.reduce_binary_power().unwrap();
        assert!(!changed);
        assert_eq!(instance.objective(), &objective);
        assert!(instance.output_objective().is_none());
    }

    #[test]
    fn reduce_binary_power_does_not_capture_constraint_only_rewrite() {
        let variable = VariableID::from(1);
        let objective = Function::from(linear!(1));
        let constraint = Constraint::equal_to_zero(Function::from(quadratic!(1, 1)));
        let mut instance = Instance::new(
            Sense::Minimize,
            objective.clone(),
            BTreeMap::from([(variable, DecisionVariable::binary())]),
            BTreeMap::from([(ConstraintID::from(1), constraint)]),
        )
        .unwrap();

        assert!(instance.reduce_binary_power().unwrap());

        assert_eq!(instance.objective(), &objective);
        assert!(instance.output_objective().is_none());
    }

    #[test]
    fn reduce_binary_power_moves_original_expression_to_output_objective() {
        let variable = VariableID::from(1);
        let objective = Function::from(linear!(1)).powi(2);
        let mut instance = Instance::new(
            Sense::Minimize,
            objective,
            BTreeMap::from([(variable, DecisionVariable::binary())]),
            BTreeMap::new(),
        )
        .unwrap();
        let Function::Expression(expression) = instance.objective() else {
            panic!("powi(2) must create an expression")
        };
        let original_instructions = crate::function::operation::instructions(expression).as_ptr();

        assert!(instance.reduce_binary_power().unwrap());

        let Function::Expression(expression) = instance.output_objective().unwrap().function()
        else {
            panic!("output objective must preserve the original expression")
        };
        assert_eq!(
            crate::function::operation::instructions(expression).as_ptr(),
            original_instructions
        );
        assert_eq!(instance.objective(), &Function::from(linear!(1)));
    }

    #[test]
    fn reduce_binary_power_preserves_instance_on_coefficient_error() {
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(VariableID::from(1), DecisionVariable::binary());

        let objective = Function::Quadratic(quadratic!(1, 1).into());
        let huge = Coefficient::try_from(f64::MAX).unwrap();
        let overflowing_constraint = Function::Quadratic(
            ((huge * quadratic!(1, 1)).unwrap() + (huge * quadratic!(1)).unwrap()).unwrap(),
        );
        let mut constraints = BTreeMap::new();
        constraints.insert(
            ConstraintID::from(1),
            Constraint {
                equality: Equality::LessThanOrEqualToZero,
                stage: CreatedData {
                    function: overflowing_constraint,
                },
            },
        );
        let mut instance =
            Instance::new(Sense::Minimize, objective, decision_variables, constraints).unwrap();
        let before_objective = instance.objective().clone();
        let before_constraints = instance.constraints().clone();

        let err = instance.reduce_binary_power().unwrap_err();

        assert_eq!(err, crate::CoefficientError::Infinite);
        assert_eq!(instance.objective(), &before_objective);
        assert_eq!(instance.constraints(), &before_constraints);
        assert!(instance.output_objective().is_none());
    }

    proptest! {
        #[test]
        fn test_instance_reduce_binary_power_idempotent(
            mut instance in Instance::arbitrary()
        ) {
            let _first = instance.reduce_binary_power().unwrap();
            let second = instance.reduce_binary_power().unwrap();
            prop_assert!(!second, "reduce_binary_power should be idempotent");
        }
    }
}