quantrs2-tytan 0.1.3

//! Compilation of symbolic expressions to QUBO/HOBO models.
//!
//! This module provides utilities for compiling symbolic expressions
//! into QUBO (Quadratic Unconstrained Binary Optimization) and
//! HOBO (Higher-Order Binary Optimization) models.

#![allow(dead_code)]

use scirs2_core::ndarray::Array;
use std::collections::{HashMap, HashSet};

#[cfg(feature = "scirs")]
use crate::scirs_stub;

#[cfg(feature = "dwave")]
use quantrs2_symengine_pure::Expression as SymEngineExpression;

#[cfg(feature = "dwave")]
type Expr = SymEngineExpression;
use thiserror::Error;

use quantrs2_anneal::QuboError;

/// Unified expression interface for examples
#[cfg(feature = "dwave")]
pub mod expr {
    use quantrs2_symengine_pure::Expression as SymEngineExpression;

    pub type Expr = SymEngineExpression;

    pub fn constant(value: f64) -> Expr {
        SymEngineExpression::from(value)
    }

    pub fn var(name: &str) -> Expr {
        SymEngineExpression::symbol(name)
    }
}

#[cfg(not(feature = "dwave"))]
pub mod expr {
    use super::SimpleExpr;

    pub type Expr = SimpleExpr;

    pub const fn constant(value: f64) -> Expr {
        SimpleExpr::constant(value)
    }

    pub fn var(name: &str) -> Expr {
        SimpleExpr::var(name)
    }
}

/// Errors that can occur during compilation
#[derive(Error, Debug)]
pub enum CompileError {
    /// Error when the expression is invalid
    #[error("Invalid expression: {0}")]
    InvalidExpression(String),

    /// Error when a term has too high a degree
    #[error("Term has degree {0}, but maximum supported is {1}")]
    DegreeTooHigh(usize, usize),

    /// Error in the underlying QUBO model
    #[error("QUBO error: {0}")]
    QuboError(#[from] QuboError),

    /// Error in Symengine operations
    #[error("Symengine error: {0}")]
    SymengineError(String),
}

/// Result type for compilation operations
pub type CompileResult<T> = Result<T, CompileError>;

// Simple expression type for when dwave feature is not enabled
#[cfg(not(feature = "dwave"))]
#[derive(Debug, Clone)]
pub enum SimpleExpr {
    /// Variable
    Var(String),
    /// Constant
    Const(f64),
    /// Addition
    Add(Box<Self>, Box<Self>),
    /// Multiplication
    Mul(Box<Self>, Box<Self>),
    /// Power
    Pow(Box<Self>, i32),
}

#[cfg(not(feature = "dwave"))]
impl SimpleExpr {
    /// Create a variable
    pub fn var(name: &str) -> Self {
        Self::Var(name.to_string())
    }

    /// Create a constant
    pub const fn constant(value: f64) -> Self {
        Self::Const(value)
    }
}

#[cfg(not(feature = "dwave"))]
impl std::ops::Add for SimpleExpr {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self::Add(Box::new(self), Box::new(rhs))
    }
}

#[cfg(not(feature = "dwave"))]
impl std::ops::Mul for SimpleExpr {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        Self::Mul(Box::new(self), Box::new(rhs))
    }
}

#[cfg(not(feature = "dwave"))]
impl std::iter::Sum for SimpleExpr {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self::Const(0.0), |acc, x| acc + x)
    }
}

/// High-level model for constraint optimization problems
#[cfg(feature = "dwave")]
#[derive(Debug, Clone)]
pub struct Model {
    /// Variables in the model
    variables: HashSet<String>,
    /// Objective function expression
    objective: Option<Expr>,
    /// Constraints
    constraints: Vec<Constraint>,
}

/// Constraint types
#[cfg(feature = "dwave")]
#[derive(Debug, Clone)]
enum Constraint {
    /// Equality constraint: sum of variables equals value
    Equality {
        name: String,
        expr: Expr,
        value: f64,
    },
    /// Inequality constraint: sum of variables <= value
    LessEqual {
        name: String,
        expr: Expr,
        value: f64,
    },
    /// At most one constraint: at most one variable can be 1
    AtMostOne { name: String, variables: Vec<Expr> },
    /// Implication constraint: if any condition is true, then result must be true
    ImpliesAny {
        name: String,
        conditions: Vec<Expr>,
        result: Expr,
    },
}

#[cfg(feature = "dwave")]
impl Default for Model {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(feature = "dwave")]
impl Model {
    /// Create a new empty model
    pub fn new() -> Self {
        Self {
            variables: HashSet::new(),
            objective: None,
            constraints: Vec::new(),
        }
    }

    /// Add a variable to the model
    pub fn add_variable(&mut self, name: &str) -> CompileResult<Expr> {
        self.variables.insert(name.to_string());
        Ok(SymEngineExpression::symbol(name))
    }

    /// Set the objective function
    pub fn set_objective(&mut self, expr: Expr) {
        self.objective = Some(expr);
    }

    /// Add constraint: exactly one of the variables must be 1
    pub fn add_constraint_eq_one(&mut self, name: &str, variables: Vec<Expr>) -> CompileResult<()> {
        // For binary variables, sum = 1 means exactly one is 1
        let sum_expr = variables
            .iter()
            .fold(Expr::from(0), |acc, v| acc + v.clone());
        self.constraints.push(Constraint::Equality {
            name: name.to_string(),
            expr: sum_expr,
            value: 1.0,
        });
        Ok(())
    }

    /// Add constraint: at most one of the variables can be 1
    pub fn add_constraint_at_most_one(
        &mut self,
        name: &str,
        variables: Vec<Expr>,
    ) -> CompileResult<()> {
        self.constraints.push(Constraint::AtMostOne {
            name: name.to_string(),
            variables,
        });
        Ok(())
    }

    /// Add constraint: if any condition is true, then result must be true
    pub fn add_constraint_implies_any(
        &mut self,
        name: &str,
        conditions: Vec<Expr>,
        result: Expr,
    ) -> CompileResult<()> {
        self.constraints.push(Constraint::ImpliesAny {
            name: name.to_string(),
            conditions,
            result,
        });
        Ok(())
    }

