gbrt_rs/io/

data_loader.rs

#![allow(unused)]

//! Data Loading and I/O Operations
//! 
//! This module provides comprehensive data loading capabilities for multiple formats
//! with built-in preprocessing support:
//! 
//! - **CSV**: Full-featured CSV parsing with headers, delimiters, missing value handling
//! - **JSON**: Structured JSON records with automatic schema inference
//! - **Parquet**: Placeholder for Parquet support (requires additional dependencies)
//! - **Categorical Encoding**: Automatic detection and encoding of categorical features
//! - **Flexible Target Selection**: Specify target column by name or use defaults
//! - **Train/Test Splitting**: Built-in splitting with shuffling and stratification support
//! - **Cross-Validation**: K-fold cross-validation split generation
//! 
//! # DataLoader Architecture
//! 
//! The [`DataLoader`] struct uses a builder pattern with [`LoadOptions`] for configuration.
//! It provides both high-level convenience methods and low-level control over the loading process.
//!

use crate::data::{
    Dataset, FeatureMatrix, DataError, FeatureMatrixError,
    preprocessing::{CategoricalEncoder, detect_categorical_columns},
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufReader, BufWriter};
use std::path::{Path, PathBuf};
use thiserror::Error;
use std::str::FromStr;
use ndarray::Array2;
use csv::ReaderBuilder;
use parquet::arrow::arrow_reader::ArrowReaderBuilder;
use parquet::file::reader::FileReader;
use std::sync::Arc;
use arrow::array::{Array, Float64Array, Int64Array, Float32Array, Int32Array};
use arrow::datatypes::DataType;

/// Errors that can occur during data loading operations.
/// 
/// This error type covers all failure modes: file I/O, parsing, validation,
/// and preprocessing errors.
#[derive(Error, Debug)]
pub enum DataLoaderError {
    /// File system I/O error.
    #[error("IO error: {0}")]
    IoError(#[from] std::io::Error),

    /// CSV parsing error.
    #[error("CSV error: {0}")]
    CsvError(#[from] csv::Error),

    /// JSON deserialization error.
    #[error("JSON error: {0}")]
    JsonError(#[from] serde_json::Error),

    /// Data validation error (shape, NaN values, etc.).
    #[error("Data error: {0}")]
    DataError(#[from] DataError),

    /// Feature matrix construction error.
    #[error("Feature matrix error: {0}")]
    FeatureMatrixError(#[from] FeatureMatrixError),

    /// Unsupported file format.
    #[error("Unsupported format: {format}")]
    UnsupportedFormat { format: String },

    /// Value parsing error with context.
    #[error("Parse error: {field} = {value} ({message})")]
    ParseError {
        field: String,
        value: String,
        message: String,
    },

    /// Data validation error with custom message.
    #[error("Validation error: {0}")]
    ValidationError(String),

    /// Requested column not found in data.
    #[error("Missing column: {0}")]
    MissingColumn(String),

    /// File format or structure is invalid.
    #[error("Invalid file: {0}")]
    InvalidFile(String),

    /// Categorical encoding or preprocessing failure.
    #[error("Preprocessing error: {0}")]
    PreprocessingError(String),

    /// Parquet file reading error.
    #[error("Parquet error: {0}")]
    ParquetError(#[from] parquet::errors::ParquetError),
}

/// Result type for data loader operations.
pub type DataLoaderResult<T> = std::result::Result<T, DataLoaderError>;

/// Supported data file formats for loading and saving.
/// 
/// Format is auto-detected from file extension when using [`load_from_file()`].
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum DataFormat {
    /// Comma-separated values with headers.
    Csv,
    /// JSON array of objects.
    Json,
    /// Apache Parquet columnar format (placeholder).
    Parquet,
    // Arrow, // Future support
}

impl std::str::FromStr for DataFormat {
    type Err = DataLoaderError;

    /// Parses DataFormat from string (file extension).
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "csv" => Ok(DataFormat::Csv),
            "json" => Ok(DataFormat::Json),
            "parquet" => Ok(DataFormat::Parquet),
            // "arrow" => Ok(DataFormat::Arrow),
            _ => Err(DataLoaderError::UnsupportedFormat {
                format: s.to_string(),
            }),
        }
    }
}

impl std::fmt::Display for DataFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            DataFormat::Csv => write!(f, "csv"),
            DataFormat::Json => write!(f, "json"),
            DataFormat::Parquet => write!(f, "parquet"),
        }
    }
}


/// Configuration options for data loading.
/// 
/// Uses builder pattern for flexible configuration of CSV/JSON parsing,
/// column selection, and preprocessing options.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LoadOptions {
    /// Field delimiter for CSV files (default: comma).
    pub delimiter: u8,
    /// Whether first row contains headers (default: true).
    pub has_headers: bool,
    /// Name of target column to predict.
    pub target_column: Option<String>,
    /// List of feature columns to use (None = all except target).
    pub feature_columns: Option<Vec<String>>,
    /// Number of rows to skip at start of file.
    pub skip_rows: usize,
    /// Maximum number of rows to read (None = all).
    pub max_rows: Option<usize>,
    /// String representation of missing values (default: "NaN").
    pub missing_value: Option<String>,
    /// File encoding (default: "utf-8").
    pub encoding: String,
}

