quantrs2-anneal 0.2.1

//! Penalty function optimization for quantum annealing
//!
//! This module implements advanced penalty function optimization techniques
//! for improving embedding quality and problem formulation in quantum annealing.
//! It includes methods for optimizing chain strengths, penalty weights, and
//! constraint handling using `SciRS2` optimization algorithms.

use crate::embedding::Embedding;
use crate::ising::{IsingModel, IsingResult, QuboModel};
use std::collections::HashMap;

/// Configuration for penalty optimization
#[derive(Debug, Clone)]
pub struct PenaltyConfig {
    /// Initial chain strength
    pub initial_chain_strength: f64,
    /// Minimum chain strength
    pub min_chain_strength: f64,
    /// Maximum chain strength
    pub max_chain_strength: f64,
    /// Chain strength scaling factor
    pub chain_strength_scale: f64,
    /// Penalty weight for constraint violations
    pub constraint_penalty: f64,
    /// Use adaptive penalty adjustment
    pub adaptive: bool,
    /// Learning rate for adaptive adjustment
    pub learning_rate: f64,
}

impl Default for PenaltyConfig {
    fn default() -> Self {
        Self {
            initial_chain_strength: 1.0,
            min_chain_strength: 0.1,
            max_chain_strength: 10.0,
            chain_strength_scale: 1.5,
            constraint_penalty: 1.0,
            adaptive: true,
            learning_rate: 0.1,
        }
    }
}

/// Statistics from penalty optimization
#[derive(Debug, Clone)]
pub struct PenaltyStats {
    /// Number of optimization iterations
    pub iterations: usize,
    /// Final chain strengths
    pub chain_strengths: HashMap<usize, f64>,
    /// Constraint violation counts
    pub violations: HashMap<String, usize>,
    /// Energy improvement
    pub energy_improvement: f64,
    /// Chain break frequency
    pub chain_break_rate: f64,
}

/// Penalty function optimizer
pub struct PenaltyOptimizer {
    config: PenaltyConfig,
    /// History of chain breaks for adaptive adjustment
    chain_break_history: HashMap<usize, Vec<bool>>,
    /// Constraint violation history
    constraint_history: HashMap<String, Vec<f64>>,
    /// Problem constraints to check violations against. Empty by default:
    /// with no constraints registered, `analyze_samples` honestly reports
    /// zero violations because there is nothing to check, not because
    /// checking was skipped.
    constraints: Vec<Constraint>,
}

impl PenaltyOptimizer {
    /// Create a new penalty optimizer
    #[must_use]
    pub fn new(config: PenaltyConfig) -> Self {
        Self {
            config,
            chain_break_history: HashMap::new(),
            constraint_history: HashMap::new(),
            constraints: Vec::new(),
        }
    }

    /// Register a problem constraint. Registered constraints are checked
    /// against every sample in [`Self::optimize_ising_penalties`] /
    /// [`Self::optimize_qubo_penalties`], and their real violation rates
    /// drive [`Self::update_constraint_penalties`].
    pub fn add_constraint(&mut self, constraint: Constraint) {
        self.constraints.push(constraint);
    }

    /// Create a penalty optimizer pre-populated with problem constraints.
    #[must_use]
    pub fn with_constraints(config: PenaltyConfig, constraints: Vec<Constraint>) -> Self {
        Self {
            config,
            chain_break_history: HashMap::new(),
            constraint_history: HashMap::new(),
            constraints,
        }
    }

    /// Optimize penalty functions for an embedded Ising model
    pub fn optimize_ising_penalties(
        &mut self,
        model: &mut IsingModel,
        embedding: &Embedding,
        samples: &[Vec<i8>],
    ) -> IsingResult<PenaltyStats> {
        let mut stats = PenaltyStats {
            iterations: 0,
            chain_strengths: HashMap::new(),
            violations: HashMap::new(),
            energy_improvement: 0.0,
            chain_break_rate: 0.0,
        };

        // Initialize chain strengths
        for (var, chain) in &embedding.chains {
            stats
                .chain_strengths
                .insert(*var, self.config.initial_chain_strength);
        }

        if self.config.adaptive {
            // Adaptive penalty optimization
            let initial_energy = self.compute_average_energy(model, samples);

            for iteration in 0..10 {
                // Max iterations
                stats.iterations = iteration + 1;

                // Analyze chain breaks and violations
                let (chain_breaks, violations) = self.analyze_samples(samples, embedding);

                // Update chain strengths based on break frequency
                self.update_chain_strengths(&mut stats.chain_strengths, &chain_breaks);

                // Update constraint penalties based on violations
                self.update_constraint_penalties(model, &violations)?;

                // Apply updated penalties
                self.apply_static_penalties(model, embedding, &stats.chain_strengths)?;

                // Check convergence
                let new_energy = self.compute_average_energy(model, samples);
                let improvement = initial_energy - new_energy;

                if improvement.abs() < 0.001 {
                    stats.energy_improvement = improvement;
                    break;
                }
            }
        } else {
            // Static penalty optimization
            self.apply_static_penalties(model, embedding, &stats.chain_strengths)?;
        }

        // Compute final statistics
        stats.chain_break_rate = self.compute_chain_break_rate(samples, embedding);

        Ok(stats)
    }

