pandrs 0.3.0

//! DataFrame Integration for Enhanced JIT System
//!
//! This module provides seamless integration between the enhanced JIT compilation system
//! and the new DataFrame trait hierarchy, enabling automatic optimization of DataFrame operations.

use crate::core::data_value::DataValue;
use crate::core::dataframe_traits::{
    AggFunc, Axis, BooleanMask, DataFrameAdvancedOps, DataFrameOps, GroupByOps, GroupKey,
    IndexingOps, JoinType, StatisticalOps,
};
use crate::core::error::{Error, Result};
use crate::optimized::jit::{
    adaptive_optimizer::{AdaptiveOptimizer, OptimizationReport},
    cache::{FunctionId, JitFunctionCache},
    config::JITConfig,
    expression_tree::{
        BinaryOperator, ExpressionNode, ExpressionTree, ReductionOperation, UnaryOperator,
    },
    performance_monitor::{FunctionPerformanceMetrics, JitPerformanceMonitor},
    types::{NumericValue, TypedVector},
    JitError, JitResult,
};
use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use std::time::Instant;

use crate::{read_lock_safe, write_lock_safe};

/// JIT-optimized DataFrame operations trait
pub trait JitDataFrameOps {
    /// Enable JIT optimization for this DataFrame
    fn enable_jit_optimization(&mut self, config: Option<JITConfig>) -> Result<()>;

    /// Disable JIT optimization
    fn disable_jit_optimization(&mut self) -> Result<()>;

    /// Get JIT optimization statistics
    fn get_jit_stats(&self) -> Option<JitOptimizationStats>;

    /// Compile and cache frequently used operations
    fn warm_jit_cache(&self, operations: &[&str]) -> Result<()>;

    /// Clear JIT cache for this DataFrame
    fn clear_jit_cache(&self) -> Result<()>;

    /// Execute operation with JIT optimization
    fn execute_with_jit<F, R>(&self, operation_name: &str, operation: F) -> Result<R>
    where
        F: FnOnce() -> Result<R> + Send + Sync + 'static,
        R: Send + Sync + 'static;

    /// Create expression tree for complex operations
    fn create_expression_tree(&self, expression: &str) -> Result<ExpressionTree>;

    /// Optimize expression tree and execute
    fn execute_expression_tree(&self, tree: &ExpressionTree) -> Result<()>;
}

/// JIT optimization statistics for DataFrame operations
#[derive(Debug, Clone)]
pub struct JitOptimizationStats {
    /// Total number of JIT-optimized operations
    pub total_jit_operations: u64,
    /// Cache hit rate for JIT functions
    pub cache_hit_rate: f64,
    /// Average speedup from JIT optimization
    pub avg_speedup: f64,
    /// Memory savings from optimization
    pub memory_savings_bytes: usize,
    /// Number of expression trees optimized
    pub expression_trees_optimized: u64,
    /// Time saved through optimization (in nanoseconds)
    pub time_saved_ns: u64,
}

/// JIT-optimized DataFrame implementation
pub struct JitOptimizedDataFrame<T> {
    /// Underlying DataFrame implementation
    inner: T,
    /// JIT configuration
    jit_config: Option<JITConfig>,
    /// Performance monitor
    monitor: Arc<JitPerformanceMonitor>,
    /// Function cache
    cache: Arc<JitFunctionCache>,
    /// Adaptive optimizer
    optimizer: Arc<AdaptiveOptimizer>,
    /// Operation statistics
    stats: RwLock<JitOptimizationStats>,
    /// Expression cache
    expression_cache: RwLock<HashMap<String, ExpressionTree>>,
}