impl Default for LoadOptions {
    fn default() -> Self {
        Self {
            delimiter: b',',
            has_headers: true,
            target_column: None,
            feature_columns: None,
            skip_rows: 0,
            max_rows: None,
            missing_value: Some("NaN".to_string()),
            encoding: "utf-8".to_string(),
        }
    }
}

impl LoadOptions {
    /// Creates new LoadOptions with default values.
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets CSV delimiter character.
    pub fn with_delimiter(mut self, delimiter: u8) -> Self {
        self.delimiter = delimiter;
        self
    }

    /// Sets whether CSV has header row.
    pub fn with_headers(mut self, has_headers: bool) -> Self {
        self.has_headers = has_headers;
        self
    }

    /// Sets target column by name.
    pub fn with_target_column(mut self, target_column: &str) -> Self {
        self.target_column = Some(target_column.to_string());
        self
    }

    /// Sets explicit feature columns (excludes others).
    pub fn with_feature_columns(mut self, feature_columns: &[&str]) -> Self {
        self.feature_columns = Some(feature_columns.iter().map(|s| s.to_string()).collect());
        self
    }

    /// Sets number of rows to skip at start.
    pub fn with_skip_rows(mut self, skip_rows: usize) -> Self {
        self.skip_rows = skip_rows;
        self
    }

    /// Sets maximum rows to read.
    pub fn with_max_rows(mut self, max_rows: usize) -> Self {
        self.max_rows = Some(max_rows);
        self
    }

    /// Sets string representation for missing values.
    pub fn with_missing_value(mut self, missing_value: &str) -> Self {
        self.missing_value = Some(missing_value.to_string());
        self
    }
}

/// Main data loader for various file formats with preprocessing capabilities.
/// 
/// [`DataLoader`] provides a unified interface for loading datasets from
/// multiple formats with built-in preprocessing, categorical encoding, and
/// validation.
/// 
/// # Usage Patterns
/// 
/// - **Simple**: Use `new()` with defaults
/// - **Configured**: Use `with_options()` for custom settings
/// - **Advanced**: Use format-specific methods for full control
pub struct DataLoader {
    /// Configuration options for loading behavior.
    options: LoadOptions,
}

impl DataLoader {
    /// Creates a new DataLoader with default options.
    /// 
    /// # Returns
    /// `Ok(DataLoader)` if creation succeeds (always succeeds currently)
    pub fn new() -> DataLoaderResult<Self> {
        Ok(Self {
            options: LoadOptions::default(),
        })
    }

    /// Creates a new DataLoader with custom options.
    ///
    /// # Parameters
    /// - `options`: LoadOptions configuration
    ///
    /// # Returns
    /// Configured DataLoader instance
    pub fn with_options(options: LoadOptions) -> Self {
        Self { options }
    }


    /// Loads dataset from file with auto-detected format.
    /// 
    /// Format is determined from file extension (.csv, .json, .parquet).
    /// 
    /// # Parameters
    /// - `path`: Path to data file
    /// 
    /// # Returns
    /// Loaded dataset
    /// 
    /// # Errors
    /// - `DataLoaderError::InvalidFile` if extension cannot be determined
    /// - `DataLoaderError::UnsupportedFormat` if format is not supported
    /// - Format-specific errors from underlying loaders
    pub fn load_file<P: AsRef<Path>>(&self, path: P) -> DataLoaderResult<Dataset> {
        load_from_file(path, &self.options)
    }

