From 8ec6afe1022dd30270d3e7b2a6c7c670bf583208 Mon Sep 17 00:00:00 2001
From: m5l14i11 <siarhei.melnik@proton.me>
Date: Thu, 3 Aug 2023 14:20:14 +0300
Subject: [PATCH] refactor

---
 ta_lib/core/src/math.rs | 24 ++++++++++++++++++++++++
 ta_lib/trend/src/wma.rs | 21 +++------------------
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/ta_lib/core/src/math.rs b/ta_lib/core/src/math.rs
index bdd3d270..9a3f5e15 100644
--- a/ta_lib/core/src/math.rs
+++ b/ta_lib/core/src/math.rs
@@ -66,6 +66,30 @@ impl Series<f64> {
         self.ew(period, |period| 1.0 / (period as f64))
     }
 
+    pub fn wma(&self, period: usize) -> Self {
+        let len = self.len();
+        let mut wma = Series::empty(len);
+
+        let weight_sum = (period * (period + 1)) as f64 / 2.0;
+
+        let mut sum = 0.0;
+
+        for i in 0..period {
+            let weight = (i + 1) as f64;
+            sum += self[i].unwrap_or(0.0) * weight;
+        }
+
+        wma[period - 1] = Some(sum / weight_sum);
+
+        for i in period..len {
+            sum += (self[i].unwrap_or(0.0) - self[i - period].unwrap_or(0.0)) * period as f64
+                - (weight_sum - period as f64);
+            wma[i] = Some(sum / weight_sum);
+        }
+
+        wma
+    }
+
     pub fn var(&self, period: usize) -> Self {
         let ma: Vec<f64> = self.ma(period).into();
 
diff --git a/ta_lib/trend/src/wma.rs b/ta_lib/trend/src/wma.rs
index b2cd8ce2..1b4a83bf 100644
--- a/ta_lib/trend/src/wma.rs
+++ b/ta_lib/trend/src/wma.rs
@@ -1,26 +1,11 @@
 use core::series::Series;
 
 pub fn wma(source: &[f64], period: usize) -> Vec<f64> {
-    let len = source.len();
-    let mut wma = Series::empty(len);
+    let source = Series::from(&source);
 
-    let weight_sum = (period * (period + 1)) as f64 / 2.0;
+    let wma = source.wma(period);
 
-    let mut sum = 0.0;
-
-    for i in 0..period {
-        let weight = (i + 1) as f64;
-        sum += source[i] * weight;
-    }
-
-    wma[period - 1] = Some(sum / weight_sum);
-
-    for i in period..len {
-        sum += (source[i] - source[i - period]) * period as f64 - (weight_sum - period as f64);
-        wma[i] = Some(sum / weight_sum);
-    }
-
-    wma.nz(Some(0.0)).into()
+    wma.into()
 }
 
 #[cfg(test)]