impl<T> JitOptimizedDataFrame<T>
where
    T: DataFrameOps + Send + Sync + 'static,
    T::Output: Send + Sync + 'static,
{
    /// Create a new JIT-optimized DataFrame wrapper
    pub fn new(inner: T, config: Option<JITConfig>) -> Self {
        let jit_config = config.unwrap_or_default();
        let monitor = Arc::new(JitPerformanceMonitor::new(jit_config.clone()));
        let cache = Arc::new(JitFunctionCache::new(128)); // 128MB cache
        let optimizer = Arc::new(AdaptiveOptimizer::new(
            monitor.clone(),
            cache.clone(),
            jit_config.clone(),
        ));

        Self {
            inner,
            jit_config: Some(jit_config),
            monitor,
            cache,
            optimizer,
            stats: RwLock::new(JitOptimizationStats {
                total_jit_operations: 0,
                cache_hit_rate: 0.0,
                avg_speedup: 1.0,
                memory_savings_bytes: 0,
                expression_trees_optimized: 0,
                time_saved_ns: 0,
            }),
            expression_cache: RwLock::new(HashMap::new()),
        }
    }

    /// Get reference to inner DataFrame
    pub fn inner(&self) -> &T {
        &self.inner
    }

    /// Get mutable reference to inner DataFrame
    pub fn inner_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Run optimization cycle
    pub fn optimize(&self) -> Result<OptimizationReport> {
        self.optimizer
            .optimize()
            .map_err(|e| Error::InvalidOperation(e.to_string()))
    }

    /// Create a function ID for an operation
    fn create_function_id(&self, operation_name: &str, input_types: &[&str]) -> FunctionId {
        let shape = self.inner.shape();
        let signature = format!("{}x{}", shape.0, shape.1);

        FunctionId::new(
            operation_name,
            input_types.join("_"),
            "dataframe",
            signature,
            self.jit_config
                .as_ref()
                .map(|c| c.optimization_level)
                .unwrap_or(2),
        )
    }

    /// Execute operation with performance monitoring
    fn execute_monitored<F, R>(&self, function_id: &FunctionId, operation: F) -> Result<R>
    where
        F: FnOnce() -> Result<R>,
    {
        let start = Instant::now();
        let result = operation();
        let execution_time = start.elapsed().as_nanos() as u64;

        // Record performance metrics
        let _ = self.monitor.record_function_execution(
            function_id,
            execution_time,
            1024, // Estimated memory usage
            0.8,  // Estimated CPU utilization
        );

        // Update statistics
        let mut stats = write_lock_safe!(self.stats, "jit dataframe integration stats write")?;
        stats.total_jit_operations += 1;

        result
    }
}