    /// Loads dataset from CSV file with automatic categorical encoding.
    /// 
    /// This is the recommended method for loading CSV data. It handles:
    /// - Automatic target column extraction
    /// - Categorical feature detection and encoding
    /// - Missing value and validation checks
    /// 
    /// # Parameters
    /// - `path`: Path to CSV file
    /// - `target_column`: Name of target column
    /// - `categorical_columns`: Optional explicit categorical column indices (None = auto-detect)
    /// - `categorical_threshold`: Ratio threshold for auto-detecting categorical columns (0.0-1.0)
    /// 
    /// # Returns
    /// Dataset with encoded features and extracted targets
    /// 
    /// # Errors
    /// - `DataLoaderError::MissingColumn` if target column not found
    /// - `DataLoaderError::ValidationError` if data is malformed
    /// - `DataLoaderError::PreprocessingError` if encoding fails
    /// - `DataLoaderError::ParseError` if target values cannot be parsed
    /// 
    /// # Categorical Detection
    /// If `categorical_columns` is `None`, automatically detects categorical columns
    /// where unique values / total samples < `categorical_threshold`.
    pub fn load_csv_with_categorical<P: AsRef<Path>>(
        &self,
        path: P,
        target_column: &str,
        categorical_columns: Option<&[usize]>,
        categorical_threshold: f64,
    ) -> DataLoaderResult<Dataset> {
        let mut reader = ReaderBuilder::new()
            .delimiter(self.options.delimiter)
            .has_headers(self.options.has_headers)
            .from_path(&path)
            .map_err(DataLoaderError::CsvError)?;
        
        let headers: Vec<String> = reader
            .headers()
            .map_err(DataLoaderError::CsvError)?
            .iter()
            .map(|s| s.to_string())
            .collect();
        
        // Find target column index
        let target_col_idx = headers.iter()
            .position(|h| h == target_column)
            .ok_or_else(|| DataLoaderError::MissingColumn(
                format!("Target column '{}' not found in CSV headers. Available columns: {:?}", 
                       target_column, headers)
            ))?;
        
        // Read all data as strings
        let mut string_data = Vec::new();
        for (row_idx, result) in reader.records().enumerate() {
            let record = result.map_err(|e| DataLoaderError::CsvError(e))?;
            let row: Vec<String> = record.iter().map(|s| s.to_string()).collect();
            
            if row.len() != headers.len() {
                return Err(DataLoaderError::ValidationError(
                    format!("Row {} has {} columns, expected {} columns. Row: {:?}", 
                           row_idx, row.len(), headers.len(), row)
                ));
            }
            string_data.push(row);
        }
        
        if string_data.is_empty() {
            return Err(DataLoaderError::ValidationError(
                "CSV file contains no data rows".to_string()
            ));
        }
        
        // Determine feature columns (exclude target)
        let n_features_total = headers.len();
        let feature_col_indices: Vec<usize> = (0..n_features_total)
            .filter(|&i| i != target_col_idx)
            .collect();
        
        // Extract feature data (excluding target)
        let feature_data: Vec<Vec<String>> = string_data.iter()
            .map(|row| {
                feature_col_indices.iter()
                    .map(|&i| row[i].clone())
                    .collect()
            })
            .collect();
        
        // Adjust categorical column indices for target exclusion
        let cat_cols_adjusted = if let Some(cols) = categorical_columns {
            cols.iter()
                .filter_map(|&col_idx| {
                    if col_idx == target_col_idx {
                        None // Skip target column
                    } else if col_idx < target_col_idx {
                        Some(col_idx)
                    } else {
                        Some(col_idx - 1)
                    }
                })
                .collect::<Vec<_>>()
        } else {
            // Auto-detect on feature columns only
            detect_categorical_columns(&feature_data, categorical_threshold)
        };
        
        // Encode categorical features
        let mut encoder = CategoricalEncoder::new();
        let encoded_features = encoder.fit_transform(&feature_data, &cat_cols_adjusted)
            .map_err(|e| DataLoaderError::PreprocessingError(e.to_string()))?;
        
        // Extract and validate targets
        let targets: Vec<f64> = string_data.iter()
            .enumerate()
            .map(|(row_idx, row)| {
                row[target_col_idx].parse::<f64>()
                    .map_err(|_| DataLoaderError::ParseError {
                        field: target_column.to_string(),
                        value: row[target_col_idx].clone(),
                        message: format!("Failed to parse target column at row {} as f64", row_idx),
                    })
            })
            .collect::<Result<Vec<_>, _>>()?;
        
        // Check for NaN/inf in targets
        for (i, &target) in targets.iter().enumerate() {
            if !target.is_finite() {
                return Err(DataLoaderError::ValidationError(
                    format!("Non-finite target value at row {}: {}", i, target)
                ));
            }
        }
        
        // Create FeatureMatrix
        let n_samples = encoded_features.len();
        let n_features = encoded_features[0].len();
        
        let flat_data: Vec<f64> = encoded_features.into_iter().flatten().collect();
        let array = Array2::from_shape_vec((n_samples, n_features), flat_data)
            .map_err(|e| DataLoaderError::ValidationError(
                format!("Failed to create feature matrix: {}", e)
            ))?;
        
        // Get feature names (excluding target)
        let feature_names: Vec<String> = feature_col_indices.iter()
            .map(|&i| headers[i].clone())
            .collect();
        
        let features = FeatureMatrix::with_feature_names(array, feature_names)?;
        
        Dataset::new(features, targets)
            .map_err(DataLoaderError::DataError)
    }

    /// Original load_csv method - delegates to categorical version with no categorical columns.
    /// 
    /// Maintains backward compatibility for simple CSV loading.    
    pub fn load_csv<P: AsRef<Path>>(&self, path: P) -> DataLoaderResult<Dataset> {
        // Default to last column as target, no categorical encoding for backward compatibility
        let target_col = self.options.target_column.as_deref().unwrap_or("");
        if target_col.is_empty() {
            load_csv(path, &self.options)
        } else {
            self.load_csv_with_categorical(path, target_col, Some(&[]), 0.0)
        }
    }

    /// Loads dataset from JSON file.
    /// 
    /// Expects JSON array of objects where each object represents a sample.
    /// 
    /// # Parameters
    /// - `path`: Path to JSON file
    /// 
    /// # Returns
    /// Loaded dataset
    /// 
    /// # Errors
    /// - `DataLoaderError::JsonError` if JSON parsing fails
    /// - `DataLoaderError::ValidationError` if data is malformed
    pub fn load_json<P: AsRef<Path>>(&self, path: P) -> DataLoaderResult<Dataset> {
        load_json(path, &self.options)
    }

    /// Loads dataset from Parquet file (placeholder).
    /// 
    /// # Note
    /// This method currently returns `UnsupportedFormat` error. To enable Parquet
    /// support, add the `parquet` crate as a dependency and implement the logic.
    pub fn load_parquet<P: AsRef<Path>>(&self, path: P) -> DataLoaderResult<Dataset> {
        load_parquet(path, &self.options)
    }

