//! Optimization passes for Ruchy code generation

use crate::ast::{BinaryOp, Literal, Param, PipelineStage, RuchyExpr};
use depyler_core::hir::HirModule;
use std::collections::{HashMap, HashSet};

/// Ruchy code optimizer
pub struct RuchyOptimizer {
    /// Optimization level (0-3)
    level: u8,

    /// Enable pipeline fusion
    enable_pipeline_fusion: bool,

    /// Enable dead code elimination
    enable_dce: bool,

    /// Enable common subexpression elimination
    enable_cse: bool,

    /// Enable function inlining
    enable_inlining: bool,
}

impl RuchyOptimizer {
    /// Creates a new optimizer with default settings
    #[must_use]
    pub fn new() -> Self {
        Self {
            level: 2,
            enable_pipeline_fusion: true,
            enable_dce: true,
            enable_cse: true,
            enable_inlining: true,
        }
    }

    /// Creates optimizer with custom configuration
    #[must_use]
    pub fn with_config(config: &crate::RuchyConfig) -> Self {
        Self {
            level: config.optimization_level,
            enable_pipeline_fusion: config.optimization_level >= 1,
            enable_dce: config.optimization_level >= 1,
            enable_cse: config.optimization_level >= 2,
            enable_inlining: config.optimization_level >= 2,
        }
    }

    /// Optimize a Ruchy AST
    pub fn optimize(&self, expr: RuchyExpr) -> anyhow::Result<RuchyExpr> {
        let mut optimized = expr;

        if self.level == 0 {
            return Ok(optimized);
        }

        // Apply optimization passes in order
        if self.enable_pipeline_fusion {
            optimized = self.fuse_pipelines(optimized);
        }

        if self.enable_cse {
            optimized = self.eliminate_common_subexpressions(optimized);
        }

        if self.enable_inlining {
            optimized = self.inline_simple_functions(optimized);
        }

        if self.enable_dce {
            optimized = self.eliminate_dead_code(optimized);
        }

        // Constant folding is always enabled at level 1+
        optimized = self.fold_constants(optimized);

        Ok(optimized)
    }

    /// Optimize HIR before transpilation
    pub fn optimize_hir(&self, hir: HirModule) -> HirModule {
        // HIR-level optimizations
        // For now, return as-is
        hir
    }

    /// Fuse adjacent pipeline operations
    fn fuse_pipelines(&self, expr: RuchyExpr) -> RuchyExpr {
        self.transform_expr(expr, &|e| match e {
            RuchyExpr::Pipeline { expr, stages } => {
                let fused_stages = self.fuse_pipeline_stages(stages);
                RuchyExpr::Pipeline {
                    expr,
                    stages: fused_stages,
                }
            }
            _ => e,
        })
    }

    /// Fuse compatible pipeline stages
    fn fuse_pipeline_stages(&self, stages: Vec<PipelineStage>) -> Vec<PipelineStage> {
        let mut fused = Vec::new();
        let mut i = 0;

        while i < stages.len() {
            if i + 1 < stages.len() {
                // Try to fuse current and next stage
                if let Some(merged) = self.try_fuse_stages(&stages[i], &stages[i + 1]) {
                    fused.push(merged);
                    i += 2;
                    continue;
                }
            }

            fused.push(stages[i].clone());
            i += 1;
        }

        fused
    }

    /// Try to fuse two pipeline stages
    fn try_fuse_stages(
        &self,
        first: &PipelineStage,
        second: &PipelineStage,
    ) -> Option<PipelineStage> {
        match (first, second) {
            // Fuse consecutive maps: map(f) |> map(g) => map(|x| g(f(x)))
            (PipelineStage::Map(f), PipelineStage::Map(g)) => {
                Some(PipelineStage::Map(Box::new(self.compose_functions(f, g))))
            }

            // Fuse consecutive filters: filter(p) |> filter(q) => filter(|x| p(x) && q(x))
            (PipelineStage::Filter(p), PipelineStage::Filter(q)) => {
                Some(PipelineStage::Filter(Box::new(self.and_predicates(p, q))))
            }

            // Fuse filter then map into flat_map with conditional
            (PipelineStage::Filter(p), PipelineStage::Map(f)) => Some(PipelineStage::FlatMap(
                Box::new(self.filter_map_fusion(p, f)),
            )),

            _ => None,
        }
    }

    /// Compose two functions
    fn compose_functions(&self, f: &RuchyExpr, g: &RuchyExpr) -> RuchyExpr {
        // Create: |x| g(f(x))
        let x_param = Param {
            name: "x".to_string(),
            typ: None,
            default: None,
        };

        let f_call = match f {
            RuchyExpr::Lambda { params, body } if params.len() == 1 => {
                // Substitute x for the parameter in f's body
                self.substitute_var(
                    &params[0].name,
                    &RuchyExpr::Identifier("x".to_string()),
                    body,
                )
            }
            _ => RuchyExpr::Call {
                func: Box::new(f.clone()),
                args: vec![RuchyExpr::Identifier("x".to_string())],
            },
        };

        let g_call = match g {
            RuchyExpr::Lambda { params, body } if params.len() == 1 => {
                // Substitute f(x) for the parameter in g's body
                self.substitute_var(&params[0].name, &f_call, body)
            }
            _ => RuchyExpr::Call {
                func: Box::new(g.clone()),
                args: vec![f_call],
            },
        };

        RuchyExpr::Lambda {
            params: vec![x_param],
            body: Box::new(g_call),
        }
    }

