semitone 0.0.30

use crate::{
    TheoryConstraint, out_of_bounds,
    rational::{InfRational, Rational},
    sat_solver::Lit,
};
use rug::Assign;
use rustc_hash::{FxHashMap, FxHashSet};
use std::{collections::BTreeMap, mem};
use tracing::trace;

pub(super) struct LraTheory {
    ints: Vec<bool>,                                      // true = integer variable, false = real variable
    reals: Vec<InfRational>,                              // Current assignments
    pub(super) lbs: Vec<(Option<Lit>, InfRational)>,      // Current lower bounds
    pub(super) ubs: Vec<(Option<Lit>, InfRational)>,      // Current upper bounds
    pub(super) lin_to_slack: FxHashMap<SparseRow, usize>, // Mapping from linear constraints to their slack variable
    pub(super) tableau: BTreeMap<usize, SparseRow>,       // Tableau: basic variable -> linear expression over non-basic variables
    pub(super) t_watches: Vec<FxHashSet<usize>>,          // For each variable, the set of tableau rows containing it
    bound_trail: Vec<BoundUpdate>,                        // Trail of bound updates for backtracking
    trail_lim: Vec<usize>,                                // Trail limits
    rational_pool: Vec<rug::Rational>,                    // Pool of rational numbers for reuse
}

impl LraTheory {
    pub(super) fn new() -> Self {
        Self {
            ints: Vec::new(),
            reals: Vec::new(),
            lbs: Vec::new(),
            ubs: Vec::new(),
            lin_to_slack: FxHashMap::default(),
            tableau: BTreeMap::new(),
            t_watches: Vec::new(),
            bound_trail: Vec::new(),
            trail_lim: Vec::new(),
            rational_pool: Vec::new(),
        }
    }

    /// Creates a new integer variable and returns its index.
    pub(super) fn mk_int(&mut self) -> usize {
        self.mk_var(true)
    }

    /// Creates a new real variable and returns its index.
    pub(super) fn mk_real(&mut self) -> usize {
        self.mk_var(false)
    }

    /// Creates a new variable (integer or real) and returns its index.
    fn mk_var(&mut self, is_int: bool) -> usize {
        let var = self.reals.len();
        trace!("Creating new variable {}{var}", if is_int { "i" } else { "r" });
        self.ints.push(is_int);
        self.reals.push(InfRational::zero());
        self.lbs.push((None, InfRational::negative_inf()));
        self.ubs.push((None, InfRational::positive_inf()));
        self.t_watches.push(FxHashSet::default());
        var
    }

    /// Checks if a variable is an integer variable.
    pub fn is_int_var(&self, var: usize) -> bool {
        self.ints[var]
    }

    /// Checks if a linear expression consists only of integer variables with integer coefficients.
    pub fn is_int_expr(&self, expr: &SparseRow) -> bool {
        expr.iter().all(|(var, coeff)| self.ints[*var] && coeff.is_integer())
    }

    #[inline]
    pub(super) fn value(&self, var: usize) -> &InfRational {
        self.reals.get(var).unwrap_or_else(|| out_of_bounds(var))
    }

    fn row_value(&self, row: &SparseRow) -> InfRational {
        row.iter().fold(InfRational::zero(), |mut acc, (v, c)| {
            acc += &self.reals[*v] * c;
            acc
        })
    }

    pub(super) fn get_or_create_slack(&mut self, vars: SparseRow) -> usize {
        if let Some(&slack) = self.lin_to_slack.get(&vars) {
            return slack;
        }

        let is_int_slack = self.is_int_expr(&vars);
        let slack = self.mk_var(is_int_slack);

        self.reals[slack] = self.row_value(&vars);

        self.tableau.insert(slack, vars.clone());
        for &var in vars.keys() {
            self.t_watches[var].insert(slack);
        }
        self.lin_to_slack.insert(vars, slack);

        slack
    }

    #[inline]
    pub(super) fn lb(&self, var: usize) -> &InfRational {
        &self.lbs.get(var).unwrap_or_else(|| out_of_bounds(var)).1
    }

    #[inline]
    pub(super) fn ub(&self, var: usize) -> &InfRational {
        &self.ubs.get(var).unwrap_or_else(|| out_of_bounds(var)).1
    }

    pub(super) fn set_lb(&mut self, lit: Option<Lit>, var: usize, new_lb: InfRational) -> Result<bool, Vec<Lit>> {
        trace!("{}{}{} ≥ {}", if let Some(l) = lit { format!("[{l}] ") } else { "".to_string() }, if self.is_int_var(var) { "i" } else { "r" }, var, new_lb);
        assert!(var < self.reals.len(), "variable index out of bounds: {var}");

        if &new_lb <= self.lb(var) {
            return Ok(false);
        }

        if &new_lb > self.ub(var) {
            return Err([lit, self.ubs[var].0].into_iter().flatten().map(|l| !l).collect());
        }

        let (c_lit, val) = self.lbs[var].clone();
        self.bound_trail.push(BoundUpdate::LowerBound { lit: c_lit, var, val });
        self.lbs[var] = (lit, new_lb.clone());

        if self.value(var) < &new_lb && !self.is_basic(var) {
            self.update(var, new_lb);
        }
        Ok(true)
    }