    /// Saves dataset to file in specified format.
    /// 
    /// # Parameters
    /// - `dataset`: Dataset to save
    /// - `path`: Output file path
    /// - `format`: DataFormat to use
    /// 
    /// # Returns
    /// `Ok(())` on successful save
    /// 
    /// # Errors
    /// - Format-specific save errors
    pub fn save_dataset<P: AsRef<Path>>(
        &self,
        dataset: &Dataset,
        path: P,
        format: DataFormat,
    ) -> DataLoaderResult<()> {
        save_dataset(dataset, path, format)
    }

    /// Splits dataset into training and testing sets.
    /// 
    /// Convenience method that delegates to [`split_data()`] function.
    /// 
    /// # Parameters
    /// - `dataset`: Dataset to split
    /// - `test_size`: Fraction of samples for test set (0.0-1.0)
    /// - `shuffle`: Whether to randomize sample order
    /// - `random_state`: Optional seed for reproducible shuffling
    /// 
    /// # Returns
    /// Tuple of (train_dataset, test_dataset)
    /// 
    /// # Errors
    /// - `DataLoaderError::ValidationError` if test_size not in (0, 1)
    /// - `DataLoaderError::ValidationError` if dataset too small
    pub fn split_data(
        &self,
        dataset: &Dataset,
        test_size: f64,
        shuffle: bool,
        random_state: Option<u64>,
    ) -> DataLoaderResult<(Dataset, Dataset)> {
        split_data(dataset, test_size, shuffle, random_state)
    }

    /// Creates k-fold cross-validation splits.
    /// 
    /// Generates `n_splits` train/test dataset pairs for cross-validation.
    /// Each fold uses a different portion as test set.
    /// 
    /// # Parameters
    /// - `dataset`: Dataset to split
    /// - `n_splits`: Number of folds (k)
    /// - `shuffle`: Whether to shuffle before splitting
    /// - `random_state`: Optional seed for reproducible shuffling
    /// 
    /// # Returns
    /// Vector of (train_dataset, test_dataset) tuples
    /// 
    /// # Errors
    /// - `DataLoaderError::ValidationError` if n_splits < 2
    /// - `DataLoaderError::ValidationError` if dataset too small
    pub fn cross_validation_splits(
        &self,
        dataset: &Dataset,
        n_splits: usize,
        shuffle: bool,
        random_state: Option<u64>,
    ) -> DataLoaderResult<Vec<(Dataset, Dataset)>> {
        cross_validation_split(dataset, n_splits, shuffle, random_state)
    }
}

impl Default for DataLoader {
    fn default() -> Self {
        Self::new().unwrap()
    }
}

// ============================================================================
// Standalone Data Loading Functions
// ============================================================================

/// Loads dataset from file with auto-detected format (CSV, JSON, Parquet).
/// 
/// # Parameters
/// - `path`: Path to data file
/// - `options`: LoadOptions configuration
/// 
/// # Returns
/// Loaded dataset
/// 
/// # Errors
/// - `DataLoaderError::InvalidFile` if extension cannot be determined
/// - `DataLoaderError::UnsupportedFormat` if format is not supported
pub fn load_from_file<P: AsRef<Path>>(
    path: P,
    options: &LoadOptions,
) -> DataLoaderResult<Dataset> {
    let path = path.as_ref();
    let extension = path
        .extension()
        .and_then(|ext| ext.to_str())
        .ok_or_else(|| DataLoaderError::InvalidFile("Cannot determine file extension".to_string()))?;

    let format = DataFormat::from_str(extension)?;

    match format {
        DataFormat::Csv => load_csv(path, options),
        DataFormat::Json => load_json(path, options),
        DataFormat::Parquet => load_parquet(path, options),
    }
}

