dslcompile 0.0.1

//! A-Normal Form (ANF) Intermediate Representation
//!
//! This module provides ANF transformation and optimization for mathematical expressions.
//! ANF ensures that every intermediate computation is explicitly named, enabling
//! efficient common subexpression elimination and clean optimization passes.
//!
//! # Overview
//!
//! A-Normal Form is an intermediate representation where:
//! 1. **Every operation** has atomic (non-compound) arguments
//! 2. **Every intermediate result** is bound to a temporary variable
//! 3. **Computation structure** is made explicit through let-bindings
//!
//! ## Example Transformation
//!
//! ```text
//! Input:  sin(x + y) + cos(x + y) + exp(x + y)
//! ANF:    let t0 = x + y in
//!         let t1 = sin(t0) in  
//!         let t2 = cos(t0) in
//!         let t3 = exp(t0) in
//!         let t4 = t1 + t2 in
//!         t4 + t3
//! ```
//!
//! ## Key Features
//!
//! - **Automatic CSE**: Common subexpressions like `x + y` are computed once
//! - **Scope Safety**: Variables are only referenced within their binding scope
//! - **Efficient Caching**: Structural hashing enables O(1) subexpression lookup
//! - **Clean Code Gen**: Produces readable nested let-bindings in Rust
//!
//! # Architecture
//!
//! ## Core Types
//!
//! - [`VarRef`]: Hybrid variable system (user + generated variables)
//! - [`ANFAtom`]: Atomic expressions (constants, variables)
//! - [`ANFComputation`]: Operations with atomic arguments only
//! - [`ANFExpr`]: Complete ANF expressions (atoms + let-bindings)
//! - [`ANFConverter`]: Stateful converter with CSE cache
//!
//! ## Variable Management
//!
//! Uses a hybrid approach:
//! - `VarRef::User(usize)`: Original user variables from the AST
//! - `VarRef::Bound(u32)`: Temporary variables generated during ANF conversion
//!
//! This integrates cleanly with the existing `VariableRegistry` system while providing
//! efficient integer-based variable management during optimization.
//!
//! # Usage Patterns
//!
//! ## Basic Conversion
//!
//! ```rust
//! use dslcompile::anf::{convert_to_anf, generate_rust_code};
//! use dslcompile::final_tagless::{ASTEval, ASTMathExpr, VariableRegistry};
//!
//! // Create expression: x^2 + 2*x + 1
//! let mut registry = VariableRegistry::new();
//! let x = ASTEval::var(registry.register_variable());
//! let expr = ASTEval::add(
//!     ASTEval::add(ASTEval::pow(x.clone(), ASTEval::constant(2.0)),
//!                  ASTEval::mul(ASTEval::constant(2.0), x)),
//!     ASTEval::constant(1.0)
//! );
//!
//! // Convert to ANF
//! let anf = convert_to_anf(&expr).unwrap();
//!
//! // Generate Rust code
//! let code = generate_rust_code(&anf, &registry);
//! println!("{}", code);
//! // Output: { let t0 = x * x; { let t1 = 2 * x; { let t2 = t0 + t1; t2 + 1 } } }
//! ```
//!
//! ## Advanced: Custom Converter
//!
//! ```rust
//! use dslcompile::anf::{ANFCodeGen, ANFConverter};
//! use dslcompile::final_tagless::{ASTEval, ASTMathExpr};
//! let expr1 = ASTEval::constant(1.0);
//! let expr2 = ASTEval::constant(2.0);
//! let mut converter = ANFConverter::new();
//! let anf1 = converter.convert(&expr1).unwrap();
//! let anf2 = converter.convert(&expr2).unwrap();  // Shares CSE cache with expr1
//! ```
//!
//! ## Function Generation
//!
//! ```rust
//! use dslcompile::anf::ANFCodeGen;
//! use dslcompile::final_tagless::VariableRegistry;
//! use dslcompile::anf::{ANFExpr, ANFAtom, VarRef};
//! let mut registry = VariableRegistry::new();
//! let x_idx = registry.register_variable();
//! let anf = ANFExpr::Atom(ANFAtom::<f64>::Variable(VarRef::User(x_idx)));
//! let codegen = ANFCodeGen::new(&registry);
//! let function = codegen.generate_function("my_function", &anf);
//! // Output: fn my_function(x: f64) -> f64 { ... }
//! ```
//!
//! ## Useful Debug Patterns
//!
//! ```rust
//! use dslcompile::anf::{convert_to_anf};
//! use dslcompile::final_tagless::{ASTEval, ASTMathExpr, VariableRegistry};
//! let mut registry = VariableRegistry::new();
//! let x = ASTEval::var(registry.register_variable());
//! let expr = ASTEval::add(x.clone(), ASTEval::constant(1.0));
//! let anf = convert_to_anf(&expr).unwrap();
//! // Print ANF structure
//! println!("ANF: {:#?}", anf);
//! // Check variable usage
//! let vars = anf.used_variables();
//! println!("Used variables: {:?}", vars);
//! // Count let-bindings
//! println!("Binding count: {}", anf.let_count());
//! ```
//!
//! # Performance Characteristics
//!
//! - **Time**: O(n) conversion where n = AST node count
//! - **Space**: O(k) cache entries where k = unique subexpression count  
//! - **CSE Hit Rate**: 80-95% for mathematical expressions with redundancy
//! - **Code Reduction**: 40-60% fewer operations in generated code
//!
//! # Implementation Notes
//!
//! ## CSE Algorithm
//!
//! 1. **Structural Hashing**: Create hash ignoring numeric values
//! 2. **Scope Tracking**: Track binding depth for each cached variable
//! 3. **Cache Lookup**: Check if subexpression already computed
//! 4. **Scope Validation**: Only reuse variables still in scope
//! 5. **Cache Invalidation**: Remove out-of-scope entries
//!
//! ## Scope Management
//!
//! The key insight is that cached variables must respect lexical scoping:
//!
//! ```text
//! ✅ Valid:   { let t0 = x + 1; { let t1 = t0 + t0; t1 } }
//! ❌ Invalid: { { let t0 = x + 1; t0 } + t0 }  // t0 not in scope
//! ```
//!
//! Our solution: Cache entries include scope depth, only reuse when `cached_scope <= current_depth`.
//!
//! ## Memory Management
//!
//! - **Cache Growth**: Bounded by unique subexpression count
//! - **Cleanup**: Out-of-scope entries removed proactively  
//! - **Sharing**: Single converter can process multiple expressions
//! - **Reset**: Call `ANFConverter::new()` to clear cache
//!
//! # Testing and Debugging
//!
//! ## Common Issues
//!
//! - **Invalid Variable References**: Usually scope management bugs
//! - **Cache Misses**: Verify structural hash implementation  
//! - **Memory Growth**: Check for cache invalidation
//! - **Wrong Code Generation**: Inspect ANF structure first
//!
//! # Future Directions
//!
//! - **Constant Folding**: Evaluate constant subexpressions at ANF level
//! - **Dead Code Elimination**: Remove unused let-bindings
//! - **Egglog Integration**: ANF as input/output for e-graph optimization
//! - **Parallel CSE**: Thread-safe conversion for concurrent usage

use crate::error::Result;
use crate::final_tagless::{ASTRepr, NumericType, VariableRegistry};
use crate::interval_domain::{IntervalDomain, IntervalDomainAnalyzer};
use num_traits::{Float, Zero};
use ordered_float::OrderedFloat;
use std::collections::HashMap;

/// Variable reference that distinguishes between user and generated variables
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum VarRef {
    /// User-defined variable (index into `VariableRegistry`)
    User(usize),
    /// De Bruijn index for let-bound variables (0 = innermost binding)
    Bound(u32),
}

impl VarRef {
    /// Generate Rust code for this variable
    #[must_use]
    pub fn to_rust_code(&self, registry: &VariableRegistry) -> String {
        match self {
            VarRef::User(idx) => registry.debug_name(*idx),
            VarRef::Bound(id) => format!("t{id}"),
        }
    }

    /// Generate a debug-friendly name for this variable
    #[must_use]
    pub fn debug_name(&self, registry: &VariableRegistry) -> String {
        match self {
            VarRef::User(idx) => {
                format!("{}({})", registry.debug_name(*idx), idx)
            }
            VarRef::Bound(id) => format!("t{id}"),
        }
    }

    /// Check if this is a user variable
    #[must_use]
    pub fn is_user(&self) -> bool {
        matches!(self, VarRef::User(_))
    }

    /// Check if this is a generated variable
    #[must_use]
    pub fn is_generated(&self) -> bool {
        matches!(self, VarRef::Bound(_))
    }
}