    /// Compile the model to a CompiledModel
    pub fn compile(&self) -> CompileResult<CompiledModel> {
        // Build the final expression with penalty terms
        let mut final_expr = self.objective.clone().unwrap_or_else(|| Expr::from(0));

        // Default penalty weight
        let penalty_weight = 10.0;

        // Add penalty terms for constraints
        for constraint in &self.constraints {
            match constraint {
                Constraint::Equality { expr, value, .. } => {
                    // (expr - value)^2 penalty
                    let diff = expr.clone() - Expr::from(*value);
                    final_expr = final_expr + Expr::from(penalty_weight) * diff.clone() * diff;
                }
                #[cfg(feature = "dwave")]
                Constraint::LessEqual { expr, value, .. } => {
                    // max(0, expr - value)^2 penalty
                    // For simplicity, we'll use a quadratic penalty
                    let excess = expr.clone() - Expr::from(*value);
                    final_expr = final_expr + Expr::from(penalty_weight) * excess.clone() * excess;
                }
                Constraint::AtMostOne { variables, .. } => {
                    // Penalty: sum(xi * xj) for all i < j
                    for i in 0..variables.len() {
                        for j in (i + 1)..variables.len() {
                            final_expr = final_expr
                                + Expr::from(penalty_weight)
                                    * variables[i].clone()
                                    * variables[j].clone();
                        }
                    }
                }
                Constraint::ImpliesAny {
                    conditions, result, ..
                } => {
                    // If any condition is true, result must be true
                    // Penalty: (max(conditions) - result)^2 where max is approximated by sum
                    let conditions_sum = conditions
                        .iter()
                        .fold(Expr::from(0), |acc, c| acc + c.clone());
                    // Penalty when conditions_sum > 0 and result = 0
                    final_expr = final_expr
                        + Expr::from(penalty_weight)
                            * conditions_sum
                            * (Expr::from(1) - result.clone());
                }
            }
        }

        // Use the standard compiler
        let mut compiler = Compile::new(final_expr);
        let ((qubo_matrix, var_map), offset) = compiler.get_qubo()?;

        Ok(CompiledModel {
            qubo_matrix,
            var_map,
            offset,
            constraints: self.constraints.clone(),
        })
    }
}

/// Compiled model ready for sampling
#[cfg(feature = "dwave")]
#[derive(Debug, Clone)]
pub struct CompiledModel {
    /// QUBO matrix
    pub qubo_matrix: Array<f64, scirs2_core::ndarray::Ix2>,
    /// Variable name to index mapping
    pub var_map: HashMap<String, usize>,
    /// Constant offset
    pub offset: f64,
    /// Original constraints (for analysis)
    constraints: Vec<Constraint>,
}

#[cfg(feature = "dwave")]
impl CompiledModel {
    /// Convert to QUBO format
    pub fn to_qubo(&self) -> quantrs2_anneal::ising::QuboModel {
        use quantrs2_anneal::ising::QuboModel;

        let mut qubo = QuboModel::new(self.var_map.len());

        // Set the offset
        qubo.offset = self.offset;

        // Set all the QUBO coefficients
        for i in 0..self.qubo_matrix.nrows() {
            for j in i..self.qubo_matrix.ncols() {
                let value = self.qubo_matrix[[i, j]];
                if value.abs() > 1e-10 {
                    if i == j {
                        // Diagonal term (linear)
                        // SAFETY: index i is derived from matrix dimensions which match QuboModel size
                        qubo.set_linear(i, value)
                            .expect("index within bounds from matrix dimensions");
                    } else {
                        // Off-diagonal term (quadratic)
                        // SAFETY: indices i,j are derived from matrix dimensions which match QuboModel size
                        qubo.set_quadratic(i, j, value)
                            .expect("indices within bounds from matrix dimensions");
                    }
                }
            }
        }

        qubo
    }

    /// Count the number of constraint violations for a given variable assignment.
    ///
    /// Returns the count of constraints that the assignment violates.
    /// `assignments` maps variable names to their binary values (true = 1, false = 0).
    pub fn count_constraint_violations(&self, assignments: &HashMap<String, bool>) -> usize {
        let float_vals: HashMap<String, f64> = assignments
            .iter()
            .map(|(k, &v)| (k.clone(), if v { 1.0 } else { 0.0 }))
            .collect();

        let mut violations = 0usize;

        for constraint in &self.constraints {
            let violated = match constraint {
                Constraint::Equality { expr, value, .. } => match expr.eval(&float_vals) {
                    Ok(result) => (result - value).abs() > 1e-6,
                    Err(_) => false,
                },
                Constraint::LessEqual { expr, value, .. } => match expr.eval(&float_vals) {
                    Ok(result) => result > value + 1e-6,
                    Err(_) => false,
                },
                Constraint::AtMostOne { variables, .. } => {
                    let count: f64 = variables
                        .iter()
                        .filter_map(|v| v.eval(&float_vals).ok())
                        .filter(|&val| val > 0.5)
                        .count() as f64;
                    count > 1.0 + 1e-6
                }
                Constraint::ImpliesAny {
                    conditions, result, ..
                } => {
                    let any_condition_true = conditions
                        .iter()
                        .any(|c| c.eval(&float_vals).map(|val| val > 0.5).unwrap_or(false));
                    if any_condition_true {
                        match result.eval(&float_vals) {
                            Ok(val) => val < 0.5,
                            Err(_) => false,
                        }
                    } else {
                        false
                    }
                }
            };
            if violated {
                violations += 1;
            }
        }

        violations
    }

    /// Return the total number of constraints in this model.
    pub fn num_constraints(&self) -> usize {
        self.constraints.len()
    }
}

/// High-level model for constraint optimization problems (non-dwave version)
#[cfg(not(feature = "dwave"))]
#[derive(Debug, Clone)]
pub struct Model {
    /// Variables in the model
    variables: HashSet<String>,
    /// Objective function expression
    objective: Option<SimpleExpr>,
    /// Constraints
    constraints: Vec<Constraint>,
}