/// Loads dataset from CSV file (legacy implementation).
/// 
/// This is the core CSV loading function used by DataLoader. It handles:
/// - Custom delimiters and headers
/// - Missing value substitution
/// - Target and feature column selection
/// - Target parsing and validation
/// 
/// # Parameters
/// - `path`: Path to CSV file
/// - `options`: LoadOptions configuration
/// 
/// # Returns
/// Dataset with extracted features and targets
/// 
/// # Errors
/// - All CSV-related errors from csv crate
/// - `DataLoaderError::MissingColumn` if target or features not found
/// - `DataLoaderError::ValidationError` for malformed data
pub fn load_csv<P: AsRef<Path>>(
    path: P,
    options: &LoadOptions,
) -> DataLoaderResult<Dataset> {
    let file = File::open(path)?;
    let mut reader = csv::ReaderBuilder::new()
        .delimiter(options.delimiter)
        .has_headers(options.has_headers)
        .from_reader(BufReader::new(file));

    // Skip rows if specified
    if options.skip_rows > 0 {
        for _ in 0..options.skip_rows {
            if reader.records().next().is_none() {
                break;
            }
        }
    }

    let headers: Vec<String> = if options.has_headers {
        reader
            .headers()?
            .iter()
            .map(|s| s.to_string())
            .collect()
    } else {
        // Generate default headers
        let first_record = reader.records().next();
        if let Some(record) = first_record {
            (0..record?.len())
                .map(|i| format!("feature_{}", i))
                .collect()
        } else {
            return Err(DataLoaderError::ValidationError("Empty CSV file".to_string()));
        }
    };

    let mut data = Vec::new();
    let mut targets = Vec::new();
    let mut record_count = 0;

    // Determine target column index
    let target_index = if let Some(ref target_col) = options.target_column {
        headers
            .iter()
            .position(|h| h == target_col)
            .ok_or_else(|| DataLoaderError::MissingColumn(target_col.clone()))?
    } else {
        // Use last column as target by default
        headers.len() - 1
    };

    // Determine feature columns
    let feature_indices: Vec<usize> = if let Some(ref feature_cols) = options.feature_columns {
        feature_cols
            .iter()
            .map(|col| {
                headers
                    .iter()
                    .position(|h| h == col)
                    .ok_or_else(|| DataLoaderError::MissingColumn(col.clone()))
            })
            .collect::<Result<Vec<_>, _>>()?
    } else {
        // Use all columns except target
        (0..headers.len())
            .filter(|&i| i != target_index)
            .collect()
    };

    for result in reader.records() {
        if let Some(max_rows) = options.max_rows {
            if record_count >= max_rows {
                break;
            }
        }

        let record = result?;
        let mut row = Vec::with_capacity(feature_indices.len());

        // Extract features
        for &idx in &feature_indices {
            let value_str = record.get(idx).unwrap_or("");
            let value = if let Some(ref missing) = options.missing_value {
                if value_str == missing {
                    f64::NAN
                } else {
                    value_str.parse().unwrap_or(f64::NAN)
                }
            } else {
                value_str.parse().unwrap_or(f64::NAN)
            };
            row.push(value);
        }

        // Extract target
        let target_str = record.get(target_index).unwrap_or("");
        let target = if let Some(ref missing) = options.missing_value {
            if target_str == missing {
                f64::NAN
            } else {
                target_str.parse().unwrap_or(f64::NAN)
            }
        } else {
            target_str.parse().unwrap_or(f64::NAN)
        };

        // Skip rows with NaN targets
        if !target.is_nan() {
            data.push(row);
            targets.push(target);
            record_count += 1;
        }
    }

    if data.is_empty() {
        return Err(DataLoaderError::ValidationError("No valid data found".to_string()));
    }

    // Convert to FeatureMatrix
    let n_samples = data.len();
    let n_features = feature_indices.len();
    let flat_data: Vec<f64> = data.into_iter().flatten().collect();
    let array = ndarray::Array2::from_shape_vec((n_samples, n_features), flat_data)
        .map_err(|e| DataLoaderError::ValidationError(e.to_string()))?;

    let feature_names: Vec<String> = feature_indices
        .iter()
        .map(|&idx| headers[idx].clone())
        .collect();

    let features = FeatureMatrix::with_feature_names(array, feature_names)?;
    Dataset::new(features, targets)
        .map_err(DataLoaderError::DataError)
}

/// Loads dataset from JSON file.
/// 
/// Expects JSON array of objects where each object represents a sample.
/// 
/// # Parameters
/// - `path`: Path to JSON file
/// - `options`: LoadOptions configuration
/// 
/// # Returns
/// Loaded dataset
/// 
/// # Errors
/// - `DataLoaderError::JsonError` if JSON parsing fails
/// - `DataLoaderError::ValidationError` if data is malformed
/// 
/// # JSON Format
/// ```json
/// [
///   {"feature1": 1.0, "feature2": 2.0, "target": 0.0},
///   {"feature1": 3.0, "feature2": 4.0, "target": 1.0}
/// ]
/// ```
pub fn load_json<P: AsRef<Path>>(
    path: P,
    options: &LoadOptions,
) -> DataLoaderResult<Dataset> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);

    #[derive(Deserialize)]
    struct JsonRecord {
        #[serde(flatten)]
        values: HashMap<String, serde_json::Value>,
    }

    let records: Vec<JsonRecord> = serde_json::from_reader(reader)?;

    if records.is_empty() {
        return Err(DataLoaderError::ValidationError("Empty JSON file".to_string()));
    }

    // Determine target column
    let target_column = if let Some(ref target_col) = options.target_column {
        target_col.clone()
    } else {
        // Try to find a common target column name
        let possible_targets = ["target", "label", "y", "class"];
        possible_targets
            .iter()
            .find(|&&col| records[0].values.contains_key(col))
            .map(|s| s.to_string())
            .unwrap_or_else(|| {
                // Use the last column as target
                records[0]
                    .values
                    .keys()
                    .last()
                    .unwrap()
                    .to_string()
            })
    };

    // Determine feature columns
    let feature_columns: Vec<String> = if let Some(ref feature_cols) = options.feature_columns {
        feature_cols.clone()
    } else {
        records[0]
            .values
            .keys()
            .filter(|&k| k != &target_column)
            .cloned()
            .collect()
    };

    let mut data = Vec::new();
    let mut targets = Vec::new();
    let mut record_count = 0;

    for record in records {
        if let Some(max_rows) = options.max_rows {
            if record_count >= max_rows {
                break;
            }
        }

        let mut row = Vec::with_capacity(feature_columns.len());

        // Extract features
        for col in &feature_columns {
            let value = record.values.get(col).and_then(|v| v.as_f64()).unwrap_or(f64::NAN);
            row.push(value);
        }

        // Extract target
        let target = record
            .values
            .get(&target_column)
            .and_then(|v| v.as_f64())
            .unwrap_or(f64::NAN);

        // Skip rows with NaN targets
        if !target.is_nan() {
            data.push(row);
            targets.push(target);
            record_count += 1;
        }
    }

    if data.is_empty() {
        return Err(DataLoaderError::ValidationError("No valid data found".to_string()));
    }

    // Convert to FeatureMatrix
    let n_samples = data.len();
    let n_features = feature_columns.len();
    let flat_data: Vec<f64> = data.into_iter().flatten().collect();
    let array = ndarray::Array2::from_shape_vec((n_samples, n_features), flat_data)
        .map_err(|e| DataLoaderError::ValidationError(e.to_string()))?;

    let features = FeatureMatrix::with_feature_names(array, feature_columns)?;
    Dataset::new(features, targets)
        .map_err(DataLoaderError::DataError)
}


