//! Comprehensive tests for graph optimization and visualization features
//!
//! This test module demonstrates the new graph optimization and visualization
//! capabilities added to the scirs2-autograd system.

use scirs2_autograd as ag;
use scirs2_autograd::optimization::{
    memory_optimization::{MemoryOptimizationConfig, MemoryOptimizer},
    ConstantFolder, ExpressionSimplifier, GraphOptimizer, OptimizationConfig, OptimizationLevel,
};
use scirs2_autograd::tensor_ops as T;
use scirs2_autograd::visualization::{OutputFormat, VisualizationConfig};
use scirs2_core::ndarray::{Array, IxDyn};

/// Test graph visualization functionality
#[cfg(test)]
mod visualization_tests {
    use super::*;

    #[test]
    fn test_visualization_config_creation() {
        let config = VisualizationConfig::default();
        let _ = config;

        // Test different output formats
        let formats = [
            OutputFormat::Dot,
            OutputFormat::Text,
            OutputFormat::Json,
            OutputFormat::Mermaid,
        ];

        for format in formats {
            let config = VisualizationConfig {
                format,
                ..Default::default()
            };
            let _ = config;
        }
    }

    #[test]
    fn test_visualization_config_options() {
        let mut config = VisualizationConfig::default();
        assert!(config.show_shapes);
        assert!(config.show_operations);
        assert!(!config.show_gradients);

        // Test configuration modifications
        config.show_gradients = true;
        config.show_values = true;
        config.max_nodes = Some(50);

        assert!(config.show_gradients);
        assert!(config.show_values);
        assert_eq!(config.max_nodes, Some(50));
    }

    #[test]
    fn test_graph_to_dot_api() {
        // Test the new graph_to_dot API with a real graph
        ag::run(|ctx: &mut ag::Context<f64>| {
            let x = ctx.placeholder("x", &[3]);
            let y = x * 2.0;
            let loss = T::reduction::sum_all(y);

            let dot = ag::visualization::graph_to_dot(&loss, ctx);
            assert!(dot.contains("digraph computation_graph"));
            assert!(dot.contains("->"));

            let summary = ag::visualization::graph_summary(&loss, ctx);
            assert!(summary.contains("Total nodes:"));

            let json = ag::visualization::graph_to_json(&loss, ctx);
            assert!(json.contains("\"nodes\""));

            let mermaid = ag::visualization::graph_to_mermaid(&loss, ctx);
            assert!(mermaid.contains("graph BT"));

            let stats = ag::visualization::GraphStats::from_tensor(&loss, ctx);
            assert!(stats.total_nodes > 0);
        });
    }
}

/// Test graph optimization functionality
#[cfg(test)]
mod optimization_tests {
    use super::*;

    #[test]
    fn test_graph_optimizer_creation() {
        let _optimizer = GraphOptimizer::<f32>::new();

        // Test with different optimization levels
        let _basic = GraphOptimizer::<f32>::with_level(OptimizationLevel::Basic);
        let _standard = GraphOptimizer::<f32>::with_level(OptimizationLevel::Standard);
        let _aggressive = GraphOptimizer::<f32>::with_level(OptimizationLevel::Aggressive);
        let _none = GraphOptimizer::<f32>::with_level(OptimizationLevel::None);

        // Test with custom config
        let config = OptimizationConfig {
            constant_folding: true,
            cse: true,
            expression_simplification: false,
            dead_code_elimination: true,
            operation_fusion: false,
            memory_optimization: true,
            max_passes: 3,
            level: OptimizationLevel::Standard,
        };
        let _custom_optimizer = GraphOptimizer::<f32>::with_config(config);
    }

    #[test]
    fn test_optimization_levels() {
        let none_config = OptimizationLevel::None.config();
        assert!(!none_config.constant_folding);
        assert!(!none_config.cse);
        assert_eq!(none_config.max_passes, 0);

        let basic_config = OptimizationLevel::Basic.config();
        assert!(basic_config.constant_folding);
        assert!(basic_config.dead_code_elimination);
        assert!(!basic_config.operation_fusion);

        let aggressive_config = OptimizationLevel::Aggressive.config();
        assert!(aggressive_config.operation_fusion);
        assert!(aggressive_config.memory_optimization);
        assert_eq!(aggressive_config.max_passes, 10);
    }

    #[test]
    fn test_constant_folding() {
        let mut folder = ConstantFolder::<f32>::new();

        // Test that we can create and use the constant folder
        let _default_folder = ConstantFolder::<f32>::default();

        // Test cache operations
        folder.clear_cache();

        // Test constant checking (these will return false without real nodes)
        let dummy_tensor_id: ag::graph::TensorID = 42;
        assert!(!folder.is_constant(dummy_tensor_id));
        assert_eq!(folder.get_constant_value(dummy_tensor_id), None);
    }