/// Constraint types (non-dwave version)
#[cfg(not(feature = "dwave"))]
#[derive(Debug, Clone)]
enum Constraint {
    /// Equality constraint: sum of variables equals value
    Equality {
        name: String,
        expr: SimpleExpr,
        value: f64,
    },
    /// At most one constraint: at most one variable can be 1
    AtMostOne {
        name: String,
        variables: Vec<SimpleExpr>,
    },
    /// Implication constraint: if any condition is true, then result must be true
    ImpliesAny {
        name: String,
        conditions: Vec<SimpleExpr>,
        result: SimpleExpr,
    },
}

#[cfg(not(feature = "dwave"))]
impl Default for Model {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(not(feature = "dwave"))]
impl Model {
    /// Create a new empty model
    pub fn new() -> Self {
        Self {
            variables: HashSet::new(),
            objective: None,
            constraints: Vec::new(),
        }
    }

    /// Add a variable to the model
    pub fn add_variable(&mut self, name: &str) -> CompileResult<SimpleExpr> {
        self.variables.insert(name.to_string());
        Ok(SimpleExpr::var(name))
    }

    /// Set the objective function
    pub fn set_objective(&mut self, expr: SimpleExpr) {
        self.objective = Some(expr);
    }

    /// Add constraint: exactly one of the variables must be 1
    pub fn add_constraint_eq_one(
        &mut self,
        name: &str,
        variables: Vec<SimpleExpr>,
    ) -> CompileResult<()> {
        let sum_expr = variables.into_iter().sum();
        self.constraints.push(Constraint::Equality {
            name: name.to_string(),
            expr: sum_expr,
            value: 1.0,
        });
        Ok(())
    }

    /// Add constraint: at most one of the variables can be 1
    pub fn add_constraint_at_most_one(
        &mut self,
        name: &str,
        variables: Vec<SimpleExpr>,
    ) -> CompileResult<()> {
        self.constraints.push(Constraint::AtMostOne {
            name: name.to_string(),
            variables,
        });
        Ok(())
    }

    /// Add constraint: if any condition is true, then result must be true
    pub fn add_constraint_implies_any(
        &mut self,
        name: &str,
        conditions: Vec<SimpleExpr>,
        result: SimpleExpr,
    ) -> CompileResult<()> {
        self.constraints.push(Constraint::ImpliesAny {
            name: name.to_string(),
            conditions,
            result,
        });
        Ok(())
    }

    /// Compile the model to a CompiledModel
    pub fn compile(&self) -> CompileResult<CompiledModel> {
        // Build QUBO directly from constraints
        let mut qubo_terms: HashMap<(String, String), f64> = HashMap::new();
        let mut offset = 0.0;
        let penalty_weight = 10.0;

        // Process objective if present
        if let Some(ref obj) = self.objective {
            self.add_expr_to_qubo(obj, 1.0, &mut qubo_terms, &mut offset)?;
        }

        // Process constraints
        for constraint in &self.constraints {
            match constraint {
                Constraint::Equality { expr, value, .. } => {
                    // (expr - value)^2 penalty
                    // Expand: expr^2 - 2*expr*value + value^2
                    self.add_expr_squared_to_qubo(
                        expr,
                        penalty_weight,
                        &mut qubo_terms,
                        &mut offset,
                    )?;
                    self.add_expr_to_qubo(
                        expr,
                        -2.0 * penalty_weight * value,
                        &mut qubo_terms,
                        &mut offset,
                    )?;
                    offset += penalty_weight * value * value;
                }
                Constraint::AtMostOne { variables, .. } => {
                    // Penalty: sum(xi * xj) for all i < j
                    for i in 0..variables.len() {
                        for j in (i + 1)..variables.len() {
                            if let (SimpleExpr::Var(vi), SimpleExpr::Var(vj)) =
                                (&variables[i], &variables[j])
                            {
                                let key = if vi < vj {
                                    (vi.clone(), vj.clone())
                                } else {
                                    (vj.clone(), vi.clone())
                                };
                                *qubo_terms.entry(key).or_insert(0.0) += penalty_weight;
                            }
                        }
                    }
                }
                Constraint::ImpliesAny {
                    conditions, result, ..
                } => {
                    // Penalty: sum(conditions) * (1 - result)
                    for cond in conditions {
                        if let (SimpleExpr::Var(c), SimpleExpr::Var(r)) = (cond, result) {
                            let key = if c < r {
                                (c.clone(), r.clone())
                            } else {
                                (r.clone(), c.clone())
                            };
                            *qubo_terms.entry(key).or_insert(0.0) -= penalty_weight;
                        }
                        // Also add linear term for condition
                        if let SimpleExpr::Var(c) = cond {
                            *qubo_terms.entry((c.clone(), c.clone())).or_insert(0.0) +=
                                penalty_weight;
                        }
                    }
                }
            }
        }

        // Convert to matrix form
        let all_vars: HashSet<String> = qubo_terms
            .keys()
            .flat_map(|(v1, v2)| vec![v1.clone(), v2.clone()])
            .collect();
        let mut sorted_vars: Vec<String> = all_vars.into_iter().collect();
        sorted_vars.sort();

        let var_map: HashMap<String, usize> = sorted_vars
            .iter()
            .enumerate()
            .map(|(i, v)| (v.clone(), i))
            .collect();

        let n = var_map.len();
        let mut matrix = Array::zeros((n, n));

        for ((v1, v2), coeff) in qubo_terms {
            let i = var_map[&v1];
            let j = var_map[&v2];
            if i == j {
                matrix[[i, i]] += coeff;
            } else {
                matrix[[i, j]] += coeff / 2.0;
                matrix[[j, i]] += coeff / 2.0;
            }
        }

        Ok(CompiledModel {
            qubo_matrix: matrix,
            var_map,
            offset,
            constraints: self.constraints.clone(),
        })
    }