/// Loads a dataset from an Apache Parquet file.
///
/// This function reads Parquet data using the Arrow integration, which provides
/// efficient columnar access and batch processing capabilities. It automatically
/// handles schema inference, type conversion, and validation.
///
/// # Target Column Selection
///
/// The target column is determined by the following priority:
/// 1. `options.target_column` if explicitly specified
/// 2. First matching column from `["target", "label", "y", "class"]`
/// 3. The last column in the file (fallback)
///
/// # Feature Column Selection
///
/// - If `options.feature_columns` is specified, uses only those columns
/// - Otherwise, uses all columns except the target column
///
/// # Type Support
///
/// Converts the following Arrow types to `f64`:
/// - `Float64` → direct conversion
/// - `Float32` → cast to `f64`
/// - `Int64` → cast to `f64`
/// - `Int32` → cast to `f64`
/// - Other types → `NaN` (skipped if in target, preserved if in features based on `missing_value`)
///
/// # Parameters
/// - `path`: Path to the Parquet file
/// - `options`: Configuration for loading behavior (target column, missing values, etc.)
///
/// # Returns
/// `Ok(Dataset)` containing features and targets, or `Err(DataLoaderError)` if loading fails.
///
/// # Errors
/// - `DataLoaderError::IoError` if file cannot be opened
/// - `DataLoaderError::ParquetError` if Parquet reading fails
/// - `DataLoaderError::MissingColumn` if configured target column doesn't exist
/// - `DataLoaderError::ValidationError` if file is empty or malformed
/// - `DataLoaderError::PreprocessingError` if data conversion fails
/// - `DataLoaderError::DataError` if dataset validation fails
///
/// # Memory Efficiency
///
/// Reads data in batches of 8192 rows to limit memory usage on large files.
/// Respects `options.max_rows` if set for additional memory control.
///
/// # Row Filtering
///
/// Rows are skipped if:
/// - Target value is not finite (NaN or infinite)
/// - Any feature value is not finite AND `missing_value` is not configured
pub fn load_parquet<P: AsRef<Path>>(
    path: P,
    options: &LoadOptions,
) -> DataLoaderResult<Dataset> {
    let file = File::open(path)?;
    let builder = ArrowReaderBuilder::try_new(file)?;
    let schema = builder.schema().clone();
    
    let field_names: Vec<String> = schema.fields().iter()
        .map(|f| f.name().clone())
        .collect();
    
    if field_names.is_empty() {
        return Err(DataLoaderError::ValidationError("Parquet file has no columns".to_string()));
    }
    
    // Determine target column
    let target_column = if let Some(ref target_col) = options.target_column {
        target_col.clone()
    } else {
        // Try common target names
        let possible_targets = ["target", "label", "y", "class"];
        possible_targets
            .iter()
            .find(|&&col| field_names.contains(&col.to_string()))
            .map(|s| s.to_string())
            .unwrap_or_else(|| field_names.last().unwrap().clone())
    };
    
    let target_idx = field_names.iter()
        .position(|h| h == &target_column)
        .ok_or_else(|| {
            DataLoaderError::MissingColumn(
                format!("Target column '{}' not found in Parquet file. Available columns: {:?}", 
                       target_column, field_names)
            )
        })?;
    
    // Determine feature columns
    let feature_names: Vec<String> = if let Some(ref feature_cols) = options.feature_columns {
        feature_cols.clone()
    } else {
        field_names.iter()
            .filter(|&name| name != &target_column)
            .cloned()
            .collect()
    };
    
    // Map feature names to indices
    let feature_indices: Vec<usize> = feature_names.iter()
        .map(|name| {
            field_names.iter()
                .position(|f| f == name)
                .ok_or_else(|| DataLoaderError::MissingColumn(name.clone()))
        })
        .collect::<Result<Vec<_>, _>>()?;
    
    // Read batches
    let reader = builder.build()?;
    let mut data: Vec<Vec<f64>> = Vec::new();
    let mut targets: Vec<f64> = Vec::new();
    
    for batch_result in reader {
        let batch = batch_result.map_err(|e| DataLoaderError::PreprocessingError(e.to_string()))?;
        
        // Check if we've reached max rows
        if let Some(max_rows) = options.max_rows {
            if data.len() >= max_rows {
                break;
            }
        }
        
        let num_rows = batch.num_rows();
        
        // Extract target column
        let target_column = batch.column(target_idx);
        
        // Process each row
        for row_idx in 0..num_rows {
            if let Some(max_rows) = options.max_rows {
                if data.len() >= max_rows {
                    break;
                }
            }
            
            // Extract features
            let mut feature_row = Vec::with_capacity(feature_indices.len());
            for &feature_idx in &feature_indices {
                let array = batch.column(feature_idx);
                let value = extract_numeric_value(array, row_idx)?;
                
                // Handle missing values
                let final_value = if !value.is_finite() {
                    options.missing_value.as_ref()
                        .and_then(|s| s.parse::<f64>().ok())
                        .unwrap_or(f64::NAN)
                } else {
                    value
                };
                feature_row.push(final_value);
            }
            
            // Extract target
            let target_value = extract_numeric_value(target_column, row_idx)?;
            
            // Skip rows with NaN targets
            if target_value.is_finite() {
                // Check for non-finite feature values
                let has_non_finite = feature_row.iter().any(|&x| !x.is_finite());
                if !has_non_finite {
                    data.push(feature_row);
                    targets.push(target_value);
                }
            }
        }
    }
    
    if data.is_empty() {
        return Err(DataLoaderError::ValidationError(
            "No valid data found in Parquet file".to_string()
        ));
    }
    
    // Create FeatureMatrix
    let n_samples = data.len();
    let n_features = feature_names.len();
    
    let flat_data: Vec<f64> = data.into_iter().flatten().collect();
    let array = Array2::from_shape_vec((n_samples, n_features), flat_data)
        .map_err(|e| DataLoaderError::ValidationError(
            format!("Failed to create feature matrix: {}", e)
        ))?;
    
    let features = FeatureMatrix::with_feature_names(array, feature_names)?;
    Dataset::new(features, targets)
        .map_err(DataLoaderError::DataError)
}