    pub(super) fn set_ub(&mut self, lit: Option<Lit>, var: usize, new_ub: InfRational) -> Result<bool, Vec<Lit>> {
        trace!("{}{}{} ≤ {}", if let Some(l) = lit { format!("[{l}] ") } else { "".to_string() }, if self.is_int_var(var) { "i" } else { "r" }, var, new_ub);
        assert!(var < self.reals.len(), "variable index out of bounds: {var}");

        if &new_ub >= self.ub(var) {
            return Ok(false);
        }

        if &new_ub < self.lb(var) {
            return Err([lit, self.lbs[var].0].into_iter().flatten().map(|l| !l).collect());
        }

        let (c_lit, val) = self.ubs[var].clone();
        self.bound_trail.push(BoundUpdate::UpperBound { lit: c_lit, var, val });
        self.ubs[var] = (lit, new_ub.clone());

        if self.value(var) > &new_ub && !self.is_basic(var) {
            self.update(var, new_ub);
        }
        Ok(true)
    }

    fn is_basic(&self, var: usize) -> bool {
        self.tableau.contains_key(&var)
    }

    fn update(&mut self, var: usize, new_value: InfRational) {
        assert!(var < self.reals.len(), "variable index out of bounds: {var}");
        assert!(!self.is_basic(var), "cannot directly update a basic variable");
        assert!(&new_value >= self.lb(var) && &new_value <= self.ub(var), "new value must be within bounds");

        let old_value = self.reals[var].clone();
        let delta_var = &new_value - &old_value;

        if delta_var.is_zero() {
            return;
        }

        let watched_rows: Vec<usize> = self.t_watches[var].iter().copied().collect();
        for row_var in watched_rows {
            let coeff = self.tableau[&row_var].get(&var).expect("watched variable must occur in tableau row").clone();
            let delta = &delta_var * &coeff;
            self.reals[row_var] += delta;
        }

        self.reals[var] = new_value;
    }

    pub fn check(&mut self) -> Result<(), Vec<Lit>> {
        loop {
            // Find a basic variable outside its bounds (is_below indicates a lower-bound violation)
            let var = self.tableau.keys().find_map(|&var| {
                let val = self.value(var);
                if val < self.lb(var) {
                    Some((var, self.lb(var).clone(), true))
                } else if val > self.ub(var) {
                    Some((var, self.ub(var).clone(), false))
                } else {
                    None
                }
            });

            let Some((leaving, target_val, is_below)) = var else {
                return Ok(());
            };

            // Find a pivot variable based on the required direction
            let entering = self.tableau[&leaving].iter().find_map(|(v, coeff)| {
                let moves_up = is_below == coeff.is_positive();
                let can_move = if moves_up { self.value(*v) < self.ub(*v) } else { self.value(*v) > self.lb(*v) };
                can_move.then_some(v)
            });

            if let Some(entering) = entering {
                self.pivot_and_update(*entering, leaving, target_val);
            } else {
                // Build the conflict clause compactly
                let mut conflict = Vec::new();
                for (vr, coeff) in &self.tableau[&leaving] {
                    let guard = if is_below == coeff.is_positive() { self.ubs[*vr].0 } else { self.lbs[*vr].0 };
                    if let Some(l) = guard {
                        conflict.push(!l);
                    }
                }

                let self_guard = if is_below { self.lbs[leaving].0 } else { self.ubs[leaving].0 };
                if let Some(l) = self_guard {
                    conflict.push(!l);
                }

                self.minimize_conflict(&mut conflict);
                return Err(conflict);
            }
        }
    }

    /// Minimizes the conflict by sorting and removing duplicates.
    fn minimize_conflict(&self, conflict: &mut Vec<Lit>) {
        if conflict.len() <= 1 {
            return;
        }
        conflict.sort_unstable();
        conflict.dedup();
    }

    fn pivot(&mut self, entering: usize, leaving: usize) {
        assert!(entering < self.reals.len(), "variable index out of bounds: {entering}");
        assert!(leaving < self.reals.len(), "variable index out of bounds: {leaving}");
        assert!(self.is_basic(leaving), "leaving variable must be basic");
        assert!(!self.is_basic(entering), "entering variable must be non-basic");

        let leaving_row_vars: Vec<usize> = self.tableau[&leaving].keys().copied().collect();
        for var in leaving_row_vars {
            self.t_watches[var].remove(&leaving);
        }

        let mut new_row = self.tableau.remove(&leaving).expect("leaving variable must have a tableau row");
        let pivot_coeff = new_row.remove(&entering).expect("entering variable must occur in leaving row");
        assert!(!pivot_coeff.is_zero(), "pivot coefficient must be non-zero");

        let inv_pivot = (-pivot_coeff.clone()).recip();
        for coeff in new_row.values_mut() {
            *coeff *= &inv_pivot;
        }

        new_row.insert(leaving, pivot_coeff.recip());

        let affected_rows: Vec<usize> = mem::take(&mut self.t_watches[entering]).into_iter().collect();

        for row_var in affected_rows {
            if row_var == leaving {
                continue;
            }
            let Some(row) = self.tableau.get_mut(&row_var) else {
                continue;
            };
            let Some(coeff_entering) = row.remove(&entering) else {
                continue;
            };

            row.add_scaled(&new_row, &coeff_entering, &mut self.t_watches, row_var, &mut self.rational_pool);
        }

        for v in new_row.keys().copied() {
            self.t_watches[v].insert(entering);
        }

        self.tableau.insert(entering, new_row);
    }