    /// Add expression to QUBO terms
    fn add_expr_to_qubo(
        &self,
        expr: &SimpleExpr,
        coeff: f64,
        terms: &mut HashMap<(String, String), f64>,
        offset: &mut f64,
    ) -> CompileResult<()> {
        match expr {
            SimpleExpr::Var(name) => {
                *terms.entry((name.clone(), name.clone())).or_insert(0.0) += coeff;
            }
            SimpleExpr::Const(val) => {
                *offset += coeff * val;
            }
            SimpleExpr::Add(left, right) => {
                self.add_expr_to_qubo(left, coeff, terms, offset)?;
                self.add_expr_to_qubo(right, coeff, terms, offset)?;
            }
            SimpleExpr::Mul(left, right) => {
                if let (SimpleExpr::Var(v1), SimpleExpr::Var(v2)) = (left.as_ref(), right.as_ref())
                {
                    let key = if v1 < v2 {
                        (v1.clone(), v2.clone())
                    } else {
                        (v2.clone(), v1.clone())
                    };
                    *terms.entry(key).or_insert(0.0) += coeff;
                } else if let (SimpleExpr::Const(c), var) | (var, SimpleExpr::Const(c)) =
                    (left.as_ref(), right.as_ref())
                {
                    self.add_expr_to_qubo(var, coeff * c, terms, offset)?;
                }
            }
            SimpleExpr::Pow(base, exp) => {
                if *exp == 2 && matches!(base.as_ref(), SimpleExpr::Var(_)) {
                    // x^2 = x for binary variables
                    self.add_expr_to_qubo(base, coeff, terms, offset)?;
                }
            }
        }
        Ok(())
    }

    /// Add expression squared to QUBO terms
    fn add_expr_squared_to_qubo(
        &self,
        expr: &SimpleExpr,
        coeff: f64,
        terms: &mut HashMap<(String, String), f64>,
        offset: &mut f64,
    ) -> CompileResult<()> {
        // For simplicity, only handle simple cases
        match expr {
            SimpleExpr::Var(name) => {
                // x^2 = x for binary
                *terms.entry((name.clone(), name.clone())).or_insert(0.0) += coeff;
            }
            SimpleExpr::Add(left, right) => {
                // (a + b)^2 = a^2 + 2ab + b^2
                self.add_expr_squared_to_qubo(left, coeff, terms, offset)?;
                self.add_expr_squared_to_qubo(right, coeff, terms, offset)?;
                // Cross term
                if let (SimpleExpr::Var(v1), SimpleExpr::Var(v2)) = (left.as_ref(), right.as_ref())
                {
                    let key = if v1 < v2 {
                        (v1.clone(), v2.clone())
                    } else {
                        (v2.clone(), v1.clone())
                    };
                    *terms.entry(key).or_insert(0.0) += 2.0 * coeff;
                }
            }
            _ => {}
        }
        Ok(())
    }
}

/// Compiled model ready for sampling (non-dwave version)
#[cfg(not(feature = "dwave"))]
#[derive(Debug, Clone)]
pub struct CompiledModel {
    /// QUBO matrix
    pub qubo_matrix: Array<f64, scirs2_core::ndarray::Ix2>,
    /// Variable name to index mapping
    pub var_map: HashMap<String, usize>,
    /// Constant offset
    pub offset: f64,
    /// Original constraints (for analysis)
    constraints: Vec<Constraint>,
}

#[cfg(not(feature = "dwave"))]
impl CompiledModel {
    /// Convert to QUBO format
    pub fn to_qubo(&self) -> quantrs2_anneal::ising::QuboModel {
        use quantrs2_anneal::ising::QuboModel;

        let mut qubo = QuboModel::new(self.var_map.len());

        // Set the offset
        qubo.offset = self.offset;

        // Set all the QUBO coefficients
        for i in 0..self.qubo_matrix.nrows() {
            for j in i..self.qubo_matrix.ncols() {
                let value = self.qubo_matrix[[i, j]];
                if value.abs() > 1e-10 {
                    if i == j {
                        // Diagonal term (linear)
                        // SAFETY: index i is derived from matrix dimensions which match QuboModel size
                        qubo.set_linear(i, value)
                            .expect("index within bounds from matrix dimensions");
                    } else {
                        // Off-diagonal term (quadratic)
                        // SAFETY: indices i,j are derived from matrix dimensions which match QuboModel size
                        qubo.set_quadratic(i, j, value)
                            .expect("indices within bounds from matrix dimensions");
                    }
                }
            }
        }

        qubo
    }
}

/// Compiler for converting symbolic expressions to QUBO models
///
/// This struct provides methods for converting symbolic expressions
/// to QUBO models, which can then be solved using quantum annealing.
#[cfg(feature = "dwave")]
pub struct Compile {
    /// The symbolic expression to compile
    expr: Expr,
}

#[cfg(feature = "dwave")]
impl Compile {
    /// Create a new compiler with the given expression
    pub fn new<T: Into<Expr>>(expr: T) -> Self {
        Self { expr: expr.into() }
    }

    /// Compile the expression to a QUBO model
    ///
    /// This method compiles the symbolic expression to a QUBO model,
    /// which can then be passed to a sampler for solving.
    ///
    /// # Returns
    ///
    /// A tuple containing:
    /// - A tuple with the QUBO matrix and a mapping of variable names to indices
    /// - An offset value that should be added to all energy values
    pub fn get_qubo(
        &self,
    ) -> CompileResult<(
        (
            Array<f64, scirs2_core::ndarray::Ix2>,
            HashMap<String, usize>,
        ),
        f64,
    )> {
        #[cfg(feature = "scirs")]
        {
            self.get_qubo_scirs()
        }
        #[cfg(not(feature = "scirs"))]
        {
            self.get_qubo_standard()
        }
    }

