Documentation
extern crate nalgebra as na;
use crate::BoxErr;

/// Simple Kalman Filter
/// Returns kalman filter for single series
pub fn simple_kalman_filter(series: &Vec<f64>) -> (Vec<f64>, Vec<f64>) {

  // Initialize state estimate and estimate uncertainty
  let mut x = 0.0;
  let mut p = 0.0;

  // Process noise
  let q: f64 = 1e-5;
  
  // Measurement noise
  let r: f64 = 1e-2;

  // Vectors to store state estimates and uncertainties
  let mut xs = vec![0.0; series.len()];
  let mut ps = vec![0.0; series.len()];

  for (i, z) in series.iter().enumerate() {
      // Prediction
      let x_pred = x;
      let p_pred = p + q;

      // Kalman gain
      let k: f64 = p_pred / (p_pred + r);

      // Update
      x = x_pred + k * (z - x_pred);
      p = (1.0 - k) * p_pred;

      // Store estimates
      xs[i] = x;
      ps[i] = p;
  }

  // State and uncertainties respectfully
  (xs, ps)
}


/// Dynamic Kalman Filter
/// Executes a Kalman Filter for finding the dynamic hedge ratio
pub fn dynamic_hedge_kalman_filter(x: &Vec<f64>, y: &Vec<f64>) -> Result<Vec<f64>, BoxErr> {
  let size: usize = x.len();
  let q: f64 = 10e-6;  // process variance
  let r: f64 = 0.1f64.powf(2.0);  // estimate of measurement variance

  // The initial state parameters
  let mut xhat: Vec<f64> = vec![0.0; size];  // a posteri estimate of x
  let mut p: Vec<f64> = vec![0.0; size];  // a posteri error estimate
  let mut xhatminus: Vec<f64> = vec![0.0; size];  // a priori estimate of x
  let mut pminus: Vec<f64> = vec![0.0; size];  // a priori error estimate
  let mut k: Vec<f64> = vec![0.0; size];  // Kalman gain
  
  // Initial guesses for x and p
  xhat[0] = 0.0;
  p[0] = 1.0;

  for i in 1..size {
    // Time update
    xhatminus[i] = xhat[i - 1];
    pminus[i] = p[i - 1] + q;

    // Measurement update
    k[i] = pminus[i] / (pminus[i] + r);
    xhat[i] = xhatminus[i] + k[i] * (y[i] - xhatminus[i]);
    p[i] = (1.0 - k[i]) * pminus[i];
  }

  Ok(xhat)  // return the dynamic hedge ratio
}

#[cfg(test)]
mod tests {
  use super::*;
  use crate::get_test_data;
  use crate::metrics::spread_dynamic;

  #[test]
  fn tests_simple_kalman_filter() {
    let (series_1, series_2) = get_test_data();
    let dummy_spread: Vec<f64> = spread_dynamic(&series_1, &series_2).unwrap();
    let res = simple_kalman_filter(&dummy_spread);
    dbg!(res);
  }

  #[test]
  fn tests_dynamic_hedge_kalman_filter() {
    let (series_1, series_2) = get_test_data();
    let res = dynamic_hedge_kalman_filter(&series_1, &series_2);
    dbg!(res);
  }
}