    #[test]
    fn test_expression_simplification() {
        let mut simplifier = ExpressionSimplifier::<f32>::new();
        let _default_simplifier = ExpressionSimplifier::<f32>::default();

        // Test cache operations
        simplifier.clear_cache();

        // Test that the simplifier has rules loaded
        // (We can't test actual simplification without real nodes)
    }

    #[test]
    fn test_memory_optimization() {
        let optimizer = MemoryOptimizer::<f32>::new();

        let config = MemoryOptimizationConfig {
            enable_gradient_checkpointing: true,
            enable_memory_pooling: false,
            enable_in_place_operations: true,
            enable_tensor_reuse: false,
            enable_lifetime_optimization: true,
            checkpoint_memory_threshold: 2048,
            checkpoint_compute_threshold: 1.5,
            pool_frequency_threshold: 3,
            max_memory_usage: Some(1024 * 1024),
        };

        let _custom_optimizer = MemoryOptimizer::<f32>::with_config(config);

        // Test that analysis can be retrieved (will be None initially)
        assert!(optimizer.get_analysis().is_none());
    }

    #[test]
    fn test_optimization_report() {
        use scirs2_autograd::optimization::OptimizationReport;

        let mut report = OptimizationReport::new();
        assert_eq!(report.total_optimizations(), 0);
        assert!(!report.has_optimizations());

        report.constant_folding_applied = 5;
        report.cse_applied = 3;
        report.expressions_simplified = 2;

        assert_eq!(report.total_optimizations(), 10);
        assert!(report.has_optimizations());

        // Test report printing (should not panic)
        report.print_summary();
    }

    #[test]
    fn test_memory_optimization_report() {
        use scirs2_autograd::optimization::memory_optimization::MemoryOptimizationReport;

        let mut report = MemoryOptimizationReport::new();
        assert_eq!(report.total_optimizations(), 0);

        report.gradient_checkpoints_added = 3;
        report.in_place_operations_applied = 7;
        report.tensors_reused = 5;

        assert_eq!(report.total_optimizations(), 15);

        // Test report printing
        report.print_summary();
    }

    #[test]
    fn test_public_optimization_api() {
        // use scirs2_autograd::optimization; // Would be used with a proper graph

        // Note: These functions exist and can be called, but we skip actual testing
        // since creating a valid graph requires more setup
        // let dummy_graph: &mut ag::graph::Graph<f32> = ...; // would need proper initialization

        // These functions exist and can be tested with a proper graph
        // With a proper graph, these would work:
        // let _optimize_result = optimization::optimize_graph(graph);
        // let _level_result = optimization::optimize_graph_with_level(graph, OptimizationLevel::Basic);
        // let config = OptimizationConfig::default();
        // let _config_result = optimization::optimize_graph_with_config(graph, config);
        // let _const_fold_result = optimization::apply_constant_folding(graph);
        // let _dce_result = optimization::apply_dead_code_elimination(graph);
        // let _cse_result = optimization::apply_cse(graph);
    }
}

/// Test integration of optimization and visualization
#[cfg(test)]
mod integration_tests {
    use super::*;

    #[test]
    fn test_optimization_and_visualization_workflow() {
        // Test that we can create both optimizers and visualizers
        let _optimizer = GraphOptimizer::<f32>::with_level(OptimizationLevel::Standard);
        let _memory_optimizer = MemoryOptimizer::<f32>::new();

        // Test configuration combinations
        let opt_config = OptimizationConfig {
            constant_folding: true,
            cse: true,
            expression_simplification: true,
            dead_code_elimination: true,
            operation_fusion: false,
            memory_optimization: true,
            max_passes: 5,
            level: OptimizationLevel::Standard,
        };

        let vis_config = VisualizationConfig {
            show_shapes: true,
            show_operations: true,
            show_gradients: true,
            max_nodes: Some(200),
            format: OutputFormat::Dot,
            show_values: false,
        };

        let mem_config = MemoryOptimizationConfig {
            enable_gradient_checkpointing: true,
            enable_memory_pooling: true,
            enable_in_place_operations: true,
            enable_tensor_reuse: true,
            enable_lifetime_optimization: true,
            checkpoint_memory_threshold: 1024,
            checkpoint_compute_threshold: 2.0,
            pool_frequency_threshold: 5,
            max_memory_usage: None,
        };

        let _combined_optimizer = GraphOptimizer::<f32>::with_config(opt_config);
        let _ = vis_config;
        let _combined_memory_optimizer = MemoryOptimizer::<f32>::with_config(mem_config);
    }

