use crate::util::dataframe_utils::check_window_size;
use polars::prelude::*;
pub fn calculate_ema(df: &DataFrame, column: &str, window: usize) -> PolarsResult<Series> {
check_window_size(df, window, "EMA")?;
let series = df.column(column)?.f64()?.clone().into_series();
let series_ca = series.f64()?;
let alpha = 2.0 / (window as f64 + 1.0);
let mut ema_values = Vec::with_capacity(series.len());
let mut sma_sum = 0.0;
let mut valid_sma_count = 0;
for i in 0..window {
let val = series_ca.get(i).unwrap_or(0.0);
if !val.is_nan() {
sma_sum += val;
valid_sma_count += 1;
}
if i < window - 1 {
ema_values.push(f64::NAN);
}
}
let initial_ema = if valid_sma_count > 0 {
sma_sum / valid_sma_count as f64
} else {
f64::NAN
};
ema_values.push(initial_ema);
let mut prev_ema = initial_ema;
for i in window..series.len() {
let price = series_ca.get(i).unwrap_or(0.0);
if price.is_nan() {
ema_values.push(f64::NAN);
continue;
}
let ema = alpha * price + (1.0 - alpha) * prev_ema;
ema_values.push(ema);
prev_ema = ema;
}
Ok(Series::new("ema".into(), ema_values))
}