ommx 2.6.1

use super::*;
use crate::{
    substitute_acyclic, substitute_one_via_acyclic, Function, Substitute, SubstitutionError,
    VariableID,
};

impl Substitute for Instance {
    type Output = Self;

    fn substitute_acyclic(
        mut self,
        acyclic: &crate::AcyclicAssignments,
    ) -> Result<Self::Output, crate::SubstitutionError> {
        if acyclic.is_empty() {
            return Ok(self);
        }

        for (var_id, function) in acyclic.iter() {
            if !self.decision_variables.contains_key(var_id) {
                return Err(SubstitutionError::UndefinedSubstitutionVariable { variable: *var_id });
            }
            for required_id in function.required_ids() {
                if !self.decision_variables.contains_key(&required_id) {
                    return Err(SubstitutionError::UndefinedSubstitutionVariable {
                        variable: required_id,
                    });
                }
            }
        }

        // Get the set of variables being substituted
        let substituted_variables: std::collections::BTreeSet<VariableID> =
            acyclic.iter().map(|(var_id, _)| *var_id).collect();

        // Identify constraint IDs that depend on substituted variables
        let mut affected_constraint_ids = std::collections::BTreeSet::new();

        // Check active constraints only.
        // Removed constraints are not checked here; they will be substituted
        // when restored via `restore_constraint`. Constraint hints for removed
        // constraints are discarded when the constraint is removed.
        for (constraint_id, constraint) in &self.constraints {
            let required_ids = constraint.required_ids();
            if !required_ids.is_disjoint(&substituted_variables) {
                affected_constraint_ids.insert(*constraint_id);
            }
        }

        // Apply substitution to the objective function
        substitute_acyclic(&mut self.objective, acyclic)?;

        // Apply substitution only to affected active constraints.
        // Removed constraints are not substituted here; they will be substituted
        // when restored via `restore_constraint`.
        for constraint_id in &affected_constraint_ids {
            if let Some(constraint) = self.constraints.get_mut(constraint_id) {
                substitute_acyclic(&mut constraint.function, acyclic)?;
            }
        }

        // Apply substitution to the existing decision_variable_dependency
        substitute_acyclic(&mut self.decision_variable_dependency, acyclic)?;

        // Remove constraint hints that reference affected constraints
        //
        // FIXME
        // ------
        // - Currently, we remove all affected constraint hints, but some can still be valid.
        //   - e.g. an one-hot constraint x1 + x2 = 1 is still one-hot after substituting x1 = x3 + x4.
        self.constraint_hints
            .one_hot_constraints
            .retain(|hint| !affected_constraint_ids.contains(&hint.id));
        self.constraint_hints.sos1_constraints.retain(|hint| {
            !affected_constraint_ids.contains(&hint.binary_constraint_id)
                && hint
                    .big_m_constraint_ids
                    .is_disjoint(&affected_constraint_ids)
        });

        Ok(self)
    }

    fn substitute_one(
        self,
        assigned: VariableID,
        f: &Function,
    ) -> Result<Self::Output, SubstitutionError> {
        substitute_one_via_acyclic(self, assigned, f)
    }
}

impl Substitute for ParametricInstance {
    type Output = Self;

    fn substitute_acyclic(
        mut self,
        acyclic: &crate::AcyclicAssignments,
    ) -> Result<Self::Output, crate::SubstitutionError> {
        if acyclic.is_empty() {
            return Ok(self);
        }

        let substituted_variables: std::collections::BTreeSet<VariableID> =
            acyclic.iter().map(|(var_id, _)| *var_id).collect();

        for (var_id, function) in acyclic.iter() {
            if self.parameters.contains_key(var_id) {
                return Err(SubstitutionError::ParameterSubstitution { parameter: *var_id });
            }
            if !self.decision_variables.contains_key(var_id) {
                return Err(SubstitutionError::UndefinedSubstitutionVariable { variable: *var_id });
            }
            for required_id in function.required_ids() {
                if !self.decision_variables.contains_key(&required_id)
                    && !self.parameters.contains_key(&required_id)
                {
                    return Err(SubstitutionError::UndefinedSubstitutionVariable {
                        variable: required_id,
                    });
                }
            }
        }

        let mut affected_constraint_ids = std::collections::BTreeSet::new();
        for (constraint_id, constraint) in &self.constraints {
            let required_ids = constraint.required_ids();
            if !required_ids.is_disjoint(&substituted_variables) {
                affected_constraint_ids.insert(*constraint_id);
            }
        }

        substitute_acyclic(&mut self.objective, acyclic)?;

        for constraint_id in &affected_constraint_ids {
            if let Some(constraint) = self.constraints.get_mut(constraint_id) {
                substitute_acyclic(&mut constraint.function, acyclic)?;
            }
        }

        substitute_acyclic(&mut self.decision_variable_dependency, acyclic)?;

        self.constraint_hints
            .one_hot_constraints
            .retain(|hint| !affected_constraint_ids.contains(&hint.id));
        self.constraint_hints.sos1_constraints.retain(|hint| {
            !affected_constraint_ids.contains(&hint.binary_constraint_id)
                && hint
                    .big_m_constraint_ids
                    .is_disjoint(&affected_constraint_ids)
        });

        Ok(self)
    }