    #[test]
    fn test_pattern_matching_and_simplification() {
        use scirs2_autograd::optimization::SimplificationPattern;

        // Test simplification patterns
        let pattern = SimplificationPattern::AddZero;
        assert_eq!(pattern, SimplificationPattern::AddZero);

        // Operation properties testing would go here when expression_simplification module is completed
        // Example: assert!(is_commutative("Add"));
        // Example: assert!(is_associative("Add"));
        // Example: assert!(has_identity("Add"));
        // Example: assert_eq!(get_identity::<f32>("Add"), Some(0.0));
    }

    #[test]
    fn test_memory_analysis_and_lifetime() {
        use scirs2_autograd::optimization::memory_optimization::{
            MemoryAnalysis, MemoryPoolManager, TensorLifetime, TensorLifetimeAnalyzer,
        };

        // Test memory analysis
        let mut analysis = MemoryAnalysis::new();
        analysis.total_memory_allocated = 2048;
        analysis.peak_memory_usage = 1536;
        analysis.num_allocations = 20;
        analysis.num_deallocations = 18;

        assert_eq!(analysis.memory_efficiency(), 0.75);
        assert_eq!(analysis.allocation_balance(), 2);

        // Test tensor lifetimes
        let lifetime1 = TensorLifetime {
            allocation_time: 0,
            deallocation_time: 10,
            size: 100,
            peak_usage: 100,
        };

        let lifetime2 = TensorLifetime {
            allocation_time: 5,
            deallocation_time: 15,
            size: 200,
            peak_usage: 200,
        };

        assert!(lifetime1.overlaps_with(&lifetime2));
        assert_eq!(lifetime1.duration(), 10);

        // Test lifetime analyzer
        let _analyzer = TensorLifetimeAnalyzer::<f32>::new();

        // Test memory pool manager
        let mut manager = MemoryPoolManager::<f32>::new();
        let buffer = manager.get_buffer(100);
        assert_eq!(buffer.len(), 100);

        manager.return_buffer(buffer);
        assert_eq!(manager.get_stats().buffers_returned, 1);

        let buffer2 = manager.get_buffer(100);
        assert_eq!(buffer2.len(), 100);
        assert_eq!(manager.get_stats().pool_hits, 1);
    }

    // Temporarily disabled - depends on graph_rewriting module which is under development
    // #[test]
    // fn test_fusion_patterns_and_graph_rewriting() {
    //     // Test would go here when graph_rewriting module is completed
    // }

    #[test]
    fn test_comprehensive_optimization_pipeline() {
        // Test that we can create a comprehensive optimization pipeline

        // 1. Create optimizers
        let graph_optimizer = GraphOptimizer::<f32>::with_level(OptimizationLevel::Aggressive);
        let memory_optimizer = MemoryOptimizer::<f32>::with_config(MemoryOptimizationConfig {
            enable_gradient_checkpointing: true,
            enable_memory_pooling: true,
            enable_in_place_operations: true,
            enable_tensor_reuse: true,
            enable_lifetime_optimization: true,
            ..Default::default()
        });

        // 2. Create visualization config
        let vis_config = VisualizationConfig {
            show_shapes: true,
            show_operations: true,
            show_gradients: true,
            format: OutputFormat::Json,
            ..Default::default()
        };

        // 3. Test that components can work together
        // In a real scenario, this would involve:
        // - Visualizing the original graph
        // - Applying optimizations
        // - Visualizing the optimized graph
        // - Analyzing memory usage
        // - Generating reports

        // For now, just test that all components exist and can be created
        let _pipeline_components = (graph_optimizer, memory_optimizer, vis_config);
    }
}

/// Test the newly implemented features from previous Advanced sessions
#[cfg(test)]
mod advanced_feature_integration_tests {
    use super::*;

    #[test]
    fn test_custom_activations_with_optimization() {
        // Test that custom activations work with the optimization system
        ag::run(|ctx: &mut ag::Context<f32>| {
            let x = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[3]), vec![-1.0, 0.0, 1.0])
                    .expect("Test: operation failed"),
                ctx,
            );

            // Apply custom activations
            let swish_result = T::custom_activation(&x, "swish");
            let gelu_result = T::custom_activation(&x, "gelu");