    /// Updates the values of the variables after a pivot operation.
    fn pivot_and_update(&mut self, entering: usize, leaving: usize, new_value: InfRational) {
        assert!(entering < self.reals.len(), "variable index out of bounds: {entering}");
        assert!(leaving < self.reals.len(), "variable index out of bounds: {leaving}");
        assert!(self.is_basic(leaving), "leaving variable must be basic");
        assert!(!self.is_basic(entering), "entering variable must be non-basic");
        assert!(&new_value >= self.lb(leaving) && &new_value <= self.ub(leaving), "new value for leaving variable must be within bounds");

        let pivot_coeff = self.tableau[&leaving].get(&entering).expect("entering variable must occur in leaving row").clone();
        assert!(!pivot_coeff.is_zero(), "pivot coefficient must be non-zero");

        let theta = (&new_value - self.value(leaving)) / &pivot_coeff;

        self.reals[leaving] = new_value;
        self.reals[entering] += &theta;

        let affected_rows: Vec<usize> = self.t_watches[entering].iter().copied().collect();
        for row_var in affected_rows {
            if row_var == leaving {
                continue;
            }
            let Some(row) = self.tableau.get(&row_var) else {
                continue;
            };
            let Some(coeff) = row.get(&entering) else {
                continue;
            };

            self.reals[row_var] += &theta * coeff;
        }

        self.pivot(entering, leaving);
    }

    /// Checks if all integer variables have integer values.
    pub(super) fn check_ints(&self) -> Result<(), (usize, InfRational)> {
        if let Some((var, val)) = self.ints.iter().enumerate().filter(|(_, is_int)| **is_int).find_map(|(var, _)| {
            let val = self.value(var);
            (!val.infinitesimal_part().is_zero() || !val.rational_part().is_integer()).then_some((var, val))
        }) {
            return Err((var, val.clone()));
        }
        Ok(())
    }

    pub(super) fn new_gomory_cut(&self) -> Option<(SparseRow, rug::Rational, Vec<TheoryConstraint>)> {
        for (b_var, row) in &self.tableau {
            if !self.ints[*b_var] {
                continue;
            }

            let val = self.value(*b_var);
            if !val.infinitesimal_part().is_zero() {
                continue;
            }

            let rat_part = val.rational_part();
            let Rational::Finite(rat_val) = rat_part else {
                continue;
            };

            if rat_val.is_integer() {
                continue;
            }

            let mut floor_val = rat_val.clone();
            floor_val.floor_mut();
            let f0 = rat_val.clone() - floor_val;

            if f0.is_zero() || f0 == 1 {
                continue;
            }

            let mut valid_row = true;
            let mut cut_row = SparseRow::new();
            let mut cut_rhs = rug::Rational::from(1);
            let one_minus_f0 = rug::Rational::from(1) - &f0;

            let mut explanation = Vec::new();

            for (nb_var, nb_coeff) in row.iter() {
                let nb_val = self.value(*nb_var);

                if !nb_val.infinitesimal_part().is_zero() {
                    valid_row = false;
                    break;
                }

                let lb = self.lb(*nb_var);
                let ub = self.ub(*nb_var);

                let at_lower = nb_val == lb;
                let at_upper = nb_val == ub;

                if !at_lower && !at_upper {
                    valid_row = false;
                    break;
                }

                let nb_rat_part = nb_val.rational_part();
                let Rational::Finite(nb_rat_val) = nb_rat_part else {
                    valid_row = false;
                    break;
                };

                if at_lower {
                    if self.lbs[*nb_var].0.is_some() {
                        explanation.push(TheoryConstraint::LraLb(*nb_var, lb.clone()));
                    }
                } else if self.ubs[*nb_var].0.is_some() {
                    explanation.push(TheoryConstraint::LraUb(*nb_var, ub.clone()));
                }

                let is_int_var = self.ints[*nb_var];
                let cut_coeff = if at_lower {
                    if !is_int_var {
                        if *nb_coeff >= 0 { nb_coeff.clone() / &one_minus_f0 } else { -nb_coeff.clone() / &f0 }
                    } else {
                        let mut floor_aj = nb_coeff.clone();
                        floor_aj.floor_mut();
                        let fj = nb_coeff.clone() - floor_aj;
                        if fj <= one_minus_f0 { fj / &one_minus_f0 } else { (rug::Rational::from(1) - fj) / &f0 }
                    }
                } else {
                    if !is_int_var {
                        if *nb_coeff >= 0 { nb_coeff.clone() / &f0 } else { -nb_coeff.clone() / &one_minus_f0 }
                    } else {
                        let mut floor_aj = nb_coeff.clone();
                        floor_aj.floor_mut();
                        let fj = nb_coeff.clone() - floor_aj;
                        if fj <= f0 { fj / &f0 } else { (rug::Rational::from(1) - fj) / &one_minus_f0 }
                    }
                };

                if !cut_coeff.is_zero() {
                    let term = rug::Rational::from(&cut_coeff * nb_rat_val);
                    if at_lower {
                        cut_rhs += term;
                        cut_row.add_coeff(*nb_var, &cut_coeff);
                    } else {
                        cut_rhs -= term;
                        let neg_coeff = -cut_coeff;
                        cut_row.add_coeff(*nb_var, &neg_coeff);
                    }
                }
            }

            if valid_row && !cut_row.is_empty() {
                return Some((cut_row, cut_rhs, explanation));
            }
        }
        None
    }