    /// Standard QUBO compilation without SciRS2
    fn get_qubo_standard(
        &self,
    ) -> CompileResult<(
        (
            Array<f64, scirs2_core::ndarray::Ix2>,
            HashMap<String, usize>,
        ),
        f64,
    )> {
        // Expand the expression to simplify
        let expr = self.expr.expand();

        // Replace all second-degree terms (x^2 and x*x) with x, since x^2 = x for binary variables
        // Do this BEFORE degree checking so that x^2 terms correctly appear as degree-1 after reduction
        let expr = replace_squared_terms(&expr)?;

        // Extract the coefficients and variables
        let (coeffs, offset) = extract_coefficients(&expr)?;

        // Check the actual degree using the extracted coefficient map (reliable, symbolic-expression-agnostic)
        let max_degree = coeffs.keys().map(|vars| vars.len()).max().unwrap_or(0);
        if max_degree > 2 {
            return Err(CompileError::DegreeTooHigh(max_degree, 2));
        }

        // Convert to a QUBO matrix
        let (matrix, var_map) = build_qubo_matrix(&coeffs)?;

        Ok(((matrix, var_map), offset))
    }

    /// QUBO compilation with SciRS2 optimization
    #[cfg(feature = "scirs")]
    fn get_qubo_scirs(
        &self,
    ) -> CompileResult<(
        (
            Array<f64, scirs2_core::ndarray::Ix2>,
            HashMap<String, usize>,
        ),
        f64,
    )> {
        // Get standard result
        let ((matrix, var_map), offset) = self.get_qubo_standard()?;

        // Apply SciRS2 enhancements
        let enhanced_matrix = crate::scirs_stub::enhance_qubo_matrix(&matrix);

        Ok(((enhanced_matrix, var_map), offset))
    }

    /// Compile the expression to a HOBO model
    ///
    /// This method compiles the symbolic expression to a Higher-Order Binary Optimization model,
    /// which can handle terms of degree higher than 2.
    ///
    /// # Returns
    ///
    /// A tuple containing:
    /// - A tuple with the HOBO tensor and a mapping of variable names to indices
    /// - An offset value that should be added to all energy values
    pub fn get_hobo(
        &self,
    ) -> CompileResult<(
        (
            Array<f64, scirs2_core::ndarray::IxDyn>,
            HashMap<String, usize>,
        ),
        f64,
    )> {
        // Expand the expression to simplify
        let mut expr = self.expr.expand();

        // Calculate highest degree (dimension of the tensor)
        let max_degree = calc_highest_degree(&expr)?;

        // Replace all squared terms (x^2) with x, since x^2 = x for binary variables
        let mut expr = replace_squared_terms(&expr)?;

        // Expand again to collect like terms
        let mut expr = expr.expand();

        // Extract the coefficients and variables
        let (coeffs, offset) = extract_coefficients(&expr)?;

        // Build the HOBO tensor
        let (tensor, var_map) = build_hobo_tensor(&coeffs, max_degree)?;

        Ok(((tensor, var_map), offset))
    }
}

// Helper function to calculate the highest degree in the expression
#[cfg(feature = "dwave")]
fn calc_highest_degree(expr: &Expr) -> CompileResult<usize> {
    // If the expression is a single variable, it's degree 1
    if expr.is_symbol() {
        return Ok(1);
    }

    // If it's a number constant, degree is 0
    if expr.is_number() {
        return Ok(0);
    }

    // If it's a negation, recursively calculate the degree of the inner expression
    if expr.is_neg() {
        // SAFETY: is_neg() check guarantees as_neg() will succeed
        let inner = expr.as_neg().expect("is_neg() was true");
        return calc_highest_degree(&inner);
    }

    // If it's a power operation (like x^2)
    if expr.is_pow() {
        // SAFETY: is_pow() check guarantees as_pow() will succeed
        let (base, exp) = expr.as_pow().expect("is_pow() was true");

        // If the base is a symbol and exponent is a number
        if base.is_symbol() && exp.is_number() {
            let exp_val = match exp.to_f64() {
                Some(n) => n,
                None => {
                    return Err(CompileError::InvalidExpression(
                        "Invalid exponent".to_string(),
                    ))
                }
            };

            // Check if exponent is a positive integer
            if exp_val.is_sign_positive() && exp_val.fract() == 0.0 {
                return Ok(exp_val as usize);
            }
        }

        // For other power expressions, recursively calculate the degree
        let base_degree = calc_highest_degree(&base)?;
        let exp_degree = if exp.is_number() {
            match exp.to_f64() {
                Some(n) => {
                    if n.is_sign_positive() && n.fract() == 0.0 {
                        n as usize
                    } else {
                        0 // Non-integer or negative exponents don't contribute to degree
                    }
                }
                None => 0,
            }
        } else {
            0 // Non-constant exponents don't contribute to degree
        };

        return Ok(base_degree * exp_degree);
    }

    // If it's a product (like x*y or x*x)
    if expr.is_mul() {
        let mut total_degree = 0;
        // SAFETY: is_mul() check guarantees as_mul() will succeed
        for factor in expr.as_mul().expect("is_mul() was true") {
            total_degree += calc_highest_degree(&factor)?;
        }
        return Ok(total_degree);
    }

    // If it's a sum (like x + y)
    if expr.is_add() {
        let mut max_degree = 0;
        // SAFETY: is_add() check guarantees as_add() will succeed
        for term in expr.as_add().expect("is_add() was true") {
            let term_degree = calc_highest_degree(&term)?;
            max_degree = std::cmp::max(max_degree, term_degree);
        }
        return Ok(max_degree);
    }

    // Check for other compound expressions by trying to parse them
    let expr_str = format!("{expr}");
    if expr_str.contains('+') || expr_str.contains('-') {
        // It's a sum-like expression but not recognized as ADD
        // Parse the string to find the highest degree term
        // This is a workaround for symengine type detection issues
        let mut max_degree = 0;

        // Split by + and - (keeping the sign)
        let parts: Vec<&str> = expr_str.split(['+', '-']).collect();

        for part in parts {
            let part = part.trim();
            if part.is_empty() {
                continue;
            }

            // Count degree based on what the term contains
            let degree = if part.contains("**") || part.contains('^') {
                // Power term like x**2 or y**3
                // Extract the exponent
                let exp_str = part
                    .split("**")
                    .nth(1)
                    .or_else(|| part.split('^').nth(1))
                    .unwrap_or("2")
                    .trim();
                exp_str.parse::<usize>().unwrap_or(2)
            } else if part.contains('*') {
                // Product term - count the number of variables
                let factors: Vec<&str> = part.split('*').collect();
                let mut var_count = 0;
                for factor in factors {
                    let factor = factor.trim();
                    // Check if it's a variable (not a number)
                    if !factor.is_empty() && factor.parse::<f64>().is_err() {
                        var_count += 1;
                    }
                }
                var_count
            } else if part.parse::<f64>().is_err() && !part.is_empty() {
                // Single variable
                1
            } else {
                // Constant
                0
            };

            max_degree = std::cmp::max(max_degree, degree);
        }

        return Ok(max_degree);
    }

    // Default case - for simplicity, we'll say degree is 0
    // but for a complete implementation, we'd need to handle all cases
    Err(CompileError::InvalidExpression(format!(
        "Can't determine degree of: {expr}"
    )))
}