            // Test that these can be evaluated
            let _swish_output = swish_result.eval(ctx).expect("Test: operation failed");
            let _gelu_output = gelu_result.eval(ctx).expect("Test: operation failed");
        });
    }

    #[test]
    fn test_performance_operations_with_memory_optimization() {
        ag::run(|ctx: &mut ag::Context<f32>| {
            // Test SIMD operations
            let a = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[4]), vec![1.0, 2.0, 3.0, 4.0])
                    .expect("Test: operation failed"),
                ctx,
            );
            let b = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[4]), vec![2.0, 3.0, 4.0, 5.0])
                    .expect("Test: operation failed"),
                ctx,
            );

            // SIMD operations
            let simd_add = T::simd_add(&a, &b);
            let simd_mul = T::simd_mul(&a, &b);

            // Memory optimized operations
            let inplace_result = T::inplace_add(&a, &b);

            // Test evaluation
            let _add_result = simd_add.eval(ctx).expect("Test: operation failed");
            let _mul_result = simd_mul.eval(ctx).expect("Test: operation failed");
            let _inplace_result = inplace_result.eval(ctx).expect("Test: operation failed");
        });
    }

    #[test]
    fn test_graph_enhancements_with_visualization() {
        ag::run(|ctx: &mut ag::Context<f32>| {
            let x = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[2]), vec![1.0, 2.0]).expect("Test: operation failed"),
                ctx,
            );

            // Test enhanced graph operations
            let cached_result = T::cached_op(&x, "square");
            let checkpointed = T::checkpoint(&cached_result);

            // Test conditional operations
            let condition = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[1]), vec![1.0]).expect("Test: operation failed"),
                ctx,
            );
            let true_branch = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[2]), vec![10.0, 20.0])
                    .expect("Test: operation failed"),
                ctx,
            );
            let false_branch = T::convert_to_tensor(
                Array::from_shape_vec(IxDyn(&[2]), vec![30.0, 40.0])
                    .expect("Test: operation failed"),
                ctx,
            );

            let conditional_result = T::conditional(
                &condition,
                &true_branch,
                &false_branch,
                T::PredicateType::GreaterThanZero,
            );

            // Test evaluation
            let _cached_output = cached_result.eval(ctx).expect("Test: operation failed");
            let _checkpointed_output = checkpointed.eval(ctx).expect("Test: operation failed");
            let _conditional_output = conditional_result
                .eval(ctx)
                .expect("Test: operation failed");
        });
    }

    #[test]
    fn test_efficient_operations_integration() {
        ag::run(|ctx: &mut ag::Context<f32>| {
            // Test efficient tensor creation
            let efficient_zeros = T::efficient_zeros(&[10, 10], ctx);
            let efficient_ones = T::efficient_ones(&[10, 10], ctx);

            // Test efficient operations
            let reshaped = T::efficient_reshape_withshape(&efficient_zeros, &[100]);

            // Test slice operations
            let slices = vec![T::SliceRange::new(Some(0), Some(50), Some(1))];
            let sliced = T::efficient_slice(&reshaped, &slices);

            // Test concatenation
            let concat_result = T::efficient_concat(&[&efficient_ones, &efficient_ones], 0);

            // Test evaluation
            let _zeros_output = efficient_zeros.eval(ctx).expect("Test: operation failed");
            let _ones_output = efficient_ones.eval(ctx).expect("Test: operation failed");
            let _reshaped_output = reshaped.eval(ctx).expect("Test: operation failed");
            let _sliced_output = sliced.eval(ctx).expect("Test: operation failed");
            let _concat_output = concat_result.eval(ctx).expect("Test: operation failed");
        });
    }

    #[test]
    fn test_complete_advanced_feature_integration() {
        // Test all Advanced features working together
        ag::run(|ctx: &mut ag::Context<f32>| {
            // 1. Create efficient tensors
            let input = T::efficient_ones(&[32, 64], ctx);

            // 2. Apply custom activation
            let activated = T::custom_activation(&input, "swish");

            // 3. Use SIMD operations
            let scaled = T::simd_mul(&activated, &activated);

            // 4. Apply memory optimization
            let optimized = T::inplace_add(&scaled, &input);

            // 5. Use efficient reshape
            let reshaped = T::efficient_reshape_withshape(&optimized, &[2048]);

            // 6. Apply caching
            let cached = T::cached_op(&reshaped, "identity");

            // 7. Apply checkpointing
            let checkpointed = T::smart_checkpoint(&cached, 100000);

            // 8. Use conditional operations
            let condition = T::efficient_ones(&[1], ctx);
            let alternative = T::efficient_zeros(&[2048], ctx);
            let final_result = T::conditional(
                &condition,
                &checkpointed,
                &alternative,
                T::PredicateType::GreaterThanZero,
            );

            // Test that the entire pipeline evaluates successfully
            let output = final_result.eval(ctx).expect("Test: operation failed");
            assert_eq!(output.len(), 2048);
            assert!(output.iter().all(|&x| x.is_finite()));
        });
    }
}