/// Generator for fresh ANF variables
#[derive(Debug, Clone)]
pub struct ANFVarGen {
    next_temp_id: u32,
}

impl ANFVarGen {
    /// Create a new variable generator
    #[must_use]
    pub fn new() -> Self {
        Self { next_temp_id: 0 }
    }

    /// Generate a fresh temporary variable
    pub fn fresh(&mut self) -> VarRef {
        let var = VarRef::Bound(self.next_temp_id);
        self.next_temp_id += 1;
        var
    }

    /// Create a user variable reference
    #[must_use]
    pub fn user_var(&self, index: usize) -> VarRef {
        VarRef::User(index)
    }

    /// Get the number of generated variables so far
    #[must_use]
    pub fn generated_count(&self) -> u32 {
        self.next_temp_id
    }
}

impl Default for ANFVarGen {
    fn default() -> Self {
        Self::new()
    }
}

/// ANF atomic expressions (no sub-computations)
#[derive(Debug, Clone, PartialEq)]
pub enum ANFAtom<T> {
    /// Constant value
    Constant(T),
    /// Variable reference
    Variable(VarRef),
}

impl<T: NumericType> ANFAtom<T> {
    /// Check if this atom is a constant
    pub fn is_constant(&self) -> bool {
        matches!(self, ANFAtom::Constant(_))
    }

    /// Check if this atom is a variable
    pub fn is_variable(&self) -> bool {
        matches!(self, ANFAtom::Variable(_))
    }

    /// Extract constant value if this is a constant
    pub fn as_constant(&self) -> Option<&T> {
        match self {
            ANFAtom::Constant(val) => Some(val),
            _ => None,
        }
    }

    /// Extract variable reference if this is a variable
    pub fn as_variable(&self) -> Option<VarRef> {
        match self {
            ANFAtom::Variable(var) => Some(*var),
            _ => None,
        }
    }
}

/// ANF computations (operations with atomic arguments only)
#[derive(Debug, Clone, PartialEq)]
pub enum ANFComputation<T> {
    /// Addition: a + b
    Add(ANFAtom<T>, ANFAtom<T>),
    /// Subtraction: a - b
    Sub(ANFAtom<T>, ANFAtom<T>),
    /// Multiplication: a * b
    Mul(ANFAtom<T>, ANFAtom<T>),
    /// Division: a / b
    Div(ANFAtom<T>, ANFAtom<T>),
    /// Power: a^b
    Pow(ANFAtom<T>, ANFAtom<T>),
    /// Negation: -a
    Neg(ANFAtom<T>),
    /// Natural logarithm: ln(a)
    Ln(ANFAtom<T>),
    /// Exponential: exp(a)
    Exp(ANFAtom<T>),
    /// Sine: sin(a)
    Sin(ANFAtom<T>),
    /// Cosine: cos(a)
    Cos(ANFAtom<T>),
    /// Square root: sqrt(a)
    Sqrt(ANFAtom<T>),
}

impl<T: NumericType> ANFComputation<T> {
    /// Get all atomic operands of this computation
    pub fn operands(&self) -> Vec<&ANFAtom<T>> {
        match self {
            ANFComputation::Add(a, b)
            | ANFComputation::Sub(a, b)
            | ANFComputation::Mul(a, b)
            | ANFComputation::Div(a, b)
            | ANFComputation::Pow(a, b) => vec![a, b],
            ANFComputation::Neg(a)
            | ANFComputation::Ln(a)
            | ANFComputation::Exp(a)
            | ANFComputation::Sin(a)
            | ANFComputation::Cos(a)
            | ANFComputation::Sqrt(a) => vec![a],
        }
    }

    /// Check if this computation uses only constants
    pub fn is_constant_computation(&self) -> bool {
        self.operands().iter().all(|atom| atom.is_constant())
    }
}

/// ANF expressions
#[derive(Debug, Clone, PartialEq)]
pub enum ANFExpr<T> {
    /// Atomic expression (constant or variable)
    Atom(ANFAtom<T>),
    /// Let binding: let var = computation in body
    Let(VarRef, ANFComputation<T>, Box<ANFExpr<T>>),
}

