dslcompile 0.0.1

//! Native egglog Integration with Domain Analysis
//!
//! This module implements domain-aware symbolic optimization using egglog's
//! native abstract interpretation capabilities, following the approach
//! demonstrated in the Herbie case study.
//!
//! ## Key Features
//!
//! - **Interval Analysis**: Native egglog lattice-based interval tracking
//! - **Domain-Safe Rules**: Conditional rewrite rules gated on domain analysis
//! - **Multiple Analyses**: Composable interval and "not-equals" analyses
//! - **Semi-Naive Evaluation**: Leverages egglog's efficient evaluation strategy
//!
//! ## Architecture
//!
//! Unlike our previous Rust wrapper approach, this uses egglog directly:
//! 1. **Native egglog Program**: Mathematical rules + analysis rules in egglog syntax
//! 2. **Interval Lattice**: Built-in egglog support for interval domains
//! 3. **Conditional Rules**: Rules that only fire when domain constraints are satisfied
//! 4. **Extraction**: Cost-based extraction with domain-aware cost functions

use crate::error::{DSLCompileError, Result};
use crate::final_tagless::ASTRepr;
use std::collections::HashMap;

#[cfg(feature = "optimization")]
use egglog::EGraph;

/// Native egglog optimizer with built-in domain analysis
#[cfg(feature = "optimization")]
pub struct NativeEgglogOptimizer {
    /// The egglog `EGraph` with mathematical and analysis rules
    egraph: EGraph,
    /// Variable counter for generating unique names
    var_counter: usize,
    /// Cache for expression mappings
    expr_cache: HashMap<String, ASTRepr<f64>>,
}

#[cfg(feature = "optimization")]
impl NativeEgglogOptimizer {
    /// Create a new native egglog optimizer with domain analysis
    pub fn new() -> Result<Self> {
        let mut egraph = EGraph::default();

        // Load the native egglog program with domain analysis
        let program = Self::create_domain_aware_program();

        egraph.parse_and_run_program(None, &program).map_err(|e| {
            DSLCompileError::Optimization(format!(
                "Failed to initialize native egglog with domain analysis: {e}"
            ))
        })?;

        Ok(Self {
            egraph,
            var_counter: 0,
            expr_cache: HashMap::new(),
        })
    }