    /// Combine two predicates with AND
    fn and_predicates(&self, p: &RuchyExpr, q: &RuchyExpr) -> RuchyExpr {
        let x_param = Param {
            name: "x".to_string(),
            typ: None,
            default: None,
        };

        let p_call = self.apply_lambda_or_call(p, "x");
        let q_call = self.apply_lambda_or_call(q, "x");

        RuchyExpr::Lambda {
            params: vec![x_param],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(p_call),
                op: BinaryOp::And,
                right: Box::new(q_call),
            }),
        }
    }

    /// Fuse filter and map into flat_map
    fn filter_map_fusion(&self, predicate: &RuchyExpr, mapper: &RuchyExpr) -> RuchyExpr {
        let x_param = Param {
            name: "x".to_string(),
            typ: None,
            default: None,
        };

        let pred_call = self.apply_lambda_or_call(predicate, "x");
        let map_call = self.apply_lambda_or_call(mapper, "x");

        RuchyExpr::Lambda {
            params: vec![x_param],
            body: Box::new(RuchyExpr::If {
                condition: Box::new(pred_call),
                then_branch: Box::new(RuchyExpr::List(vec![map_call])),
                else_branch: Some(Box::new(RuchyExpr::List(vec![]))),
            }),
        }
    }

    /// Apply lambda or create function call
    fn apply_lambda_or_call(&self, func: &RuchyExpr, arg_name: &str) -> RuchyExpr {
        match func {
            RuchyExpr::Lambda { params, body } if params.len() == 1 => self.substitute_var(
                &params[0].name,
                &RuchyExpr::Identifier(arg_name.to_string()),
                body,
            ),
            _ => RuchyExpr::Call {
                func: Box::new(func.clone()),
                args: vec![RuchyExpr::Identifier(arg_name.to_string())],
            },
        }
    }

    /// Substitute a variable in an expression
    #[allow(clippy::only_used_in_recursion)]
    fn substitute_var(&self, var: &str, replacement: &RuchyExpr, expr: &RuchyExpr) -> RuchyExpr {
        match expr {
            RuchyExpr::Identifier(name) if name == var => replacement.clone(),

            RuchyExpr::Binary { left, op, right } => RuchyExpr::Binary {
                left: Box::new(self.substitute_var(var, replacement, left)),
                op: *op,
                right: Box::new(self.substitute_var(var, replacement, right)),
            },

            RuchyExpr::Call { func, args } => RuchyExpr::Call {
                func: Box::new(self.substitute_var(var, replacement, func)),
                args: args
                    .iter()
                    .map(|arg| self.substitute_var(var, replacement, arg))
                    .collect(),
            },

            _ => expr.clone(),
        }
    }

    /// Eliminate common subexpressions
    fn eliminate_common_subexpressions(&self, expr: RuchyExpr) -> RuchyExpr {
        let mut cse_map = HashMap::new();
        let mut next_temp_id = 0;

        self.cse_transform(expr, &mut cse_map, &mut next_temp_id)
    }

    /// CSE transformation helper
    fn cse_transform(
        &self,
        expr: RuchyExpr,
        cse_map: &mut HashMap<String, String>,
        next_id: &mut usize,
    ) -> RuchyExpr {
        // Hash the expression
        let expr_hash = self.hash_expr(&expr);

        // Check if we've seen this expression before
        if let Some(temp_var) = cse_map.get(&expr_hash) {
            return RuchyExpr::Identifier(temp_var.clone());
        }

        // Process the expression
        match expr {
            RuchyExpr::Binary { left, op, right }
                if self.is_pure_expr(&left) && self.is_pure_expr(&right) =>
            {
                let left_transformed = self.cse_transform(*left, cse_map, next_id);
                let right_transformed = self.cse_transform(*right, cse_map, next_id);

                let new_expr = RuchyExpr::Binary {
                    left: Box::new(left_transformed),
                    op,
                    right: Box::new(right_transformed),
                };

                // Create temporary variable for complex expressions to enable common subexpression elimination
                if self.is_complex_expr(&new_expr) {
                    let temp_var = format!("_cse_{}", next_id);
                    *next_id += 1;
                    cse_map.insert(self.hash_expr(&new_expr), temp_var.clone());
                    RuchyExpr::Identifier(temp_var)
                } else {
                    new_expr
                }
            }

            RuchyExpr::Block(exprs) => {
                let transformed = exprs
                    .into_iter()
                    .map(|e| self.cse_transform(e, cse_map, next_id))
                    .collect();
                RuchyExpr::Block(transformed)
            }

            _ => expr,
        }
    }

    /// Hash an expression for CSE
    fn hash_expr(&self, expr: &RuchyExpr) -> String {
        // Simple string representation for hashing
        format!("{:?}", expr)
    }

    /// Check if expression is pure (no side effects)
    #[allow(clippy::only_used_in_recursion)]
    fn is_pure_expr(&self, expr: &RuchyExpr) -> bool {
        match expr {
            RuchyExpr::Literal(_) | RuchyExpr::Identifier(_) => true,
            RuchyExpr::Binary { left, right, .. } => {
                self.is_pure_expr(left) && self.is_pure_expr(right)
            }
            RuchyExpr::Unary { operand, .. } => self.is_pure_expr(operand),
            _ => false,
        }
    }

    /// Check if expression is complex enough to CSE
    fn is_complex_expr(&self, expr: &RuchyExpr) -> bool {
        matches!(expr, RuchyExpr::Binary { .. } | RuchyExpr::Call { .. })
    }

    /// Inline simple functions
    fn inline_simple_functions(&self, expr: RuchyExpr) -> RuchyExpr {
        let mut inline_map = HashMap::new();

        // First pass: collect inline candidates
        self.collect_inline_candidates(&expr, &mut inline_map);

        // Second pass: perform inlining
        self.perform_inlining(expr, &inline_map)
    }

    /// Collect functions that can be inlined
    fn collect_inline_candidates(
        &self,
        expr: &RuchyExpr,
        inline_map: &mut HashMap<String, RuchyExpr>,
    ) {
        match expr {
            RuchyExpr::Function {
                name, params, body, ..
            } => {
                if self.should_inline(params, body) {
                    inline_map.insert(name.clone(), (**body).clone());
                }
            }

            RuchyExpr::Block(exprs) => {
                for e in exprs {
                    self.collect_inline_candidates(e, inline_map);
                }
            }

            _ => {}
        }
    }

    /// Check if function should be inlined
    fn should_inline(&self, params: &[Param], body: &RuchyExpr) -> bool {
        // Inline if:
        // - No parameters or single parameter
        // - Body is simple (not recursive, small size)
        params.len() <= 1 && self.expr_size(body) < 10
    }

    /// Calculate expression size for inlining heuristic
    #[allow(clippy::only_used_in_recursion)]
    fn expr_size(&self, expr: &RuchyExpr) -> usize {
        match expr {
            RuchyExpr::Literal(_) | RuchyExpr::Identifier(_) => 1,
            RuchyExpr::Binary { left, right, .. } => {
                1 + self.expr_size(left) + self.expr_size(right)
            }
            RuchyExpr::Call { args, .. } => {
                2 + args.iter().map(|a| self.expr_size(a)).sum::<usize>()
            }
            RuchyExpr::Block(exprs) => exprs.iter().map(|e| self.expr_size(e)).sum(),
            _ => 5, // Default weight for complex expressions
        }
    }

    /// Perform function inlining
    fn perform_inlining(
        &self,
        expr: RuchyExpr,
        inline_map: &HashMap<String, RuchyExpr>,
    ) -> RuchyExpr {
        self.transform_expr(expr, &|e| match e {
            RuchyExpr::Call { func, args } => {
                if let RuchyExpr::Identifier(name) = func.as_ref() {
                    if let Some(body) = inline_map.get(name) {
                        // Inline the function body
                        return body.clone();
                    }
                }
                RuchyExpr::Call { func, args }
            }
            _ => e,
        })
    }

    /// Eliminate dead code
    fn eliminate_dead_code(&self, expr: RuchyExpr) -> RuchyExpr {
        let mut used_vars = HashSet::new();

        // First pass: collect used variables
        self.collect_used_vars(&expr, &mut used_vars);

        // Second pass: remove unused definitions
        self.remove_unused_defs(expr, &used_vars)
    }

    /// Collect used variables
    #[allow(clippy::only_used_in_recursion)]
    fn collect_used_vars(&self, expr: &RuchyExpr, used: &mut HashSet<String>) {
        match expr {
            RuchyExpr::Identifier(name) => {
                used.insert(name.clone());
            }

            RuchyExpr::Binary { left, right, .. } => {
                self.collect_used_vars(left, used);
                self.collect_used_vars(right, used);
            }

            RuchyExpr::Call { func, args } => {
                self.collect_used_vars(func, used);
                for arg in args {
                    self.collect_used_vars(arg, used);
                }
            }

            RuchyExpr::Block(exprs) => {
                for e in exprs {
                    self.collect_used_vars(e, used);
                }
            }

            _ => {}
        }
    }

    /// Remove unused definitions
    fn remove_unused_defs(&self, expr: RuchyExpr, used_vars: &HashSet<String>) -> RuchyExpr {
        match expr {
            RuchyExpr::Let {
                name,
                value,
                body,
                is_mutable,
            } => {
                if used_vars.contains(&name) {
                    RuchyExpr::Let {
                        name,
                        value,
                        body: Box::new(self.remove_unused_defs(*body, used_vars)),
                        is_mutable,
                    }
                } else {
                    // Skip unused let binding
                    self.remove_unused_defs(*body, used_vars)
                }
            }

            RuchyExpr::Block(exprs) => {
                let filtered: Vec<_> = exprs
                    .into_iter()
                    .filter_map(|e| {
                        let transformed = self.remove_unused_defs(e, used_vars);
                        // Remove empty blocks
                        if matches!(transformed, RuchyExpr::Block(ref v) if v.is_empty()) {
                            None
                        } else {
                            Some(transformed)
                        }
                    })
                    .collect();

                RuchyExpr::Block(filtered)
            }

            _ => expr,
        }
    }

    /// Fold constant expressions
    fn fold_constants(&self, expr: RuchyExpr) -> RuchyExpr {
        self.transform_expr(expr, &|e| match e {
            RuchyExpr::Binary { left, op, right } => {
                match (left.as_ref(), right.as_ref()) {
                    (RuchyExpr::Literal(l1), RuchyExpr::Literal(l2)) => {
                        if let Some(result) = self.fold_binary_op(l1, op, l2) {
                            return RuchyExpr::Literal(result);
                        }
                    }
                    _ => {}
                }
                RuchyExpr::Binary { left, op, right }
            }

            RuchyExpr::Unary { op, operand } => {
                if let RuchyExpr::Literal(lit) = operand.as_ref() {
                    if let Some(result) = self.fold_unary_op(op, lit) {
                        return RuchyExpr::Literal(result);
                    }
                }
                RuchyExpr::Unary { op, operand }
            }

            _ => e,
        })
    }

    /// Fold binary operation on literals
    fn fold_binary_op(&self, left: &Literal, op: BinaryOp, right: &Literal) -> Option<Literal> {
        match (left, right) {
            (Literal::Integer(a), Literal::Integer(b)) => Some(match op {
                BinaryOp::Add => Literal::Integer(a + b),
                BinaryOp::Subtract => Literal::Integer(a - b),
                BinaryOp::Multiply => Literal::Integer(a * b),
                BinaryOp::Divide if *b != 0 => Literal::Integer(a / b),
                BinaryOp::Modulo if *b != 0 => Literal::Integer(a % b),
                BinaryOp::Equal => Literal::Bool(a == b),
                BinaryOp::NotEqual => Literal::Bool(a != b),
                BinaryOp::Less => Literal::Bool(a < b),
                BinaryOp::LessEqual => Literal::Bool(a <= b),
                BinaryOp::Greater => Literal::Bool(a > b),
                BinaryOp::GreaterEqual => Literal::Bool(a >= b),
                _ => return None,
            }),

            (Literal::Float(a), Literal::Float(b)) => Some(match op {
                BinaryOp::Add => Literal::Float(a + b),
                BinaryOp::Subtract => Literal::Float(a - b),
                BinaryOp::Multiply => Literal::Float(a * b),
                BinaryOp::Divide if *b != 0.0 => Literal::Float(a / b),
                _ => return None,
            }),

            (Literal::Bool(a), Literal::Bool(b)) => Some(match op {
                BinaryOp::And => Literal::Bool(*a && *b),
                BinaryOp::Or => Literal::Bool(*a || *b),
                BinaryOp::Equal => Literal::Bool(a == b),
                BinaryOp::NotEqual => Literal::Bool(a != b),
                _ => return None,
            }),

            _ => None,
        }
    }

    /// Fold unary operation on literal
    fn fold_unary_op(&self, op: crate::ast::UnaryOp, lit: &Literal) -> Option<Literal> {
        match (op, lit) {
            (crate::ast::UnaryOp::Negate, Literal::Integer(n)) => Some(Literal::Integer(-n)),
            (crate::ast::UnaryOp::Negate, Literal::Float(f)) => Some(Literal::Float(-f)),
            (crate::ast::UnaryOp::Not, Literal::Bool(b)) => Some(Literal::Bool(!b)),
            _ => None,
        }
    }

    /// Transform expression with a function
    fn transform_expr<F>(&self, expr: RuchyExpr, transform_fn: &F) -> RuchyExpr
    where
        F: Fn(RuchyExpr) -> RuchyExpr,
    {
        let transformed = match expr {
            RuchyExpr::Binary { left, op, right } => RuchyExpr::Binary {
                left: Box::new(self.transform_expr(*left, transform_fn)),
                op,
                right: Box::new(self.transform_expr(*right, transform_fn)),
            },

            RuchyExpr::Unary { op, operand } => RuchyExpr::Unary {
                op,
                operand: Box::new(self.transform_expr(*operand, transform_fn)),
            },

            RuchyExpr::Call { func, args } => RuchyExpr::Call {
                func: Box::new(self.transform_expr(*func, transform_fn)),
                args: args
                    .into_iter()
                    .map(|arg| self.transform_expr(arg, transform_fn))
                    .collect(),
            },

            RuchyExpr::Block(exprs) => RuchyExpr::Block(
                exprs
                    .into_iter()
                    .map(|e| self.transform_expr(e, transform_fn))
                    .collect(),
            ),

            RuchyExpr::If {
                condition,
                then_branch,
                else_branch,
            } => RuchyExpr::If {
                condition: Box::new(self.transform_expr(*condition, transform_fn)),
                then_branch: Box::new(self.transform_expr(*then_branch, transform_fn)),
                else_branch: else_branch.map(|e| Box::new(self.transform_expr(*e, transform_fn))),
            },

            RuchyExpr::Pipeline { expr, stages } => RuchyExpr::Pipeline {
                expr: Box::new(self.transform_expr(*expr, transform_fn)),
                stages, // Stages are already optimized in fuse_pipelines
            },

            _ => expr,
        };

        transform_fn(transformed)
    }
}

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

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

    #[test]
    fn test_constant_folding() {
        let optimizer = RuchyOptimizer::new();

        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
        };

        let result = optimizer.fold_constants(expr);

        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(5)));
    }

    #[test]
    fn test_pipeline_fusion() {
        let optimizer = RuchyOptimizer::new();

        // Create a pipeline with two consecutive maps
        let pipeline = RuchyExpr::Pipeline {
            expr: Box::new(RuchyExpr::List(vec![])),
            stages: vec![
                PipelineStage::Map(Box::new(RuchyExpr::Lambda {
                    params: vec![Param {
                        name: "x".to_string(),
                        typ: None,
                        default: None,
                    }],
                    body: Box::new(RuchyExpr::Binary {
                        left: Box::new(RuchyExpr::Identifier("x".to_string())),
                        op: BinaryOp::Add,
                        right: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
                    }),
                })),
                PipelineStage::Map(Box::new(RuchyExpr::Lambda {
                    params: vec![Param {
                        name: "y".to_string(),
                        typ: None,
                        default: None,
                    }],
                    body: Box::new(RuchyExpr::Binary {
                        left: Box::new(RuchyExpr::Identifier("y".to_string())),
                        op: BinaryOp::Multiply,
                        right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
                    }),
                })),
            ],
        };

        let result = optimizer.fuse_pipelines(pipeline);

        // Should have fused the two maps into one
        if let RuchyExpr::Pipeline { stages, .. } = result {
            assert_eq!(stages.len(), 1);
        } else {
            panic!("Expected pipeline");
        }
    }

    // Test optimizer creation and configuration
    #[test]
    fn test_optimizer_new() {
        let optimizer = RuchyOptimizer::new();
        assert_eq!(optimizer.level, 2);
        assert!(optimizer.enable_pipeline_fusion);
        assert!(optimizer.enable_dce);
        assert!(optimizer.enable_cse);
        assert!(optimizer.enable_inlining);
    }

    #[test]
    fn test_optimizer_default() {
        let optimizer = RuchyOptimizer::default();
        assert_eq!(optimizer.level, 2);
    }

    #[test]
    fn test_optimizer_with_config_level_0() {
        let config = crate::RuchyConfig {
            optimization_level: 0,
            ..Default::default()
        };
        let optimizer = RuchyOptimizer::with_config(&config);
        assert_eq!(optimizer.level, 0);
        assert!(!optimizer.enable_pipeline_fusion);
        assert!(!optimizer.enable_dce);
        assert!(!optimizer.enable_cse);
        assert!(!optimizer.enable_inlining);
    }

    #[test]
    fn test_optimizer_with_config_level_1() {
        let config = crate::RuchyConfig {
            optimization_level: 1,
            ..Default::default()
        };
        let optimizer = RuchyOptimizer::with_config(&config);
        assert_eq!(optimizer.level, 1);
        assert!(optimizer.enable_pipeline_fusion);
        assert!(optimizer.enable_dce);
        assert!(!optimizer.enable_cse);
        assert!(!optimizer.enable_inlining);
    }

    #[test]
    fn test_optimizer_with_config_level_2() {
        let config = crate::RuchyConfig {
            optimization_level: 2,
            ..Default::default()
        };
        let optimizer = RuchyOptimizer::with_config(&config);
        assert_eq!(optimizer.level, 2);
        assert!(optimizer.enable_pipeline_fusion);
        assert!(optimizer.enable_dce);
        assert!(optimizer.enable_cse);
        assert!(optimizer.enable_inlining);
    }

    #[test]
    fn test_optimizer_with_config_level_3() {
        let config = crate::RuchyConfig {
            optimization_level: 3,
            ..Default::default()
        };
        let optimizer = RuchyOptimizer::with_config(&config);
        assert_eq!(optimizer.level, 3);
        assert!(optimizer.enable_pipeline_fusion);
        assert!(optimizer.enable_dce);
        assert!(optimizer.enable_cse);
        assert!(optimizer.enable_inlining);
    }

    // Test optimize method
    #[test]
    fn test_optimize_level_0_no_changes() {
        let config = crate::RuchyConfig {
            optimization_level: 0,
            ..Default::default()
        };
        let optimizer = RuchyOptimizer::with_config(&config);

        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
        };

        let result = optimizer.optimize(expr.clone()).unwrap();
        // At level 0, expression should be unchanged (no constant folding)
        assert_eq!(result, expr);
    }

    #[test]
    fn test_optimize_full_pipeline() {
        let optimizer = RuchyOptimizer::new();

        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            op: BinaryOp::Multiply,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };

        // The optimize() method applies multiple passes in order, which may result in CSE
        // extracting the expression. The important thing is that it doesn't error.
        let result = optimizer.optimize(expr);
        assert!(result.is_ok());
    }

    // Test constant folding operations
    #[test]
    fn test_fold_binary_integer_subtract() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            op: BinaryOp::Subtract,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(7)));
    }

    #[test]
    fn test_fold_binary_integer_multiply() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(4))),
            op: BinaryOp::Multiply,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(20)));
    }

    #[test]
    fn test_fold_binary_integer_divide() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(20))),
            op: BinaryOp::Divide,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(4))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(5)));
    }

    #[test]
    fn test_fold_binary_integer_divide_by_zero() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(20))),
            op: BinaryOp::Divide,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(0))),
        };
        let result = optimizer.fold_constants(expr.clone());
        // Division by zero should not fold
        assert_eq!(result, expr);
    }

    #[test]
    fn test_fold_binary_integer_modulo() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(17))),
            op: BinaryOp::Modulo,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(2)));
    }

    #[test]
    fn test_fold_binary_integer_modulo_by_zero() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(17))),
            op: BinaryOp::Modulo,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(0))),
        };
        let result = optimizer.fold_constants(expr.clone());
        // Modulo by zero should not fold
        assert_eq!(result, expr);
    }

    #[test]
    fn test_fold_binary_integer_comparison_equal() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
            op: BinaryOp::Equal,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_integer_comparison_not_equal() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
            op: BinaryOp::NotEqual,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_integer_comparison_less() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
            op: BinaryOp::Less,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_integer_comparison_less_equal() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
            op: BinaryOp::LessEqual,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_integer_comparison_greater() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            op: BinaryOp::Greater,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_integer_comparison_greater_equal() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            op: BinaryOp::GreaterEqual,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    // Float constant folding
    #[test]
    fn test_fold_binary_float_add() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Float(2.5))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Float(3.5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Float(6.0)));
    }

    #[test]
    fn test_fold_binary_float_subtract() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Float(10.0))),
            op: BinaryOp::Subtract,
            right: Box::new(RuchyExpr::Literal(Literal::Float(4.0))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Float(6.0)));
    }

    #[test]
    fn test_fold_binary_float_multiply() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Float(3.0))),
            op: BinaryOp::Multiply,
            right: Box::new(RuchyExpr::Literal(Literal::Float(4.0))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Float(12.0)));
    }

    #[test]
    fn test_fold_binary_float_divide() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Float(10.0))),
            op: BinaryOp::Divide,
            right: Box::new(RuchyExpr::Literal(Literal::Float(2.0))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Float(5.0)));
    }

    #[test]
    fn test_fold_binary_float_divide_by_zero() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Float(10.0))),
            op: BinaryOp::Divide,
            right: Box::new(RuchyExpr::Literal(Literal::Float(0.0))),
        };
        let result = optimizer.fold_constants(expr.clone());
        // Division by zero should not fold
        assert_eq!(result, expr);
    }

    // Boolean constant folding
    #[test]
    fn test_fold_binary_bool_and_true() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
            op: BinaryOp::And,
            right: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_bool_and_false() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
            op: BinaryOp::And,
            right: Box::new(RuchyExpr::Literal(Literal::Bool(false))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(false)));
    }

    #[test]
    fn test_fold_binary_bool_or_true() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Bool(false))),
            op: BinaryOp::Or,
            right: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_bool_or_false() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Bool(false))),
            op: BinaryOp::Or,
            right: Box::new(RuchyExpr::Literal(Literal::Bool(false))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(false)));
    }

    #[test]
    fn test_fold_binary_bool_equal() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
            op: BinaryOp::Equal,
            right: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    #[test]
    fn test_fold_binary_bool_not_equal() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
            op: BinaryOp::NotEqual,
            right: Box::new(RuchyExpr::Literal(Literal::Bool(false))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(true)));
    }

    // Unary constant folding
    #[test]
    fn test_fold_unary_negate_integer() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Unary {
            op: UnaryOp::Negate,
            operand: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(-5)));
    }

    #[test]
    fn test_fold_unary_negate_float() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Unary {
            op: UnaryOp::Negate,
            operand: Box::new(RuchyExpr::Literal(Literal::Float(3.14))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Float(-3.14)));
    }

    #[test]
    fn test_fold_unary_not_bool() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Unary {
            op: UnaryOp::Not,
            operand: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
        };
        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Bool(false)));
    }

    // Test is_pure_expr
    #[test]
    fn test_is_pure_literal() {
        let optimizer = RuchyOptimizer::new();
        assert!(optimizer.is_pure_expr(&RuchyExpr::Literal(Literal::Integer(5))));
    }

    #[test]
    fn test_is_pure_identifier() {
        let optimizer = RuchyOptimizer::new();
        assert!(optimizer.is_pure_expr(&RuchyExpr::Identifier("x".to_string())));
    }

    #[test]
    fn test_is_pure_binary() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Identifier("x".to_string())),
        };
        assert!(optimizer.is_pure_expr(&expr));
    }

    #[test]
    fn test_is_pure_unary() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Unary {
            op: UnaryOp::Negate,
            operand: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
        };
        assert!(optimizer.is_pure_expr(&expr));
    }

    #[test]
    fn test_is_not_pure_call() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Call {
            func: Box::new(RuchyExpr::Identifier("print".to_string())),
            args: vec![],
        };
        assert!(!optimizer.is_pure_expr(&expr));
    }

    // Test is_complex_expr
    #[test]
    fn test_is_complex_binary() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
        };
        assert!(optimizer.is_complex_expr(&expr));
    }

    #[test]
    fn test_is_complex_call() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Call {
            func: Box::new(RuchyExpr::Identifier("foo".to_string())),
            args: vec![],
        };
        assert!(optimizer.is_complex_expr(&expr));
    }

    #[test]
    fn test_is_not_complex_literal() {
        let optimizer = RuchyOptimizer::new();
        assert!(!optimizer.is_complex_expr(&RuchyExpr::Literal(Literal::Integer(5))));
    }

    // Test expr_size
    #[test]
    fn test_expr_size_literal() {
        let optimizer = RuchyOptimizer::new();
        assert_eq!(
            optimizer.expr_size(&RuchyExpr::Literal(Literal::Integer(5))),
            1
        );
    }

    #[test]
    fn test_expr_size_identifier() {
        let optimizer = RuchyOptimizer::new();
        assert_eq!(
            optimizer.expr_size(&RuchyExpr::Identifier("x".to_string())),
            1
        );
    }

    #[test]
    fn test_expr_size_binary() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
        };
        assert_eq!(optimizer.expr_size(&expr), 3); // 1 + 1 + 1
    }

    #[test]
    fn test_expr_size_call() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Call {
            func: Box::new(RuchyExpr::Identifier("foo".to_string())),
            args: vec![
                RuchyExpr::Literal(Literal::Integer(1)),
                RuchyExpr::Literal(Literal::Integer(2)),
            ],
        };
        assert_eq!(optimizer.expr_size(&expr), 4); // 2 + 1 + 1
    }

    #[test]
    fn test_expr_size_block() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Block(vec![
            RuchyExpr::Literal(Literal::Integer(1)),
            RuchyExpr::Literal(Literal::Integer(2)),
        ]);
        assert_eq!(optimizer.expr_size(&expr), 2); // 1 + 1
    }

    #[test]
    fn test_expr_size_complex() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::If {
            condition: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
            then_branch: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            else_branch: None,
        };
        assert_eq!(optimizer.expr_size(&expr), 5); // Default weight for complex
    }

    // Test should_inline
    #[test]
    fn test_should_inline_small_function() {
        let optimizer = RuchyOptimizer::new();
        let params = vec![Param {
            name: "x".to_string(),
            typ: None,
            default: None,
        }];
        let body = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Identifier("x".to_string())),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
        };
        assert!(optimizer.should_inline(&params, &body));
    }

    #[test]
    fn test_should_not_inline_many_params() {
        let optimizer = RuchyOptimizer::new();
        let params = vec![
            Param {
                name: "x".to_string(),
                typ: None,
                default: None,
            },
            Param {
                name: "y".to_string(),
                typ: None,
                default: None,
            },
        ];
        let body = RuchyExpr::Literal(Literal::Integer(1));
        assert!(!optimizer.should_inline(&params, &body));
    }

    // Test hash_expr
    #[test]
    fn test_hash_expr_same() {
        let optimizer = RuchyOptimizer::new();
        let expr1 = RuchyExpr::Literal(Literal::Integer(5));
        let expr2 = RuchyExpr::Literal(Literal::Integer(5));
        assert_eq!(optimizer.hash_expr(&expr1), optimizer.hash_expr(&expr2));
    }

    #[test]
    fn test_hash_expr_different() {
        let optimizer = RuchyOptimizer::new();
        let expr1 = RuchyExpr::Literal(Literal::Integer(5));
        let expr2 = RuchyExpr::Literal(Literal::Integer(6));
        assert_ne!(optimizer.hash_expr(&expr1), optimizer.hash_expr(&expr2));
    }

    // Test substitute_var
    #[test]
    fn test_substitute_var_identifier() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Identifier("x".to_string());
        let replacement = RuchyExpr::Literal(Literal::Integer(5));
        let result = optimizer.substitute_var("x", &replacement, &expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(5)));
    }

    #[test]
    fn test_substitute_var_no_match() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Identifier("y".to_string());
        let replacement = RuchyExpr::Literal(Literal::Integer(5));
        let result = optimizer.substitute_var("x", &replacement, &expr);
        assert_eq!(result, RuchyExpr::Identifier("y".to_string()));
    }

    #[test]
    fn test_substitute_var_binary() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Identifier("x".to_string())),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Identifier("x".to_string())),
        };
        let replacement = RuchyExpr::Literal(Literal::Integer(5));
        let result = optimizer.substitute_var("x", &replacement, &expr);

        if let RuchyExpr::Binary { left, right, .. } = result {
            assert_eq!(*left, RuchyExpr::Literal(Literal::Integer(5)));
            assert_eq!(*right, RuchyExpr::Literal(Literal::Integer(5)));
        } else {
            panic!("Expected binary");
        }
    }

    #[test]
    fn test_substitute_var_call() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Call {
            func: Box::new(RuchyExpr::Identifier("f".to_string())),
            args: vec![RuchyExpr::Identifier("x".to_string())],
        };
        let replacement = RuchyExpr::Literal(Literal::Integer(5));
        let result = optimizer.substitute_var("x", &replacement, &expr);

        if let RuchyExpr::Call { args, .. } = result {
            assert_eq!(args[0], RuchyExpr::Literal(Literal::Integer(5)));
        } else {
            panic!("Expected call");
        }
    }

    // Test filter fusion
    #[test]
    fn test_fuse_consecutive_filters() {
        let optimizer = RuchyOptimizer::new();

        let p = RuchyExpr::Lambda {
            params: vec![Param {
                name: "x".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("x".to_string())),
                op: BinaryOp::Greater,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(0))),
            }),
        };

        let q = RuchyExpr::Lambda {
            params: vec![Param {
                name: "y".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("y".to_string())),
                op: BinaryOp::Less,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            }),
        };

        let stage1 = PipelineStage::Filter(Box::new(p));
        let stage2 = PipelineStage::Filter(Box::new(q));

        let result = optimizer.try_fuse_stages(&stage1, &stage2);
        assert!(result.is_some());
        assert!(matches!(result, Some(PipelineStage::Filter(_))));
    }

    #[test]
    fn test_fuse_filter_map() {
        let optimizer = RuchyOptimizer::new();

        let pred = RuchyExpr::Lambda {
            params: vec![Param {
                name: "x".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("x".to_string())),
                op: BinaryOp::Greater,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(0))),
            }),
        };

        let mapper = RuchyExpr::Lambda {
            params: vec![Param {
                name: "y".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("y".to_string())),
                op: BinaryOp::Multiply,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            }),
        };

        let stage1 = PipelineStage::Filter(Box::new(pred));
        let stage2 = PipelineStage::Map(Box::new(mapper));

        let result = optimizer.try_fuse_stages(&stage1, &stage2);
        assert!(result.is_some());
        assert!(matches!(result, Some(PipelineStage::FlatMap(_))));
    }

    #[test]
    fn test_cannot_fuse_incompatible() {
        let optimizer = RuchyOptimizer::new();

        let stage1 = PipelineStage::Reduce(Box::new(RuchyExpr::Identifier("add".to_string())));
        let stage2 = PipelineStage::Map(Box::new(RuchyExpr::Identifier("f".to_string())));

        let result = optimizer.try_fuse_stages(&stage1, &stage2);
        assert!(result.is_none());
    }

    // Test apply_lambda_or_call
    #[test]
    fn test_apply_lambda_with_lambda() {
        let optimizer = RuchyOptimizer::new();
        let func = RuchyExpr::Lambda {
            params: vec![Param {
                name: "a".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("a".to_string())),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            }),
        };

        let result = optimizer.apply_lambda_or_call(&func, "x");
        // Should substitute a with x in body
        if let RuchyExpr::Binary { left, .. } = result {
            assert_eq!(*left, RuchyExpr::Identifier("x".to_string()));
        } else {
            panic!("Expected binary");
        }
    }

    #[test]
    fn test_apply_lambda_or_call_non_lambda() {
        let optimizer = RuchyOptimizer::new();
        let func = RuchyExpr::Identifier("foo".to_string());
        let result = optimizer.apply_lambda_or_call(&func, "x");

        if let RuchyExpr::Call { func: f, args } = result {
            assert_eq!(*f, RuchyExpr::Identifier("foo".to_string()));
            assert_eq!(args.len(), 1);
            assert_eq!(args[0], RuchyExpr::Identifier("x".to_string()));
        } else {
            panic!("Expected call");
        }
    }

    // Test dead code elimination
    #[test]
    fn test_collect_used_vars_identifier() {
        let optimizer = RuchyOptimizer::new();
        let mut used = HashSet::new();
        let expr = RuchyExpr::Identifier("x".to_string());
        optimizer.collect_used_vars(&expr, &mut used);
        assert!(used.contains("x"));
    }

    #[test]
    fn test_collect_used_vars_binary() {
        let optimizer = RuchyOptimizer::new();
        let mut used = HashSet::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Identifier("x".to_string())),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Identifier("y".to_string())),
        };
        optimizer.collect_used_vars(&expr, &mut used);
        assert!(used.contains("x"));
        assert!(used.contains("y"));
    }

    #[test]
    fn test_collect_used_vars_call() {
        let optimizer = RuchyOptimizer::new();
        let mut used = HashSet::new();
        let expr = RuchyExpr::Call {
            func: Box::new(RuchyExpr::Identifier("f".to_string())),
            args: vec![
                RuchyExpr::Identifier("a".to_string()),
                RuchyExpr::Identifier("b".to_string()),
            ],
        };
        optimizer.collect_used_vars(&expr, &mut used);
        assert!(used.contains("f"));
        assert!(used.contains("a"));
        assert!(used.contains("b"));
    }

    #[test]
    fn test_collect_used_vars_block() {
        let optimizer = RuchyOptimizer::new();
        let mut used = HashSet::new();
        let expr = RuchyExpr::Block(vec![
            RuchyExpr::Identifier("x".to_string()),
            RuchyExpr::Identifier("y".to_string()),
        ]);
        optimizer.collect_used_vars(&expr, &mut used);
        assert!(used.contains("x"));
        assert!(used.contains("y"));
    }

    #[test]
    fn test_remove_unused_let() {
        let optimizer = RuchyOptimizer::new();
        let used = HashSet::new();
        // x is not used
        let expr = RuchyExpr::Let {
            name: "x".to_string(),
            value: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
            body: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            is_mutable: false,
        };

        let result = optimizer.remove_unused_defs(expr, &used);
        // The let binding should be removed, leaving just the body
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(10)));
    }

    #[test]
    fn test_keep_used_let() {
        let optimizer = RuchyOptimizer::new();
        let mut used = HashSet::new();
        used.insert("x".to_string());

        let expr = RuchyExpr::Let {
            name: "x".to_string(),
            value: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
            body: Box::new(RuchyExpr::Identifier("x".to_string())),
            is_mutable: false,
        };

        let result = optimizer.remove_unused_defs(expr.clone(), &used);
        // The let binding should remain
        if let RuchyExpr::Let { name, .. } = result {
            assert_eq!(name, "x");
        } else {
            panic!("Expected Let");
        }
    }

    // Test CSE
    #[test]
    fn test_cse_simple() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Block(vec![
            RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            },
            RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            },
        ]);

        let result = optimizer.eliminate_common_subexpressions(expr);
        // Both expressions should be CSE'd to the same temp var
        if let RuchyExpr::Block(exprs) = result {
            assert_eq!(exprs.len(), 2);
        } else {
            panic!("Expected block");
        }
    }

    // Test compose_functions
    #[test]
    fn test_compose_functions_lambdas() {
        let optimizer = RuchyOptimizer::new();

        let f = RuchyExpr::Lambda {
            params: vec![Param {
                name: "a".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("a".to_string())),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            }),
        };

        let g = RuchyExpr::Lambda {
            params: vec![Param {
                name: "b".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("b".to_string())),
                op: BinaryOp::Multiply,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            }),
        };

        let result = optimizer.compose_functions(&f, &g);

        if let RuchyExpr::Lambda { params, .. } = result {
            assert_eq!(params.len(), 1);
            assert_eq!(params[0].name, "x");
        } else {
            panic!("Expected lambda");
        }
    }

    #[test]
    fn test_compose_functions_non_lambda() {
        let optimizer = RuchyOptimizer::new();
        let f = RuchyExpr::Identifier("f".to_string());
        let g = RuchyExpr::Identifier("g".to_string());

        let result = optimizer.compose_functions(&f, &g);

        if let RuchyExpr::Lambda { params, body } = result {
            assert_eq!(params.len(), 1);
            assert_eq!(params[0].name, "x");
            // Body should be g(f(x))
            if let RuchyExpr::Call { func, args } = *body {
                assert_eq!(*func, RuchyExpr::Identifier("g".to_string()));
                assert_eq!(args.len(), 1);
                // First arg should be f(x)
                if let RuchyExpr::Call {
                    func: inner_func, ..
                } = &args[0]
                {
                    assert_eq!(**inner_func, RuchyExpr::Identifier("f".to_string()));
                } else {
                    panic!("Expected inner call");
                }
            } else {
                panic!("Expected outer call");
            }
        } else {
            panic!("Expected lambda");
        }
    }

    // Test and_predicates
    #[test]
    fn test_and_predicates() {
        let optimizer = RuchyOptimizer::new();

        let p = RuchyExpr::Lambda {
            params: vec![Param {
                name: "a".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("a".to_string())),
                op: BinaryOp::Greater,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(0))),
            }),
        };

        let q = RuchyExpr::Lambda {
            params: vec![Param {
                name: "b".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("b".to_string())),
                op: BinaryOp::Less,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
            }),
        };

        let result = optimizer.and_predicates(&p, &q);

        if let RuchyExpr::Lambda { params, body } = result {
            assert_eq!(params.len(), 1);
            assert_eq!(params[0].name, "x");
            if let RuchyExpr::Binary { op, .. } = *body {
                assert_eq!(op, BinaryOp::And);
            } else {
                panic!("Expected binary And");
            }
        } else {
            panic!("Expected lambda");
        }
    }

    // Test filter_map_fusion
    #[test]
    fn test_filter_map_fusion() {
        let optimizer = RuchyOptimizer::new();

        let pred = RuchyExpr::Lambda {
            params: vec![Param {
                name: "a".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("a".to_string())),
                op: BinaryOp::Greater,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(0))),
            }),
        };

        let mapper = RuchyExpr::Lambda {
            params: vec![Param {
                name: "b".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Identifier("b".to_string())),
                op: BinaryOp::Multiply,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
            }),
        };

        let result = optimizer.filter_map_fusion(&pred, &mapper);

        if let RuchyExpr::Lambda { params, body } = result {
            assert_eq!(params.len(), 1);
            assert_eq!(params[0].name, "x");
            if let RuchyExpr::If {
                else_branch: Some(_),
                ..
            } = *body
            {
                // Expected structure
            } else {
                panic!("Expected if with else");
            }
        } else {
            panic!("Expected lambda");
        }
    }

    // Test transform_expr
    #[test]
    fn test_transform_expr_identity() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Literal(Literal::Integer(5));
        let result = optimizer.transform_expr(expr.clone(), &|e| e);
        assert_eq!(result, expr);
    }

    #[test]
    fn test_transform_expr_binary() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            op: BinaryOp::Add,
            right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
        };
        let result = optimizer.transform_expr(expr, &|e| {
            if let RuchyExpr::Literal(Literal::Integer(n)) = e {
                RuchyExpr::Literal(Literal::Integer(n * 10))
            } else {
                e
            }
        });

        if let RuchyExpr::Binary { left, right, .. } = result {
            assert_eq!(*left, RuchyExpr::Literal(Literal::Integer(10)));
            assert_eq!(*right, RuchyExpr::Literal(Literal::Integer(20)));
        } else {
            panic!("Expected binary");
        }
    }

    #[test]
    fn test_transform_expr_if() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::If {
            condition: Box::new(RuchyExpr::Literal(Literal::Bool(true))),
            then_branch: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            else_branch: Some(Box::new(RuchyExpr::Literal(Literal::Integer(2)))),
        };

        let result = optimizer.transform_expr(expr, &|e| e);
        assert!(matches!(result, RuchyExpr::If { .. }));
    }

    #[test]
    fn test_transform_expr_pipeline() {
        let optimizer = RuchyOptimizer::new();
        let expr = RuchyExpr::Pipeline {
            expr: Box::new(RuchyExpr::List(vec![])),
            stages: vec![PipelineStage::Map(Box::new(RuchyExpr::Identifier(
                "f".to_string(),
            )))],
        };

        let result = optimizer.transform_expr(expr, &|e| e);
        assert!(matches!(result, RuchyExpr::Pipeline { .. }));
    }

    // Test inline candidates collection
    #[test]
    fn test_collect_inline_candidates_function() {
        let optimizer = RuchyOptimizer::new();
        let mut inline_map = HashMap::new();

        let expr = RuchyExpr::Function {
            name: "small_fn".to_string(),
            params: vec![Param {
                name: "x".to_string(),
                typ: None,
                default: None,
            }],
            body: Box::new(RuchyExpr::Identifier("x".to_string())),
            return_type: None,
            is_async: false,
        };

        optimizer.collect_inline_candidates(&expr, &mut inline_map);
        assert!(inline_map.contains_key("small_fn"));
    }

    #[test]
    fn test_collect_inline_candidates_block() {
        let optimizer = RuchyOptimizer::new();
        let mut inline_map = HashMap::new();

        let expr = RuchyExpr::Block(vec![RuchyExpr::Function {
            name: "inner_fn".to_string(),
            params: vec![],
            body: Box::new(RuchyExpr::Literal(Literal::Integer(1))),
            return_type: None,
            is_async: false,
        }]);

        optimizer.collect_inline_candidates(&expr, &mut inline_map);
        assert!(inline_map.contains_key("inner_fn"));
    }

    // Test full optimization flow
    #[test]
    fn test_optimize_hir_passthrough() {
        let optimizer = RuchyOptimizer::new();
        let hir = HirModule {
            functions: vec![],
            classes: vec![],
            imports: vec![],
            type_aliases: vec![],
            constants: vec![],
            protocols: vec![],
        };

        let result = optimizer.optimize_hir(hir.clone());
        assert_eq!(result.functions.len(), hir.functions.len());
    }

    #[test]
    fn test_full_optimize_with_all_passes() {
        let optimizer = RuchyOptimizer::new();

        // Create a complex expression that triggers multiple passes
        let expr = RuchyExpr::Block(vec![
            RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
            },
            RuchyExpr::Pipeline {
                expr: Box::new(RuchyExpr::List(vec![
                    RuchyExpr::Literal(Literal::Integer(1)),
                    RuchyExpr::Literal(Literal::Integer(2)),
                ])),
                stages: vec![PipelineStage::Map(Box::new(RuchyExpr::Lambda {
                    params: vec![Param {
                        name: "x".to_string(),
                        typ: None,
                        default: None,
                    }],
                    body: Box::new(RuchyExpr::Identifier("x".to_string())),
                }))],
            },
        ]);

        let result = optimizer.optimize(expr).unwrap();
        // Should successfully optimize without error
        assert!(matches!(result, RuchyExpr::Block(_)));
    }

    #[test]
    fn test_inline_simple_functions() {
        let optimizer = RuchyOptimizer::new();

        // Create a block with a function and a call to it
        let expr = RuchyExpr::Block(vec![
            RuchyExpr::Function {
                name: "double".to_string(),
                params: vec![Param {
                    name: "x".to_string(),
                    typ: None,
                    default: None,
                }],
                body: Box::new(RuchyExpr::Binary {
                    left: Box::new(RuchyExpr::Identifier("x".to_string())),
                    op: BinaryOp::Multiply,
                    right: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
                }),
                return_type: None,
                is_async: false,
            },
            RuchyExpr::Call {
                func: Box::new(RuchyExpr::Identifier("double".to_string())),
                args: vec![RuchyExpr::Literal(Literal::Integer(5))],
            },
        ]);

        let result = optimizer.inline_simple_functions(expr);
        // The function should be inlined
        if let RuchyExpr::Block(exprs) = result {
            assert_eq!(exprs.len(), 2);
        } else {
            panic!("Expected block");
        }
    }

    #[test]
    fn test_eliminate_dead_code() {
        let optimizer = RuchyOptimizer::new();

        // Create expression with unused variable
        let expr = RuchyExpr::Block(vec![
            RuchyExpr::Let {
                name: "unused".to_string(),
                value: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
                body: Box::new(RuchyExpr::Let {
                    name: "used".to_string(),
                    value: Box::new(RuchyExpr::Literal(Literal::Integer(10))),
                    body: Box::new(RuchyExpr::Identifier("used".to_string())),
                    is_mutable: false,
                }),
                is_mutable: false,
            },
        ]);

        let result = optimizer.eliminate_dead_code(expr);
        // The unused let should be eliminated
        assert!(matches!(result, RuchyExpr::Block(_)));
    }

    #[test]
    fn test_fuse_pipeline_stages_single_stage() {
        let optimizer = RuchyOptimizer::new();
        let stages = vec![PipelineStage::Map(Box::new(RuchyExpr::Identifier(
            "f".to_string(),
        )))];

        let result = optimizer.fuse_pipeline_stages(stages);
        assert_eq!(result.len(), 1);
    }

    #[test]
    fn test_fuse_pipeline_stages_unfusable() {
        let optimizer = RuchyOptimizer::new();
        let stages = vec![
            PipelineStage::Reduce(Box::new(RuchyExpr::Identifier("add".to_string()))),
            PipelineStage::Call("collect".to_string(), vec![]),
        ];

        let result = optimizer.fuse_pipeline_stages(stages);
        // Should remain unfused
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_nested_constant_folding() {
        let optimizer = RuchyOptimizer::new();

        // (2 + 3) * (4 + 5)
        let expr = RuchyExpr::Binary {
            left: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Literal(Literal::Integer(2))),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(3))),
            }),
            op: BinaryOp::Multiply,
            right: Box::new(RuchyExpr::Binary {
                left: Box::new(RuchyExpr::Literal(Literal::Integer(4))),
                op: BinaryOp::Add,
                right: Box::new(RuchyExpr::Literal(Literal::Integer(5))),
            }),
        };

        let result = optimizer.fold_constants(expr);
        assert_eq!(result, RuchyExpr::Literal(Literal::Integer(45)));
    }
}