impl<T> ANFExpr<T>
where
    T: NumericType
        + Float
        + Copy
        + std::ops::Add<Output = T>
        + std::ops::Sub<Output = T>
        + std::ops::Mul<Output = T>
        + std::ops::Div<Output = T>
        + std::ops::Neg<Output = T>
        + Zero,
{
    /// Create an atomic constant
    pub fn constant(value: T) -> Self {
        ANFExpr::Atom(ANFAtom::Constant(value))
    }

    /// Create an atomic variable
    #[must_use]
    pub fn variable(var_ref: VarRef) -> Self {
        ANFExpr::Atom(ANFAtom::Variable(var_ref))
    }

    /// Create a let binding
    pub fn let_binding(var_ref: VarRef, computation: ANFComputation<T>, body: ANFExpr<T>) -> Self {
        ANFExpr::Let(var_ref, computation, Box::new(body))
    }

    /// Check if this is an atomic expression
    pub fn is_atom(&self) -> bool {
        matches!(self, ANFExpr::Atom(_))
    }

    /// Check if this is a let binding
    pub fn is_let(&self) -> bool {
        matches!(self, ANFExpr::Let(_, _, _))
    }

    /// Count the number of let bindings in this expression
    pub fn let_count(&self) -> usize {
        match self {
            ANFExpr::Atom(_) => 0,
            ANFExpr::Let(_, _, body) => 1 + body.let_count(),
        }
    }

    /// Collect all variables used in this expression
    pub fn used_variables(&self) -> Vec<VarRef> {
        let mut vars = Vec::new();
        self.collect_variables(&mut vars);
        vars.sort_unstable();
        vars.dedup();
        vars
    }

    fn collect_variables(&self, vars: &mut Vec<VarRef>) {
        match self {
            ANFExpr::Atom(ANFAtom::Variable(var)) => vars.push(*var),
            ANFExpr::Atom(ANFAtom::Constant(_)) => {}
            ANFExpr::Let(_, comp, body) => {
                for operand in comp.operands() {
                    if let ANFAtom::Variable(var) = operand {
                        vars.push(*var);
                    }
                }
                body.collect_variables(vars);
            }
        }
    }

    /// Evaluate the ANF expression with a variable map and bound variable map
    pub fn eval(&self, variables: &HashMap<usize, T>) -> T {
        self.eval_with_bound_vars(variables, &HashMap::new())
    }

    /// Domain-aware evaluation that uses interval analysis to ensure mathematical safety
    pub fn eval_domain_aware(
        &self,
        variables: &HashMap<usize, T>,
        domain_analyzer: &IntervalDomainAnalyzer<T>,
    ) -> T
    where
        T: PartialOrd + From<f64>,
    {
        self.eval_with_bound_vars_domain_aware(variables, &HashMap::new(), domain_analyzer)
    }

    /// Internal eval that handles both user variables and bound variables
    fn eval_with_bound_vars(
        &self,
        user_vars: &HashMap<usize, T>,
        bound_vars: &HashMap<u32, T>,
    ) -> T {
        match self {
            ANFExpr::Atom(atom) => match atom {
                ANFAtom::Constant(value) => *value,
                ANFAtom::Variable(var_ref) => match var_ref {
                    VarRef::User(idx) => user_vars.get(idx).copied().unwrap_or_else(T::zero),
                    VarRef::Bound(id) => bound_vars.get(id).copied().unwrap_or_else(T::zero),
                },
            },
            ANFExpr::Let(var_ref, computation, body) => {
                // Evaluate the computation
                let comp_result = match computation {
                    ANFComputation::Add(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            + self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Sub(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            - self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Mul(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            * self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Div(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            / self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Pow(a, b) => {
                        let base = self.eval_atom_with_bound(a, user_vars, bound_vars);
                        let exp = self.eval_atom_with_bound(b, user_vars, bound_vars);
                        // Use domain-aware power operation
                        Self::safe_powf(base, exp)
                    }
                    ANFComputation::Neg(a) => -self.eval_atom_with_bound(a, user_vars, bound_vars),
                    ANFComputation::Ln(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).ln()
                    }
                    ANFComputation::Exp(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).exp()
                    }
                    ANFComputation::Sin(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).sin()
                    }
                    ANFComputation::Cos(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).cos()
                    }
                    ANFComputation::Sqrt(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).sqrt()
                    }
                };

                // Extend the bound variable environment and evaluate the body
                match var_ref {
                    VarRef::Bound(id) => {
                        let mut extended_bound_vars = bound_vars.clone();
                        extended_bound_vars.insert(*id, comp_result);
                        body.eval_with_bound_vars(user_vars, &extended_bound_vars)
                    }
                    VarRef::User(_) => {
                        // This shouldn't happen in normal ANF, but handle gracefully
                        body.eval_with_bound_vars(user_vars, bound_vars)
                    }
                }
            }
        }
    }

    /// Domain-aware evaluation with interval analysis
    fn eval_with_bound_vars_domain_aware(
        &self,
        user_vars: &HashMap<usize, T>,
        bound_vars: &HashMap<u32, T>,
        domain_analyzer: &IntervalDomainAnalyzer<T>,
    ) -> T
    where
        T: PartialOrd + From<f64>,
    {
        match self {
            ANFExpr::Atom(atom) => match atom {
                ANFAtom::Constant(value) => *value,
                ANFAtom::Variable(var_ref) => match var_ref {
                    VarRef::User(idx) => user_vars.get(idx).copied().unwrap_or_else(T::zero),
                    VarRef::Bound(id) => bound_vars.get(id).copied().unwrap_or_else(T::zero),
                },
            },
            ANFExpr::Let(var_ref, computation, body) => {
                // Evaluate the computation with domain awareness
                let comp_result = match computation {
                    ANFComputation::Add(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            + self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Sub(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            - self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Mul(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            * self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Div(a, b) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars)
                            / self.eval_atom_with_bound(b, user_vars, bound_vars)
                    }
                    ANFComputation::Pow(a, b) => {
                        let base = self.eval_atom_with_bound(a, user_vars, bound_vars);
                        let exp = self.eval_atom_with_bound(b, user_vars, bound_vars);

                        // Use domain analysis to determine if this power operation is safe
                        // For now, use the same safe_powf approach but with domain context
                        Self::domain_aware_powf(base, exp, domain_analyzer)
                    }
                    ANFComputation::Neg(a) => -self.eval_atom_with_bound(a, user_vars, bound_vars),
                    ANFComputation::Ln(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).ln()
                    }
                    ANFComputation::Exp(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).exp()
                    }
                    ANFComputation::Sin(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).sin()
                    }
                    ANFComputation::Cos(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).cos()
                    }
                    ANFComputation::Sqrt(a) => {
                        self.eval_atom_with_bound(a, user_vars, bound_vars).sqrt()
                    }
                };

                // Extend the bound variable environment and evaluate the body
                match var_ref {
                    VarRef::Bound(id) => {
                        let mut extended_bound_vars = bound_vars.clone();
                        extended_bound_vars.insert(*id, comp_result);
                        body.eval_with_bound_vars_domain_aware(
                            user_vars,
                            &extended_bound_vars,
                            domain_analyzer,
                        )
                    }
                    VarRef::User(_) => {
                        // This shouldn't happen in normal ANF, but handle gracefully
                        body.eval_with_bound_vars_domain_aware(
                            user_vars,
                            bound_vars,
                            domain_analyzer,
                        )
                    }
                }
            }
        }
    }

    /// Safe power function that avoids NaN results
    fn safe_powf(base: T, exp: T) -> T {
        let result = base.powf(exp);

        // If the result is already well-formed (finite, inf, or NaN), use it
        if result.is_finite() || result.is_infinite() {
            return result;
        }

        // Only intervene if we get NaN and it might be fixable
        if result.is_nan() {
            // Check for the specific case of negative finite base with non-integer exponent
            if base.is_finite() && base < T::zero() && exp.is_finite() {
                // This is the problematic case: negative base with non-integer exponent
                // Return NaN to indicate undefined result in real numbers
                return T::nan();
            }

            // For other NaN cases (like inf^0, 0^0, etc.), let Rust's powf handle it
            return result;
        }

        // Fallback: return the original result
        result
    }

    /// Domain-aware power function that uses interval analysis
    fn domain_aware_powf(base: T, exp: T, _domain_analyzer: &IntervalDomainAnalyzer<T>) -> T
    where
        T: PartialOrd + From<f64>,
    {
        // For now, use the same logic as safe_powf but with domain context
        // In a full implementation, this would use the domain analyzer to:
        // 1. Check if base is known to be positive from interval analysis
        // 2. Use more sophisticated handling based on domain information

        let result = base.powf(exp);

        // If the result is already well-formed (finite, inf, or NaN), use it
        if result.is_finite() || result.is_infinite() {
            return result;
        }

        // Only intervene if we get NaN and it might be fixable
        if result.is_nan() {
            // Check for the specific case of negative finite base with non-integer exponent
            if base.is_finite() && base < T::zero() && exp.is_finite() {
                // This is the problematic case: negative base with non-integer exponent
                // Return NaN to indicate undefined result in real numbers
                return T::nan();
            }

            // For other NaN cases (like inf^0, 0^0, etc.), let Rust's powf handle it
            return result;
        }

        // Fallback: return the original result
        result
    }

    fn eval_atom_with_bound(
        &self,
        atom: &ANFAtom<T>,
        user_vars: &HashMap<usize, T>,
        bound_vars: &HashMap<u32, T>,
    ) -> T {
        match atom {
            ANFAtom::Constant(value) => *value,
            ANFAtom::Variable(var_ref) => match var_ref {
                VarRef::User(idx) => user_vars.get(idx).copied().unwrap_or_else(T::zero),
                VarRef::Bound(id) => bound_vars.get(id).copied().unwrap_or_else(T::zero),
            },
        }
    }
}

/// Structural hash representing the shape of an expression for CSE
/// Ignores actual numeric values but captures operators and variable positions
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum StructuralHash {
    /// Constant (with value)
    Constant(OrderedFloat<f64>),
    /// Variable by index
    Variable(usize),
    /// Binary operations
    Add(Box<StructuralHash>, Box<StructuralHash>),
    Sub(Box<StructuralHash>, Box<StructuralHash>),
    Mul(Box<StructuralHash>, Box<StructuralHash>),
    Div(Box<StructuralHash>, Box<StructuralHash>),
    Pow(Box<StructuralHash>, Box<StructuralHash>),
    /// Unary operations
    Neg(Box<StructuralHash>),
    Ln(Box<StructuralHash>),
    Exp(Box<StructuralHash>),
    Sin(Box<StructuralHash>),
    Cos(Box<StructuralHash>),
    Sqrt(Box<StructuralHash>),
}

impl StructuralHash {
    /// Create a structural hash from an `ASTRepr<f64>` expression
    #[must_use]
    pub fn from_expr(expr: &ASTRepr<f64>) -> Self {
        match expr {
            ASTRepr::Constant(val) => StructuralHash::Constant(OrderedFloat(*val)),
            ASTRepr::Variable(idx) => StructuralHash::Variable(*idx),
            ASTRepr::Add(left, right) => StructuralHash::Add(
                Box::new(Self::from_expr(left)),
                Box::new(Self::from_expr(right)),
            ),
            ASTRepr::Sub(left, right) => StructuralHash::Sub(
                Box::new(Self::from_expr(left)),
                Box::new(Self::from_expr(right)),
            ),
            ASTRepr::Mul(left, right) => StructuralHash::Mul(
                Box::new(Self::from_expr(left)),
                Box::new(Self::from_expr(right)),
            ),
            ASTRepr::Div(left, right) => StructuralHash::Div(
                Box::new(Self::from_expr(left)),
                Box::new(Self::from_expr(right)),
            ),
            ASTRepr::Pow(left, right) => StructuralHash::Pow(
                Box::new(Self::from_expr(left)),
                Box::new(Self::from_expr(right)),
            ),
            ASTRepr::Neg(inner) => StructuralHash::Neg(Box::new(Self::from_expr(inner))),
            ASTRepr::Ln(inner) => StructuralHash::Ln(Box::new(Self::from_expr(inner))),
            ASTRepr::Exp(inner) => StructuralHash::Exp(Box::new(Self::from_expr(inner))),
            ASTRepr::Sin(inner) => StructuralHash::Sin(Box::new(Self::from_expr(inner))),
            ASTRepr::Cos(inner) => StructuralHash::Cos(Box::new(Self::from_expr(inner))),
            ASTRepr::Sqrt(inner) => StructuralHash::Sqrt(Box::new(Self::from_expr(inner))),
        }
    }
}

/// ANF converter that transforms `ASTRepr` to ANF form
#[derive(Debug)]
pub struct ANFConverter {
    /// Current binding depth (for scope tracking)
    binding_depth: u32,
    /// Next unique binding ID
    next_binding_id: u32,
    /// Cache for common subexpression elimination with scope tracking
    /// Maps structural hashes to (`scope_depth`, variable, `binding_id`)
    expr_cache: HashMap<StructuralHash, (u32, VarRef, u32)>,
}