    fn substitute_one(
        self,
        assigned: VariableID,
        f: &Function,
    ) -> Result<Self::Output, SubstitutionError> {
        substitute_one_via_acyclic(self, assigned, f)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        coeff,
        constraint::Equality,
        constraint_hints::{OneHot, Sos1},
        linear, DecisionVariable, Sense,
    };
    use std::collections::BTreeMap;

    #[test]
    fn test_instance_substitute() {
        // Create decision variables
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(
            VariableID::from(1),
            DecisionVariable::continuous(VariableID::from(1)),
        );
        decision_variables.insert(
            VariableID::from(2),
            DecisionVariable::continuous(VariableID::from(2)),
        );
        decision_variables.insert(
            VariableID::from(3),
            DecisionVariable::continuous(VariableID::from(3)),
        );

        // Create a simple instance: minimize x1 + 2*x2, subject to x1 + x2 <= 10
        let objective = Function::from(linear!(1) + coeff!(2.0) * linear!(2));
        let constraint_function = Function::from(linear!(1) + linear!(2) + coeff!(-10.0));

        let mut constraints = BTreeMap::new();
        let constraint = Constraint {
            id: ConstraintID::from(1),
            function: constraint_function,
            equality: Equality::LessThanOrEqualToZero,
            name: None,
            subscripts: Vec::new(),
            parameters: Default::default(),
            description: None,
        };
        constraints.insert(ConstraintID::from(1), constraint);
        let _constraint_hints = ConstraintHints::default();

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

        // Substitute x1 with x3 + 1
        let substitution = Function::from(linear!(3) + coeff!(1.0));
        let result = instance
            .substitute_one(VariableID::from(1), &substitution)
            .unwrap();

        // Check that the decision_variable_dependency contains the assignment x1 <- x3 + 1
        assert_eq!(result.decision_variable_dependency.len(), 1);
        assert!(result
            .decision_variable_dependency
            .get(&VariableID::from(1))
            .is_some());
    }

    #[test]
    fn test_instance_substitute_undefined_rhs_fails() {
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(
            VariableID::from(0),
            DecisionVariable::continuous(VariableID::from(0)),
        );

        let instance = Instance::new(
            Sense::Minimize,
            Function::from(linear!(0)),
            decision_variables,
            BTreeMap::new(),
        )
        .unwrap();

        let err = instance
            .substitute_one(VariableID::from(0), &Function::from(linear!(999)))
            .unwrap_err();
        assert!(matches!(
            err,
            SubstitutionError::UndefinedSubstitutionVariable { variable }
                if variable == VariableID::from(999)
        ));
    }

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

        let mut parameters = BTreeMap::new();
        parameters.insert(
            VariableID::from(100),
            crate::v1::Parameter {
                id: 100,
                ..Default::default()
            },
        );

        let objective = Function::from(linear!(0) + linear!(100));
        let parametric = ParametricInstance::new(
            Sense::Minimize,
            objective,
            decision_variables,
            parameters,
            BTreeMap::new(),
        )
        .unwrap();

        let substituted = parametric
            .substitute_one(
                VariableID::from(0),
                &Function::from(linear!(1) + linear!(100)),
            )
            .unwrap();

        assert_eq!(substituted.decision_variable_dependency.len(), 1);
        assert!(substituted
            .decision_variable_dependency
            .get(&VariableID::from(0))
            .is_some());