// Helper function to replace squared terms with linear terms
#[cfg(feature = "dwave")]
fn replace_squared_terms(expr: &Expr) -> CompileResult<Expr> {
    // For binary variables, x^2 = x since x ∈ {0,1}

    // If the expression is a symbol or number, just return it
    if expr.is_symbol() || expr.is_number() {
        return Ok(expr.clone());
    }

    // If it's a negation, recursively process the inner expression
    if expr.is_neg() {
        // SAFETY: is_neg() check guarantees as_neg() will succeed
        let inner = expr.as_neg().expect("is_neg() was true");
        let new_inner = replace_squared_terms(&inner)?;
        return Ok(-new_inner);
    }

    // If it's a power operation (like x^2)
    if expr.is_pow() {
        // SAFETY: is_pow() check guarantees as_pow() will succeed
        let (base, exp) = expr.as_pow().expect("is_pow() was true");

        // If the base is a symbol and exponent is 2, replace with base
        if base.is_symbol() && exp.is_number() {
            let exp_val = match exp.to_f64() {
                Some(n) => n,
                None => {
                    return Err(CompileError::InvalidExpression(
                        "Invalid exponent".to_string(),
                    ))
                }
            };

            // Check if exponent is 2 (for higher exponents we'd need to recurse)
            if exp_val == 2.0 {
                return Ok(base);
            }
        }

        // For other power expressions, recursively replace
        let new_base = replace_squared_terms(&base)?;
        return Ok(new_base.pow(&exp));
    }

    // If it's a product (like x*y or x*x)
    if expr.is_mul() {
        let mut new_terms = Vec::new();
        // SAFETY: is_mul() check guarantees as_mul() will succeed
        for factor in expr.as_mul().expect("is_mul() was true") {
            new_terms.push(replace_squared_terms(&factor)?);
        }

        // Check for x*x pattern (same symbol multiplied by itself)
        // For binary variables, x*x = x
        if new_terms.len() == 2 {
            if let (Some(name1), Some(name2)) = (new_terms[0].as_symbol(), new_terms[1].as_symbol())
            {
                if name1 == name2 {
                    // x*x = x for binary variables
                    return Ok(new_terms.remove(0));
                }
            }
        }

        // Combine the terms back into a product (without identity element)
        if new_terms.is_empty() {
            return Ok(Expr::from(1));
        }
        let mut result = new_terms.remove(0);
        for term in new_terms {
            result = result * term;
        }
        return Ok(result);
    }

    // If it's a sum (like x + y)
    if expr.is_add() {
        let mut new_terms = Vec::new();
        // SAFETY: is_add() check guarantees as_add() will succeed
        for term in expr.as_add().expect("is_add() was true") {
            new_terms.push(replace_squared_terms(&term)?);
        }

        // Combine the terms back into a sum (without identity element)
        if new_terms.is_empty() {
            return Ok(Expr::from(0));
        }
        let mut result = new_terms.remove(0);
        for term in new_terms {
            result = result + term;
        }
        return Ok(result);
    }

    // For any other type of expression, just return it unchanged
    Ok(expr.clone())
}

// Helper function to extract coefficients and variables from the expression
#[cfg(feature = "dwave")]
fn extract_coefficients(expr: &Expr) -> CompileResult<(HashMap<Vec<String>, f64>, f64)> {
    let mut coeffs = HashMap::new();
    let mut offset = 0.0;

    // Process expression as a sum of terms
    if expr.is_add() {
        // SAFETY: is_add() check guarantees as_add() will succeed
        for term in expr.as_add().expect("is_add() was true") {
            let (term_coeffs, term_offset) = extract_term_coefficients(&term)?;

            // Merge coefficients
            for (vars, coeff) in term_coeffs {
                *coeffs.entry(vars).or_insert(0.0) += coeff;
            }

            // Add constant terms to offset
            offset += term_offset;
        }
    } else {
        // Check if it's a sum-like expression that wasn't detected as ADD
        let expr_str = format!("{expr}");
        if expr_str.contains('+') || expr_str.contains('-') {
            // Use regex to split properly maintaining signs
            // This is a more robust workaround for symengine type detection issues
            use regex::Regex;
            // SAFETY: Static regex pattern is known to be valid at compile time
            let re = Regex::new(r"([+-]?)([^+-]+)").expect("static regex pattern is valid");

            for caps in re.captures_iter(&expr_str) {
                let sign = caps.get(1).map_or("", |m| m.as_str());
                let term = caps.get(2).map_or("", |m| m.as_str()).trim();

                if term.is_empty() {
                    continue;
                }

                let sign_mult = if sign == "-" { -1.0 } else { 1.0 };

                // Handle x**2 or x^2 (becomes just x for binary)
                if term.contains("**") || term.contains('^') {
                    let base = if term.contains("**") {
                        term.split("**").next().unwrap_or(term)
                    } else {
                        term.split('^').next().unwrap_or(term)
                    }
                    .trim();

                    // Extract coefficient if present (e.g., "10*x^2" -> coeff=10, base="x")
                    let (coeff_mult, var_name) = if base.contains('*') {
                        let parts: Vec<&str> = base.split('*').collect();
                        if parts.len() == 2 {
                            if let Ok(num) = parts[0].trim().parse::<f64>() {
                                (num, parts[1].trim().to_string())
                            } else if let Ok(num) = parts[1].trim().parse::<f64>() {
                                (num, parts[0].trim().to_string())
                            } else {
                                (1.0, base.to_string())
                            }
                        } else {
                            (1.0, base.to_string())
                        }
                    } else {
                        (1.0, base.to_string())
                    };

                    let vars = vec![var_name.clone()];
                    *coeffs.entry(vars).or_insert(0.0) += sign_mult * coeff_mult;
                } else if term.contains('*') {
                    // Handle multiplication: could be "x*y", "2*x", "x*2", "x*y*z", etc.
                    let parts: Vec<&str> = term.split('*').collect();
                    let mut coeff = sign_mult;
                    let mut vars = Vec::new();

                    for part in parts {
                        let part = part.trim();
                        if let Ok(num) = part.parse::<f64>() {
                            coeff *= num;
                        } else {
                            // It's a variable
                            vars.push(part.to_string());
                        }
                    }

                    // Sort variables for consistent ordering
                    vars.sort();
                    *coeffs.entry(vars).or_insert(0.0) += coeff;
                } else if let Ok(num) = term.parse::<f64>() {
                    // Constant term
                    offset += sign_mult * num;
                } else {
                    // Single variable with coefficient 1
                    let vars = vec![term.to_string()];
                    *coeffs.entry(vars).or_insert(0.0) += sign_mult;
                }
            }
            return Ok((coeffs, offset));
        }

        // Only process as a single term if we haven't processed it as ADD yet
        if coeffs.is_empty() {
            // Process a single term
            let (term_coeffs, term_offset) = extract_term_coefficients(expr)?;

            // Merge coefficients
            for (vars, coeff) in term_coeffs {
                *coeffs.entry(vars).or_insert(0.0) += coeff;
            }

            // Add constant terms to offset
            offset += term_offset;
        }
    }

    Ok((coeffs, offset))
}