    /// Create the native egglog program with domain analysis
    /// This follows the Herbie paper's approach for interval analysis
    fn create_domain_aware_program() -> String {
        r"
; ========================================
; CORE DATATYPES
; ========================================

(datatype Math
  (Num f64)
  (Var String)
  (Add Math Math)
  (Mul Math Math)
  (Neg Math)
  (Pow Math Math)
  (Ln Math)
  (Exp Math)
  (Sin Math)
  (Cos Math)
  (Sqrt Math))

; ========================================
; INTERVAL DOMAIN FOR ABSTRACT INTERPRETATION
; ========================================

(datatype Interval
  (IVal f64 f64)  ; [lower, upper] bounds
  (IBot)          ; Bottom (empty interval)
  (ITop))         ; Top (all reals)

; Interval analysis function
(function ival (Math) Interval :merge (ITop))

; ========================================
; DOMAIN PREDICATES
; ========================================

; Check if an expression is provably positive
(function ival-positive (Math) bool :merge false)

; Check if an expression is provably non-negative
(function ival-nonneg (Math) bool :merge false)

; Check if an expression is provably non-zero
(function ival-nonzero (Math) bool :merge false)

; ========================================
; BASIC INTERVAL ANALYSIS RULES
; ========================================

; Constants have singleton intervals
(rule ((= e (Num ?x))) 
      ((set (ival e) (IVal ?x ?x))))

; Variables have top interval (unknown bounds)
(rule ((= e (Var ?name))) 
      ((set (ival e) (ITop))))

; Positive constants are positive
(rule ((= e (Num ?x))
       (> ?x 0.0))
      ((set (ival-positive e) true)))

; Non-negative constants are non-negative
(rule ((= e (Num ?x))
       (>= ?x 0.0))
      ((set (ival-nonneg e) true)))

; Non-zero constants are non-zero
(rule ((= e (Num ?x))
       (!= ?x 0.0))
      ((set (ival-nonzero e) true)))

; Exponential is always positive
(rule ((= e (Exp ?x)))
      ((set (ival-positive e) true)))

; Exponential is always non-negative
(rule ((= e (Exp ?x)))
      ((set (ival-nonneg e) true)))

; Exponential is always non-zero
(rule ((= e (Exp ?x)))
      ((set (ival-nonzero e) true)))

; Square root is always non-negative
(rule ((= e (Sqrt ?x)))
      ((set (ival-nonneg e) true)))

; Product of positive expressions is positive
(rule ((= e (Mul ?a ?b))
       (= (ival-positive ?a) true)
       (= (ival-positive ?b) true))
      ((set (ival-positive e) true)))

; ========================================
; BASIC MATHEMATICAL RULES
; ========================================

; Additive identity
(rewrite (Add a (Num 0.0)) a)
(rewrite (Add (Num 0.0) a) a)

; Multiplicative identity
(rewrite (Mul a (Num 1.0)) a)
(rewrite (Mul (Num 1.0) a) a)

; Multiplicative zero
(rewrite (Mul a (Num 0.0)) (Num 0.0))
(rewrite (Mul (Num 0.0) a) (Num 0.0))

; Power rules
(rewrite (Pow a (Num 0.0)) (Num 1.0))
(rewrite (Pow a (Num 1.0)) a)

; Commutativity
(rewrite (Add a b) (Add b a))
(rewrite (Mul a b) (Mul b a))

; Associativity for addition
(rewrite (Add (Add a b) c) (Add a (Add b c)))

; Associativity for multiplication  
(rewrite (Mul (Mul a b) c) (Mul a (Mul b c)))

; ========================================
; DOMAIN-AWARE TRANSCENDENTAL RULES
; ========================================

; ln(exp(x)) = x (always safe)
(rewrite (Ln (Exp x)) x)

; exp(ln(x)) = x (only if x is positive)
(rule ((= e (Exp (Ln ?x)))
       (= (ival-positive ?x) true))
      ((union e ?x)))

; ln(a * b) = ln(a) + ln(b) (only if both a and b are positive)
(rule ((= e (Ln (Mul ?a ?b)))
       (= (ival-positive ?a) true)
       (= (ival-positive ?b) true))
      ((union e (Add (Ln ?a) (Ln ?b)))))

; Special case: ln(exp(x) * exp(y)) = ln(exp(x)) + ln(exp(y)) = x + y
(rewrite (Ln (Mul (Exp a) (Exp b))) (Add a b))

; Special case: ln(exp(x) * y) = ln(exp(x)) + ln(y) = x + ln(y) (if y is positive)
(rule ((= e (Ln (Mul (Exp ?x) ?y)))
       (= (ival-positive ?y) true))
      ((union e (Add ?x (Ln ?y)))))

; Special case: ln(x * exp(y)) = ln(x) + ln(exp(y)) = ln(x) + y (if x is positive)
(rule ((= e (Ln (Mul ?x (Exp ?y))))
       (= (ival-positive ?x) true))
      ((union e (Add (Ln ?x) ?y))))

; ln(a / b) = ln(a) - ln(b) (only if both a and b are positive)
; Note: Division is represented as Mul(a, Pow(b, Neg(Num 1.0))) in canonical form
(rule ((= e (Ln (Mul ?a (Pow ?b (Neg (Num 1.0))))))
       (= (ival-positive ?a) true)
       (= (ival-positive ?b) true))
      ((union e (Add (Ln ?a) (Neg (Ln ?b))))))

; ln(x^a) = a * ln(x) (only if x is positive)
(rule ((= e (Ln (Pow ?x ?a)))
       (= (ival-positive ?x) true))
      ((union e (Mul ?a (Ln ?x)))))

; exp(a + b) = exp(a) * exp(b)
(rewrite (Exp (Add a b)) (Mul (Exp a) (Exp b)))

; exp(a - b) = exp(a) / exp(b) -> exp(a) * exp(-b) in canonical form
(rewrite (Exp (Add a (Neg b))) (Mul (Exp a) (Exp (Neg b))))

; ========================================
; DOMAIN-AWARE SQUARE ROOT RULES
; ========================================

; sqrt(0) = 0
(rewrite (Sqrt (Num 0.0)) (Num 0.0))

; sqrt(1) = 1
(rewrite (Sqrt (Num 1.0)) (Num 1.0))

; sqrt(x^2) = |x| = x (only if x is non-negative)
(rule ((= e (Sqrt (Pow ?x (Num 2.0))))
       (= (ival-nonneg ?x) true))
      ((union e ?x)))

; sqrt(x * x) = |x| = x (only if x is non-negative)
(rule ((= e (Sqrt (Mul ?x ?x)))
       (= (ival-nonneg ?x) true))
      ((union e ?x)))

; sqrt(a * b) = sqrt(a) * sqrt(b) (only if both a and b are non-negative)
(rule ((= e (Sqrt (Mul ?a ?b)))
       (= (ival-nonneg ?a) true)
       (= (ival-nonneg ?b) true))
      ((union e (Mul (Sqrt ?a) (Sqrt ?b)))))

; ========================================
; POWER SIMPLIFICATION RULES
; ========================================

; x^(a + b) = x^a * x^b (only if x is positive)
(rule ((= e (Pow ?x (Add ?a ?b)))
       (= (ival-positive ?x) true))
      ((union e (Mul (Pow ?x ?a) (Pow ?x ?b)))))

; (x^a)^b = x^(a*b) (only if x is positive)
(rule ((= e (Pow (Pow ?x ?a) ?b))
       (= (ival-positive ?x) true))
      ((union e (Pow ?x (Mul ?a ?b)))))

; (a * b)^c = a^c * b^c (only if a and b are positive)
(rule ((= e (Pow (Mul ?a ?b) ?c))
       (= (ival-positive ?a) true)
       (= (ival-positive ?b) true))
      ((union e (Mul (Pow ?a ?c) (Pow ?b ?c)))))

"
        .to_string()
    }