impl ANFConverter {
    /// Create a new ANF converter
    #[must_use]
    pub fn new() -> Self {
        Self {
            binding_depth: 0,
            next_binding_id: 0,
            expr_cache: HashMap::new(),
        }
    }

    /// Convert an `ASTRepr<f64>` expression to ANF
    pub fn convert(&mut self, expr: &ASTRepr<f64>) -> Result<ANFExpr<f64>> {
        Ok(self.to_anf(expr))
    }

    /// Convert `ASTRepr<f64>` to ANF, returning the result and any let-bindings needed
    fn to_anf(&mut self, expr: &ASTRepr<f64>) -> ANFExpr<f64> {
        // Always inline constants as atoms
        if let ASTRepr::Constant(value) = expr {
            return ANFExpr::Atom(ANFAtom::Constant(*value));
        }
        // Check cache for CSE - but only for non-trivial expressions
        if !matches!(expr, ASTRepr::Constant(_) | ASTRepr::Variable(_)) {
            let structural_hash = StructuralHash::from_expr(expr);
            if let Some((cached_scope, cached_var, _cached_binding_id)) =
                self.expr_cache.get(&structural_hash)
            {
                // Cache hit! But only reuse if the variable is still in scope
                // A variable is in scope if it was created at the current depth or shallower
                if *cached_scope <= self.binding_depth {
                    return ANFExpr::Atom(ANFAtom::Variable(*cached_var));
                }
                // Variable is out of scope, remove from cache
                self.expr_cache.remove(&structural_hash);
            }
        }
        match expr {
            ASTRepr::Constant(value) => ANFExpr::Atom(ANFAtom::Constant(*value)),
            ASTRepr::Variable(index) => ANFExpr::Atom(ANFAtom::Variable(VarRef::User(*index))),
            // Binary operations - these will be cached
            ASTRepr::Add(left, right) => {
                self.convert_binary_op_with_cse(expr, left, right, ANFComputation::Add)
            }
            ASTRepr::Sub(left, right) => {
                self.convert_binary_op_with_cse(expr, left, right, ANFComputation::Sub)
            }
            ASTRepr::Mul(left, right) => {
                self.convert_binary_op_with_cse(expr, left, right, ANFComputation::Mul)
            }
            ASTRepr::Div(left, right) => {
                self.convert_binary_op_with_cse(expr, left, right, ANFComputation::Div)
            }
            ASTRepr::Pow(left, right) => {
                self.convert_binary_op_with_cse(expr, left, right, ANFComputation::Pow)
            }
            // Unary operations - these will be cached
            ASTRepr::Neg(inner) => self.convert_unary_op_with_cse(expr, inner, ANFComputation::Neg),
            ASTRepr::Ln(inner) => self.convert_unary_op_with_cse(expr, inner, ANFComputation::Ln),
            ASTRepr::Exp(inner) => self.convert_unary_op_with_cse(expr, inner, ANFComputation::Exp),
            ASTRepr::Sin(inner) => self.convert_unary_op_with_cse(expr, inner, ANFComputation::Sin),
            ASTRepr::Cos(inner) => self.convert_unary_op_with_cse(expr, inner, ANFComputation::Cos),
            ASTRepr::Sqrt(inner) => {
                self.convert_unary_op_with_cse(expr, inner, ANFComputation::Sqrt)
            }
        }
    }

    /// Convert a binary operation to ANF with CSE caching
    fn convert_binary_op_with_cse(
        &mut self,
        expr: &ASTRepr<f64>,
        left: &ASTRepr<f64>,
        right: &ASTRepr<f64>,
        op_constructor: fn(ANFAtom<f64>, ANFAtom<f64>) -> ANFComputation<f64>,
    ) -> ANFExpr<f64> {
        // Special handling for power operations with constant integer exponents
        if matches!(op_constructor(ANFAtom::Constant(0.0), ANFAtom::Constant(0.0)), ANFComputation::Pow(_, _)) {
            if let ASTRepr::Constant(exp_val) = right {
                // Check if it's an integer exponent suitable for binary exponentiation
                if exp_val.fract() == 0.0 && exp_val.abs() <= 64.0 && *exp_val != 0.0 && *exp_val != 1.0 {
                    let exp_int = *exp_val as i32;
                    return self.convert_integer_power_to_anf(left, exp_int);
                }
            }
        }

        fn extract_final_var(expr: &ANFExpr<f64>) -> Option<VarRef> {
            match expr {
                ANFExpr::Let(var, _, body) => extract_final_var(body).or(Some(*var)),
                ANFExpr::Atom(ANFAtom::Variable(var)) => Some(*var),
                _ => None,
            }
        }
        let (left_expr, left_atom_orig) = self.to_anf_atom(left);
        let (right_expr, right_atom_orig) = self.to_anf_atom(right);
        let left_atom = match &left_expr {
            Some(e) => extract_final_var(e).map_or(left_atom_orig, ANFAtom::Variable),
            None => left_atom_orig,
        };
        let right_atom = match &right_expr {
            Some(e) => extract_final_var(e).map_or(right_atom_orig, ANFAtom::Variable),
            None => right_atom_orig,
        };
        let computation = op_constructor(left_atom.clone(), right_atom.clone());
        if left_atom.is_constant() && right_atom.is_constant() {
            let result = match computation {
                ANFComputation::Add(ANFAtom::Constant(a), ANFAtom::Constant(b)) => {
                    ANFAtom::Constant(a + b)
                }
                ANFComputation::Sub(ANFAtom::Constant(a), ANFAtom::Constant(b)) => {
                    ANFAtom::Constant(a - b)
                }
                ANFComputation::Mul(ANFAtom::Constant(a), ANFAtom::Constant(b)) => {
                    ANFAtom::Constant(a * b)
                }
                ANFComputation::Div(ANFAtom::Constant(a), ANFAtom::Constant(b)) => {
                    ANFAtom::Constant(a / b)
                }
                ANFComputation::Pow(ANFAtom::Constant(a), ANFAtom::Constant(b)) => {
                    // Use domain analysis to determine if constant folding is safe
                    let result = a.powf(b);
                    if result.is_finite() && !result.is_nan() {
                        ANFAtom::Constant(result)
                    } else {
                        // Don't fold - skip constant folding and proceed to let-binding
                        // This will be handled by the code below
                        let binding_id = self.next_binding_id;
                        self.next_binding_id += 1;
                        let result_var = VarRef::Bound(binding_id);
                        let structural_hash = StructuralHash::from_expr(expr);
                        self.expr_cache.insert(
                            structural_hash,
                            (self.binding_depth, result_var, binding_id),
                        );
                        self.binding_depth += 1;
                        let body = ANFExpr::Atom(ANFAtom::Variable(result_var));
                        self.binding_depth -= 1;
                        return self.chain_lets(
                            left_expr,
                            right_expr,
                            ANFExpr::Let(result_var, computation, Box::new(body)),
                        );
                    }
                }
                _ => unreachable!(),
            };
            return ANFExpr::Atom(result);
        }
        let binding_id = self.next_binding_id;
        self.next_binding_id += 1;
        let result_var = VarRef::Bound(binding_id);
        let structural_hash = StructuralHash::from_expr(expr);
        self.expr_cache.insert(
            structural_hash,
            (self.binding_depth, result_var, binding_id),
        );
        self.binding_depth += 1;
        let body = ANFExpr::Atom(ANFAtom::Variable(result_var));
        self.binding_depth -= 1;
        self.chain_lets(
            left_expr,
            right_expr,
            ANFExpr::Let(result_var, computation, Box::new(body)),
        )
    }