/// Extracts a numeric value from an Arrow array at a specific row index.
/// 
/// This helper function converts Arrow's native types to f64 for consistency
/// with the rest of the gradient boosting implementation.
/// 
/// # Supported Types
/// 
/// - `Float64`: Direct extraction
/// - `Float32`: Casts to f64
/// - `Int64`: Casts to f64 (may lose precision for very large values)
/// - `Int32`: Casts to f64
/// - All others: Returns `NaN`
/// 
/// # Parameters
/// - `array`: Reference to a dynamic Arrow array (`dyn Array`)
/// - `row_idx`: Row index to extract (0-based)
/// 
/// # Returns
/// `Ok(f64)` value extracted and converted, or `Err` if the operation fails.
/// 
/// # Panics
/// This function will panic if `row_idx` is out of bounds for the given array.
/// The caller should ensure bounds checking.
/// 
/// # Safety
/// Downcasting to specific array types is safe when checked against `array.data_type()`.
fn extract_numeric_value(array: &dyn Array, row_idx: usize) -> DataLoaderResult<f64> {
    
    Ok(match array.data_type() {
        DataType::Float64 => {
            let arr = array.as_any().downcast_ref::<Float64Array>().unwrap();
            arr.value(row_idx)
        },
        DataType::Float32 => {
            let arr = array.as_any().downcast_ref::<Float32Array>().unwrap();
            arr.value(row_idx) as f64
        },
        DataType::Int64 => {
            let arr = array.as_any().downcast_ref::<Int64Array>().unwrap();
            arr.value(row_idx) as f64
        },
        DataType::Int32 => {
            let arr = array.as_any().downcast_ref::<Int32Array>().unwrap();
            arr.value(row_idx) as f64
        },
        _ => {
            // Try numeric conversion for other types
            if let Some(arr) = array.as_any().downcast_ref::<Float64Array>() {
                arr.value(row_idx)
            } else if let Some(arr) = array.as_any().downcast_ref::<Float32Array>() {
                arr.value(row_idx) as f64
            } else if let Some(arr) = array.as_any().downcast_ref::<Int64Array>() {
                arr.value(row_idx) as f64
            } else if let Some(arr) = array.as_any().downcast_ref::<Int32Array>() {
                arr.value(row_idx) as f64
            } else {
                // For unsupported types, return NaN
                f64::NAN
            }
        }
    })
}

// ============================================================================
// Dataset Saving Functions
// ============================================================================

/// Saves dataset to file in specified format.
/// 
/// # Parameters
/// - `dataset`: Dataset to save
/// - `path`: Output file path
/// - `format`: DataFormat to use
/// 
/// # Returns
/// `Ok(())` on successful save
/// 
/// # Errors
/// - Format-specific save errors
/// - I/O errors
pub fn save_dataset<P: AsRef<Path>>(
    dataset: &Dataset,
    path: P,
    format: DataFormat,
) -> DataLoaderResult<()> {
    match format {
        DataFormat::Csv => save_dataset_csv(dataset, path),
        DataFormat::Json => save_dataset_json(dataset, path),
        DataFormat::Parquet => save_dataset_parquet(dataset, path),
    }
}

/// Saves dataset to CSV file.
///
/// # Parameters
/// - `dataset`: Dataset to save
/// - `path`: Output CSV file path
///
    /// # Format
/// - First row: feature_names_1, feature_names_2, ..., "target"
/// - Subsequent rows: feature values, target value
fn save_dataset_csv<P: AsRef<Path>>(dataset: &Dataset, path: P) -> DataLoaderResult<()> {
    let file = File::create(path)?;
    let mut writer = csv::Writer::from_writer(BufWriter::new(file));

    // Write headers
    let mut headers: Vec<String> = dataset.features().feature_names().to_vec();
    headers.push("target".to_string());
    writer.write_record(&headers)?;

    // Write data
    for i in 0..dataset.n_samples() {
        let sample = dataset.features().get_sample(i)?;
        let target = dataset.targets()[i];

        let mut record: Vec<String> = sample.iter().map(|x| x.to_string()).collect();
        record.push(target.to_string());

        writer.write_record(&record)?;
    }

    writer.flush()?;
    Ok(())
}

