///|
/// Unscented Kalman Filter using scaled sigma points.
///
/// UKF is useful when a model is smooth but its Jacobian is inconvenient to
/// derive.  The implementation falls back to diagonal spread when a noisy
/// covariance is not Cholesky-decomposable, keeping an edge device alive
/// while diagnostics can report the covariance issue separately.
pub struct UKF {
  mut x : Array[Double]
  mut p : Matrix
  q : Matrix
  r : Matrix
  initial_state : Array[Double]
  initial_covariance : Matrix
  mut alpha : Double
  mut beta : Double
  mut kappa : Double
  mut predicted_sigma_points : Array[Array[Double]]
  mut last_innovation : Array[Double]
  mut last_innovation_covariance : Matrix
  mut last_gain : Matrix
  mut last_nis : Double
  mut gate_threshold : Double
  mut predict_count : Int
  mut accepted_count : Int
  mut rejected_count : Int
  mut missing_count : Int
}

///|
pub fn UKF::new(
  initial_state : Array[Double],
  initial_covariance : Array[Array[Double]],
  process_noise : Array[Array[Double]],
  measurement_noise : Array[Array[Double]],
) -> UKF {
  let covariance = Matrix::from_rows(initial_covariance)
  let state = initial_state.copy()
  {
    x: state.copy(),
    p: covariance.copy(),
    q: Matrix::from_rows(process_noise),
    r: Matrix::from_rows(measurement_noise),
    initial_state: state,
    initial_covariance: covariance,
    alpha: 0.35,
    beta: 2.0,
    kappa: 0.0,
    predicted_sigma_points: [],
    last_innovation: [],
    last_innovation_covariance: Matrix::zeros(0, 0),
    last_gain: Matrix::zeros(state.length(), 0),
    last_nis: 0.0,
    gate_threshold: 9.210340371976184,
    predict_count: 0,
    accepted_count: 0,
    rejected_count: 0,
    missing_count: 0,
  }
}

///|
pub fn UKF::state(self : UKF) -> Array[Double] {
  self.x.copy()
}

///|
pub fn UKF::covariance(self : UKF) -> Matrix {
  self.p.copy()
}

///|
pub fn UKF::set_parameters(
  self : UKF,
  alpha : Double,
  beta : Double,
  kappa : Double,
) -> Unit {
  self.alpha = if alpha <= 0.0001 { 0.0001 } else { alpha }
  self.beta = beta
  self.kappa = kappa
}

///|
pub fn UKF::parameters(self : UKF) -> (Double, Double, Double) {
  (self.alpha, self.beta, self.kappa)
}

///|
pub fn UKF::set_gate_threshold(self : UKF, threshold : Double) -> Unit {
  self.gate_threshold = if threshold < 0.0 { 0.0 } else { threshold }
}

///|
pub fn UKF::gate_threshold(self : UKF) -> Double {
  self.gate_threshold
}

///|
fn UKF::weights(self : UKF) -> (Double, Double, Double) {
  let n = self.x.length().to_double()
  let lambda = self.alpha * self.alpha * (n + self.kappa) - n
  let scale = n + lambda
  if scale.abs() <= 0.000000000001 {
    (0.0, 0.5, 0.0)
  } else {
    let mean_zero = lambda / scale
    let covariance_zero = mean_zero +
      (1.0 - self.alpha * self.alpha + self.beta)
    let other = 1.0 / (2.0 * scale)
    (mean_zero, covariance_zero, other)
  }
}

///|
fn UKF::sigma_points(self : UKF) -> Array[Array[Double]] {
  let n = self.x.length()
  if n == 0 {
    return []
  }
  let (mean_zero, _, other) = self.weights()
  let scale = if other <= 0.0 { 1.0 } else { 1.0 / (2.0 * other) }
  let spread = self.p.scale(scale)
  let factor = match spread.cholesky() {
    Some(cholesky) => cholesky
    None => {
      let diagonal = Matrix::zeros(n, n)
      for i in 0.. ignore
      }
      diagonal
    }
  }
  let points : Array[Array[Double]] = []
  points.push(self.x.copy())
  for i in 0.. Array[Double] {
  if points.length() == 0 {
    return []
  }
  let dimension = points[0].length()
  let result = Array::make(dimension, 0.0)
  for j in 0.. Matrix {
  if points.length() == 0 {
    return Matrix::zeros(0, 0)
  }
  let result = Matrix::zeros(mean.length(), mean.length())
  for i in 0.. ignore
      }
    }
  }
  result
}

///|
pub fn UKF::predict(self : UKF, f : (Array[Double]) -> Array[Double]) -> Unit {
  let (mean_zero, covariance_zero, other) = self.weights()
  let points = self.sigma_points()
  let transformed : Array[Array[Double]] = []
  for point in points {
    let value = f(point)
    if value.length() == self.x.length() {
      transformed.push(value)
    }
  }
  if transformed.length() != points.length() || transformed.length() == 0 {
    return
  }
  self.x = weighted_point_mean(transformed, mean_zero, other)
  self.p = weighted_point_covariance(
      transformed,
      self.x,
      covariance_zero,
      other,
    )
    .add(self.q)
    .symmetric_part()
  self.predicted_sigma_points = transformed
  self.predict_count = self.predict_count + 1
}