    /// Optimize an expression using native egglog with domain analysis
    pub fn optimize(&mut self, expr: &ASTRepr<f64>) -> Result<ASTRepr<f64>> {
        // Convert expression to egglog format
        let egglog_expr = self.ast_to_egglog(expr)?;
        let expr_id = format!("expr_{}", self.var_counter);
        self.var_counter += 1;

        // Store original expression
        self.expr_cache.insert(expr_id.clone(), expr.clone());

        // Add expression to egglog
        let add_command = format!("(let {expr_id} {egglog_expr})");
        self.egraph
            .parse_and_run_program(None, &add_command)
            .map_err(|e| {
                DSLCompileError::Optimization(format!("Failed to add expression to egglog: {e}"))
            })?;

        // Run mathematical optimization rules
        self.egraph
            .parse_and_run_program(None, "(run 10)")
            .map_err(|e| {
                DSLCompileError::Optimization(format!("Failed to run mathematical rules: {e}"))
            })?;

        // Extract the best expression
        self.extract_best(&expr_id)
    }

    /// Get interval analysis information for an expression
    pub fn analyze_interval(&mut self, expr: &ASTRepr<f64>) -> Result<String> {
        // Convert expression to egglog format and add it
        let egglog_expr = self.ast_to_egglog(expr)?;
        let expr_id = format!("interval_expr_{}", self.var_counter);
        self.var_counter += 1;

        // Add expression to egglog
        let add_command = format!("(let {expr_id} {egglog_expr})");
        self.egraph
            .parse_and_run_program(None, &add_command)
            .map_err(|e| {
                DSLCompileError::Optimization(format!(
                    "Failed to add expression for interval analysis: {e}"
                ))
            })?;

        // Run analysis rules to compute intervals
        self.egraph
            .parse_and_run_program(None, "(run 5)")
            .map_err(|e| {
                DSLCompileError::Optimization(format!("Failed to run interval analysis: {e}"))
            })?;

        // Try to extract interval information
        // For now, we'll return a basic analysis based on the expression structure
        self.analyze_interval_heuristic(expr)
    }

    /// Check if an expression is domain-safe for a specific operation
    pub fn is_domain_safe(&mut self, expr: &ASTRepr<f64>, operation: &str) -> Result<bool> {
        match operation {
            "ln" => self.is_positive_definite(expr),
            "sqrt" => self.is_non_negative(expr),
            "div" => self.is_nonzero_denominator(expr),
            _ => Ok(true), // Conservative: assume safe for unknown operations
        }
    }