    /// Optimize penalty functions for a QUBO model
    pub fn optimize_qubo_penalties(
        &mut self,
        model: &mut QuboModel,
        embedding: &Embedding,
        samples: &[Vec<i8>],
    ) -> IsingResult<PenaltyStats> {
        // Convert QUBO to Ising for penalty optimization
        let (mut ising, offset) = model.to_ising();

        // Optimize Ising penalties
        let stats = self.optimize_ising_penalties(&mut ising, embedding, samples)?;

        // Convert back to QUBO
        *model = ising.to_qubo();

        Ok(stats)
    }

    /// Apply static penalties to the model
    fn apply_static_penalties(
        &self,
        model: &mut IsingModel,
        embedding: &Embedding,
        chain_strengths: &HashMap<usize, f64>,
    ) -> IsingResult<()> {
        // Add chain coupling terms
        for (var, chain) in &embedding.chains {
            let strength = chain_strengths
                .get(var)
                .copied()
                .unwrap_or(self.config.initial_chain_strength);

            // Add coupling between all pairs in the chain
            for i in 0..chain.len() {
                for j in (i + 1)..chain.len() {
                    model.set_coupling(chain[i], chain[j], -strength)?;
                }
            }
        }

        Ok(())
    }

    /// Analyze samples for chain breaks and constraint violations.
    ///
    /// `violations` maps each registered [`Constraint`]'s name to its real
    /// violation rate over `samples` (the fraction of samples that fail
    /// [`constraint_is_satisfied`]) — not a placeholder. Optimizers that
    /// haven't registered any constraints via [`Self::add_constraint`] /
    /// [`Self::with_constraints`] will honestly get back an empty map.
    fn analyze_samples(
        &self,
        samples: &[Vec<i8>],
        embedding: &Embedding,
    ) -> (HashMap<usize, f64>, HashMap<String, f64>) {
        let mut chain_breaks = HashMap::new();
        let mut violations = HashMap::new();

        for sample in samples {
            // Check chain integrity
            for (var, chain) in &embedding.chains {
                let mut broken = false;
                if chain.len() > 1 {
                    let first_val = sample[chain[0]];
                    for &qubit in &chain[1..] {
                        if sample[qubit] != first_val {
                            broken = true;
                            break;
                        }
                    }
                }

                let count = chain_breaks.entry(*var).or_insert(0.0);
                if broken {
                    *count += 1.0;
                }
            }
        }

        // Normalize by number of samples
        let n = samples.len() as f64;
        for count in chain_breaks.values_mut() {
            *count /= n;
        }

        // Real per-constraint violation rates over the registered constraints.
        if !samples.is_empty() {
            for constraint in &self.constraints {
                let violation_count = samples
                    .iter()
                    .filter(|sample| !constraint_is_satisfied(constraint, sample))
                    .count();
                violations.insert(
                    constraint.name.clone(),
                    violation_count as f64 / samples.len() as f64,
                );
            }
        }

        (chain_breaks, violations)
    }

    /// Update chain strengths based on break frequency
    fn update_chain_strengths(
        &self,
        chain_strengths: &mut HashMap<usize, f64>,
        chain_breaks: &HashMap<usize, f64>,
    ) {
        for (var, break_rate) in chain_breaks {
            if let Some(strength) = chain_strengths.get_mut(var) {
                if *break_rate > 0.1 {
                    // More than 10% breaks
                    // Increase chain strength
                    *strength = (*strength * self.config.chain_strength_scale)
                        .min(self.config.max_chain_strength);
                } else if *break_rate < 0.01 {
                    // Less than 1% breaks
                    // Decrease chain strength (might be too strong)
                    *strength = (*strength / self.config.chain_strength_scale)
                        .max(self.config.min_chain_strength);
                }
            }
        }
    }

