pandrs 0.4.1

//! GroupBy operations for DataFrames
//!
//! Provides pandas-compatible GroupBy functionality for aggregating data.

use crate::core::error::{Error, Result};
use crate::dataframe::base::DataFrame;
use crate::series::Series;
use std::collections::HashMap;

/// GroupBy object that holds the grouped data
pub struct DataFrameGroupBy<'a> {
    df: &'a DataFrame,
    group_columns: Vec<String>,
    /// Maps group key (as string representation) to row indices
    groups: HashMap<String, Vec<usize>>,
    /// Maps group key to the actual group values
    group_keys: HashMap<String, Vec<String>>,
}

impl<'a> DataFrameGroupBy<'a> {
    /// Create a new GroupBy object
    pub fn new(df: &'a DataFrame, by: &[&str]) -> Result<Self> {
        if by.is_empty() {
            return Err(Error::InvalidValue(
                "GroupBy requires at least one column".to_string(),
            ));
        }

        // Validate columns exist
        for col in by {
            if !df.contains_column(col) {
                return Err(Error::InvalidValue(format!(
                    "Column '{}' not found in DataFrame",
                    col
                )));
            }
        }

        let group_columns: Vec<String> = by.iter().map(|s| s.to_string()).collect();
        let mut groups: HashMap<String, Vec<usize>> = HashMap::new();
        let mut group_keys: HashMap<String, Vec<String>> = HashMap::new();

        let row_count = df.row_count();

        // Materialize each group-by column once up front instead of
        // re-fetching (and re-downcasting) it inside the row loop below --
        // the previous version called `get_column_string_values`/
        // `get_column_numeric_values` for every (row, group column) pair,
        // making `GroupBy::new` (the entry point of every pandas_compat
        // groupby operation) O(rows^2 * group_columns) instead of
        // O(rows * group_columns).
        enum KeyColumn {
            Str(Vec<String>),
            Num(Vec<f64>),
            Missing,
        }
        let key_columns: Vec<KeyColumn> = group_columns
            .iter()
            .map(|col| {
                if let Ok(values) = df.get_column_string_values(col) {
                    KeyColumn::Str(values)
                } else if let Ok(values) = df.get_column_numeric_values(col) {
                    KeyColumn::Num(values)
                } else {
                    KeyColumn::Missing
                }
            })
            .collect();

        // Build group indices
        for row_idx in 0..row_count {
            let mut key_parts: Vec<String> = Vec::with_capacity(key_columns.len());

            for col in &key_columns {
                let value = match col {
                    KeyColumn::Str(values) => values.get(row_idx).cloned().unwrap_or_default(),
                    KeyColumn::Num(values) => {
                        let v = values.get(row_idx).copied().unwrap_or(f64::NAN);
                        if v.is_nan() {
                            "NaN".to_string()
                        } else {
                            v.to_string()
                        }
                    }
                    KeyColumn::Missing => "".to_string(),
                };
                key_parts.push(value);
            }

            let key = key_parts.join("|||");
            groups
                .entry(key.clone())
                .or_insert_with(Vec::new)
                .push(row_idx);
            group_keys.entry(key).or_insert(key_parts);
        }

        Ok(Self {
            df,
            group_columns,
            groups,
            group_keys,
        })
    }

    /// Get the number of groups
    pub fn ngroups(&self) -> usize {
        self.groups.len()
    }

    /// Return this GroupBy's group keys in a stable, deterministic order.
    ///
    /// `self.groups`/`self.group_keys` are `HashMap`s, whose iteration
    /// order is not guaranteed to be the same from one run to the next;
    /// every method that walks all groups sorts through this helper first
    /// so the resulting DataFrame's row order is reproducible.
    fn sorted_group_keys(&self) -> Vec<&String> {
        let mut keys: Vec<&String> = self.groups.keys().collect();
        keys.sort();
        keys
    }