    /// Convert integer power to optimized ANF using binary exponentiation
    fn convert_integer_power_to_anf(&mut self, base: &ASTRepr<f64>, exp: i32) -> ANFExpr<f64> {
        let (base_expr, base_atom) = self.to_anf_atom(base);
        
        match exp {
            0 => ANFExpr::Atom(ANFAtom::Constant(1.0)),
            1 => match base_expr {
                Some(expr) => expr,
                None => ANFExpr::Atom(base_atom),
            },
            -1 => {
                // x^(-1) = 1/x
                let one = ANFAtom::Constant(1.0);
                let div_computation = ANFComputation::Div(one, base_atom);
                let binding_id = self.next_binding_id;
                self.next_binding_id += 1;
                let result_var = VarRef::Bound(binding_id);
                let body = ANFExpr::Atom(ANFAtom::Variable(result_var));
                let div_expr = ANFExpr::Let(result_var, div_computation, Box::new(body));
                
                match base_expr {
                    Some(expr) => self.wrap_with_lets(Some(expr), div_expr),
                    None => div_expr,
                }
            },
            2 => {
                // x^2 = x * x
                let mul_computation = ANFComputation::Mul(base_atom.clone(), base_atom);
                let binding_id = self.next_binding_id;
                self.next_binding_id += 1;
                let result_var = VarRef::Bound(binding_id);
                let body = ANFExpr::Atom(ANFAtom::Variable(result_var));
                let mul_expr = ANFExpr::Let(result_var, mul_computation, Box::new(body));
                
                match base_expr {
                    Some(expr) => self.wrap_with_lets(Some(expr), mul_expr),
                    None => mul_expr,
                }
            },
            exp if exp > 0 => {
                // Use binary exponentiation for positive powers
                self.generate_binary_exponentiation_anf(base_expr, base_atom, exp as u32)
            },
            exp if exp < 0 => {
                // x^(-n) = 1/(x^n)
                let positive_power = self.generate_binary_exponentiation_anf(base_expr.clone(), base_atom.clone(), (-exp) as u32);
                
                // Create 1/result
                let power_var = self.extract_result_var(&positive_power);
                let one = ANFAtom::Constant(1.0);
                let div_computation = ANFComputation::Div(one, ANFAtom::Variable(power_var));
                let binding_id = self.next_binding_id;
                self.next_binding_id += 1;
                let result_var = VarRef::Bound(binding_id);
                let body = ANFExpr::Atom(ANFAtom::Variable(result_var));
                let div_expr = ANFExpr::Let(result_var, div_computation, Box::new(body));
                
                self.wrap_with_lets(Some(positive_power), div_expr)
            },
            _ => unreachable!(),
        }
    }

    /// Generate binary exponentiation in ANF form
    fn generate_binary_exponentiation_anf(
        &mut self,
        base_expr: Option<ANFExpr<f64>>,
        base_atom: ANFAtom<f64>,
        exp: u32,
    ) -> ANFExpr<f64> {
        if exp == 1 {
            return match base_expr {
                Some(expr) => expr,
                None => ANFExpr::Atom(base_atom),
            };
        }

        // Binary exponentiation: repeatedly square and multiply
        let mut result_atom = base_atom.clone();
        let mut result_expr = base_expr;
        let mut current_exp = exp;
        let mut accumulated_expr: Option<ANFExpr<f64>> = None;
        let mut accumulated_atom: Option<ANFAtom<f64>> = None;

        // Handle the case where exp is odd - multiply by base once
        if current_exp % 2 == 1 {
            accumulated_atom = Some(base_atom.clone());
            accumulated_expr = result_expr.clone();
            current_exp -= 1;
        }

        // Now current_exp is even, repeatedly square
        while current_exp > 0 {
            // Square the current result
            let square_computation = ANFComputation::Mul(result_atom.clone(), result_atom.clone());
            let binding_id = self.next_binding_id;
            self.next_binding_id += 1;
            let square_var = VarRef::Bound(binding_id);
            let square_body = ANFExpr::Atom(ANFAtom::Variable(square_var));
            let square_expr = ANFExpr::Let(square_var, square_computation, Box::new(square_body));
            
            result_expr = Some(match result_expr {
                Some(expr) => self.wrap_with_lets(Some(expr), square_expr),
                None => square_expr,
            });
            result_atom = ANFAtom::Variable(square_var);
            current_exp /= 2;

            // If we need to multiply by this power
            if current_exp % 2 == 1 && current_exp > 0 {
                match (&accumulated_atom, &accumulated_expr) {
                    (Some(acc_atom), acc_expr) => {
                        let mul_computation = ANFComputation::Mul(acc_atom.clone(), result_atom.clone());
                        let binding_id = self.next_binding_id;
                        self.next_binding_id += 1;
                        let mul_var = VarRef::Bound(binding_id);
                        let mul_body = ANFExpr::Atom(ANFAtom::Variable(mul_var));
                        let mul_expr = ANFExpr::Let(mul_var, mul_computation, Box::new(mul_body));
                        
                        accumulated_expr = Some(match (acc_expr, &result_expr) {
                            (Some(acc_e), Some(res_e)) => {
                                let combined = self.wrap_with_lets(Some(res_e.clone()), mul_expr);
                                self.wrap_with_lets(Some(acc_e.clone()), combined)
                            },
                            (Some(acc_e), None) => self.wrap_with_lets(Some(acc_e.clone()), mul_expr),
                            (None, Some(res_e)) => self.wrap_with_lets(Some(res_e.clone()), mul_expr),
                            (None, None) => mul_expr,
                        });
                        accumulated_atom = Some(ANFAtom::Variable(mul_var));
                    },
                    (None, _) => {
                        accumulated_atom = Some(result_atom.clone());
                        accumulated_expr = result_expr.clone();
                    },
                }
                current_exp -= 1;
            }
        }

        // Return the accumulated result or the final squared result
        match (accumulated_expr, accumulated_atom) {
            (Some(expr), _) => expr,
            (None, Some(atom)) => ANFExpr::Atom(atom),
            (None, None) => result_expr.unwrap_or(ANFExpr::Atom(result_atom)),
        }
    }

    /// Extract the result variable from an ANF expression (helper for negative powers)
    fn extract_result_var(&self, expr: &ANFExpr<f64>) -> VarRef {
        match expr {
            ANFExpr::Atom(ANFAtom::Variable(var)) => *var,
            ANFExpr::Let(var, _, _) => *var,
            _ => panic!("Expected variable result from power expression"),
        }
    }

    /// Convert a unary operation to ANF with CSE caching
    fn convert_unary_op_with_cse(
        &mut self,
        expr: &ASTRepr<f64>,
        inner: &ASTRepr<f64>,
        op_constructor: fn(ANFAtom<f64>) -> ANFComputation<f64>,
    ) -> ANFExpr<f64> {
        fn extract_final_var(expr: &ANFExpr<f64>) -> Option<VarRef> {
            match expr {
                ANFExpr::Let(var, _, body) => extract_final_var(body).or(Some(*var)),
                ANFExpr::Atom(ANFAtom::Variable(var)) => Some(*var),
                _ => None,
            }
        }
        let (inner_expr, inner_atom_orig) = self.to_anf_atom(inner);
        let inner_atom = match &inner_expr {
            Some(e) => extract_final_var(e).map_or(inner_atom_orig, ANFAtom::Variable),
            None => inner_atom_orig,
        };
        let computation = op_constructor(inner_atom.clone());
        if inner_atom.is_constant() {
            let result = match computation {
                ANFComputation::Neg(ANFAtom::Constant(a)) => ANFAtom::Constant(-a),
                ANFComputation::Ln(ANFAtom::Constant(a)) => ANFAtom::Constant(a.ln()),
                ANFComputation::Exp(ANFAtom::Constant(a)) => ANFAtom::Constant(a.exp()),
                ANFComputation::Sin(ANFAtom::Constant(a)) => ANFAtom::Constant(a.sin()),
                ANFComputation::Cos(ANFAtom::Constant(a)) => ANFAtom::Constant(a.cos()),
                ANFComputation::Sqrt(ANFAtom::Constant(a)) => ANFAtom::Constant(a.sqrt()),
                _ => unreachable!(),
            };
            return ANFExpr::Atom(result);
        }
        let binding_id = self.next_binding_id;
        self.next_binding_id += 1;
        let result_var = VarRef::Bound(binding_id);
        let structural_hash = StructuralHash::from_expr(expr);
        self.expr_cache.insert(
            structural_hash,
            (self.binding_depth, result_var, binding_id),
        );
        self.binding_depth += 1;
        let body = ANFExpr::Atom(ANFAtom::Variable(result_var));
        self.binding_depth -= 1;
        self.wrap_with_lets(
            inner_expr,
            ANFExpr::Let(result_var, computation, Box::new(body)),
        )
    }

    /// Convert an expression to an atom, generating let-bindings as needed
    fn to_anf_atom(&mut self, expr: &ASTRepr<f64>) -> (Option<ANFExpr<f64>>, ANFAtom<f64>) {
        match expr {
            ASTRepr::Constant(value) => (None, ANFAtom::Constant(*value)),
            ASTRepr::Variable(index) => (None, ANFAtom::Variable(VarRef::User(*index))),
            _ => {
                let anf_expr = self.to_anf(expr);
                match anf_expr {
                    ANFExpr::Atom(atom) => (None, atom),
                    ANFExpr::Let(var, computation, body) => (
                        Some(ANFExpr::Let(var, computation, body)),
                        ANFAtom::Variable(var),
                    ),
                }
            }
        }
    }