// Helper function to extract coefficient and variables from a single term
#[cfg(feature = "dwave")]
fn extract_term_coefficients(term: &Expr) -> CompileResult<(HashMap<Vec<String>, f64>, f64)> {
    let mut coeffs = HashMap::new();

    // If it's a number constant, it's an offset
    if term.is_number() {
        let value = match term.to_f64() {
            Some(n) => n,
            None => {
                return Err(CompileError::InvalidExpression(
                    "Invalid number".to_string(),
                ))
            }
        };
        return Ok((coeffs, value));
    }

    // If it's an addition, recursively extract from both sides
    if term.is_add() {
        let mut offset = 0.0;
        // SAFETY: is_add() check guarantees as_add() will succeed
        for sub_term in term.as_add().expect("is_add() was true") {
            let (sub_coeffs, sub_offset) = extract_term_coefficients(&sub_term)?;
            for (vars, coeff) in sub_coeffs {
                *coeffs.entry(vars).or_insert(0.0) += coeff;
            }
            offset += sub_offset;
        }
        return Ok((coeffs, offset));
    }

    // If it's a negation, recursively extract and negate
    if term.is_neg() {
        // SAFETY: is_neg() check guarantees as_neg() will succeed
        let inner = term.as_neg().expect("is_neg() was true");
        let (inner_coeffs, inner_offset) = extract_term_coefficients(&inner)?;

        // Negate all coefficients
        for (vars, coeff) in inner_coeffs {
            coeffs.insert(vars, -coeff);
        }

        return Ok((coeffs, -inner_offset));
    }

    // If it's a symbol, it's a linear term with coefficient 1
    if term.is_symbol() {
        // SAFETY: is_symbol() check guarantees as_symbol() will succeed
        let var_name = term.as_symbol().expect("is_symbol() was true");
        let vars = vec![var_name.to_string()];
        coeffs.insert(vars, 1.0);
        return Ok((coeffs, 0.0));
    }

    // If it's a product of terms
    if term.is_mul() {
        let mut coeff = 1.0;
        let mut vars = Vec::new();

        // SAFETY: is_mul() check guarantees as_mul() will succeed
        let factors = term.as_mul().expect("is_mul() was true");
        // Use a stack to iteratively flatten nested products (handles symengine's
        // internal representation where x*y*z may appear as (* (* x y) z))
        let mut factor_stack: Vec<_> = factors.into_iter().collect();
        while let Some(factor) = factor_stack.pop() {
            if factor.is_number() {
                // Numerical factor is a coefficient
                let value = match factor.to_f64() {
                    Some(n) => n,
                    None => {
                        return Err(CompileError::InvalidExpression(
                            "Invalid number in product".to_string(),
                        ))
                    }
                };
                coeff *= value;
            } else if factor.is_symbol() {
                // Symbol is a variable
                // SAFETY: is_symbol() check guarantees as_symbol() will succeed
                let var_name = factor.as_symbol().expect("is_symbol() was true");
                vars.push(var_name.to_string());
            } else if factor.is_mul() {
                // Nested product — flatten by pushing sub-factors back onto the stack
                let sub_factors = factor.as_mul().expect("is_mul() was true");
                factor_stack.extend(sub_factors);
            } else if factor.is_pow() {
                // Power term like x^2 — for binary vars x^k = x, treat as x
                let (base, exp) = factor.as_pow().expect("is_pow() was true");
                if base.is_symbol() && exp.is_number() {
                    let exp_val = exp.to_f64().unwrap_or(0.0);
                    if exp_val.is_sign_positive() && exp_val.fract() == 0.0 && exp_val >= 1.0 {
                        // Binary variable: x^k = x for k >= 1
                        let var_name = base.as_symbol().expect("is_symbol() was true");
                        vars.push(var_name.to_string());
                    } else {
                        return Err(CompileError::InvalidExpression(format!(
                            "Unsupported power in product: {factor}"
                        )));
                    }
                } else {
                    return Err(CompileError::InvalidExpression(format!(
                        "Unsupported power term in product: {factor}"
                    )));
                }
            } else {
                // More complex factors not supported
                return Err(CompileError::InvalidExpression(format!(
                    "Unsupported term in product: {factor}"
                )));
            }
        }

        // Sort variables for consistent ordering
        vars.sort();

        if vars.is_empty() {
            // If there are no variables, it's a constant term
            return Ok((coeffs, coeff));
        }
        coeffs.insert(vars, coeff);

        return Ok((coeffs, 0.0));
    }

    // If it's a power operation (like x^2), should have been simplified earlier
    if term.is_pow() {
        return Err(CompileError::InvalidExpression(format!(
            "Unexpected power term after simplification: {term}"
        )));
    }

    // Unsupported term type
    Err(CompileError::InvalidExpression(format!(
        "Unsupported term: {term}"
    )))
}