    /// Update constraint penalties based on real, measured violation rates.
    ///
    /// For every registered constraint whose violation rate exceeds
    /// `self.config.learning_rate` (used here as an adaptive tolerance
    /// threshold), reinforces the standard quadratic penalty for
    /// `(sum_{i in variables} s_i - target)^2` on `model`: a linear bias
    /// term `-2 * weight * target` on each variable in the constraint, plus
    /// a pairwise coupling `+2 * weight` between every pair of the
    /// constraint's variables, where
    /// `weight = self.config.constraint_penalty * self.config.learning_rate * violation_rate`
    /// scales with how often the constraint was actually violated. This
    /// is the standard QUBO/Ising equality-constraint encoding; inequality
    /// constraint types (`LessEqual`/`GreaterEqual`) are approximated with
    /// the same toward-target formula rather than a full slack-variable
    /// encoding.
    fn update_constraint_penalties(
        &self,
        model: &mut IsingModel,
        violations: &HashMap<String, f64>,
    ) -> IsingResult<()> {
        for constraint in &self.constraints {
            let Some(&violation_rate) = violations.get(&constraint.name) else {
                continue;
            };
            if violation_rate <= self.config.learning_rate {
                continue;
            }

            let weight =
                self.config.constraint_penalty * self.config.learning_rate * violation_rate;
            if weight.abs() < 1e-12 {
                continue;
            }

            for &var in &constraint.variables {
                let current_bias = model.get_bias(var).unwrap_or(0.0);
                model.set_bias(
                    var,
                    (-2.0 * weight).mul_add(constraint.target, current_bias),
                )?;
            }
            for i in 0..constraint.variables.len() {
                for j in (i + 1)..constraint.variables.len() {
                    let (vi, vj) = (constraint.variables[i], constraint.variables[j]);
                    let current_coupling = model.get_coupling(vi, vj).unwrap_or(0.0);
                    model.set_coupling(vi, vj, current_coupling + 2.0 * weight)?;
                }
            }
        }

        Ok(())
    }

    /// Compute average energy of samples
    fn compute_average_energy(&self, model: &IsingModel, samples: &[Vec<i8>]) -> f64 {
        let mut total_energy = 0.0;

        for sample in samples {
            // Ignore errors for invalid samples
            if let Ok(energy) = model.energy(sample) {
                total_energy += energy;
            }
        }

        total_energy / samples.len() as f64
    }

    /// Compute chain break rate across all samples
    fn compute_chain_break_rate(&self, samples: &[Vec<i8>], embedding: &Embedding) -> f64 {
        let mut total_breaks = 0;
        let mut total_chains = 0;

        for sample in samples {
            for (_var, chain) in &embedding.chains {
                if chain.len() > 1 {
                    total_chains += 1;
                    let first_val = sample[chain[0]];
                    for &qubit in &chain[1..] {
                        if sample[qubit] != first_val {
                            total_breaks += 1;
                            break;
                        }
                    }
                }
            }
        }

        if total_chains > 0 {
            f64::from(total_breaks) / f64::from(total_chains)
        } else {
            0.0
        }
    }
}

/// Advanced penalty optimization using SciRS2-style optimization
pub struct AdvancedPenaltyOptimizer {
    /// Base penalty optimizer
    base_optimizer: PenaltyOptimizer,
    /// Use gradient-based optimization
    use_gradients: bool,
    /// Regularization parameter
    regularization: f64,
}

impl AdvancedPenaltyOptimizer {
    /// Create a new advanced penalty optimizer
    #[must_use]
    pub fn new(config: PenaltyConfig) -> Self {
        Self {
            base_optimizer: PenaltyOptimizer::new(config),
            use_gradients: true,
            regularization: 0.01,
        }
    }