    pub(super) fn push(&mut self) {
        self.trail_lim.push(self.bound_trail.len());
    }

    pub(super) fn cancel_until(&mut self, level: usize) {
        if level >= self.trail_lim.len() {
            return;
        }
        let target_len = self.trail_lim[level];

        for update in self.bound_trail.drain(target_len..).rev() {
            match update {
                BoundUpdate::LowerBound { lit, var, val } => self.lbs[var] = (lit, val),
                BoundUpdate::UpperBound { lit, var, val } => self.ubs[var] = (lit, val),
            }
        }

        self.trail_lim.truncate(level);
    }

    /// Optimizes the given objective function, returning the optimal value.
    #[allow(dead_code)]
    pub(super) fn optimize(&mut self, objective: SparseRow, maximize: bool) -> InfRational {
        let mut obj_row = self.canonicalize(objective);

        loop {
            let mut entering: Option<(usize, bool)> = None;
            let mut best_magnitude: Option<rug::Rational> = None;

            for (v, coeff) in obj_row.iter() {
                if !coeff.is_positive() && !coeff.is_negative() {
                    continue;
                }

                let dir = maximize == coeff.is_positive();

                let movable = if dir { self.value(*v) < self.ub(*v) } else { self.value(*v) > self.lb(*v) };
                if !movable {
                    continue;
                }

                let magnitude = if coeff.is_positive() { coeff.clone() } else { -coeff.clone() };

                if best_magnitude.as_ref().is_none_or(|cur| magnitude > *cur) {
                    best_magnitude = Some(magnitude);
                    entering = Some((*v, dir));
                }
            }

            let Some((entering_var, dir)) = entering else {
                return self.row_value(&obj_row);
            };

            let mut best_target: Option<InfRational> = None;
            let mut leaving: Option<usize> = None;
            let mut leaving_target: Option<InfRational> = None;

            let self_bound = if dir { self.ub(entering_var) } else { self.lb(entering_var) }.clone();

            if matches!(self_bound.rational_part(), Rational::Finite(_)) {
                best_target = Some(self_bound.clone());
                leaving_target = Some(self_bound);
            }

            let watched_rows: Vec<usize> = self.t_watches[entering_var].iter().copied().collect();
            for row_var in watched_rows {
                let coeff = self.tableau[&row_var].get(&entering_var).expect("watched variable must occur in tableau row");

                let moves_up = dir == coeff.is_positive();
                let bound = if moves_up { self.ub(row_var) } else { self.lb(row_var) }.clone();

                if !matches!(bound.rational_part(), Rational::Finite(_)) {
                    continue;
                }

                let signed_delta = (bound.clone() - self.value(row_var).clone()) / coeff;
                let target = self.value(entering_var).clone() + signed_delta;

                let is_better = match &best_target {
                    None => true,
                    Some(cur) => {
                        let better_val = if dir { target < *cur } else { target > *cur };
                        better_val || (target == *cur && leaving.is_some_and(|l| row_var < l))
                    }
                };

                if is_better {
                    best_target = Some(target);
                    leaving = Some(row_var);
                    leaving_target = Some(bound);
                }
            }

            match (leaving, leaving_target) {
                (None, None) => return if maximize { InfRational::positive_inf() } else { InfRational::negative_inf() },
                (None, Some(new_value)) => {
                    self.update(entering_var, new_value); // bound flip, obj_row remains valid
                }
                (Some(leaving_var), Some(new_value)) => {
                    self.pivot_and_update(entering_var, leaving_var, new_value);
                    let coeff = obj_row.remove(&entering_var).expect("entering variable must occur in the objective row");
                    let new_row = &self.tableau[&entering_var];
                    obj_row.add_scaled_untracked(new_row, &coeff);
                }
                (Some(_), None) => unreachable!(),
            }
        }
    }