    /// Get group sizes
    pub fn size(&self) -> Result<DataFrame> {
        let mut result = DataFrame::new();

        let mut group_col_values: Vec<Vec<String>> = vec![Vec::new(); self.group_columns.len()];
        let mut sizes: Vec<f64> = Vec::new();

        for key in self.sorted_group_keys() {
            let indices = self.groups.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its row indices", key))
            })?;
            let key_values = self.group_keys.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its key values", key))
            })?;
            for (i, val) in key_values.iter().enumerate() {
                group_col_values[i].push(val.clone());
            }
            sizes.push(indices.len() as f64);
        }

        // Add group columns
        for (i, col_name) in self.group_columns.iter().enumerate() {
            result.add_column(
                col_name.clone(),
                Series::new(group_col_values[i].clone(), Some(col_name.clone()))?,
            )?;
        }

        // Add size column
        result.add_column(
            "size".to_string(),
            Series::new(sizes, Some("size".to_string()))?,
        )?;

        Ok(result)
    }

    /// Count non-null values per group, for every non-group column.
    ///
    /// Unlike [`size`](Self::size) (which just counts *rows*), pandas'
    /// `GroupBy.count()` reports, per group and per column, how many of
    /// that group's values in that column are non-null -- a group with a
    /// missing value in one column but not another gets different counts
    /// for the two. Object (string) columns follow this crate's existing
    /// convention (see `count_valid` / `PandasCompatExt`) of treating an
    /// empty string as the missing marker.
    pub fn count(&self) -> Result<DataFrame> {
        let mut result = DataFrame::new();

        let other_cols: Vec<String> = self
            .df
            .column_names()
            .iter()
            .filter(|col| !self.group_columns.contains(*col))
            .cloned()
            .collect();

        // Materialize each non-group column once (dispatched by concrete
        // dtype) instead of re-fetching it once per group.
        enum ColData {
            Num(Vec<f64>),
            Str(Vec<String>),
        }
        let mut col_data: HashMap<String, ColData> = HashMap::new();
        for col in &other_cols {
            if self.df.is_numeric_column(col) {
                if let Ok(v) = self.df.get_column_numeric_values(col) {
                    col_data.insert(col.clone(), ColData::Num(v));
                }
            } else if let Ok(v) = self.df.get_column_string_values(col) {
                col_data.insert(col.clone(), ColData::Str(v));
            }
        }

        let mut group_col_values: Vec<Vec<String>> = vec![Vec::new(); self.group_columns.len()];
        let mut counts: HashMap<String, Vec<f64>> = other_cols
            .iter()
            .filter(|c| col_data.contains_key(c.as_str()))
            .map(|c| (c.clone(), Vec::new()))
            .collect();

        for key in self.sorted_group_keys() {
            let indices = self.groups.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its row indices", key))
            })?;
            let key_values = self.group_keys.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its key values", key))
            })?;
            for (i, val) in key_values.iter().enumerate() {
                group_col_values[i].push(val.clone());
            }

            for col in &other_cols {
                let Some(data) = col_data.get(col) else {
                    continue;
                };
                let non_null = match data {
                    ColData::Num(vals) => indices
                        .iter()
                        .filter(|&&i| vals.get(i).map(|v| !v.is_nan()).unwrap_or(false))
                        .count(),
                    ColData::Str(vals) => indices
                        .iter()
                        .filter(|&&i| vals.get(i).map(|v| !v.is_empty()).unwrap_or(false))
                        .count(),
                };
                if let Some(bucket) = counts.get_mut(col) {
                    bucket.push(non_null as f64);
                }
            }
        }

        for (i, col_name) in self.group_columns.iter().enumerate() {
            result.add_column(
                col_name.clone(),
                Series::new(group_col_values[i].clone(), Some(col_name.clone()))?,
            )?;
        }
        for col in &other_cols {
            if let Some(values) = counts.get(col) {
                result.add_column(col.clone(), Series::new(values.clone(), Some(col.clone()))?)?;
            }
        }

        Ok(result)
    }

    /// Sum numeric columns per group
    pub fn sum(&self) -> Result<DataFrame> {
        self.aggregate(|values| values.iter().filter(|v| !v.is_nan()).sum())
    }

    /// Mean of numeric columns per group
    pub fn mean(&self) -> Result<DataFrame> {
        self.aggregate(|values| {
            let valid: Vec<f64> = values.iter().filter(|v| !v.is_nan()).copied().collect();
            if valid.is_empty() {
                f64::NAN
            } else {
                valid.iter().sum::<f64>() / valid.len() as f64
            }
        })
    }

    /// Minimum of numeric columns per group
    pub fn min(&self) -> Result<DataFrame> {
        self.aggregate(|values| {
            values
                .iter()
                .filter(|v| !v.is_nan())
                .copied()
                .fold(f64::INFINITY, f64::min)
        })
    }

    /// Maximum of numeric columns per group
    pub fn max(&self) -> Result<DataFrame> {
        self.aggregate(|values| {
            values
                .iter()
                .filter(|v| !v.is_nan())
                .copied()
                .fold(f64::NEG_INFINITY, f64::max)
        })
    }

    /// Standard deviation of numeric columns per group (sample std, using n-1)
    pub fn std(&self) -> Result<DataFrame> {
        self.aggregate(|values| {
            let valid: Vec<f64> = values.iter().filter(|v| !v.is_nan()).copied().collect();
            if valid.len() <= 1 {
                f64::NAN
            } else {
                let mean = valid.iter().sum::<f64>() / valid.len() as f64;
                let variance: f64 = valid.iter().map(|v| (v - mean).powi(2)).sum::<f64>()
                    / (valid.len() - 1) as f64;
                variance.sqrt()
            }
        })
    }

    /// Variance of numeric columns per group (sample variance, using n-1)
    pub fn var(&self) -> Result<DataFrame> {
        self.aggregate(|values| {
            let valid: Vec<f64> = values.iter().filter(|v| !v.is_nan()).copied().collect();
            if valid.len() <= 1 {
                f64::NAN
            } else {
                let mean = valid.iter().sum::<f64>() / valid.len() as f64;
                valid.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / (valid.len() - 1) as f64
            }
        })
    }

    /// First value of each group
    pub fn first(&self) -> Result<DataFrame> {
        self.aggregate_first_last(true)
    }

    /// Last value of each group
    pub fn last(&self) -> Result<DataFrame> {
        self.aggregate_first_last(false)
    }

    /// Internal method to apply an aggregation function
    fn aggregate<F>(&self, agg_fn: F) -> Result<DataFrame>
    where
        F: Fn(&[f64]) -> f64,
    {
        let mut result = DataFrame::new();

        // Get numeric columns (excluding group columns)
        let numeric_cols: Vec<String> = self
            .df
            .column_names()
            .iter()
            .filter(|col| !self.group_columns.contains(*col) && self.df.is_numeric_column(col))
            .cloned()
            .collect();

        // Materialize each numeric column once instead of re-fetching (and
        // re-downcasting) it once per group -- O(groups * cols) column
        // materializations instead of O(cols).
        let mut col_values: HashMap<&str, Vec<f64>> = HashMap::with_capacity(numeric_cols.len());
        for col in &numeric_cols {
            col_values.insert(col.as_str(), self.df.get_column_numeric_values(col)?);
        }

        // Prepare group column values
        let mut group_col_values: Vec<Vec<String>> = vec![Vec::new(); self.group_columns.len()];

        // Prepare aggregated values for each numeric column
        let mut agg_values: HashMap<String, Vec<f64>> = HashMap::new();
        for col in &numeric_cols {
            agg_values.insert(col.clone(), Vec::new());
        }

        // Process each group, in a deterministic (sorted-key) order rather
        // than raw `HashMap` iteration order.
        for key in self.sorted_group_keys() {
            let indices = self.groups.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its row indices", key))
            })?;
            let key_values = self.group_keys.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its key values", key))
            })?;
            for (i, val) in key_values.iter().enumerate() {
                group_col_values[i].push(val.clone());
            }

            // Aggregate each numeric column. Always push exactly one value
            // per group per column (NaN if the column's data were somehow
            // unavailable) so the group-key columns and every aggregated
            // column stay row-aligned; the previous version pushed a group
            // key unconditionally but the aggregated value only inside an
            // `if let Ok(..)`, which -- on any future divergence between
            // that check and `numeric_cols`'s upfront filter -- would shift
            // every later group's value into the wrong row.
            for col in &numeric_cols {
                let aggregated = match col_values.get(col.as_str()) {
                    Some(all_values) => {
                        let group_values: Vec<f64> = indices
                            .iter()
                            .filter_map(|&i| all_values.get(i).copied())
                            .collect();
                        agg_fn(&group_values)
                    }
                    None => f64::NAN,
                };
                agg_values
                    .get_mut(col)
                    .ok_or_else(|| {
                        Error::InvalidValue(format!(
                            "column '{}' missing from aggregation buffer",
                            col
                        ))
                    })?
                    .push(aggregated);
            }
        }

        // Build result DataFrame
        // Add group columns
        for (i, col_name) in self.group_columns.iter().enumerate() {
            result.add_column(
                col_name.clone(),
                Series::new(group_col_values[i].clone(), Some(col_name.clone()))?,
            )?;
        }

        // Add aggregated columns
        for col in &numeric_cols {
            if let Some(values) = agg_values.get(col) {
                result.add_column(col.clone(), Series::new(values.clone(), Some(col.clone()))?)?;
            }
        }

        Ok(result)
    }

    /// Internal method for first/last aggregations
    fn aggregate_first_last(&self, first: bool) -> Result<DataFrame> {
        let mut result = DataFrame::new();

        // Get all non-group columns
        let other_cols: Vec<String> = self
            .df
            .column_names()
            .iter()
            .filter(|col| !self.group_columns.contains(*col))
            .cloned()
            .collect();

        // Materialize each column once (dispatched by concrete dtype, not
        // by whether numeric *or* string conversion happens to succeed --
        // every numeric column also renders through
        // `get_column_string_values`) instead of re-fetching it once per
        // group.
        enum ColData {
            Num(Vec<f64>),
            Str(Vec<String>),
        }
        let mut col_data: HashMap<String, ColData> = HashMap::new();
        for col in &other_cols {
            if self.df.is_numeric_column(col) {
                if let Ok(v) = self.df.get_column_numeric_values(col) {
                    col_data.insert(col.clone(), ColData::Num(v));
                }
            } else if let Ok(v) = self.df.get_column_string_values(col) {
                col_data.insert(col.clone(), ColData::Str(v));
            }
        }

        // Prepare group column values
        let mut group_col_values: Vec<Vec<String>> = vec![Vec::new(); self.group_columns.len()];

        // Prepare values for each column
        let mut numeric_values: HashMap<String, Vec<f64>> = HashMap::new();
        let mut string_values: HashMap<String, Vec<String>> = HashMap::new();
        for (col, data) in &col_data {
            match data {
                ColData::Num(_) => {
                    numeric_values.insert(col.clone(), Vec::new());
                }
                ColData::Str(_) => {
                    string_values.insert(col.clone(), Vec::new());
                }
            }
        }

        // Process each group, in a deterministic (sorted-key) order.
        for key in self.sorted_group_keys() {
            let indices = self.groups.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its row indices", key))
            })?;
            let key_values = self.group_keys.get(key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its key values", key))
            })?;
            for (i, val) in key_values.iter().enumerate() {
                group_col_values[i].push(val.clone());
            }

            let target_idx = if first {
                *indices
                    .first()
                    .ok_or_else(|| Error::InvalidValue(format!("group '{}' has no rows", key)))?
            } else {
                *indices
                    .last()
                    .ok_or_else(|| Error::InvalidValue(format!("group '{}' has no rows", key)))?
            };

            // Get first/last value for each column
            for col in &other_cols {
                match col_data.get(col) {
                    Some(ColData::Num(all_values)) => {
                        let value = all_values.get(target_idx).copied().unwrap_or(f64::NAN);
                        numeric_values
                            .get_mut(col)
                            .ok_or_else(|| {
                                Error::InvalidValue(format!(
                                    "column '{}' missing from first/last buffer",
                                    col
                                ))
                            })?
                            .push(value);
                    }
                    Some(ColData::Str(all_values)) => {
                        let value = all_values.get(target_idx).cloned().unwrap_or_default();
                        string_values
                            .get_mut(col)
                            .ok_or_else(|| {
                                Error::InvalidValue(format!(
                                    "column '{}' missing from first/last buffer",
                                    col
                                ))
                            })?
                            .push(value);
                    }
                    None => {}
                }
            }
        }

        // Build result DataFrame
        // Add group columns
        for (i, col_name) in self.group_columns.iter().enumerate() {
            result.add_column(
                col_name.clone(),
                Series::new(group_col_values[i].clone(), Some(col_name.clone()))?,
            )?;
        }

        // Add other columns (preserve order from original DataFrame)
        for col in &other_cols {
            if let Some(values) = numeric_values.get(col) {
                result.add_column(col.clone(), Series::new(values.clone(), Some(col.clone()))?)?;
            } else if let Some(values) = string_values.get(col) {
                result.add_column(col.clone(), Series::new(values.clone(), Some(col.clone()))?)?;
            }
        }

        Ok(result)
    }

    /// Apply multiple aggregations at once
    pub fn agg(&self, aggs: &[(&str, &str)]) -> Result<DataFrame> {
        let mut result = DataFrame::new();

        // Process every group in the same deterministic order for both the
        // group-key columns and every aggregated column below. The
        // previous version built `group_col_values` from one
        // `self.groups.iter()` pass and each aggregated column from
        // *another* -- internally consistent within a single run (nothing
        // mutates `self.groups` in between), but the row order itself
        // still varied from run to run with `HashMap`'s iteration order.
        let keys = self.sorted_group_keys();

        // Prepare group column values
        let mut group_col_values: Vec<Vec<String>> = vec![Vec::new(); self.group_columns.len()];
        for key in &keys {
            let key_values = self.group_keys.get(*key).ok_or_else(|| {
                Error::InvalidValue(format!("group '{}' missing its key values", key))
            })?;
            for (i, val) in key_values.iter().enumerate() {
                group_col_values[i].push(val.clone());
            }
        }

        // Add group columns
        for (i, col_name) in self.group_columns.iter().enumerate() {
            result.add_column(
                col_name.clone(),
                Series::new(group_col_values[i].clone(), Some(col_name.clone()))?,
            )?;
        }

        // Process each aggregation
        for (col, agg_name) in aggs {
            if !self.df.contains_column(col) {
                continue;
            }

            if let Ok(all_values) = self.df.get_column_numeric_values(col) {
                let mut agg_values: Vec<f64> = Vec::with_capacity(keys.len());

                for key in &keys {
                    let indices = self.groups.get(*key).ok_or_else(|| {
                        Error::InvalidValue(format!("group '{}' missing its row indices", key))
                    })?;
                    let group_values: Vec<f64> = indices
                        .iter()
                        .filter_map(|&i| all_values.get(i).copied())
                        .collect();

                    let aggregated = match *agg_name {
                        "sum" => group_values.iter().filter(|v| !v.is_nan()).sum(),
                        "mean" => {
                            let valid: Vec<f64> = group_values
                                .iter()
                                .filter(|v| !v.is_nan())
                                .copied()
                                .collect();
                            if valid.is_empty() {
                                f64::NAN
                            } else {
                                valid.iter().sum::<f64>() / valid.len() as f64
                            }
                        }
                        "min" => group_values
                            .iter()
                            .filter(|v| !v.is_nan())
                            .copied()
                            .fold(f64::INFINITY, f64::min),
                        "max" => group_values
                            .iter()
                            .filter(|v| !v.is_nan())
                            .copied()
                            .fold(f64::NEG_INFINITY, f64::max),
                        "count" => group_values.iter().filter(|v| !v.is_nan()).count() as f64,
                        "std" => {
                            let valid: Vec<f64> = group_values
                                .iter()
                                .filter(|v| !v.is_nan())
                                .copied()
                                .collect();
                            if valid.len() <= 1 {
                                f64::NAN
                            } else {
                                let mean = valid.iter().sum::<f64>() / valid.len() as f64;
                                let variance: f64 =
                                    valid.iter().map(|v| (v - mean).powi(2)).sum::<f64>()
                                        / (valid.len() - 1) as f64;
                                variance.sqrt()
                            }
                        }
                        "var" => {
                            let valid: Vec<f64> = group_values
                                .iter()
                                .filter(|v| !v.is_nan())
                                .copied()
                                .collect();
                            if valid.len() <= 1 {
                                f64::NAN
                            } else {
                                let mean = valid.iter().sum::<f64>() / valid.len() as f64;
                                valid.iter().map(|v| (v - mean).powi(2)).sum::<f64>()
                                    / (valid.len() - 1) as f64
                            }
                        }
                        "first" => group_values.first().copied().unwrap_or(f64::NAN),
                        "last" => group_values.last().copied().unwrap_or(f64::NAN),
                        _ => f64::NAN,
                    };

                    agg_values.push(aggregated);
                }

                let result_col_name = format!("{}_{}", col, agg_name);
                result.add_column(
                    result_col_name.clone(),
                    Series::new(agg_values, Some(result_col_name))?,
                )?;
            }
        }

        Ok(result)
    }
}