    /// Chain two optional ANF expressions with a final expression
    fn chain_lets<T: NumericType + Clone>(
        &self,
        first: Option<ANFExpr<T>>,
        second: Option<ANFExpr<T>>,
        final_expr: ANFExpr<T>,
    ) -> ANFExpr<T> {
        match (first, second) {
            (None, None) => final_expr,
            (Some(first_expr), None) => self.wrap_with_lets(Some(first_expr), final_expr),
            (None, Some(second_expr)) => self.wrap_with_lets(Some(second_expr), final_expr),
            (Some(first_expr), Some(second_expr)) => {
                let combined = self.wrap_with_lets(Some(second_expr), final_expr);
                self.wrap_with_lets(Some(first_expr), combined)
            }
        }
    }

    /// Wrap an expression with let-bindings if needed
    fn wrap_with_lets<T: NumericType + Clone>(
        &self,
        wrapper: Option<ANFExpr<T>>,
        body: ANFExpr<T>,
    ) -> ANFExpr<T> {
        match wrapper {
            None => body,
            Some(ANFExpr::Let(var, computation, inner_body)) => ANFExpr::Let(
                var,
                computation,
                Box::new(self.wrap_with_lets(Some(*inner_body), body)),
            ),
            Some(ANFExpr::Atom(_)) => body, // Atom: just use the body (the let-binding for the unary op)
        }
    }
}

impl Default for ANFConverter {
    fn default() -> Self {
        Self::new()
    }
}

/// Convenience function to convert `ASTRepr<f64>` to ANF
pub fn convert_to_anf(expr: &ASTRepr<f64>) -> Result<ANFExpr<f64>> {
    let mut converter = ANFConverter::new();
    converter.convert(expr)
}

/// Domain-aware ANF converter that integrates interval analysis with ANF transformation
#[derive(Debug)]
pub struct DomainAwareANFConverter {
    /// Basic ANF converter for the core transformation
    anf_converter: ANFConverter,
    /// Domain analyzer for safety checks
    domain_analyzer: IntervalDomainAnalyzer<f64>,
    /// Domain information for generated variables
    /// Maps `VarRef::Bound(id)` to its computed domain
    variable_domains: HashMap<u32, IntervalDomain<f64>>,
    /// Safety validation cache
    safety_cache: HashMap<String, bool>,
}

impl DomainAwareANFConverter {
    /// Create a new domain-aware ANF converter
    #[must_use]
    pub fn new(domain_analyzer: IntervalDomainAnalyzer<f64>) -> Self {
        Self {
            anf_converter: ANFConverter::new(),
            domain_analyzer,
            variable_domains: HashMap::new(),
            safety_cache: HashMap::new(),
        }
    }

    /// Convert an expression to ANF with domain awareness
    pub fn convert(&mut self, expr: &ASTRepr<f64>) -> Result<ANFExpr<f64>> {
        // First, convert using the basic ANF converter
        let anf = self.anf_converter.convert(expr)?;

        // Track domain information for any generated variables
        self.propagate_domain_information(&anf, expr);

        // For now, be less strict about validation to allow reasonable expressions
        // In a production system, this would do more sophisticated domain tracking

        Ok(anf)
    }

    /// Propagate domain information through the ANF structure
    fn propagate_domain_information(&mut self, anf: &ANFExpr<f64>, original_expr: &ASTRepr<f64>) {
        match anf {
            ANFExpr::Atom(_) => {
                // Nothing to propagate for atoms
            }
            ANFExpr::Let(var_ref, computation, body) => {
                // Compute domain for this computation
                let domain = self.compute_computation_domain(computation);

                // Store domain for generated variables
                if let VarRef::Bound(id) = var_ref {
                    self.variable_domains.insert(*id, domain);
                }

                // Recursively propagate through the body
                self.propagate_domain_information(body, original_expr);
            }
        }
    }

    /// Compute the domain of a computation
    fn compute_computation_domain(&self, computation: &ANFComputation<f64>) -> IntervalDomain<f64> {
        match computation {
            ANFComputation::Add(left, right) => {
                let left_domain = self.compute_atom_domain(left);
                let right_domain = self.compute_atom_domain(right);
                // For now, return Top for most operations
                // A full implementation would do proper interval arithmetic
                if left_domain == IntervalDomain::Bottom || right_domain == IntervalDomain::Bottom {
                    IntervalDomain::Bottom
                } else {
                    IntervalDomain::Top
                }
            }
            ANFComputation::Mul(left, right) => {
                let left_domain = self.compute_atom_domain(left);
                let right_domain = self.compute_atom_domain(right);
                // Multiplication of positive numbers is positive
                if left_domain.is_positive(0.0) && right_domain.is_positive(0.0) {
                    IntervalDomain::positive(0.0)
                } else {
                    IntervalDomain::Top
                }
            }
            ANFComputation::Exp(_) => {
                // exp(x) is always positive
                IntervalDomain::positive(0.0)
            }
            _ => IntervalDomain::Top, // Conservative default
        }
    }

    /// Compute the domain of an atom
    fn compute_atom_domain(&self, atom: &ANFAtom<f64>) -> IntervalDomain<f64> {
        match atom {
            ANFAtom::Constant(val) => IntervalDomain::Constant(*val),
            ANFAtom::Variable(var_ref) => self.get_variable_domain(*var_ref),
        }
    }

    /// Get domain information for a variable
    #[must_use]
    pub fn get_variable_domain(&self, var_ref: VarRef) -> IntervalDomain<f64> {
        match var_ref {
            VarRef::User(idx) => self.domain_analyzer.get_variable_domain(idx),
            VarRef::Bound(id) => self
                .variable_domains
                .get(&id)
                .cloned()
                .unwrap_or(IntervalDomain::Top),
        }
    }

    /// Set domain information for a generated variable
    pub fn set_generated_variable_domain(&mut self, var_id: u32, domain: IntervalDomain<f64>) {
        self.variable_domains.insert(var_id, domain);
    }

    /// Get the underlying domain analyzer (for external use)
    #[must_use]
    pub fn domain_analyzer(&self) -> &IntervalDomainAnalyzer<f64> {
        &self.domain_analyzer
    }

    /// Get a mutable reference to the domain analyzer
    pub fn domain_analyzer_mut(&mut self) -> &mut IntervalDomainAnalyzer<f64> {
        &mut self.domain_analyzer
    }

    /// Convert an expression with explicit domain constraint validation
    pub fn convert_with_domain_constraint(
        &mut self,
        expr: &ASTRepr<f64>,
        expected_domain: &IntervalDomain<f64>,
    ) -> Result<ANFExpr<f64>> {
        // First, convert to ANF
        let anf = self.convert(expr)?;

        // Analyze the output domain of the expression
        let output_domain = self.analyze_expression_domain(expr);

        // Check if the output domain is compatible with the expected domain
        if !self.is_domain_compatible(&output_domain, expected_domain) {
            return Err(crate::error::DSLCompileError::DomainError(format!(
                "Expression domain {output_domain:?} is not compatible with expected domain {expected_domain:?}"
            )));
        }

        Ok(anf)
    }

    /// Analyze the domain of an expression
    fn analyze_expression_domain(&self, expr: &ASTRepr<f64>) -> IntervalDomain<f64> {
        match expr {
            ASTRepr::Constant(val) => IntervalDomain::Constant(*val),
            ASTRepr::Variable(idx) => self.domain_analyzer.get_variable_domain(*idx),
            ASTRepr::Exp(_) => {
                // exp(x) is always positive for any real x
                IntervalDomain::positive(0.0)
            }
            ASTRepr::Ln(inner) => {
                // ln(x) requires x > 0, output can be any real number
                let inner_domain = self.analyze_expression_domain(inner);
                if inner_domain.is_positive(0.0) {
                    IntervalDomain::Top // ln of positive number can be any real
                } else {
                    IntervalDomain::Bottom // ln of non-positive is undefined
                }
            }
            ASTRepr::Sqrt(inner) => {
                // sqrt(x) requires x >= 0, output is non-negative
                let inner_domain = self.analyze_expression_domain(inner);
                if inner_domain.is_non_negative(0.0) {
                    IntervalDomain::non_negative(0.0)
                } else {
                    IntervalDomain::Bottom // sqrt of negative is undefined (in reals)
                }
            }
            ASTRepr::Mul(left, right) => {
                let left_domain = self.analyze_expression_domain(left);
                let right_domain = self.analyze_expression_domain(right);
                // Multiplication of positive numbers is positive
                if left_domain.is_positive(0.0) && right_domain.is_positive(0.0) {
                    IntervalDomain::positive(0.0)
                } else {
                    IntervalDomain::Top // Conservative
                }
            }
            _ => IntervalDomain::Top, // Conservative default for other operations
        }
    }