    /// Canonicalizes a row by eliminating basic variables.
    #[allow(dead_code)]
    fn canonicalize(&self, mut row: SparseRow) -> SparseRow {
        loop {
            let Some(&(basic_var, _)) = row.iter().find(|(v, _)| self.is_basic(*v)) else {
                return row;
            };
            let coeff = row.remove(&basic_var).expect("just found in row");
            let basic_row = &self.tableau[&basic_var];
            for (v, c) in basic_row.iter() {
                let mut delta = rug::Rational::from(0);
                delta.assign(c * &coeff);
                row.add_coeff(*v, &delta);
            }
        }
    }
}

pub(super) fn tighten_int_bound(mut r: rug::Rational, is_upper: bool, strict: bool) -> rug::Rational {
    if is_upper {
        if strict {
            r.ceil_mut();
            r - rug::Rational::from(1)
        } else {
            r.floor_mut();
            r
        }
    } else {
        if strict {
            r.floor_mut();
            r + rug::Rational::from(1)
        } else {
            r.ceil_mut();
            r
        }
    }
}

#[derive(Clone, Default, Debug, PartialEq, Eq, Hash)]
pub(super) struct SparseRow {
    pub terms: Vec<(usize, rug::Rational)>,
}

impl SparseRow {
    pub fn new() -> Self {
        Self { terms: Vec::new() }
    }

    pub fn is_empty(&self) -> bool {
        self.terms.is_empty()
    }

    pub fn len(&self) -> usize {
        self.terms.len()
    }

    pub fn iter(&self) -> std::slice::Iter<'_, (usize, rug::Rational)> {
        self.terms.iter()
    }

    pub fn get(&self, var: &usize) -> Option<&rug::Rational> {
        self.terms.binary_search_by_key(var, |&(v, _)| v).ok().map(|idx| &self.terms[idx].1)
    }

    pub fn insert(&mut self, var: usize, coeff: rug::Rational) {
        match self.terms.binary_search_by_key(&var, |&(v, _)| v) {
            Ok(idx) => self.terms[idx].1 = coeff,
            Err(idx) => self.terms.insert(idx, (var, coeff)),
        }
    }

    pub fn remove(&mut self, var: &usize) -> Option<rug::Rational> {
        self.terms.binary_search_by_key(var, |&(v, _)| v).ok().map(|idx| self.terms.remove(idx).1)
    }

    pub fn retain<F>(&mut self, mut f: F)
    where
        F: FnMut(&usize, &mut rug::Rational) -> bool,
    {
        self.terms.retain_mut(|(v, c)| f(v, c));
    }

    pub fn keys(&self) -> impl Iterator<Item = &usize> {
        self.terms.iter().map(|(v, _)| v)
    }

    pub fn values_mut(&mut self) -> impl Iterator<Item = &mut rug::Rational> {
        self.terms.iter_mut().map(|(_, c)| c)
    }

    pub fn add_coeff(&mut self, var: usize, delta: &rug::Rational) {
        if delta.is_zero() {
            return;
        }
        match self.terms.binary_search_by_key(&var, |&(v, _)| v) {
            Ok(idx) => {
                self.terms[idx].1 += delta;
            }
            Err(idx) => {
                self.terms.insert(idx, (var, delta.clone()));
            }
        }
    }

    #[allow(dead_code)]
    fn add_scaled_untracked(&mut self, other: &SparseRow, scale: &rug::Rational) {
        let old_terms = std::mem::take(&mut self.terms);
        let mut new_terms = Vec::with_capacity(old_terms.len());

        let mut old_iter = old_terms.into_iter().peekable();
        let mut other_iter = other.terms.iter().peekable();

        while let (Some((v1, _)), Some((v2, _))) = (old_iter.peek(), other_iter.peek()) {
            if v1 < v2 {
                new_terms.push(old_iter.next().unwrap());
            } else if v1 > v2 {
                let (v2, c2) = other_iter.next().unwrap();
                let mut delta = rug::Rational::from(0);
                delta.assign(c2 * scale);
                if !delta.is_zero() {
                    new_terms.push((*v2, delta));
                }
            } else {
                let (v1, mut c1) = old_iter.next().unwrap();
                let (_, c2) = other_iter.next().unwrap();
                let mut tmp = rug::Rational::from(0);
                tmp.assign(c2 * scale);
                c1 += tmp;
                if !c1.is_zero() {
                    new_terms.push((v1, c1));
                }
            }
        }

        new_terms.extend(old_iter);
        for (v2, c2) in other_iter {
            let mut delta = rug::Rational::from(0);
            delta.assign(c2 * scale);
            if !delta.is_zero() {
                new_terms.push((*v2, delta));
            }
        }

        self.terms = new_terms;
    }

    fn add_scaled(&mut self, other: &SparseRow, scale: &rug::Rational, watches: &mut [FxHashSet<usize>], target_row_var: usize, pool: &mut Vec<rug::Rational>) {
        let old_terms = std::mem::take(&mut self.terms);
        let mut new_terms = Vec::with_capacity(old_terms.len());

        let mut old_iter = old_terms.into_iter().peekable();
        let mut other_iter = other.terms.iter().peekable();

        while let (Some(&(v1, _)), Some(&(v2, _))) = (old_iter.peek(), other_iter.peek()) {
            if v1 < *v2 {
                new_terms.push(old_iter.next().unwrap());
            } else if v1 > *v2 {
                let (v2, c2) = other_iter.next().unwrap();

                let mut delta = pool.pop().unwrap_or_else(|| rug::Rational::from(0));
                delta.assign(c2 * scale);

                if !delta.is_zero() {
                    watches[*v2].insert(target_row_var);
                    new_terms.push((*v2, delta));
                } else {
                    pool.push(delta);
                }
            } else {
                let (v1, mut c1) = old_iter.next().unwrap();
                let (_, c2) = other_iter.next().unwrap();

                let mut tmp_delta = pool.pop().unwrap_or_else(|| rug::Rational::from(0));
                tmp_delta.assign(c2 * scale);

                c1 += &tmp_delta;
                pool.push(tmp_delta);

                if c1.is_zero() {
                    watches[v1].remove(&target_row_var);
                    pool.push(c1);
                } else {
                    new_terms.push((v1, c1));
                }
            }
        }

        new_terms.extend(old_iter);

        for (v2, c2) in other_iter {
            let mut delta = pool.pop().unwrap_or_else(|| rug::Rational::from(0));
            delta.assign(c2 * scale);

            if !delta.is_zero() {
                watches[*v2].insert(target_row_var);
                new_terms.push((*v2, delta));
            } else {
                pool.push(delta);
            }
        }

        self.terms = new_terms;
    }
}