    /// Heuristic interval analysis based on expression structure
    fn analyze_interval_heuristic(&self, expr: &ASTRepr<f64>) -> Result<String> {
        match expr {
            ASTRepr::Constant(val) => Ok(format!("[{val}, {val}] (singleton interval)")),
            ASTRepr::Variable(_) => Ok("(-∞, +∞) (unknown variable bounds)".to_string()),
            ASTRepr::Add(left, right) => {
                let left_analysis = self.analyze_interval_heuristic(left)?;
                let right_analysis = self.analyze_interval_heuristic(right)?;
                Ok(format!(
                    "Sum of intervals: {left_analysis} + {right_analysis}"
                ))
            }
            ASTRepr::Mul(left, right) => {
                let left_analysis = self.analyze_interval_heuristic(left)?;
                let right_analysis = self.analyze_interval_heuristic(right)?;
                Ok(format!(
                    "Product of intervals: {left_analysis} * {right_analysis}"
                ))
            }
            ASTRepr::Exp(_) => Ok("(0, +∞) (exponential is always positive)".to_string()),
            ASTRepr::Ln(inner) => {
                if self.is_positive_definite(inner)? {
                    Ok("(-∞, +∞) (ln of positive expression)".to_string())
                } else {
                    Ok("Domain error: ln requires positive argument".to_string())
                }
            }
            ASTRepr::Sqrt(inner) => {
                if self.is_non_negative(inner)? {
                    Ok("[0, +∞) (sqrt of non-negative expression)".to_string())
                } else {
                    Ok("Domain error: sqrt requires non-negative argument".to_string())
                }
            }
            _ => Ok("Complex expression - detailed analysis needed".to_string()),
        }
    }

    /// Check if an expression is provably positive
    fn is_positive_definite(&self, expr: &ASTRepr<f64>) -> Result<bool> {
        match expr {
            ASTRepr::Constant(val) => Ok(*val > 0.0),
            ASTRepr::Exp(_) => Ok(true), // exp(x) > 0 for all x
            ASTRepr::Mul(left, right) => {
                // Product is positive if both factors are positive or both are negative
                let left_pos = self.is_positive_definite(left)?;
                let right_pos = self.is_positive_definite(right)?;
                let left_neg = self.is_negative_definite(left)?;
                let right_neg = self.is_negative_definite(right)?;
                Ok((left_pos && right_pos) || (left_neg && right_neg))
            }
            ASTRepr::Pow(base, exp) => {
                // x^a > 0 if x > 0, or if x < 0 and a is even integer
                if self.is_positive_definite(base)? {
                    Ok(true)
                } else if let ASTRepr::Constant(exp_val) = exp.as_ref() {
                    // Check if exponent is even integer
                    if exp_val.fract() == 0.0 && (*exp_val as i64) % 2 == 0 {
                        Ok(true)
                    } else {
                        Ok(false)
                    }
                } else {
                    Ok(false)
                }
            }
            ASTRepr::Sqrt(inner) => {
                // sqrt(x) >= 0, and > 0 if x > 0
                self.is_positive_definite(inner)
            }
            _ => Ok(false), // Conservative: assume not provably positive
        }
    }

    /// Check if an expression is provably negative
    fn is_negative_definite(&self, expr: &ASTRepr<f64>) -> Result<bool> {
        match expr {
            ASTRepr::Constant(val) => Ok(*val < 0.0),
            ASTRepr::Neg(inner) => self.is_positive_definite(inner),
            _ => Ok(false), // Conservative: assume not provably negative
        }
    }

    /// Check if an expression is provably non-negative
    fn is_non_negative(&self, expr: &ASTRepr<f64>) -> Result<bool> {
        match expr {
            ASTRepr::Constant(val) => Ok(*val >= 0.0),
            ASTRepr::Exp(_) => Ok(true),  // exp(x) >= 0 for all x
            ASTRepr::Sqrt(_) => Ok(true), // sqrt(x) >= 0 by definition
            ASTRepr::Mul(left, right) => {
                // Product is non-negative if both factors have same sign
                let left_nonneg = self.is_non_negative(left)?;
                let right_nonneg = self.is_non_negative(right)?;
                let left_nonpos = self.is_non_positive(left)?;
                let right_nonpos = self.is_non_positive(right)?;
                Ok((left_nonneg && right_nonneg) || (left_nonpos && right_nonpos))
            }
            ASTRepr::Pow(base, exp) => {
                // x^a >= 0 if x >= 0, or if a is even
                if self.is_non_negative(base)? {
                    Ok(true)
                } else if let ASTRepr::Constant(exp_val) = exp.as_ref() {
                    // Check if exponent is even integer
                    if exp_val.fract() == 0.0 && (*exp_val as i64) % 2 == 0 {
                        Ok(true)
                    } else {
                        Ok(false)
                    }
                } else {
                    Ok(false)
                }
            }
            ASTRepr::Add(left, right) => {
                // Sum is non-negative if both terms are non-negative
                Ok(self.is_non_negative(left)? && self.is_non_negative(right)?)
            }
            _ => Ok(false), // Conservative: assume not provably non-negative
        }
    }