///|
pub fn UKF::predict_with_control(
  self : UKF,
  f : (Array[Double], Array[Double]) -> Array[Double],
  control : Array[Double],
) -> Unit {
  let (mean_zero, covariance_zero, other) = self.weights()
  let points = self.sigma_points()
  let transformed : Array[Array[Double]] = []
  for point in points {
    let value = f(point, control.copy())
    if value.length() == self.x.length() {
      transformed.push(value)
    }
  }
  if transformed.length() != points.length() || transformed.length() == 0 {
    return
  }
  self.x = weighted_point_mean(transformed, mean_zero, other)
  self.p = weighted_point_covariance(
      transformed,
      self.x,
      covariance_zero,
      other,
    )
    .add(self.q)
    .symmetric_part()
  self.predicted_sigma_points = transformed
  self.predict_count = self.predict_count + 1
}

///|
pub fn UKF::update(
  self : UKF,
  z : Array[Double],
  h : (Array[Double]) -> Array[Double],
) -> UpdateResult {
  self.update_gated(z, h, self.gate_threshold)
}

///|
pub fn UKF::update_gated(
  self : UKF,
  z : Array[Double],
  h : (Array[Double]) -> Array[Double],
  threshold : Double,
) -> UpdateResult {
  let points = if self.predicted_sigma_points.length() == 0 {
    self.sigma_points()
  } else {
    self.predicted_sigma_points
  }
  if points.length() == 0 {
    self.rejected_count = self.rejected_count + 1
    return InvalidMeasurement
  }
  let (mean_zero, covariance_zero, other) = self.weights()
  let measurement_points : Array[Array[Double]] = []
  for point in points {
    let value = h(point)
    if vector_is_finite(value) {
      measurement_points.push(value)
    }
  }
  if measurement_points.length() != points.length() ||
    measurement_points.length() == 0 ||
    measurement_points[0].length() != z.length() {
    self.rejected_count = self.rejected_count + 1
    return InvalidMeasurement
  }
  if !vector_is_finite(z) {
    self.rejected_count = self.rejected_count + 1
    return InvalidMeasurement
  }
  let predicted_measurement = weighted_point_mean(
    measurement_points, mean_zero, other,
  )
  let innovation = vector_sub(z, predicted_measurement)
  let innovation_covariance = weighted_point_covariance(
      measurement_points, predicted_measurement, covariance_zero, other,
    )
    .add(self.r)
    .symmetric_part()
  match innovation_covariance.inverse() {
    None => {
      self.rejected_count = self.rejected_count + 1
      SingularInnovation
    }
    Some(inverse_covariance) => {
      let nis = vector_dot(
        innovation,
        inverse_covariance.multiply_vector(innovation),
      )
      self.last_innovation = innovation.copy()
      self.last_innovation_covariance = innovation_covariance.copy()
      self.last_nis = nis
      let safe_threshold = if threshold < 0.0 { 0.0 } else { threshold }
      if nis > safe_threshold {
        self.rejected_count = self.rejected_count + 1
        RejectedByGate
      } else {
        let cross = Matrix::zeros(self.x.length(), z.length())
        for i in 0.. ignore
            }
          }
        }
        let gain = cross.multiply(inverse_covariance)
        self.x = vector_add(self.x, gain.multiply_vector(innovation))
        self.p = self.p
          .sub(gain.multiply(innovation_covariance).multiply(gain.transpose()))
          .symmetric_part()
        self.last_gain = gain
        self.accepted_count = self.accepted_count + 1
        Accepted
      }
    }
  }
}

///|
pub fn UKF::update_missing(self : UKF) -> UpdateResult {
  self.missing_count = self.missing_count + 1
  MissingMeasurement
}

///|
pub fn UKF::innovation(self : UKF) -> Array[Double] {
  self.last_innovation.copy()
}

///|
pub fn UKF::innovation_covariance(self : UKF) -> Matrix {
  self.last_innovation_covariance.copy()
}

///|
pub fn UKF::kalman_gain(self : UKF) -> Matrix {
  self.last_gain.copy()
}

///|
pub fn UKF::normalized_innovation_squared(self : UKF) -> Double {
  self.last_nis
}

///|
pub fn UKF::predict_count(self : UKF) -> Int {
  self.predict_count
}

///|
pub fn UKF::accepted_count(self : UKF) -> Int {
  self.accepted_count
}

///|
pub fn UKF::rejected_count(self : UKF) -> Int {
  self.rejected_count
}

///|
pub fn UKF::missing_count(self : UKF) -> Int {
  self.missing_count
}

///|
pub fn UKF::reset(self : UKF) -> Unit {
  self.x = self.initial_state.copy()
  self.p = self.initial_covariance.copy()
  self.predicted_sigma_points = []
  self.last_innovation = []
  self.last_innovation_covariance = Matrix::zeros(0, 0)
  self.last_gain = Matrix::zeros(self.x.length(), 0)
  self.last_nis = 0.0
  self.predict_count = 0
  self.accepted_count = 0
  self.rejected_count = 0
  self.missing_count = 0
}

///|
pub fn UKF::filter(
  self : UKF,
  measurements : Array[Array[Double]],
  f : (Array[Double]) -> Array[Double],
  h : (Array[Double]) -> Array[Double],
) -> Array[Array[Double]] {
  let states : Array[Array[Double]] = []
  for measurement in measurements {
    self.predict(f)
    self.update(measurement, h) |> ignore
    states.push(self.state())
  }
  states
}