/// Extension trait to add groupby method to DataFrame (multi-column support)
pub trait PandasGroupByExt {
    /// Group DataFrame by one or more columns
    ///
    /// # Example
    /// ```ignore
    /// use pandrs::dataframe::pandas_compat::PandasGroupByExt;
    ///
    /// let result = df.groupby_multi(&["category"]).expect("test should succeed").sum().expect("test should succeed");
    /// ```
    fn groupby_multi(&self, by: &[&str]) -> Result<DataFrameGroupBy>;
}

impl PandasGroupByExt for DataFrame {
    fn groupby_multi(&self, by: &[&str]) -> Result<DataFrameGroupBy> {
        DataFrameGroupBy::new(self, by)
    }
}

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

    fn create_test_df() -> DataFrame {
        let mut df = DataFrame::new();
        df.add_column(
            "category".to_string(),
            Series::new(
                vec![
                    "A".to_string(),
                    "B".to_string(),
                    "A".to_string(),
                    "B".to_string(),
                    "A".to_string(),
                ],
                Some("category".to_string()),
            )
            .expect("test should succeed"),
        )
        .expect("test should succeed");
        df.add_column(
            "value".to_string(),
            Series::new(
                vec![10.0, 20.0, 30.0, 40.0, 50.0],
                Some("value".to_string()),
            )
            .expect("test should succeed"),
        )
        .expect("test should succeed");
        df.add_column(
            "score".to_string(),
            Series::new(vec![1.0, 2.0, 3.0, 4.0, 5.0], Some("score".to_string()))
                .expect("test should succeed"),
        )
        .expect("test should succeed");
        df
    }

    #[test]
    fn test_groupby_sum() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .sum()
            .expect("test should succeed");

        assert_eq!(result.row_count(), 2);

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        // Find indices for A and B
        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A: 10 + 30 + 50 = 90
        assert_eq!(values[a_idx], 90.0);
        // B: 20 + 40 = 60
        assert_eq!(values[b_idx], 60.0);
    }

    #[test]
    fn test_groupby_mean() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .mean()
            .expect("test should succeed");

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A: (10 + 30 + 50) / 3 = 30
        assert_eq!(values[a_idx], 30.0);
        // B: (20 + 40) / 2 = 30
        assert_eq!(values[b_idx], 30.0);
    }

    #[test]
    fn test_groupby_min() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .min()
            .expect("test should succeed");

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A: min(10, 30, 50) = 10
        assert_eq!(values[a_idx], 10.0);
        // B: min(20, 40) = 20
        assert_eq!(values[b_idx], 20.0);
    }

    #[test]
    fn test_groupby_max() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .max()
            .expect("test should succeed");

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A: max(10, 30, 50) = 50
        assert_eq!(values[a_idx], 50.0);
        // B: max(20, 40) = 40
        assert_eq!(values[b_idx], 40.0);
    }

    #[test]
    fn test_groupby_count() {
        // `count()` reports non-null counts *per column*, not row counts
        // (that's `size()`) -- previously `count()` literally called
        // `size()`, so it returned one "size" column instead of a
        // per-column non-null tally. With no actual nulls in this fixture
        // the numbers happen to match `size()`'s, but the *shape* of the
        // result (one count column per original non-group column, not a
        // single "size" column) is what this test now asserts; see
        // `test_groupby_count_skips_nulls` for the behavior that actually
        // distinguishes `count()` from `size()`.
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .count()
            .expect("test should succeed");

        assert!(!result.contains_column("size"));
        assert!(result.contains_column("value"));
        assert!(result.contains_column("score"));

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let sizes = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A: 3 rows, all non-null
        assert_eq!(sizes[a_idx], 3.0);
        // B: 2 rows, all non-null
        assert_eq!(sizes[b_idx], 2.0);
    }

    #[test]
    fn test_groupby_std() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .std()
            .expect("test should succeed");

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");

        // A: std of [10, 30, 50] with sample std (n-1)
        // mean = 30, variance = ((10-30)^2 + (30-30)^2 + (50-30)^2) / 2 = 400
        // std = 20
        assert!((values[a_idx] - 20.0).abs() < 0.001);
    }

    #[test]
    fn test_groupby_first() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .first()
            .expect("test should succeed");

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A first: 10
        assert_eq!(values[a_idx], 10.0);
        // B first: 20
        assert_eq!(values[b_idx], 20.0);
    }

    #[test]
    fn test_groupby_last() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .last()
            .expect("test should succeed");

        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");
        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        let b_idx = cats
            .iter()
            .position(|c| c == "B")
            .expect("test should succeed");

        // A last: 50
        assert_eq!(values[a_idx], 50.0);
        // B last: 40
        assert_eq!(values[b_idx], 40.0);
    }

    #[test]
    fn test_groupby_multiple_columns() {
        let mut df = DataFrame::new();
        df.add_column(
            "cat1".to_string(),
            Series::new(
                vec![
                    "A".to_string(),
                    "A".to_string(),
                    "B".to_string(),
                    "B".to_string(),
                ],
                Some("cat1".to_string()),
            )
            .expect("test should succeed"),
        )
        .expect("test should succeed");
        df.add_column(
            "cat2".to_string(),
            Series::new(
                vec![
                    "X".to_string(),
                    "Y".to_string(),
                    "X".to_string(),
                    "Y".to_string(),
                ],
                Some("cat2".to_string()),
            )
            .expect("test should succeed"),
        )
        .expect("test should succeed");
        df.add_column(
            "value".to_string(),
            Series::new(vec![1.0, 2.0, 3.0, 4.0], Some("value".to_string()))
                .expect("test should succeed"),
        )
        .expect("test should succeed");

        let result = df
            .groupby_multi(&["cat1", "cat2"])
            .expect("test should succeed")
            .sum()
            .expect("test should succeed");

        // Should have 4 groups: (A,X), (A,Y), (B,X), (B,Y)
        assert_eq!(result.row_count(), 4);
    }

    #[test]
    fn test_groupby_with_nan() {
        let mut df = DataFrame::new();
        df.add_column(
            "category".to_string(),
            Series::new(
                vec!["A".to_string(), "A".to_string(), "A".to_string()],
                Some("category".to_string()),
            )
            .expect("test should succeed"),
        )
        .expect("test should succeed");
        df.add_column(
            "value".to_string(),
            Series::new(vec![10.0, f64::NAN, 30.0], Some("value".to_string()))
                .expect("test should succeed"),
        )
        .expect("test should succeed");

        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .sum()
            .expect("test should succeed");

        let values = result
            .get_column_numeric_values("value")
            .expect("test should succeed");
        // Sum should ignore NaN: 10 + 30 = 40
        assert_eq!(values[0], 40.0);
    }

    #[test]
    fn test_groupby_agg() {
        let df = create_test_df();
        let result = df
            .groupby_multi(&["category"])
            .expect("test should succeed")
            .agg(&[("value", "sum"), ("value", "mean"), ("score", "max")])
            .expect("test should succeed");

        assert!(result.contains_column("value_sum"));
        assert!(result.contains_column("value_mean"));
        assert!(result.contains_column("score_max"));

        let value_sums = result
            .get_column_numeric_values("value_sum")
            .expect("test should succeed");
        let cats = result
            .get_column_string_values("category")
            .expect("test should succeed");

        let a_idx = cats
            .iter()
            .position(|c| c == "A")
            .expect("test should succeed");
        assert_eq!(value_sums[a_idx], 90.0);
    }

    #[test]
    fn test_ngroups() {
        let df = create_test_df();
        let gb = df
            .groupby_multi(&["category"])
            .expect("test should succeed");
        assert_eq!(gb.ngroups(), 2);
    }
}