    /// Check if an expression is provably non-positive
    fn is_non_positive(&self, expr: &ASTRepr<f64>) -> Result<bool> {
        match expr {
            ASTRepr::Constant(val) => Ok(*val <= 0.0),
            ASTRepr::Neg(inner) => self.is_non_negative(inner),
            _ => Ok(false), // Conservative: assume not provably non-positive
        }
    }

    /// Check if a denominator expression is provably non-zero
    fn is_nonzero_denominator(&self, expr: &ASTRepr<f64>) -> Result<bool> {
        match expr {
            ASTRepr::Constant(val) => Ok(*val != 0.0),
            ASTRepr::Exp(_) => Ok(true), // exp(x) != 0 for all x
            ASTRepr::Sqrt(inner) => {
                // sqrt(x) != 0 if x > 0
                self.is_positive_definite(inner)
            }
            _ => Ok(false), // Conservative: assume might be zero
        }
    }

    /// Convert `ASTRepr` to egglog s-expression
    pub fn ast_to_egglog(&self, expr: &ASTRepr<f64>) -> Result<String> {
        match expr {
            ASTRepr::Constant(value) => {
                // Always format f64 with decimal point for egglog compatibility
                if value.fract() == 0.0 {
                    Ok(format!("(Num {value:.1})"))
                } else {
                    Ok(format!("(Num {value})"))
                }
            }
            ASTRepr::Variable(index) => Ok(format!("(Var \"x{index}\")")),
            ASTRepr::Add(left, right) => {
                let left_s = self.ast_to_egglog(left)?;
                let right_s = self.ast_to_egglog(right)?;
                Ok(format!("(Add {left_s} {right_s})"))
            }
            ASTRepr::Sub(left, right) => {
                // Convert Sub to Add + Neg for canonical form
                let left_s = self.ast_to_egglog(left)?;
                let right_s = self.ast_to_egglog(right)?;
                Ok(format!("(Add {left_s} (Neg {right_s}))"))
            }
            ASTRepr::Mul(left, right) => {
                let left_s = self.ast_to_egglog(left)?;
                let right_s = self.ast_to_egglog(right)?;
                Ok(format!("(Mul {left_s} {right_s})"))
            }
            ASTRepr::Div(left, right) => {
                // Convert Div to Mul + Pow(-1) for canonical form
                let left_s = self.ast_to_egglog(left)?;
                let right_s = self.ast_to_egglog(right)?;
                Ok(format!("(Mul {left_s} (Pow {right_s} (Neg (Num 1.0))))"))
            }
            ASTRepr::Pow(base, exp) => {
                let base_s = self.ast_to_egglog(base)?;
                let exp_s = self.ast_to_egglog(exp)?;
                Ok(format!("(Pow {base_s} {exp_s})"))
            }
            ASTRepr::Neg(inner) => {
                let inner_s = self.ast_to_egglog(inner)?;
                Ok(format!("(Neg {inner_s})"))
            }
            ASTRepr::Ln(inner) => {
                let inner_s = self.ast_to_egglog(inner)?;
                Ok(format!("(Ln {inner_s})"))
            }
            ASTRepr::Exp(inner) => {
                let inner_s = self.ast_to_egglog(inner)?;
                Ok(format!("(Exp {inner_s})"))
            }
            ASTRepr::Sin(inner) => {
                let inner_s = self.ast_to_egglog(inner)?;
                Ok(format!("(Sin {inner_s})"))
            }
            ASTRepr::Cos(inner) => {
                let inner_s = self.ast_to_egglog(inner)?;
                Ok(format!("(Cos {inner_s})"))
            }
            ASTRepr::Sqrt(inner) => {
                let inner_s = self.ast_to_egglog(inner)?;
                Ok(format!("(Sqrt {inner_s})"))
            }
        }
    }

    /// Extract the best expression using egglog's extraction
    fn extract_best(&mut self, expr_id: &str) -> Result<ASTRepr<f64>> {
        // Use egglog's extraction to get the best (lowest cost) expression
        let extract_command = format!("(extract {expr_id})");

        // Run the extraction command
        let extract_result = self
            .egraph
            .parse_and_run_program(None, &extract_command)
            .map_err(|e| {
                DSLCompileError::Optimization(format!(
                    "Failed to extract optimized expression: {e}"
                ))
            })?;

        // Convert Vec<String> to a single string for parsing
        let output_string = extract_result.join("\n");

        // Parse the extraction result
        // For now, we'll try to parse the output and convert back to ASTRepr
        // If extraction fails, fall back to the original expression
        match self.parse_egglog_output(&output_string) {
            Ok(optimized) => Ok(optimized),
            Err(_) => {
                // Extraction parsing failed, return original expression
                self.expr_cache.get(expr_id).cloned().ok_or_else(|| {
                    DSLCompileError::Optimization("Expression not found in cache".to_string())
                })
            }
        }
    }

