theory_core 0.1.0

//! Z3 SMT Solver Integration
//!
//! Provides real SMT solving capabilities using the Z3 solver.
//! This module is only available when the `z3-solver` feature is enabled.

#[cfg(feature = "z3-solver")]
use z3::{Config, Context, Optimize, SatResult, Solver, ast};

use crate::ir::{BinaryOp, Constraint, ConstraintIR, Expr, VarId};
use crate::surrogate::{SolutionResult, SolutionValue};
use anyhow::{Result, anyhow};
use std::collections::HashMap;

/// Z3-based SMT solver backend
pub struct Z3Backend {
    #[cfg(feature = "z3-solver")]
    context: Context,
    timeout_ms: u32,
}

impl Z3Backend {
    /// Create a new Z3 backend
    pub fn new(timeout_ms: u32) -> Result<Self> {
        #[cfg(feature = "z3-solver")]
        {
            let mut cfg = Config::new();
            cfg.set_timeout_msec(timeout_ms as u64);
            let context = Context::new(&cfg);
            Ok(Self {
                context,
                timeout_ms,
            })
        }

        #[cfg(not(feature = "z3-solver"))]
        {
            Err(anyhow!(
                "Z3 solver not available. Enable the 'z3-solver' feature to use Z3."
            ))
        }
    }

    /// Solve constraints using Z3
    #[cfg(feature = "z3-solver")]
    pub fn solve(&self, constraints: &ConstraintIR) -> Result<SolutionResult> {
        let solver = Solver::new(&self.context);

        // Create Z3 variables
        let mut var_map: HashMap<VarId, ast::Int> = HashMap::new();
        for var_id in constraints.get_variable_ids() {
            let var_name = format!("v_{}", var_id.0);
            let z3_var = ast::Int::new_const(&self.context, var_name);
            var_map.insert(var_id.clone(), z3_var);
        }

        // Translate constraints to Z3
        for constraint in constraints.get_constraints() {
            let z3_constraint = self.translate_constraint(constraint, &var_map)?;
            solver.assert(&z3_constraint);
        }

        // Solve
        match solver.check() {
            SatResult::Sat => {
                let model = solver
                    .get_model()
                    .ok_or_else(|| anyhow!("No model available"))?;

                // Extract solution
                let mut assignment = HashMap::new();
                for (var_id, z3_var) in &var_map {
                    if let Some(value) = model.eval(z3_var, true) {
                        let value_str = value.to_string();
                        // Try to parse as integer
                        if let Ok(int_val) = value_str.parse::<i64>() {
                            assignment.insert(var_id.to_string(), SolutionValue::Integer(int_val));
                        } else if let Ok(float_val) = value_str.parse::<f64>() {
                            assignment.insert(var_id.to_string(), SolutionValue::Real(float_val));
                        } else {
                            // Fallback to string representation
                            assignment.insert(var_id.to_string(), SolutionValue::Integer(0));
                        }
                    }
                }

                Ok(SolutionResult {
                    satisfiable: true,
                    assignment,
                    confidence: 1.0,        // Z3 solutions are exact
                    estimated_time_ms: 0.0, // TODO: Track actual solve time
                    method: "Z3-SMT".to_string(),
                    winning_strategy: Some("SMT".to_string()),
                    winning_backend: Some("Z3".to_string()),
                    winning_branch_id: None,
                    variable_ranges: None,
                })
            }
            SatResult::Unsat => Err(anyhow!("Constraints are unsatisfiable")),
            SatResult::Unknown => Err(anyhow!("Z3 could not determine satisfiability")),
        }
    }

    /// Translate a constraint to Z3 format
    #[cfg(feature = "z3-solver")]
    fn translate_constraint(
        &self,
        constraint: &Constraint,
        var_map: &HashMap<VarId, ast::Int>,
    ) -> Result<ast::Bool> {
        match constraint {
            Constraint::LessEqual { lhs, rhs } => {
                let lhs_ast = self.translate_expr(lhs, var_map)?;
                let rhs_ast = self.translate_expr(rhs, var_map)?;
                Ok(lhs_ast.le(&rhs_ast))
            }
            Constraint::Equal { lhs, rhs } => {
                let lhs_ast = self.translate_expr(lhs, var_map)?;
                let rhs_ast = self.translate_expr(rhs, var_map)?;
                Ok(lhs_ast._eq(&rhs_ast))
            }
            Constraint::NotEqual { lhs, rhs } => {
                let lhs_ast = self.translate_expr(lhs, var_map)?;
                let rhs_ast = self.translate_expr(rhs, var_map)?;
                Ok(lhs_ast._eq(&rhs_ast).not())
            }
        }
    }

    /// Translate an expression to Z3 format
    #[cfg(feature = "z3-solver")]
    fn translate_expr(&self, expr: &Expr, var_map: &HashMap<VarId, ast::Int>) -> Result<ast::Int> {
        match expr {
            Expr::Var(var_id) => var_map
                .get(var_id)
                .cloned()
                .ok_or_else(|| anyhow!("Variable {:?} not found in variable map", var_id)),
            Expr::Const(value) => {
                // Z3 integers are exact, so convert float to i64
                Ok(ast::Int::from_i64(&self.context, *value as i64))
            }
            Expr::Binary { op, lhs, rhs } => {
                let lhs_ast = self.translate_expr(lhs, var_map)?;
                let rhs_ast = self.translate_expr(rhs, var_map)?;

                match op {
                    BinaryOp::Add => Ok(lhs_ast + rhs_ast),
                    BinaryOp::Sub => Ok(lhs_ast - rhs_ast),
                    BinaryOp::Mul => Ok(lhs_ast * rhs_ast),
                    BinaryOp::Div => {
                        // Integer division in Z3
                        Ok(lhs_ast / rhs_ast)
                    }
                }
            }
        }
    }

    /// Stub implementation when Z3 is not available
    #[cfg(not(feature = "z3-solver"))]
    pub fn solve(&self, _constraints: &ConstraintIR) -> Result<SolutionResult> {
        Err(anyhow!(
            "Z3 solver not available. Enable the 'z3-solver' feature to use Z3."
        ))
    }
}

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

    #[test]
    #[cfg(feature = "z3-solver")]
    fn test_z3_simple_constraint() {
        use crate::ir::{Constraint, ConstraintIR, Expr, VarId};

        let backend = Z3Backend::new(5000).unwrap();

        // Create a simple constraint: x + y = 10, x = 3
        let x = VarId(0);
        let y = VarId(1);

        let mut ir = ConstraintIR::new();
        // x = 3
        ir.add_constraint(Constraint::Equal {
            lhs: Expr::Var(x.clone()),
            rhs: Expr::Const(3.0),
        });
        // y = 7 (implicitly from x + y = 10)
        ir.add_constraint(Constraint::Equal {
            lhs: Expr::Binary {
                op: BinaryOp::Add,
                lhs: Box::new(Expr::Var(x.clone())),
                rhs: Box::new(Expr::Var(y.clone())),
            },
            rhs: Expr::Const(10.0),
        });

        let result = backend.solve(&ir);
        assert!(result.is_ok());

        let solution = result.unwrap();
        assert_eq!(solution.confidence, 1.0);
        assert!(solution.assignment.contains_key("v_0"));
        assert!(solution.assignment.contains_key("v_1"));
    }
}