    /// Optimize penalties using gradient descent
    pub fn optimize_with_gradients(
        &mut self,
        model: &mut IsingModel,
        embedding: &Embedding,
        samples: &[Vec<i8>],
        max_iterations: usize,
    ) -> IsingResult<PenaltyStats> {
        let mut chain_strengths: HashMap<usize, f64> = embedding
            .chains
            .keys()
            .map(|&var| (var, self.base_optimizer.config.initial_chain_strength))
            .collect();

        let mut best_energy = f64::INFINITY;
        let mut best_strengths = chain_strengths.clone();

        for iteration in 0..max_iterations {
            // Compute gradients with respect to chain strengths
            let gradients = self.compute_gradients(model, embedding, samples, &chain_strengths)?;

            // Update chain strengths using gradient descent
            for (var, strength) in &mut chain_strengths {
                if let Some(&grad) = gradients.get(var) {
                    let new_strength = self
                        .base_optimizer
                        .config
                        .learning_rate
                        .mul_add(-grad, *strength);
                    *strength = new_strength
                        .max(self.base_optimizer.config.min_chain_strength)
                        .min(self.base_optimizer.config.max_chain_strength);
                }
            }

            // Apply penalties and evaluate
            self.base_optimizer
                .apply_static_penalties(model, embedding, &chain_strengths)?;
            let energy = self.base_optimizer.compute_average_energy(model, samples);

            if energy < best_energy {
                best_energy = energy;
                best_strengths = chain_strengths.clone();
            }

            // Check convergence
            if iteration > 0 && (best_energy - energy).abs() < 1e-6 {
                break;
            }
        }

        Ok(PenaltyStats {
            iterations: max_iterations,
            chain_strengths: best_strengths,
            violations: HashMap::new(),
            energy_improvement: 0.0,
            chain_break_rate: self
                .base_optimizer
                .compute_chain_break_rate(samples, embedding),
        })
    }

    /// Compute gradients of the objective with respect to chain strengths
    fn compute_gradients(
        &self,
        model: &IsingModel,
        embedding: &Embedding,
        samples: &[Vec<i8>],
        chain_strengths: &HashMap<usize, f64>,
    ) -> IsingResult<HashMap<usize, f64>> {
        let mut gradients = HashMap::new();
        let epsilon = 0.01;

        // Numerical gradient computation
        for (var, &current_strength) in chain_strengths {
            // Forward difference
            let mut strengths_plus = chain_strengths.clone();
            strengths_plus.insert(*var, current_strength + epsilon);

            let energy_plus = {
                let mut model_copy = model.clone();
                self.base_optimizer.apply_static_penalties(
                    &mut model_copy,
                    embedding,
                    &strengths_plus,
                )?;
                self.base_optimizer
                    .compute_average_energy(&model_copy, samples)
            };

            // Backward difference
            let mut strengths_minus = chain_strengths.clone();
            strengths_minus.insert(*var, current_strength - epsilon);

            let energy_minus = {
                let mut model_copy = model.clone();
                self.base_optimizer.apply_static_penalties(
                    &mut model_copy,
                    embedding,
                    &strengths_minus,
                )?;
                self.base_optimizer
                    .compute_average_energy(&model_copy, samples)
            };

            // Gradient with regularization
            let gradient = self.regularization.mul_add(
                current_strength,
                (energy_plus - energy_minus) / (2.0 * epsilon),
            );
            gradients.insert(*var, gradient);
        }

        Ok(gradients)
    }
}

/// Penalty optimization for constrained problems
pub struct ConstraintPenaltyOptimizer {
    /// Constraint definitions
    constraints: Vec<Constraint>,
    /// Penalty weights for each constraint
    penalty_weights: HashMap<String, f64>,
    /// Violation tolerance
    tolerance: f64,
}

/// Represents a constraint in the optimization problem
#[derive(Debug, Clone)]
pub struct Constraint {
    /// Constraint name
    pub name: String,
    /// Variables involved in the constraint
    pub variables: Vec<usize>,
    /// Constraint type
    pub constraint_type: ConstraintType,
    /// Target value
    pub target: f64,
}

/// Types of constraints
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConstraintType {
    /// Equality constraint (sum = target)
    Equality,
    /// Less than or equal constraint (sum <= target)
    LessEqual,
    /// Greater than or equal constraint (sum >= target)
    GreaterEqual,
    /// Exactly one constraint (exactly one variable is 1)
    ExactlyOne,
    /// At most one constraint (at most one variable is 1)
    AtMostOne,
}