    /// Parse egglog output back to `ASTRepr`
    fn parse_egglog_output(&self, output: &str) -> Result<ASTRepr<f64>> {
        // Remove any whitespace and newlines
        let cleaned = output.trim();

        // Parse the s-expression recursively
        self.parse_sexpr(cleaned)
    }

    /// Parse a single s-expression
    fn parse_sexpr(&self, s: &str) -> Result<ASTRepr<f64>> {
        let s = s.trim();

        if !s.starts_with('(') || !s.ends_with(')') {
            return Err(DSLCompileError::Optimization(format!(
                "Invalid s-expression: {s}"
            )));
        }

        // Remove outer parentheses
        let inner = &s[1..s.len() - 1];

        // Split into tokens
        let tokens = self.tokenize_sexpr(inner)?;

        if tokens.is_empty() {
            return Err(DSLCompileError::Optimization(
                "Empty s-expression".to_string(),
            ));
        }

        match tokens[0].as_str() {
            "Num" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Num requires exactly one argument".to_string(),
                    ));
                }
                let value = tokens[1].parse::<f64>().map_err(|_| {
                    DSLCompileError::Optimization(format!("Invalid number: {}", tokens[1]))
                })?;
                Ok(ASTRepr::Constant(value))
            }
            "Var" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Var requires exactly one argument".to_string(),
                    ));
                }
                // Parse variable name like "x0", "x1", etc.
                let var_name = tokens[1].trim_matches('"');
                if let Some(index_str) = var_name.strip_prefix('x') {
                    let index = index_str.parse::<usize>().map_err(|_| {
                        DSLCompileError::Optimization(format!(
                            "Invalid variable index: {index_str}"
                        ))
                    })?;
                    Ok(ASTRepr::Variable(index))
                } else {
                    Err(DSLCompileError::Optimization(format!(
                        "Invalid variable name: {var_name}"
                    )))
                }
            }
            "Add" => {
                if tokens.len() != 3 {
                    return Err(DSLCompileError::Optimization(
                        "Add requires exactly two arguments".to_string(),
                    ));
                }
                let left = self.parse_sexpr(&tokens[1])?;
                let right = self.parse_sexpr(&tokens[2])?;
                Ok(ASTRepr::Add(Box::new(left), Box::new(right)))
            }
            "Mul" => {
                if tokens.len() != 3 {
                    return Err(DSLCompileError::Optimization(
                        "Mul requires exactly two arguments".to_string(),
                    ));
                }
                let left = self.parse_sexpr(&tokens[1])?;
                let right = self.parse_sexpr(&tokens[2])?;
                Ok(ASTRepr::Mul(Box::new(left), Box::new(right)))
            }
            "Neg" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Neg requires exactly one argument".to_string(),
                    ));
                }
                let inner = self.parse_sexpr(&tokens[1])?;
                Ok(ASTRepr::Neg(Box::new(inner)))
            }
            "Pow" => {
                if tokens.len() != 3 {
                    return Err(DSLCompileError::Optimization(
                        "Pow requires exactly two arguments".to_string(),
                    ));
                }
                let base = self.parse_sexpr(&tokens[1])?;
                let exp = self.parse_sexpr(&tokens[2])?;
                Ok(ASTRepr::Pow(Box::new(base), Box::new(exp)))
            }
            "Ln" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Ln requires exactly one argument".to_string(),
                    ));
                }
                let inner = self.parse_sexpr(&tokens[1])?;
                Ok(ASTRepr::Ln(Box::new(inner)))
            }
            "Exp" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Exp requires exactly one argument".to_string(),
                    ));
                }
                let inner = self.parse_sexpr(&tokens[1])?;
                Ok(ASTRepr::Exp(Box::new(inner)))
            }
            "Sin" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Sin requires exactly one argument".to_string(),
                    ));
                }
                let inner = self.parse_sexpr(&tokens[1])?;
                Ok(ASTRepr::Sin(Box::new(inner)))
            }
            "Cos" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Cos requires exactly one argument".to_string(),
                    ));
                }
                let inner = self.parse_sexpr(&tokens[1])?;
                Ok(ASTRepr::Cos(Box::new(inner)))
            }
            "Sqrt" => {
                if tokens.len() != 2 {
                    return Err(DSLCompileError::Optimization(
                        "Sqrt requires exactly one argument".to_string(),
                    ));
                }
                let inner = self.parse_sexpr(&tokens[1])?;
                Ok(ASTRepr::Sqrt(Box::new(inner)))
            }
            _ => Err(DSLCompileError::Optimization(format!(
                "Unknown operation: {}",
                tokens[0]
            ))),
        }
    }

    /// Tokenize an s-expression into its components
    fn tokenize_sexpr(&self, s: &str) -> Result<Vec<String>> {
        let mut tokens = Vec::new();
        let mut current_token = String::new();
        let mut paren_depth = 0;
        let mut in_string = false;
        let chars = s.chars().peekable();

        for ch in chars {
            match ch {
                '"' => {
                    in_string = !in_string;
                    current_token.push(ch);
                }
                '(' if !in_string => {
                    if paren_depth == 0 && !current_token.is_empty() {
                        // We're starting a new nested expression, save the current token first
                        tokens.push(current_token.trim().to_string());
                        current_token.clear();
                    }
                    current_token.push(ch);
                    paren_depth += 1;
                }
                ')' if !in_string => {
                    current_token.push(ch);
                    paren_depth -= 1;
                    if paren_depth == 0 {
                        // We've completed a nested expression
                        tokens.push(current_token.trim().to_string());
                        current_token.clear();
                    }
                }
                ' ' | '\t' | '\n' if !in_string && paren_depth == 0 => {
                    if !current_token.is_empty() {
                        tokens.push(current_token.trim().to_string());
                        current_token.clear();
                    }
                }
                _ => {
                    current_token.push(ch);
                }
            }
        }

        if !current_token.is_empty() {
            tokens.push(current_token.trim().to_string());
        }

        Ok(tokens)
    }
}