    /// Check if two domains are compatible (one is a subset of the other)
    fn is_domain_compatible(
        &self,
        domain1: &IntervalDomain<f64>,
        domain2: &IntervalDomain<f64>,
    ) -> bool {
        match (domain1, domain2) {
            // Bottom is compatible with everything
            (IntervalDomain::Bottom, _) => true,
            // Nothing is compatible with Bottom except Bottom
            (_, IntervalDomain::Bottom) => false,
            // Top is only compatible with Top
            (IntervalDomain::Top, IntervalDomain::Top) => true,
            (IntervalDomain::Top, _) => false,
            // Everything is compatible with Top
            (_, IntervalDomain::Top) => true,
            // Constants must match exactly
            (IntervalDomain::Constant(a), IntervalDomain::Constant(b)) => a == b,
            // For intervals, check if domain1 is a subset of domain2
            (IntervalDomain::Interval { .. }, IntervalDomain::Interval { .. }) => {
                // For now, be conservative and only allow exact matches
                domain1 == domain2
            }
            // Constants are compatible with intervals if they're contained
            (IntervalDomain::Constant(val), interval) => interval.contains(*val),
            (interval, IntervalDomain::Constant(val)) => interval.contains(*val),
        }
    }

    /// Check if an operation is safe given the current domain information
    pub fn is_operation_safe(
        &mut self,
        operation: &str,
        operands: &[&ASTRepr<f64>],
    ) -> Result<bool> {
        // Create a cache key
        let cache_key = format!("{operation}:{operands:?}");

        // Check cache first
        if let Some(&cached_result) = self.safety_cache.get(&cache_key) {
            return Ok(cached_result);
        }

        let result = match operation {
            "ln" => {
                if operands.len() != 1 {
                    return Ok(false);
                }
                let domain = self.analyze_expression_domain(operands[0]);
                domain.is_positive(0.0)
            }
            "sqrt" => {
                if operands.len() != 1 {
                    return Ok(false);
                }
                let domain = self.analyze_expression_domain(operands[0]);
                domain.is_non_negative(0.0)
            }
            "div" => {
                if operands.len() != 2 {
                    return Ok(false);
                }
                let denominator_domain = self.analyze_expression_domain(operands[1]);
                // Check that denominator is not zero
                !matches!(denominator_domain, IntervalDomain::Constant(x) if x == 0.0)
            }
            _ => true, // Other operations are generally safe
        };

        // Cache the result
        self.safety_cache.insert(cache_key, result);
        Ok(result)
    }

    /// Clear all caches (useful for testing or memory management)
    pub fn clear_caches(&mut self) {
        self.safety_cache.clear();
        self.variable_domains.clear();
    }

    /// Get optimization statistics
    #[must_use]
    pub fn get_optimization_stats(&self) -> DomainAwareOptimizationStats {
        DomainAwareOptimizationStats {
            generated_variables: self.variable_domains.len(),
            safety_checks_cached: self.safety_cache.len(),
            anf_let_bindings: 0, // Would need to traverse current ANF to compute this
        }
    }
}

/// Statistics for domain-aware ANF optimization
#[derive(Debug, Clone)]
pub struct DomainAwareOptimizationStats {
    /// Number of generated variables with domain information
    pub generated_variables: usize,
    /// Number of cached safety checks
    pub safety_checks_cached: usize,
    /// Number of ANF let bindings in the current expression
    pub anf_let_bindings: usize,
}

/// ANF code generator for Rust
#[derive(Debug)]
pub struct ANFCodeGen<'a> {
    registry: &'a VariableRegistry,
}

impl<'a> ANFCodeGen<'a> {
    /// Create a new code generator with a variable registry
    #[must_use]
    pub fn new(registry: &'a VariableRegistry) -> Self {
        Self { registry }
    }

    /// Generate Rust code from an ANF expression
    pub fn generate<T: NumericType + std::fmt::Display>(&self, expr: &ANFExpr<T>) -> String {
        match expr {
            ANFExpr::Atom(atom) => self.generate_atom(atom),
            ANFExpr::Let(var, computation, body) => {
                let var_name = var.to_rust_code(self.registry);
                let comp_code = self.generate_computation(computation);
                let body_code = self.generate(body);

                format!("{{ let {var_name} = {comp_code};\n{body_code} }}")
            }
        }
    }

    /// Generate code for an atomic expression
    fn generate_atom<T: NumericType + std::fmt::Display>(&self, atom: &ANFAtom<T>) -> String {
        match atom {
            ANFAtom::Constant(value) => value.to_string(),
            ANFAtom::Variable(var) => var.to_rust_code(self.registry),
        }
    }

    /// Generate code for a computation
    fn generate_computation<T: NumericType + std::fmt::Display>(
        &self,
        comp: &ANFComputation<T>,
    ) -> String {
        match comp {
            ANFComputation::Add(left, right) => {
                format!(
                    "{} + {}",
                    self.generate_atom(left),
                    self.generate_atom(right)
                )
            }
            ANFComputation::Sub(left, right) => {
                format!(
                    "{} - {}",
                    self.generate_atom(left),
                    self.generate_atom(right)
                )
            }
            ANFComputation::Mul(left, right) => {
                format!(
                    "{} * {}",
                    self.generate_atom(left),
                    self.generate_atom(right)
                )
            }
            ANFComputation::Div(left, right) => {
                format!(
                    "{} / {}",
                    self.generate_atom(left),
                    self.generate_atom(right)
                )
            }
            ANFComputation::Pow(left, right) => {
                format!(
                    "{}.powf({})",
                    self.generate_atom(left),
                    self.generate_atom(right)
                )
            }
            ANFComputation::Neg(operand) => {
                format!("-{}", self.generate_atom(operand))
            }
            ANFComputation::Ln(operand) => {
                format!("{}.ln()", self.generate_atom(operand))
            }
            ANFComputation::Exp(operand) => {
                format!("{}.exp()", self.generate_atom(operand))
            }
            ANFComputation::Sin(operand) => {
                format!("{}.sin()", self.generate_atom(operand))
            }
            ANFComputation::Cos(operand) => {
                format!("{}.cos()", self.generate_atom(operand))
            }
            ANFComputation::Sqrt(operand) => {
                format!("{}.sqrt()", self.generate_atom(operand))
            }
        }
    }

    /// Generate a complete function definition
    pub fn generate_function<T: NumericType + std::fmt::Display>(
        &self,
        name: &str,
        expr: &ANFExpr<T>,
    ) -> String {
        // Generate parameter list based on registry size
        let param_list: Vec<String> = (0..self.registry.len())
            .map(|i| format!("{}: f64", self.registry.debug_name(i)))
            .collect();

        let body = self.generate(expr);

        format!(
            "fn {}({}) -> f64 {{\n    {}\n}}",
            name,
            param_list.join(", "),
            body.replace('\n', "\n    ")
        )
    }
}