/// Check whether `sample` satisfies `constraint`. Shared by
/// [`ConstraintPenaltyOptimizer::check_constraint`] and
/// [`PenaltyOptimizer::analyze_samples`] so both real constraint-checking
/// paths in this module agree on the same semantics.
fn constraint_is_satisfied(constraint: &Constraint, sample: &[i8]) -> bool {
    let sum: i8 = constraint
        .variables
        .iter()
        .map(|&var| sample.get(var).copied().unwrap_or(0))
        .sum();

    match constraint.constraint_type {
        ConstraintType::Equality => (f64::from(sum) - constraint.target).abs() < 1e-6,
        ConstraintType::LessEqual => f64::from(sum) <= constraint.target,
        ConstraintType::GreaterEqual => f64::from(sum) >= constraint.target,
        ConstraintType::ExactlyOne => sum == 1,
        ConstraintType::AtMostOne => sum <= 1,
    }
}

impl ConstraintPenaltyOptimizer {
    /// Create a new constraint penalty optimizer
    #[must_use]
    pub fn new(tolerance: f64) -> Self {
        Self {
            constraints: Vec::new(),
            penalty_weights: HashMap::new(),
            tolerance,
        }
    }

    /// Add a constraint
    pub fn add_constraint(&mut self, constraint: Constraint) {
        let default_weight = 1.0;
        self.penalty_weights
            .insert(constraint.name.clone(), default_weight);
        self.constraints.push(constraint);
    }

    /// Optimize penalty weights for constraints
    pub fn optimize_penalties(
        &mut self,
        samples: &[Vec<i8>],
        max_iterations: usize,
    ) -> HashMap<String, f64> {
        for _ in 0..max_iterations {
            // Analyze constraint violations
            let violations = self.analyze_constraint_violations(samples);

            // Update penalty weights based on violations
            for (constraint_name, violation_rate) in violations {
                if let Some(weight) = self.penalty_weights.get_mut(&constraint_name) {
                    if violation_rate > self.tolerance {
                        *weight *= 1.5; // Increase penalty
                    } else if violation_rate < self.tolerance / 10.0 {
                        *weight *= 0.8; // Decrease penalty
                    }
                }
            }
        }

        self.penalty_weights.clone()
    }

    /// Analyze constraint violations in samples
    fn analyze_constraint_violations(&self, samples: &[Vec<i8>]) -> HashMap<String, f64> {
        let mut violations = HashMap::new();

        for constraint in &self.constraints {
            let mut violation_count = 0;

            for sample in samples {
                if !self.check_constraint(constraint, sample) {
                    violation_count += 1;
                }
            }

            let violation_rate = f64::from(violation_count) / samples.len() as f64;
            violations.insert(constraint.name.clone(), violation_rate);
        }

        violations
    }