/// Saves dataset to JSON file.
///
/// # Parameters
/// - `dataset`: Dataset to save
/// - `path`: Output JSON file path
///
/// # Format
/// JSON array of objects, each containing all features and target:
/// ```json
/// [
///   {"feature1": 1.0, "feature2": 2.0, "target": 0.0},
///   {"feature1": 3.0, "feature2": 4.0, "target": 1.0}
/// ]
/// ```
fn save_dataset_json<P: AsRef<Path>>(dataset: &Dataset, path: P) -> DataLoaderResult<()> {
    use serde_json::Value;

    let mut records = Vec::new();

    for i in 0..dataset.n_samples() {
        let sample = dataset.features().get_sample(i)?;
        let target = dataset.targets()[i];

        let mut record = serde_json::Map::new();
        
        // Add features
        for (j, &value) in sample.iter().enumerate() {
            let feature_name = dataset.features().feature_names()[j].clone();
            record.insert(feature_name, Value::from(value));
        }
        
        // Add target
        record.insert("target".to_string(), Value::from(target));
        
        records.push(Value::Object(record));
    }

    let file = File::create(path)?;
    serde_json::to_writer_pretty(BufWriter::new(file), &records)?;
    Ok(())
}

/// Saves dataset to Parquet file (placeholder).
///
/// # Note
/// This method currently returns `UnsupportedFormat` error.
fn save_dataset_parquet<P: AsRef<Path>>(_dataset: &Dataset, _path: P) -> DataLoaderResult<()> {
    // Placeholder implementation
    Err(DataLoaderError::UnsupportedFormat {
        format: "parquet".to_string(),
    })
}

// ============================================================================
// Splitting and Cross-Validation Functions
// ============================================================================

/// Splits dataset into training and testing sets.
/// 
/// # Parameters
/// - `dataset`: Dataset to split
/// - `test_size`: Fraction of samples for test set (0.0-1.0)
/// - `shuffle`: Whether to randomize sample order before splitting
/// - `random_state`: Optional seed for reproducible shuffling
/// 
/// # Returns
/// Tuple of (train_dataset, test_dataset)
/// 
/// # Errors
/// - `DataLoaderError::ValidationError` if test_size not in (0, 1)
/// - `DataLoaderError::ValidationError` if dataset too small
pub fn split_data(
    dataset: &Dataset,
    test_size: f64,
    shuffle: bool,
    random_state: Option<u64>,
) -> DataLoaderResult<(Dataset, Dataset)> {
    if !(0.0..1.0).contains(&test_size) {
        return Err(DataLoaderError::ValidationError(
            "test_size must be between 0 and 1".to_string(),
        ));
    }

    dataset
        .train_test_split(test_size, shuffle)
        .map_err(DataLoaderError::DataError)
}

/// Creates k-fold cross-validation splits.
/// 
/// Generates `n_splits` train/test dataset pairs for cross-validation.
/// Each fold uses a different, non-overlapping portion as test set.
/// 
/// # Parameters
/// - `dataset`: Dataset to split
/// - `n_splits`: Number of folds (k)
/// - `shuffle`: Whether to shuffle before splitting
/// - `random_state`: Optional seed for reproducible shuffling
/// 
/// # Returns
/// Vector of (train_dataset, test_dataset) tuples, length = n_splits
/// 
    /// # Errors
/// - `DataLoaderError::ValidationError` if n_splits < 2
/// - `DataLoaderError::ValidationError` if dataset too small
pub fn cross_validation_split(
    dataset: &Dataset,
    n_splits: usize,
    shuffle: bool,
    _random_state: Option<u64>,
) -> DataLoaderResult<Vec<(Dataset, Dataset)>> {
    if n_splits < 2 {
        return Err(DataLoaderError::ValidationError(
            "n_splits must be at least 2".to_string(),
        ));
    }

    if dataset.n_samples() < n_splits {
        return Err(DataLoaderError::ValidationError(
            "Number of samples must be at least n_splits".to_string(),
        ));
    }

    let mut splits = Vec::new();
    let fold_size = dataset.n_samples() / n_splits;

    // Create indices
    let mut indices: Vec<usize> = (0..dataset.n_samples()).collect();
    if shuffle {
        use rand::seq::SliceRandom;
        use rand::thread_rng;
        indices.shuffle(&mut thread_rng());
    }

    for fold in 0..n_splits {
        let test_start = fold * fold_size;
        let test_end = if fold == n_splits - 1 {
            dataset.n_samples()
        } else {
            (fold + 1) * fold_size
        };

        let test_indices: Vec<usize> = indices[test_start..test_end].to_vec();
        let train_indices: Vec<usize> = indices[0..test_start]
            .iter()
            .chain(&indices[test_end..])
            .cloned()
            .collect();

        let train_dataset = dataset.select_samples(&train_indices)?;
        let test_dataset = dataset.select_samples(&test_indices)?;

        splits.push((train_dataset, test_dataset));
    }

    Ok(splits)
}