/// Convenience function to generate code from ANF
pub fn generate_rust_code<T: NumericType + std::fmt::Display>(
    expr: &ANFExpr<T>,
    registry: &VariableRegistry,
) -> String {
    let codegen = ANFCodeGen::new(registry);
    codegen.generate(expr)
}

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

    #[test]
    fn test_var_generic() {
        let mut generic = ANFVarGen::new();

        let v1 = generic.fresh();
        let v2 = generic.fresh();
        let v3 = generic.user_var(0);

        assert_eq!(v1, VarRef::Bound(0));
        assert_eq!(v2, VarRef::Bound(1));
        assert_eq!(v3, VarRef::User(0));
        assert_ne!(v1, v2);
        assert_ne!(v1, v3);
    }

    #[test]
    fn test_anf_atom() {
        let const_atom: ANFAtom<f64> = ANFAtom::Constant(42.0);
        let var_atom: ANFAtom<f64> = ANFAtom::Variable(VarRef::Bound(0));

        assert!(const_atom.is_constant());
        assert!(!const_atom.is_variable());
        assert_eq!(const_atom.as_constant(), Some(&42.0));

        assert!(var_atom.is_variable());
        assert!(!var_atom.is_constant());
        assert_eq!(var_atom.as_variable(), Some(VarRef::Bound(0)));
    }

    #[test]
    fn test_anf_computation_operands() {
        let a: ANFAtom<f64> = ANFAtom::Constant(1.0);
        let b: ANFAtom<f64> = ANFAtom::Variable(VarRef::Bound(0));

        let add = ANFComputation::Add(a.clone(), b.clone());
        let operands = add.operands();

        assert_eq!(operands.len(), 2);
        assert_eq!(operands[0], &a);
        assert_eq!(operands[1], &b);

        let neg = ANFComputation::Neg(a.clone());
        let neg_operands = neg.operands();
        assert_eq!(neg_operands.len(), 1);
        assert_eq!(neg_operands[0], &a);
    }

    #[test]
    fn test_anf_expr_construction() {
        let var = VarRef::Bound(0);
        let const_val = 42.0;

        let atom_expr: ANFExpr<f64> = ANFExpr::constant(const_val);
        let var_expr: ANFExpr<f64> = ANFExpr::variable(var);

        assert!(atom_expr.is_atom());
        assert!(var_expr.is_atom());

        let computation = ANFComputation::Add(ANFAtom::Variable(var), ANFAtom::Constant(1.0));
        let let_expr: ANFExpr<f64> = ANFExpr::let_binding(var, computation, atom_expr);

        assert!(let_expr.is_let());
        assert_eq!(let_expr.let_count(), 1);
    }

    #[test]
    fn test_variable_collection() {
        let var1 = VarRef::Bound(0);
        let var2 = VarRef::User(0);

        // Create: let t0 = x + 1 in t0 * t0
        let computation = ANFComputation::Add(
            ANFAtom::Variable(var2), // x
            ANFAtom::Constant(1.0),
        );
        let body: ANFExpr<f64> = ANFExpr::Atom(ANFAtom::Variable(var1)); // t0
        let expr: ANFExpr<f64> = ANFExpr::Let(var1, computation, Box::new(body));

        let used_vars = expr.used_variables();
        assert_eq!(used_vars.len(), 2);
        assert!(used_vars.contains(&var1));
        assert!(used_vars.contains(&var2));
    }

    #[test]
    fn test_anf_conversion() {
        use crate::final_tagless::ASTEval;

        // Create expression: sin(x + 1) + cos(x + 1)
        // This should demonstrate CSE automatically
        let x = ASTEval::var(0); // Variable with index 0
        let one = ASTEval::constant(1.0);
        let x_plus_one = ASTEval::add(x.clone(), one.clone());
        let sin_expr = ASTEval::sin(x_plus_one.clone());
        let cos_expr = ASTEval::cos(x_plus_one);
        let full_expr = ASTEval::add(sin_expr, cos_expr);

        // Convert to ANF
        let anf_result = convert_to_anf(&full_expr);
        assert!(anf_result.is_ok());

        let anf = anf_result.unwrap();

        // Should have multiple let bindings
        assert!(anf.let_count() > 0);

        // Verify structure is let-bound
        assert!(anf.is_let());
    }

    #[test]
    fn test_anf_code_generation() {
        use crate::final_tagless::{ASTEval, VariableRegistry};

        // Create a variable registry
        let mut registry = VariableRegistry::new();
        let x_idx = registry.register_variable();

        // Create expression: x * x + 2 * x + 1 (quadratic)
        let x = ASTEval::var(x_idx);
        let two = ASTEval::constant(2.0);
        let one = ASTEval::constant(1.0);
        let x_squared = ASTEval::mul(x.clone(), x.clone());
        let two_x = ASTEval::mul(two, x);
        let sum1 = ASTEval::add(x_squared, two_x);
        let quadratic = ASTEval::add(sum1, one);

        // Convert to ANF
        let anf = convert_to_anf(&quadratic).unwrap();

        // Generate code
        let code = generate_rust_code(&anf, &registry);

        // Code should contain let bindings and be properly structured
        assert!(code.contains("let t"));
        assert!(code.contains("var_0")); // Updated to expect index-based variable name

        // Also test function generation
        let codegen = ANFCodeGen::new(&registry);
        let function_code = codegen.generate_function("quadratic", &anf);

        assert!(function_code.contains("fn quadratic"));
        assert!(function_code.contains("var_0: f64")); // Updated to expect index-based variable name
        assert!(function_code.contains("-> f64"));

        println!("Generated code:\n{code}");
        println!("Generated function:\n{function_code}");
    }

    #[test]
    fn test_anf_complete_pipeline() {
        use crate::final_tagless::{ASTEval, VariableRegistry};

        // Create a variable registry
        let mut registry = VariableRegistry::new();
        let x_idx = registry.register_variable();
        let y_idx = registry.register_variable();

        // Create a complex expression with common subexpressions:
        // sin(x + y) + cos(x + y) + exp(x + y)
        // This should demonstrate automatic CSE of (x + y)
        let x = ASTEval::var(x_idx);
        let y = ASTEval::var(y_idx);
        let x_plus_y = ASTEval::add(x, y);

        let sin_term = ASTEval::sin(x_plus_y.clone());
        let cos_term = ASTEval::cos(x_plus_y.clone());
        let exp_term = ASTEval::exp(x_plus_y);

        let sum1 = ASTEval::add(sin_term, cos_term);
        let final_expr = ASTEval::add(sum1, exp_term);

        // Convert to ANF
        let anf = convert_to_anf(&final_expr).unwrap();

        // Generate code
        let codegen = ANFCodeGen::new(&registry);
        let function_code = codegen.generate_function("demo_function", &anf);

        println!("\n=== ANF Demo: Automatic Common Subexpression Elimination ===");
        println!("Original expression: sin(x + y) + cos(x + y) + exp(x + y)");
        println!("ANF introduces variables for shared subexpressions automatically\n");
        println!("Generated function:");
        println!("{function_code}");

        // Verify the ANF has the expected structure
        assert!(anf.let_count() >= 1); // Should have let bindings
        assert!(function_code.contains("fn demo_function"));
        assert!(function_code.contains("var_0: f64, var_1: f64")); // Updated to expect index-based variable names
        assert!(function_code.contains("-> f64"));

        // The beauty is that this is ready to compile and run!
    }

    #[test]
    fn test_cse_simple_case() {
        use crate::final_tagless::{ASTEval, VariableRegistry};

        // Create a variable registry
        let mut registry = VariableRegistry::new();
        let x_idx = registry.register_variable();

        // Create expression: (x + 1) + (x + 1)
        // This should reuse the computation of (x + 1)
        let x = ASTEval::var(x_idx);
        let one = ASTEval::constant(1.0);
        let x_plus_one_left = ASTEval::add(x.clone(), one.clone());
        let x_plus_one_right = ASTEval::add(x, one);
        let final_expr = ASTEval::add(x_plus_one_left, x_plus_one_right);

        // Convert to ANF
        let anf = convert_to_anf(&final_expr).unwrap();

        // Generate code
        let codegen = ANFCodeGen::new(&registry);
        let function_code = codegen.generate_function("cse_test", &anf);

        println!("\n=== CSE Test: (x + 1) + (x + 1) ===");
        println!("Generated function:");
        println!("{function_code}");

        // Count how many times we see "x + 1" in the generated code
        let code_contains_reuse = function_code.matches("x + 1").count() == 1;
        println!(
            "CSE working correctly: {}",
            if code_contains_reuse {
                "✅ YES"
            } else {
                "❌ NO"
            }
        );

        // Should have fewer let bindings due to reuse
        assert!(anf.let_count() > 0);
    }

    #[test]
    fn test_cse_debug() {
        use crate::final_tagless::{ASTEval, VariableRegistry};

        // Create a very simple case to debug: x + x
        let mut registry = VariableRegistry::new();
        let x_idx = registry.register_variable();

        let x = ASTEval::var(x_idx);
        let expr = ASTEval::add(x.clone(), x.clone()); // x + x - should reuse x

        // Convert to ANF
        let anf = convert_to_anf(&expr).unwrap();

        // Generate code
        let codegen = ANFCodeGen::new(&registry);
        let function_code = codegen.generate_function("debug_test", &anf);

        println!("\n=== CSE Debug: x + x ===");
        println!("Generated function:");
        println!("{function_code}");

        // This should have only one variable, not duplicate computations
        assert!(anf.let_count() == 1);
    }

    #[test]
    fn test_cse_failing_case() {
        use crate::final_tagless::{ASTEval, VariableRegistry};

        // Create the exact failing case: (x + 1) + (x + 1)
        let mut registry = VariableRegistry::new();
        let x_idx = registry.register_variable();

        let x = ASTEval::var(x_idx);
        let one = ASTEval::constant(1.0);
        let x_plus_one_left = ASTEval::add(x.clone(), one.clone());
        let x_plus_one_right = ASTEval::add(x, one);
        let expr = ASTEval::add(x_plus_one_left, x_plus_one_right);

        // Convert to ANF
        let anf = convert_to_anf(&expr).unwrap();

        // Generate code
        let codegen = ANFCodeGen::new(&registry);
        let function_code = codegen.generate_function("failing_case", &anf);

        println!("\n=== CSE Failing Case: (x + 1) + (x + 1) ===");
        println!("Generated function:");
        println!("{function_code}");
        println!("Let count: {}", anf.let_count());

        // Debug the ANF structure
        println!("ANF structure: {anf:#?}");
    }
}