    /// Check if a sample satisfies a constraint
    fn check_constraint(&self, constraint: &Constraint, sample: &[i8]) -> bool {
        constraint_is_satisfied(constraint, sample)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_penalty_optimizer_creation() {
        let config = PenaltyConfig::default();
        let optimizer = PenaltyOptimizer::new(config);
        assert!(optimizer.chain_break_history.is_empty());
    }

    #[test]
    fn test_chain_break_analysis() {
        let config = PenaltyConfig::default();
        let optimizer = PenaltyOptimizer::new(config);

        // Create a simple embedding
        let mut embedding = Embedding::new();
        embedding.chains.insert(0, vec![0, 1]);
        embedding.chains.insert(1, vec![2, 3]);

        // Create samples with some chain breaks
        let samples = vec![
            vec![1, 1, -1, -1], // No breaks
            vec![1, -1, 1, 1],  // Break in chain 0
            vec![1, 1, 1, -1],  // Break in chain 1
        ];

        let rate = optimizer.compute_chain_break_rate(&samples, &embedding);
        assert!(rate > 0.0 && rate < 1.0);
    }

    #[test]
    fn test_constraint_checking() {
        let mut optimizer = ConstraintPenaltyOptimizer::new(0.1);

        // Add an equality constraint
        optimizer.add_constraint(Constraint {
            name: "sum_equals_one".to_string(),
            variables: vec![0, 1, 2],
            constraint_type: ConstraintType::Equality,
            target: 1.0,
        });

        // Test samples
        let sample1 = vec![1, 0, 0]; // Satisfies
        let sample2 = vec![1, 1, 0]; // Violates

        assert!(optimizer.check_constraint(&optimizer.constraints[0], &sample1));
        assert!(!optimizer.check_constraint(&optimizer.constraints[0], &sample2));
    }

    #[test]
    fn analyze_samples_reports_real_constraint_violation_rates() {
        let config = PenaltyConfig::default();
        let mut optimizer = PenaltyOptimizer::new(config);
        optimizer.add_constraint(Constraint {
            name: "balanced_pair".to_string(),
            variables: vec![0, 1],
            constraint_type: ConstraintType::Equality,
            target: 0.0, // For +-1 spins, sum==0 means one spin up and one down.
        });

        let embedding = Embedding::new();
        let samples = vec![
            vec![1, -1],  // sum=0 -> satisfies
            vec![1, 1],   // sum=2 -> violates
            vec![-1, -1], // sum=-2 -> violates
            vec![-1, 1],  // sum=0 -> satisfies
        ];

        let (_chain_breaks, violations) = optimizer.analyze_samples(&samples, &embedding);
        let rate = violations
            .get("balanced_pair")
            .copied()
            .expect("registered constraint must be tracked, not silently dropped");
        assert!(
            (rate - 0.5).abs() < 1e-9,
            "expected the real 50% violation rate, got {rate}"
        );
    }

    #[test]
    fn analyze_samples_reports_no_violations_when_no_constraints_are_registered() {
        // Honest zero: nothing was registered, so nothing is violated -- this
        // must be distinguishable in intent from the old placeholder (which
        // never checked anything regardless of registration).
        let config = PenaltyConfig::default();
        let optimizer = PenaltyOptimizer::new(config);
        let embedding = Embedding::new();
        let samples = vec![vec![1, 1], vec![-1, -1]];

        let (_chain_breaks, violations) = optimizer.analyze_samples(&samples, &embedding);
        assert!(violations.is_empty());
    }

    #[test]
    fn update_constraint_penalties_actually_modifies_the_model_when_violated() {
        let config = PenaltyConfig {
            learning_rate: 0.05,
            constraint_penalty: 2.0,
            ..PenaltyConfig::default()
        };

        let mut optimizer = PenaltyOptimizer::new(config);
        optimizer.add_constraint(Constraint {
            name: "sum_equals_one".to_string(),
            variables: vec![0, 1],
            constraint_type: ConstraintType::Equality,
            target: 1.0,
        });

        let mut model = IsingModel::new(2);
        let embedding = Embedding::new();
        // Every sample violates the constraint (sum=2, not 1): a 100%
        // violation rate, well above the learning-rate tolerance.
        let samples = vec![vec![1, 1], vec![1, 1], vec![1, 1], vec![1, 1]];

        let bias_before = model.get_bias(0).expect("bias should be readable");
        let coupling_before = model
            .get_coupling(0, 1)
            .expect("coupling should be readable");

        let (_chain_breaks, violations) = optimizer.analyze_samples(&samples, &embedding);
        optimizer
            .update_constraint_penalties(&mut model, &violations)
            .expect("penalty update should succeed");

        let bias_after = model.get_bias(0).expect("bias should be readable");
        let coupling_after = model
            .get_coupling(0, 1)
            .expect("coupling should be readable");

        // With a non-zero target, the toward-target quadratic penalty
        // contributes both a linear bias term (`-2*weight*target`) and a
        // pairwise coupling term (`+2*weight`); both must actually move.
        assert_ne!(
            bias_after, bias_before,
            "a real, heavily-violated constraint must actually adjust the model's bias"
        );
        assert_ne!(
            coupling_after, coupling_before,
            "a real, heavily-violated constraint must actually adjust the model's coupling"
        );
    }

    #[test]
    fn update_constraint_penalties_leaves_model_untouched_when_within_tolerance() {
        let config = PenaltyConfig {
            learning_rate: 0.5, // High tolerance: a 25% violation rate should be ignored.
            constraint_penalty: 2.0,
            ..PenaltyConfig::default()
        };

        let mut optimizer = PenaltyOptimizer::new(config);
        optimizer.add_constraint(Constraint {
            name: "balanced_pair".to_string(),
            variables: vec![0, 1],
            constraint_type: ConstraintType::Equality,
            target: 0.0,
        });

        let mut model = IsingModel::new(2);
        let embedding = Embedding::new();
        let samples = vec![vec![1, -1], vec![1, -1], vec![1, -1], vec![1, 1]];

        let (_chain_breaks, violations) = optimizer.analyze_samples(&samples, &embedding);
        optimizer
            .update_constraint_penalties(&mut model, &violations)
            .expect("penalty update should succeed");

        assert_eq!(model.get_bias(0).expect("bias should be readable"), 0.0);
        assert_eq!(
            model
                .get_coupling(0, 1)
                .expect("coupling should be readable"),
            0.0
        );
    }
}