use super::plus_dm::calculate_plus_dm;
use crate::util::dataframe_utils::check_window_size;
use polars::prelude::*;
pub fn calculate_plus_di(df: &DataFrame, window: usize) -> PolarsResult<Series> {
check_window_size(df, window, "+DI")?;
let plus_dm = calculate_plus_dm(df, window)?;
let high = df.column("high")?.f64()?;
let low = df.column("low")?.f64()?;
let close_prev = df.column("close")?.f64()?.shift(1);
let mut tr_values = Vec::with_capacity(df.height());
tr_values.push(high.get(0).unwrap_or(0.0) - low.get(0).unwrap_or(0.0));
for i in 1..df.height() {
let h = high.get(i).unwrap_or(0.0);
let l = low.get(i).unwrap_or(0.0);
let cp = close_prev.get(i).unwrap_or(0.0);
let tr = (h - l).max((h - cp).abs()).max((l - cp).abs());
tr_values.push(tr);
}
let tr_series = Series::new("tr".into(), tr_values);
let atr = tr_series.rolling_mean(RollingOptionsFixedWindow {
window_size: window,
min_periods: window,
center: false,
weights: None,
fn_params: None,
})?;
let mut plus_di_values = Vec::with_capacity(df.height());
for i in 0..df.height() {
let plus_dm_val = plus_dm.f64()?.get(i).unwrap_or(0.0);
let atr_val = atr.f64()?.get(i).unwrap_or(1.0);
if atr_val > 0.0 {
plus_di_values.push((100.0 * plus_dm_val) / atr_val);
} else {
plus_di_values.push(0.0);
}
}
Ok(Series::new("plus_di".into(), plus_di_values))
}