impl<'a> IntoIterator for &'a SparseRow {
    type Item = &'a (usize, rug::Rational);
    type IntoIter = std::slice::Iter<'a, (usize, rug::Rational)>;

    fn into_iter(self) -> Self::IntoIter {
        self.terms.iter()
    }
}

enum BoundUpdate {
    LowerBound { lit: Option<Lit>, var: usize, val: InfRational },
    UpperBound { lit: Option<Lit>, var: usize, val: InfRational },
}

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

    fn add_test_row(theory: &mut LraTheory, basic_var: usize, terms: &[(usize, i32)]) {
        let mut row = SparseRow::new();
        for &(var, coeff) in terms {
            row.insert(var, rug::Rational::from(coeff));
            theory.t_watches[var].insert(basic_var);
        }
        theory.tableau.insert(basic_var, row);
    }

    fn build_row(terms: &[(usize, i32)]) -> SparseRow {
        let mut row = SparseRow::new();
        for &(var, coeff) in terms {
            row.insert(var, rug::Rational::from(coeff));
        }
        row
    }

    #[test]
    fn test_pure_pivot_algebra() {
        let mut lra = LraTheory::new();

        let x = lra.mk_real(); // 0
        let y = lra.mk_real(); // 1
        let s = lra.mk_real(); // 2 (slack)

        add_test_row(&mut lra, s, &[(x, 2), (y, -3)]);

        assert!(lra.is_basic(s));
        assert!(!lra.is_basic(x));
        assert!(lra.t_watches[x].contains(&s));

        lra.pivot(y, s);

        assert!(!lra.is_basic(s));
        assert!(lra.is_basic(y));
        assert!(!lra.t_watches[y].contains(&s), "y should not be watching s after pivoting");
        assert!(lra.t_watches[x].contains(&y), "x should now be watched by y after pivoting");
        assert!(lra.t_watches[s].contains(&y), "s should now be watched by y after pivoting");

        let row_y = lra.tableau.get(&y).expect("y must have a row in the tableau");
        assert_eq!(row_y.get(&x).unwrap(), &rug::Rational::from((2, 3)));
        assert_eq!(row_y.get(&s).unwrap(), &rug::Rational::from((-1, 3)));
    }

    #[test]
    fn test_pivot_and_update_maintains_equality() {
        let mut lra = LraTheory::new();

        let x = lra.mk_real();
        let y = lra.mk_real();
        let s = lra.mk_real();

        add_test_row(&mut lra, s, &[(x, 1), (y, 1)]);

        lra.reals[x] = InfRational::from(5);
        lra.reals[y] = InfRational::from(3);
        lra.reals[s] = InfRational::from(8);

        lra.set_lb(None, s, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, s, InfRational::from(10)).expect("setting upper bound should succeed");

        lra.pivot_and_update(y, s, InfRational::from(6));

        assert_eq!(lra.value(s), &InfRational::from(6));
        assert_eq!(lra.value(x), &InfRational::from(5));
        assert_eq!(lra.value(y), &InfRational::from(1), "y should have absorbed the delta of -2");
    }

    #[test]
    fn test_check_resolves_out_of_bounds() {
        let mut lra = LraTheory::new();

        let x = lra.mk_real();
        let y = lra.mk_real();
        let s = lra.mk_real();

        add_test_row(&mut lra, s, &[(x, 1), (y, 1)]); // s = x + y

        lra.set_lb(None, x, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, x, InfRational::from(10)).expect("setting upper bound should succeed");
        lra.set_lb(None, y, InfRational::from(-10)).expect("setting lower bound should succeed");
        lra.set_ub(None, y, InfRational::from(10)).expect("setting upper bound should succeed");
        lra.set_ub(None, s, InfRational::from(5)).expect("setting upper bound should succeed");

        lra.set_lb(None, x, InfRational::from(6)).expect("setting lower bound should succeed");

        assert!(lra.value(s) > lra.ub(s));

        let result = lra.check();

        assert!(result.is_ok(), "check should resolve the out-of-bounds situation");

        assert_eq!(lra.value(s), &InfRational::from(5));
        assert_eq!(lra.value(y), &InfRational::from(-1));

        assert!(!lra.is_basic(s));
        assert!(lra.value(s) <= lra.ub(s));
    }

    #[test]
    fn test_check_detects_conflict() {
        let mut lra = LraTheory::new();

        let x = lra.mk_real();
        let y = lra.mk_real();
        let s = lra.mk_real();

        add_test_row(&mut lra, s, &[(x, 1), (y, 1)]);

        lra.set_lb(Some(Lit::new(1, false)), x, InfRational::from(3)).expect("setting lower bound should succeed");
        assert_eq!(lra.value(s), &InfRational::from(3));

        lra.set_lb(Some(Lit::new(2, false)), y, InfRational::from(4)).expect("setting lower bound should succeed");
        assert_eq!(lra.value(s), &InfRational::from(7));

        lra.set_ub(Some(Lit::new(3, false)), s, InfRational::from(5)).expect("setting upper bound should succeed");

        let result = lra.check();

        assert!(result.is_err());
        let conflict = result.unwrap_err();

        assert!(conflict.contains(&Lit::new(1, true)));
        assert!(conflict.contains(&Lit::new(2, true)));
        assert!(conflict.contains(&Lit::new(3, true)));
    }

    #[test]
    fn test_sparse_row_add_scaled_cancellation() {
        let rational_pool = &mut Vec::new();
        // row1 = 2*v0 + 3*v1 - 1*v3
        let mut row1 = build_row(&[(0, 2), (1, 3), (3, -1)]);
        // row2 = 1*v1 + 4*v2 + 2*v3
        let row2 = build_row(&[(1, 1), (2, 4), (3, 2)]);

        let mut watches = vec![FxHashSet::default(); 4];
        let target_row = 10;

        watches[0].insert(target_row);
        watches[1].insert(target_row);
        watches[3].insert(target_row);

        let scale = rug::Rational::from(-3);
        row1.add_scaled(&row2, &scale, &mut watches, target_row, rational_pool);

        assert_eq!(row1.len(), 3, "The row should have exactly 3 active terms");
        assert_eq!(row1.get(&0), Some(&rug::Rational::from(2)));
        assert_eq!(row1.get(&1), None, "v1 should have been removed from the row");
        assert_eq!(row1.get(&2), Some(&rug::Rational::from(-12)));
        assert_eq!(row1.get(&3), Some(&rug::Rational::from(-7)));

        assert!(watches[0].contains(&target_row), "The watch for v0 should remain unchanged");
        assert!(!watches[1].contains(&target_row), "The watch for v1 should have been removed because the coefficient became zero");
        assert!(watches[2].contains(&target_row), "The watch for v2 should have been added dynamically");
        assert!(watches[3].contains(&target_row), "The watch for v3 should remain unchanged after the coefficient update");
    }

    #[test]
    fn test_tableau_pivot_nested_substitution() {
        let mut lra = LraTheory::new();

        let x = lra.mk_real(); // 0
        let y = lra.mk_real(); // 1
        let z = lra.mk_real(); // 2

        let s1 = lra.mk_real(); // 3 (slack 1)
        let s2 = lra.mk_real(); // 4 (slack 2)

        // s1 = 2x + 1y - 1z
        add_test_row(&mut lra, s1, &[(x, 2), (y, 1), (z, -1)]);
        // s2 = 1x - 1y + 2z
        add_test_row(&mut lra, s2, &[(x, 1), (y, -1), (z, 2)]);

        lra.pivot(x, s2);

        let row_s1 = lra.tableau.get(&s1).expect("s1 should still be present in the tableau");

        assert_eq!(row_s1.get(&s2), Some(&rug::Rational::from(2)));
        assert_eq!(row_s1.get(&y), Some(&rug::Rational::from(3)));
        assert_eq!(row_s1.get(&z), Some(&rug::Rational::from(-5)));

        assert_eq!(row_s1.get(&x), None, "x is now basic, so it cannot appear in s1");

        // Verify that the reverse dependencies (watches) remain consistent.
        assert!(!lra.t_watches[x].contains(&s1), "s1 should no longer watch x");
        assert!(lra.t_watches[s2].contains(&s1), "s1 should now watch s2");
        assert!(lra.t_watches[y].contains(&s1));
        assert!(lra.t_watches[z].contains(&s1));
    }

    #[test]
    fn test_optimize_maximize_simple_bound() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0

        lra.set_lb(None, x, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, x, InfRational::from(10)).expect("setting upper bound should succeed");

        let obj = build_row(&[(x, 1)]);
        let result = lra.optimize(obj, true);

        assert_eq!(result, InfRational::from(10));
        assert_eq!(lra.value(x), &InfRational::from(10));
    }

    #[test]
    fn test_optimize_minimize_simple_bound() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0

        lra.set_lb(None, x, InfRational::from(-5)).expect("setting lower bound should succeed");
        lra.set_ub(None, x, InfRational::from(10)).expect("setting upper bound should succeed");

        let obj = build_row(&[(x, 1)]);
        let result = lra.optimize(obj, false);

        assert_eq!(result, InfRational::from(-5));
        assert_eq!(lra.value(x), &InfRational::from(-5));
    }

    #[test]
    fn test_optimize_unbounded() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0

        lra.set_lb(None, x, InfRational::from(0)).expect("setting lower bound should succeed");
        // No upper bound: remains +inf by default

        let obj = build_row(&[(x, 1)]);
        let result = lra.optimize(obj, true);

        assert_eq!(result, InfRational::positive_inf());
    }

    #[test]
    fn test_optimize_requires_pivot() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0
        let y = lra.mk_real(); // 1
        let s = lra.mk_real(); // 2, basic: s = x + y

        add_test_row(&mut lra, s, &[(x, 1), (y, 1)]);

        lra.set_lb(None, x, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, x, InfRational::from(100)).expect("setting upper bound should succeed");
        lra.set_lb(None, y, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, y, InfRational::from(100)).expect("setting upper bound should succeed");
        lra.set_lb(None, s, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, s, InfRational::from(5)).expect("setting upper bound should succeed");

        let obj = build_row(&[(x, 1)]);
        let result = lra.optimize(obj, true);

        assert_eq!(result, InfRational::from(5));
        assert_eq!(lra.value(x), &InfRational::from(5));
        assert_eq!(lra.value(s), &InfRational::from(5));
        assert_eq!(lra.value(y), &InfRational::from(0), "y should have absorbed the delta to maintain s = x + y");
        assert!(lra.is_basic(x), "x should be basic after pivoting");
        assert!(!lra.is_basic(s), "s should no longer be basic after pivoting");
    }

    #[test]
    fn test_optimize_objective_on_basic_variable() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0
        let y = lra.mk_real(); // 1
        let s = lra.mk_real(); // 2, basic: s = x + 2y

        add_test_row(&mut lra, s, &[(x, 1), (y, 2)]);

        lra.set_lb(None, x, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, x, InfRational::from(3)).expect("setting upper bound should succeed");
        lra.set_lb(None, y, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, y, InfRational::from(3)).expect("setting upper bound should succeed");

        let obj = build_row(&[(s, 1)]);
        let result = lra.optimize(obj, true);

        assert_eq!(result, InfRational::from(9)); // x=3, y=3 => s = 3 + 2*3 = 9
        assert_eq!(lra.value(x), &InfRational::from(3));
        assert_eq!(lra.value(y), &InfRational::from(3));
    }

    #[test]
    fn test_optimize_degenerate_tie_picks_lowest_index() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0
        let s1 = lra.mk_real(); // 1, basic: s1 = x
        let s2 = lra.mk_real(); // 2, basic: s2 = x

        add_test_row(&mut lra, s1, &[(x, 1)]);
        add_test_row(&mut lra, s2, &[(x, 1)]);

        lra.set_lb(None, x, InfRational::from(0)).expect("setting lower bound should succeed");
        lra.set_ub(None, s1, InfRational::from(5)).expect("setting upper bound should succeed");
        lra.set_ub(None, s2, InfRational::from(5)).expect("setting upper bound should succeed");

        let obj = build_row(&[(x, 1)]);
        let result = lra.optimize(obj, true);

        assert_eq!(result, InfRational::from(5));
        assert_eq!(lra.value(x), &InfRational::from(5));
        assert_eq!(lra.value(s1), &InfRational::from(5));
        assert_eq!(lra.value(s2), &InfRational::from(5));
    }

    #[test]
    fn test_canonicalize_substitutes_basic_variables() {
        let mut lra = LraTheory::new();
        let x = lra.mk_real(); // 0
        let y = lra.mk_real(); // 1
        let s = lra.mk_real(); // 2, basic: s = x + 2y

        add_test_row(&mut lra, s, &[(x, 1), (y, 2)]);

        let row = build_row(&[(s, 3)]);
        let canon = lra.canonicalize(row);

        assert_eq!(canon.get(&s), None, "s should be eliminated from the canonicalized row");
        assert_eq!(canon.get(&x), Some(&rug::Rational::from(3)));
        assert_eq!(canon.get(&y), Some(&rug::Rational::from(6)));
    }
}