/// Fallback implementation when optimization feature is not enabled
#[cfg(not(feature = "optimization"))]
pub struct NativeEgglogOptimizer;

#[cfg(not(feature = "optimization"))]
impl NativeEgglogOptimizer {
    pub fn new() -> Result<Self> {
        Ok(Self)
    }

    pub fn optimize(&mut self, expr: &ASTRepr<f64>) -> Result<ASTRepr<f64>> {
        Ok(expr.clone())
    }
}

/// Helper function to create and use the native egglog optimizer
pub fn optimize_with_native_egglog(expr: &ASTRepr<f64>) -> Result<ASTRepr<f64>> {
    let mut optimizer = NativeEgglogOptimizer::new()?;
    optimizer.optimize(expr)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::final_tagless::{ASTEval, ASTMathExpr};

    #[test]
    fn test_native_egglog_creation() {
        let result = NativeEgglogOptimizer::new();
        if let Err(e) = &result {
            println!("Error creating NativeEgglogOptimizer: {e}");
        }
        assert!(result.is_ok());
    }

    #[test]
    fn test_ast_to_egglog_conversion() {
        let optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test basic conversions
        let num = ASTRepr::Constant(42.0);
        let egglog_str = optimizer.ast_to_egglog(&num).unwrap();
        assert_eq!(egglog_str, "(Num 42.0)");

        let var = ASTRepr::Variable(0);
        let egglog_str = optimizer.ast_to_egglog(&var).unwrap();
        assert_eq!(egglog_str, "(Var \"x0\")");

        let add = ASTRepr::Add(
            Box::new(ASTRepr::Variable(0)),
            Box::new(ASTRepr::Constant(1.0)),
        );
        let egglog_str = optimizer.ast_to_egglog(&add).unwrap();
        assert_eq!(egglog_str, "(Add (Var \"x0\") (Num 1.0))");
    }

    #[test]
    fn test_canonical_form_conversion() {
        let optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test conversion of canonical form (Sub -> Add + Neg)
        let sub = ASTRepr::Sub(
            Box::new(ASTRepr::Variable(0)),
            Box::new(ASTRepr::Constant(1.0)),
        );
        let egglog_str = optimizer.ast_to_egglog(&sub).unwrap();
        assert_eq!(egglog_str, "(Add (Var \"x0\") (Neg (Num 1.0)))");

        // Test conversion of canonical form (Div -> Mul + Pow)
        let div = ASTRepr::Div(
            Box::new(ASTRepr::Variable(0)),
            Box::new(ASTRepr::Constant(2.0)),
        );
        let egglog_str = optimizer.ast_to_egglog(&div).unwrap();
        assert_eq!(
            egglog_str,
            "(Mul (Var \"x0\") (Pow (Num 2.0) (Neg (Num 1.0))))"
        );
    }

    #[test]
    fn test_basic_optimization() {
        let expr = ASTEval::add(ASTEval::var(0), ASTEval::constant(0.0));
        let result = optimize_with_native_egglog(&expr);

        #[cfg(feature = "optimization")]
        {
            // Should run without error
            assert!(result.is_ok());
        }

        #[cfg(not(feature = "optimization"))]
        {
            // Should return unchanged
            assert!(result.is_ok());
        }
    }

    #[test]
    fn test_domain_aware_optimization() {
        let mut optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test domain-safe ln(exp(x)) = x (should always work)
        let safe_expr = ASTRepr::Ln(Box::new(ASTRepr::Exp(Box::new(ASTRepr::Variable(0)))));

        let result = optimizer.optimize(&safe_expr);
        assert!(result.is_ok());

        // Test potentially unsafe exp(ln(x)) (depends on domain analysis)
        let potentially_unsafe =
            ASTRepr::Exp(Box::new(ASTRepr::Ln(Box::new(ASTRepr::Variable(0)))));

        let result = optimizer.optimize(&potentially_unsafe);
        assert!(result.is_ok());
    }

    #[test]
    fn test_interval_analysis() {
        let mut optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test interval analysis on a constant
        let constant_expr = ASTRepr::Constant(5.0);
        let interval_info = optimizer.analyze_interval(&constant_expr);
        assert!(interval_info.is_ok());

        // Test interval analysis on a variable
        let var_expr = ASTRepr::Variable(0);
        let interval_info = optimizer.analyze_interval(&var_expr);
        assert!(interval_info.is_ok());

        // Test interval analysis on a complex expression
        let complex_expr = ASTRepr::Add(
            Box::new(ASTRepr::Constant(2.0)),
            Box::new(ASTRepr::Mul(
                Box::new(ASTRepr::Variable(0)),
                Box::new(ASTRepr::Constant(3.0)),
            )),
        );
        let interval_info = optimizer.analyze_interval(&complex_expr);
        assert!(interval_info.is_ok());
    }

    #[test]
    fn test_domain_safety_checks() {
        let mut optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test safety check for ln operation
        let positive_constant = ASTRepr::Constant(5.0);
        let is_safe = optimizer.is_domain_safe(&positive_constant, "ln");
        assert!(is_safe.is_ok());

        // Test safety check for division
        let nonzero_expr = ASTRepr::Add(
            Box::new(ASTRepr::Variable(0)),
            Box::new(ASTRepr::Constant(1.0)),
        );
        let is_safe = optimizer.is_domain_safe(&nonzero_expr, "div");
        assert!(is_safe.is_ok());

        // Test safety check for sqrt
        let sqrt_expr = ASTRepr::Pow(
            Box::new(ASTRepr::Variable(0)),
            Box::new(ASTRepr::Constant(2.0)),
        );
        let is_safe = optimizer.is_domain_safe(&sqrt_expr, "sqrt");
        assert!(is_safe.is_ok());
    }

    #[test]
    fn test_domain_aware_ln_rules() {
        let mut optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test ln(a * b) = ln(a) + ln(b) with positive constants
        let ln_product = ASTRepr::Ln(Box::new(ASTRepr::Mul(
            Box::new(ASTRepr::Constant(2.0)),
            Box::new(ASTRepr::Constant(3.0)),
        )));

        let result = optimizer.optimize(&ln_product);
        assert!(result.is_ok());

        // The optimization should work since both constants are positive
        // In a full implementation, we'd check that it actually transformed to ln(2) + ln(3)
    }

    #[test]
    fn test_sqrt_domain_awareness() {
        let mut optimizer = NativeEgglogOptimizer::new().unwrap();

        // Test sqrt(x^2) = x with domain analysis
        let sqrt_square = ASTRepr::Sqrt(Box::new(ASTRepr::Pow(
            Box::new(ASTRepr::Variable(0)),
            Box::new(ASTRepr::Constant(2.0)),
        )));

        let result = optimizer.optimize(&sqrt_square);
        assert!(result.is_ok());

        // The rule should only apply if x is known to be non-negative
        // For variables, this would typically not apply without additional constraints
    }
}