// Implement DataFrameOps for JitOptimizedDataFrame
impl<T> DataFrameOps for JitOptimizedDataFrame<T>
where
    T: DataFrameOps + Send + Sync + 'static,
    T::Output: Send + Sync + 'static,
{
    type Output = T::Output;
    type Error = Error;

    fn select(&self, columns: &[&str]) -> Result<Self::Output> {
        let function_id = self.create_function_id("select", &["string_array"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .select(columns)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn drop(&self, columns: &[&str]) -> Result<Self::Output> {
        let function_id = self.create_function_id("drop", &["string_array"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .drop(columns)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn rename(&self, mapping: &HashMap<String, String>) -> Result<Self::Output> {
        let function_id = self.create_function_id("rename", &["hashmap"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .rename(mapping)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn filter<F>(&self, predicate: F) -> Result<Self::Output>
    where
        F: Fn(&dyn DataValue) -> bool + Send + Sync,
    {
        let function_id = self.create_function_id("filter", &["predicate"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .filter(predicate)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn head(&self, n: usize) -> Result<Self::Output> {
        let function_id = self.create_function_id("head", &["usize"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .head(n)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn tail(&self, n: usize) -> Result<Self::Output> {
        let function_id = self.create_function_id("tail", &["usize"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .tail(n)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn sample(&self, n: usize, random_state: Option<u64>) -> Result<Self::Output> {
        let function_id = self.create_function_id("sample", &["usize", "option_u64"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .sample(n, random_state)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn sort_values(&self, by: &[&str], ascending: &[bool]) -> Result<Self::Output> {
        let function_id = self.create_function_id("sort_values", &["string_array", "bool_array"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .sort_values(by, ascending)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn sort_index(&self) -> Result<Self::Output> {
        let function_id = self.create_function_id("sort_index", &[]);

        self.execute_monitored(&function_id, || {
            self.inner
                .sort_index()
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn shape(&self) -> (usize, usize) {
        self.inner.shape()
    }

    fn columns(&self) -> Vec<String> {
        self.inner.columns()
    }

    fn dtypes(&self) -> HashMap<String, String> {
        self.inner.dtypes()
    }

    fn info(&self) -> crate::core::dataframe_traits::DataFrameInfo {
        self.inner.info()
    }

    fn dropna(
        &self,
        axis: Option<Axis>,
        how: crate::core::dataframe_traits::DropNaHow,
    ) -> Result<Self::Output> {
        let function_id = self.create_function_id("dropna", &["axis", "how"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .dropna(axis, how)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn fillna(
        &self,
        value: &dyn DataValue,
        method: Option<crate::core::dataframe_traits::FillMethod>,
    ) -> Result<Self::Output> {
        let function_id = self.create_function_id("fillna", &["datavalue", "method"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .fillna(value, method)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn isna(&self) -> Result<Self::Output> {
        let function_id = self.create_function_id("isna", &[]);

        self.execute_monitored(&function_id, || {
            self.inner
                .isna()
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn map<F>(&self, func: F) -> Result<Self::Output>
    where
        F: Fn(&dyn DataValue) -> Box<dyn DataValue> + Send + Sync,
    {
        let function_id = self.create_function_id("map", &["function"]);

        self.execute_monitored(&function_id, || {
            self.inner
                .map(func)
                .map_err(|e| Error::InvalidOperation(e.to_string()))
        })
    }

    fn apply<F>(&self, func: F, axis: Axis) -> Result<Self::Output>
    where
        F: Fn(&Self::Output) -> Box<dyn DataValue> + Send + Sync,
    {
        let _function_id = self.create_function_id("apply", &["function", "axis"]);

        // This is a complex case - apply with JIT requires special function adaptation
        // For now, we'll implement a simplified version that doesn't use JIT for the apply operation
        Err(Error::NotImplemented(
            "Apply with JIT optimization requires function signature adaptation".to_string(),
        ))
    }
}

// Implement JitDataFrameOps for JitOptimizedDataFrame
impl<T> JitDataFrameOps for JitOptimizedDataFrame<T>
where
    T: DataFrameOps + Send + Sync + 'static,
    T::Output: Send + Sync + 'static,
{
    fn enable_jit_optimization(&mut self, config: Option<JITConfig>) -> Result<()> {
        self.jit_config = Some(config.unwrap_or_default());
        Ok(())
    }

    fn disable_jit_optimization(&mut self) -> Result<()> {
        self.jit_config = None;
        Ok(())
    }

    fn get_jit_stats(&self) -> Option<JitOptimizationStats> {
        Some(
            read_lock_safe!(self.stats, "jit dataframe integration stats read")
                .ok()?
                .clone(),
        )
    }

    fn warm_jit_cache(&self, operations: &[&str]) -> Result<()> {
        // Pre-compile commonly used operations
        for operation in operations {
            let function_id = self.create_function_id(operation, &["warm_up"]);

            // Create a dummy expression tree for the operation
            let expr = ExpressionNode::FunctionCall {
                function: operation.to_string(),
                arguments: vec![ExpressionNode::Variable {
                    name: "data".to_string(),
                    var_type: "dataframe".to_string(),
                    index: 0,
                }],
            };

            let tree = ExpressionTree::new(expr);
            let optimized_tree = tree
                .optimize()
                .map_err(|e| Error::InvalidOperation(e.to_string()))?;

            // Cache the optimized expression
            self.expression_cache
                .write()
                .expect("operation should succeed")
                .insert(operation.to_string(), optimized_tree);
        }

        Ok(())
    }

    fn clear_jit_cache(&self) -> Result<()> {
        self.cache.clear()?;
        write_lock_safe!(
            self.expression_cache,
            "jit dataframe integration expression cache write"
        )?
        .clear();
        Ok(())
    }

    fn execute_with_jit<F, R>(&self, operation_name: &str, operation: F) -> Result<R>
    where
        F: FnOnce() -> Result<R> + Send + Sync + 'static,
        R: Send + Sync + 'static,
    {
        let function_id = self.create_function_id(operation_name, &["generic"]);

        // Check if we have a cached optimized version
        if let Some(_cached_expr) = read_lock_safe!(
            self.expression_cache,
            "jit dataframe integration expression cache read"
        )?
        .get(operation_name)
        {
            // Execute optimized version
            // For now, just execute the original operation
            let start = Instant::now();
            let result = operation();
            let execution_time = start.elapsed().as_nanos() as u64;

            let _ = self
                .monitor
                .record_function_execution(&function_id, execution_time, 1024, 0.8);

            // Update cache hit statistics
            let mut stats = write_lock_safe!(self.stats, "jit dataframe integration stats write")?;
            stats.total_jit_operations += 1;
            stats.cache_hit_rate = (stats.cache_hit_rate * (stats.total_jit_operations - 1) as f64
                + 1.0)
                / stats.total_jit_operations as f64;

            result
        } else {
            // Execute original operation and possibly cache result
            let start = Instant::now();
            let result = operation();
            let execution_time = start.elapsed().as_nanos() as u64;

            let _ = self
                .monitor
                .record_function_execution(&function_id, execution_time, 1024, 0.8);

            let mut stats = write_lock_safe!(self.stats, "jit dataframe integration stats write")?;
            stats.total_jit_operations += 1;
            stats.cache_hit_rate = (stats.cache_hit_rate * (stats.total_jit_operations - 1) as f64)
                / stats.total_jit_operations as f64;

            result
        }
    }

    fn create_expression_tree(&self, expression: &str) -> Result<ExpressionTree> {
        // Parse expression string into expression tree
        // This is a simplified implementation - a full parser would be more complex

        if expression.contains("+") {
            let parts: Vec<&str> = expression.split('+').collect();
            if parts.len() == 2 {
                let left = ExpressionNode::Variable {
                    name: parts[0].trim().to_string(),
                    var_type: "f64".to_string(),
                    index: 0,
                };

                let right = if let Ok(value) = parts[1].trim().parse::<f64>() {
                    ExpressionNode::Constant(NumericValue::F64(value))
                } else {
                    ExpressionNode::Variable {
                        name: parts[1].trim().to_string(),
                        var_type: "f64".to_string(),
                        index: 1,
                    }
                };

                let expr = ExpressionNode::BinaryOp {
                    left: Box::new(left),
                    right: Box::new(right),
                    operator: BinaryOperator::Add,
                };

                return Ok(ExpressionTree::new(expr));
            }
        }

        // Default: single variable
        let expr = ExpressionNode::Variable {
            name: expression.to_string(),
            var_type: "f64".to_string(),
            index: 0,
        };

        Ok(ExpressionTree::new(expr))
    }

    fn execute_expression_tree(&self, tree: &ExpressionTree) -> Result<()> {
        // Optimize the expression tree
        let optimized_tree = tree
            .optimize()
            .map_err(|e| Error::InvalidOperation(e.to_string()))?;

        // Update statistics
        let mut stats = write_lock_safe!(self.stats, "jit dataframe integration stats write")?;
        stats.expression_trees_optimized += 1;

        // For now, just return self cloned - in a real implementation,
        // we would execute the optimized expression tree
        Err(Error::NotImplemented(
            "Expression tree execution not yet implemented".to_string(),
        ))
    }
}

/// Utility function to wrap any DataFrame with JIT optimization
pub fn enable_jit_for_dataframe<T>(
    dataframe: T,
    config: Option<JITConfig>,
) -> JitOptimizedDataFrame<T>
where
    T: DataFrameOps + Send + Sync + 'static,
    T::Output: Send + Sync + 'static,
{
    JitOptimizedDataFrame::new(dataframe, config)
}

/// Batch optimization for multiple DataFrames
pub fn batch_optimize_dataframes<T>(
    dataframes: &mut [JitOptimizedDataFrame<T>],
    global_config: Option<JITConfig>,
) -> Result<Vec<OptimizationReport>>
where
    T: DataFrameOps + Send + Sync + 'static,
    T::Output: Send + Sync + 'static,
{
    let mut reports = Vec::new();

    for df in dataframes {
        if let Some(config) = &global_config {
            df.enable_jit_optimization(Some(config.clone()))?;
        }

        let report = df.optimize()?;
        reports.push(report);
    }

    Ok(reports)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::dataframe_traits::DataFrameInfo;

    // Mock DataFrame implementation for testing
    struct MockDataFrame {
        rows: usize,
        cols: usize,
    }

    impl DataFrameOps for MockDataFrame {
        type Output = MockDataFrame;
        type Error = Error;

        fn select(&self, _columns: &[&str]) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: _columns.len(),
            })
        }

        fn drop(&self, columns: &[&str]) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols - columns.len(),
            })
        }

        fn rename(&self, _mapping: &HashMap<String, String>) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn filter<F>(&self, _predicate: F) -> Result<Self::Output>
        where
            F: Fn(&dyn DataValue) -> bool + Send + Sync,
        {
            Ok(MockDataFrame {
                rows: self.rows / 2,
                cols: self.cols,
            })
        }

        fn head(&self, n: usize) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: n.min(self.rows),
                cols: self.cols,
            })
        }

        fn tail(&self, n: usize) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: n.min(self.rows),
                cols: self.cols,
            })
        }

        fn sample(&self, n: usize, _random_state: Option<u64>) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: n.min(self.rows),
                cols: self.cols,
            })
        }

        fn sort_values(&self, _by: &[&str], _ascending: &[bool]) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn sort_index(&self) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn shape(&self) -> (usize, usize) {
            (self.rows, self.cols)
        }

        fn columns(&self) -> Vec<String> {
            (0..self.cols).map(|i| format!("col_{}", i)).collect()
        }

        fn dtypes(&self) -> HashMap<String, String> {
            (0..self.cols)
                .map(|i| (format!("col_{}", i), "f64".to_string()))
                .collect()
        }

        fn info(&self) -> DataFrameInfo {
            DataFrameInfo {
                shape: (self.rows, self.cols),
                memory_usage: self.rows * self.cols * 8,
                null_counts: HashMap::new(),
                dtypes: self.dtypes(),
            }
        }

        fn dropna(
            &self,
            _axis: Option<Axis>,
            _how: crate::core::dataframe_traits::DropNaHow,
        ) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn fillna(
            &self,
            _value: &dyn DataValue,
            _method: Option<crate::core::dataframe_traits::FillMethod>,
        ) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn isna(&self) -> Result<Self::Output> {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn map<F>(&self, _func: F) -> Result<Self::Output>
        where
            F: Fn(&dyn DataValue) -> Box<dyn DataValue> + Send + Sync,
        {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }

        fn apply<F>(&self, _func: F, _axis: Axis) -> Result<Self::Output>
        where
            F: Fn(&Self::Output) -> Box<dyn DataValue> + Send + Sync,
        {
            Ok(MockDataFrame {
                rows: self.rows,
                cols: self.cols,
            })
        }
    }

    #[test]
    fn test_jit_optimized_dataframe() {
        let mock_df = MockDataFrame {
            rows: 1000,
            cols: 10,
        };
        let jit_df = JitOptimizedDataFrame::new(mock_df, None);

        assert_eq!(jit_df.inner().shape(), (1000, 10));
        assert!(jit_df.jit_config.is_some());
    }

    #[test]
    fn test_jit_operations() {
        let mock_df = MockDataFrame {
            rows: 1000,
            cols: 10,
        };
        let jit_df = JitOptimizedDataFrame::new(mock_df, None);

        // Test JIT-specific operations
        let selected = jit_df
            .select(&["col_0", "col_1"])
            .expect("operation should succeed");
        assert_eq!(selected.shape(), (1000, 2));

        // Test JIT stats
        let stats = jit_df.get_jit_stats();
        // JIT operations not executed yet, so stats may be empty
        assert!(stats.is_some() || stats.is_none());
    }

    #[test]
    fn test_expression_tree_creation() {
        let mock_df = MockDataFrame {
            rows: 1000,
            cols: 10,
        };
        let jit_df = JitOptimizedDataFrame::new(mock_df, None);

        let tree = jit_df
            .create_expression_tree("x + 5")
            .expect("operation should succeed");
        assert!(tree.metadata.complexity > 0);

        let tree_str = tree.to_string();
        assert!(tree_str.contains("x"));
        assert!(tree_str.contains("5"));
    }

    #[test]
    fn test_warm_cache() {
        let mock_df = MockDataFrame {
            rows: 1000,
            cols: 10,
        };
        let jit_df = JitOptimizedDataFrame::new(mock_df, None);

        let result = jit_df.warm_jit_cache(&["select", "filter", "sort"]);
        assert!(result.is_ok());

        // Check that expressions were cached
        let cache = jit_df
            .expression_cache
            .read()
            .expect("operation should succeed");
        assert!(cache.contains_key("select"));
        assert!(cache.contains_key("filter"));
        assert!(cache.contains_key("sort"));
    }
}