use common_error::DaftResult;

use crate::{
    array::{
        DataArray,
        ops::{DaftStddevAggable, GroupIndices},
    },
    datatypes::Float64Type,
    utils::stats,
};

impl DaftStddevAggable for DataArray<Float64Type> {
    type Output = DaftResult<Self>;

    fn stddev(&self) -> Self::Output {
        let stats = stats::calculate_stats(self)?;
        let values = self.into_iter().flatten();
        let stddev = stats::calculate_stddev(stats, values);
        Ok(Self::from_iter(
            self.field().clone(),
            std::iter::once(stddev),
        ))
    }

    fn grouped_stddev(&self, groups: &GroupIndices) -> Self::Output {
        let grouped_stddevs_iter = stats::grouped_stats(self, groups)?.map(|(stats, group)| {
            let values = group.iter().filter_map(|&index| self.get(index as _));
            stats::calculate_stddev(stats, values)
        });

        Ok(Self::from_iter(self.field().clone(), grouped_stddevs_iter))
    }
}