// Helper function to build the QUBO matrix
#[allow(dead_code)]
fn build_qubo_matrix(
    coeffs: &HashMap<Vec<String>, f64>,
) -> CompileResult<(
    Array<f64, scirs2_core::ndarray::Ix2>,
    HashMap<String, usize>,
)> {
    // Collect all unique variable names
    let mut all_vars = HashSet::new();
    for vars in coeffs.keys() {
        for var in vars {
            all_vars.insert(var.clone());
        }
    }

    // Convert to a sorted vector
    let mut sorted_vars: Vec<String> = all_vars.into_iter().collect();
    sorted_vars.sort();

    // Create the variable-to-index mapping
    let var_map: HashMap<String, usize> = sorted_vars
        .iter()
        .enumerate()
        .map(|(i, var)| (var.clone(), i))
        .collect();

    // Size of the matrix
    let n = var_map.len();

    // Create an empty matrix
    let mut matrix = Array::zeros((n, n));

    // Fill the matrix with coefficients
    for (vars, &coeff) in coeffs {
        match vars.len() {
            0 => {
                // Should never happen since constants are handled in offset
            }
            1 => {
                // Linear term: var * coeff
                // SAFETY: var_map was built from the same variables in coeffs
                let i = *var_map
                    .get(&vars[0])
                    .expect("variable exists in var_map built from coeffs");
                matrix[[i, i]] += coeff;
            }
            2 => {
                // Quadratic term: var1 * var2 * coeff
                // SAFETY: var_map was built from the same variables in coeffs
                let i = *var_map
                    .get(&vars[0])
                    .expect("variable exists in var_map built from coeffs");
                let j = *var_map
                    .get(&vars[1])
                    .expect("variable exists in var_map built from coeffs");

                // QUBO format requires i <= j
                if i == j {
                    // Diagonal term
                    matrix[[i, i]] += coeff;
                } else {
                    // Off-diagonal term - store full coefficient in upper triangular, zero in lower
                    if i <= j {
                        matrix[[i, j]] += coeff;
                    } else {
                        matrix[[j, i]] += coeff;
                    }
                }
            }
            _ => {
                // Higher-order terms are not supported in QUBO
                return Err(CompileError::DegreeTooHigh(vars.len(), 2));
            }
        }
    }

    Ok((matrix, var_map))
}

// Helper function to build the HOBO tensor
#[allow(dead_code)]
fn build_hobo_tensor(
    coeffs: &HashMap<Vec<String>, f64>,
    max_degree: usize,
) -> CompileResult<(
    Array<f64, scirs2_core::ndarray::IxDyn>,
    HashMap<String, usize>,
)> {
    // Collect all unique variable names
    let mut all_vars = HashSet::new();
    for vars in coeffs.keys() {
        for var in vars {
            all_vars.insert(var.clone());
        }
    }

    // Convert to a sorted vector
    let mut sorted_vars: Vec<String> = all_vars.into_iter().collect();
    sorted_vars.sort();

    // Create the variable-to-index mapping
    let var_map: HashMap<String, usize> = sorted_vars
        .iter()
        .enumerate()
        .map(|(i, var)| (var.clone(), i))
        .collect();

    // Size of each dimension
    let n = var_map.len();

    // Create shape vector for the tensor
    let shape: Vec<usize> = vec![n; max_degree];

    // Create an empty tensor
    let mut tensor = Array::zeros(scirs2_core::ndarray::IxDyn(&shape));

    // Fill the tensor with coefficients
    for (vars, &coeff) in coeffs {
        let degree = vars.len();

        if degree == 0 {
            // Should never happen since constants are handled in offset
            continue;
        }

        if degree > max_degree {
            return Err(CompileError::DegreeTooHigh(degree, max_degree));
        }

        // Convert variable names to indices
        // SAFETY: var_map was built from the same variables in coeffs
        let mut indices: Vec<usize> = vars
            .iter()
            .map(|var| {
                *var_map
                    .get(var)
                    .expect("variable exists in var_map built from coeffs")
            })
            .collect();

        // Sort indices (canonical ordering)
        indices.sort_unstable();

        // Pad indices to match tensor order if necessary
        while indices.len() < max_degree {
            indices.insert(0, indices[0]); // Padding with first index
        }

        // Set the coefficient in the tensor
        let idx = scirs2_core::ndarray::IxDyn(&indices);
        tensor[idx] += coeff;
    }

    Ok((tensor, var_map))
}

/// Special compiler for problems with one-hot constraints
///
/// This is a specialized compiler that is optimized for problems
/// with one-hot constraints, common in many optimization problems.
#[cfg(feature = "dwave")]
pub struct PieckCompile {
    /// The symbolic expression to compile
    expr: Expr,
    /// Whether to show verbose output
    verbose: bool,
}

#[cfg(feature = "dwave")]
impl PieckCompile {
    /// Create a new Pieck compiler with the given expression
    pub fn new<T: Into<Expr>>(expr: T, verbose: bool) -> Self {
        Self {
            expr: expr.into(),
            verbose,
        }
    }

    /// Compile the expression to a QUBO model optimized for one-hot constraints
    pub fn get_qubo(
        &self,
    ) -> CompileResult<(
        (
            Array<f64, scirs2_core::ndarray::Ix2>,
            HashMap<String, usize>,
        ),
        f64,
    )> {
        // Implementation will compile the expression using specialized techniques
        // For now, call the regular compiler
        Compile::new(self.expr.clone()).get_qubo()
    }
}