        let mut parameter_values = crate::v1::Parameters::default();
        parameter_values.entries.insert(100, 2.0);
        let instance = substituted.with_parameters(parameter_values).unwrap();
        let state = crate::v1::State::from_iter([(1, 3.0)]);
        let value = instance
            .objective()
            .evaluate(&state, crate::ATol::default())
            .unwrap();
        assert_eq!(value, 7.0);
    }

    #[test]
    fn test_parametric_instance_substitute_parameter_target_fails() {
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(
            VariableID::from(0),
            DecisionVariable::continuous(VariableID::from(0)),
        );

        let mut parameters = BTreeMap::new();
        parameters.insert(
            VariableID::from(100),
            crate::v1::Parameter {
                id: 100,
                ..Default::default()
            },
        );

        let parametric = ParametricInstance::new(
            Sense::Minimize,
            Function::from(linear!(0) + linear!(100)),
            decision_variables,
            parameters,
            BTreeMap::new(),
        )
        .unwrap();

        let err = parametric
            .substitute_one(VariableID::from(100), &Function::from(linear!(0)))
            .unwrap_err();
        assert!(matches!(
            err,
            SubstitutionError::ParameterSubstitution { parameter }
                if parameter == VariableID::from(100)
        ));
    }

    #[test]
    fn test_parametric_instance_substitute_undefined_rhs_fails() {
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(
            VariableID::from(0),
            DecisionVariable::continuous(VariableID::from(0)),
        );

        let parametric = ParametricInstance::new(
            Sense::Minimize,
            Function::from(linear!(0)),
            decision_variables,
            BTreeMap::new(),
            BTreeMap::new(),
        )
        .unwrap();

        let err = parametric
            .substitute_one(VariableID::from(0), &Function::from(linear!(999)))
            .unwrap_err();
        assert!(matches!(
            err,
            SubstitutionError::UndefinedSubstitutionVariable { variable }
                if variable == VariableID::from(999)
        ));
    }

    #[test]
    fn test_constraint_hints_removal_on_substitute() {
        // Create decision variables
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(
            VariableID::from(1),
            DecisionVariable::binary(VariableID::from(1)),
        );
        decision_variables.insert(
            VariableID::from(2),
            DecisionVariable::binary(VariableID::from(2)),
        );
        decision_variables.insert(
            VariableID::from(3),
            DecisionVariable::binary(VariableID::from(3)),
        );

        // Create constraint that depends on x1: x1 + x2 == 1 (constraint_id=1)
        let mut constraints = BTreeMap::new();
        let constraint1_function = Function::from(linear!(1) + linear!(2) + coeff!(-1.0));
        let constraint1 = Constraint {
            id: ConstraintID::from(1),
            function: constraint1_function,
            equality: Equality::EqualToZero,
            name: None,
            subscripts: Vec::new(),
            parameters: Default::default(),
            description: None,
        };
        constraints.insert(ConstraintID::from(1), constraint1);

        // Create constraint that doesn't depend on x1: x2 + x3 == 1 (constraint_id=2)
        let constraint2_function = Function::from(linear!(2) + linear!(3) + coeff!(-1.0));
        let constraint2 = Constraint {
            id: ConstraintID::from(2),
            function: constraint2_function,
            equality: Equality::EqualToZero,
            name: None,
            subscripts: Vec::new(),
            parameters: Default::default(),
            description: None,
        };
        constraints.insert(ConstraintID::from(2), constraint2);

        // Create constraint hints
        let one_hot_for_constraint1 = OneHot {
            id: ConstraintID::from(1),
            variables: [VariableID::from(1), VariableID::from(2)]
                .into_iter()
                .collect(),
        };
        let one_hot_for_constraint2 = OneHot {
            id: ConstraintID::from(2),
            variables: [VariableID::from(2), VariableID::from(3)]
                .into_iter()
                .collect(),
        };
        let sos1 = Sos1 {
            binary_constraint_id: ConstraintID::from(1),
            big_m_constraint_ids: [ConstraintID::from(2)].into_iter().collect(),
            variables: [
                VariableID::from(1),
                VariableID::from(2),
                VariableID::from(3),
            ]
            .into_iter()
            .collect(),
        };

        let constraint_hints = ConstraintHints {
            one_hot_constraints: vec![one_hot_for_constraint1, one_hot_for_constraint2],
            sos1_constraints: vec![sos1],
        };

        // Create objective
        let objective = Function::from(linear!(1) + linear!(2) + linear!(3));

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

        // Before substitution, verify we have 2 OneHot constraints and 1 SOS1 constraint
        assert_eq!(instance.constraint_hints.one_hot_constraints.len(), 2);
        assert_eq!(instance.constraint_hints.sos1_constraints.len(), 1);

        // Substitute x1 with a constant: x1 = 1
        let substitution = Function::from(coeff!(1.0));
        let result = instance
            .substitute_one(VariableID::from(1), &substitution)
            .unwrap();

        // After substitution:
        // - OneHot for constraint1 should be removed (constraint1 depends on x1)
        // - OneHot for constraint2 should remain (constraint2 doesn't depend on x1)
        // - SOS1 should be removed (it references constraint1 which depends on x1)
        assert_eq!(result.constraint_hints.one_hot_constraints.len(), 1);
        assert_eq!(
            result.constraint_hints.one_hot_constraints[0].id,
            ConstraintID::from(2)
        );
        assert_eq!(result.constraint_hints.sos1_constraints.len(), 0);
    }

    #[test]
    fn test_constraint_hints_removal_with_removed_constraints() {
        // Create decision variables
        let mut decision_variables = BTreeMap::new();
        decision_variables.insert(
            VariableID::from(1),
            DecisionVariable::binary(VariableID::from(1)),
        );
        decision_variables.insert(
            VariableID::from(2),
            DecisionVariable::binary(VariableID::from(2)),
        );
        decision_variables.insert(
            VariableID::from(3),
            DecisionVariable::binary(VariableID::from(3)),
        );

        // Create active constraint: x1 + x2 == 1 (constraint_id=1)
        let mut constraints = BTreeMap::new();
        let constraint1_function = Function::from(linear!(1) + linear!(2) + coeff!(-1.0));
        let constraint1 = Constraint {
            id: ConstraintID::from(1),
            function: constraint1_function,
            equality: Equality::EqualToZero,
            name: None,
            subscripts: Vec::new(),
            parameters: Default::default(),
            description: None,
        };
        constraints.insert(ConstraintID::from(1), constraint1);

        // Create removed constraint that depends on x1: x1 + x3 == 1 (constraint_id=2)
        let mut removed_constraints = BTreeMap::new();
        let removed_constraint_function = Function::from(linear!(1) + linear!(3) + coeff!(-1.0));
        let removed_constraint_inner = Constraint {
            id: ConstraintID::from(2),
            function: removed_constraint_function,
            equality: Equality::EqualToZero,
            name: None,
            subscripts: Vec::new(),
            parameters: Default::default(),
            description: None,
        };
        let removed_constraint = RemovedConstraint {
            constraint: removed_constraint_inner,
            removed_reason: "test".to_string(),
            removed_reason_parameters: Default::default(),
        };
        removed_constraints.insert(ConstraintID::from(2), removed_constraint);

        // Create constraint hints - OneHot for removed constraint should also be removed
        let one_hot_for_active = OneHot {
            id: ConstraintID::from(1),
            variables: [VariableID::from(1), VariableID::from(2)]
                .into_iter()
                .collect(),
        };
        let one_hot_for_removed = OneHot {
            id: ConstraintID::from(2),
            variables: [VariableID::from(1), VariableID::from(3)]
                .into_iter()
                .collect(),
        };

        let constraint_hints = ConstraintHints {
            one_hot_constraints: vec![one_hot_for_active, one_hot_for_removed],
            sos1_constraints: vec![],
        };

        // Create objective
        let objective = Function::from(linear!(1) + linear!(2) + linear!(3));

        // Create instance directly since new() only accepts active constraints
        let mut instance =
            Instance::new(Sense::Minimize, objective, decision_variables, constraints).unwrap();

        // Manually set removed constraints and constraint hints
        instance.removed_constraints = removed_constraints;
        instance.constraint_hints = constraint_hints;

        // Before substitution, verify we have 2 OneHot constraints
        assert_eq!(instance.constraint_hints.one_hot_constraints.len(), 2);

        // Substitute x1 with a constant: x1 = 1
        let substitution = Function::from(coeff!(1.0));
        let result = instance
            .substitute_one(VariableID::from(1), &substitution)
            .unwrap();

        // After substitution:
        // Only the OneHot constraint for the active constraint is removed.
        // The hint for the removed constraint remains (it should have been
        // discarded when the constraint was removed, not during substitution).
        assert_eq!(result.constraint_hints.one_hot_constraints.len(), 1);
        assert_eq!(
            result.constraint_hints.one_hot_constraints[0].id,
            ConstraintID::from(2)
        );
    }
}