///|
/// A one-dimensional estimate and its uncertainty.
pub struct Estimate1D {
  value : Double
  variance : Double
} derive(Debug)

///|
pub fn Estimate1D::value(self : Estimate1D) -> Double {
  self.value
}

///|
pub fn Estimate1D::variance(self : Estimate1D) -> Double {
  self.variance
}

///|
/// A production-ready scalar Kalman filter for cheap sensor smoothing.
pub struct Kalman1D {
  mut x : Double
  mut p : Double
  mut q : Double
  mut r : Double
  initial_state : Double
  initial_uncertainty : Double
  mut last_innovation : Double
  mut last_innovation_variance : Double
  mut last_gain : Double
  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 Kalman1D::new(
  initial_state : Double,
  initial_uncertainty : Double,
  process_noise : Double,
  measurement_noise : Double,
) -> Kalman1D {
  {
    x: initial_state,
    p: if initial_uncertainty < 0.0 {
      0.0
    } else {
      initial_uncertainty
    },
    q: if process_noise < 0.0 {
      0.0
    } else {
      process_noise
    },
    r: if measurement_noise <= 0.0 {
      0.000000000001
    } else {
      measurement_noise
    },
    initial_state,
    initial_uncertainty: if initial_uncertainty < 0.0 {
      0.0
    } else {
      initial_uncertainty
    },
    last_innovation: 0.0,
    last_innovation_variance: 0.0,
    last_gain: 0.0,
    last_nis: 0.0,
    gate_threshold: 9.210340371976184,
    predict_count: 0,
    accepted_count: 0,
    rejected_count: 0,
    missing_count: 0,
  }
}

///|
pub fn Kalman1D::predict(self : Kalman1D, u : Double) -> Unit {
  if u.is_nan() || u.is_inf() {
    return
  }
  self.x = self.x + u
  self.p = self.p + self.q
  self.predict_count = self.predict_count + 1
}

///|
pub fn Kalman1D::predict_without_control(self : Kalman1D) -> Unit {
  self.predict(0.0)
}

///|
pub fn Kalman1D::update(self : Kalman1D, z : Double) -> Unit {
  self.update_gated(z, self.gate_threshold) |> ignore
}

///|
pub fn Kalman1D::update_if_valid(self : Kalman1D, z : Double) -> UpdateResult {
  self.update_gated(z, self.gate_threshold)
}

///|
pub fn Kalman1D::update_gated(
  self : Kalman1D,
  z : Double,
  threshold : Double,
) -> UpdateResult {
  if z.is_nan() || z.is_inf() {
    self.rejected_count = self.rejected_count + 1
    return InvalidMeasurement
  }
  let innovation = z - self.x
  let innovation_variance = self.p + self.r
  if innovation_variance <= 0.000000000001 {
    self.rejected_count = self.rejected_count + 1
    return SingularInnovation
  }
  let nis = innovation * innovation / innovation_variance
  self.last_innovation = innovation
  self.last_innovation_variance = innovation_variance
  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 gain = self.p / innovation_variance
    self.x = self.x + gain * innovation
    self.p = (1.0 - gain) * self.p
    self.last_gain = gain
    self.accepted_count = self.accepted_count + 1
    Accepted
  }
}

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

///|
pub fn Kalman1D::state(self : Kalman1D) -> Double {
  self.x
}

///|
pub fn Kalman1D::uncertainty(self : Kalman1D) -> Double {
  self.p
}

///|
pub fn Kalman1D::estimate(self : Kalman1D) -> Estimate1D {
  { value: self.x, variance: self.p }
}

///|
pub fn Kalman1D::innovation(self : Kalman1D) -> Double {
  self.last_innovation
}

///|
pub fn Kalman1D::innovation_variance(self : Kalman1D) -> Double {
  self.last_innovation_variance
}

///|
pub fn Kalman1D::kalman_gain(self : Kalman1D) -> Double {
  self.last_gain
}

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

///|
pub fn Kalman1D::set_process_noise(
  self : Kalman1D,
  process_noise : Double,
) -> Unit {
  self.q = if process_noise < 0.0 { 0.0 } else { process_noise }
}

///|
pub fn Kalman1D::process_noise(self : Kalman1D) -> Double {
  self.q
}

///|
pub fn Kalman1D::set_measurement_noise(
  self : Kalman1D,
  measurement_noise : Double,
) -> Unit {
  self.r = if measurement_noise <= 0.0 {
    0.000000000001
  } else {
    measurement_noise
  }
}

///|
pub fn Kalman1D::measurement_noise(self : Kalman1D) -> Double {
  self.r
}

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

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

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

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

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

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

///|
pub fn Kalman1D::reset(self : Kalman1D) -> Unit {
  self.x = self.initial_state
  self.p = self.initial_uncertainty
  self.last_innovation = 0.0
  self.last_innovation_variance = 0.0
  self.last_gain = 0.0
  self.last_nis = 0.0
  self.predict_count = 0
  self.accepted_count = 0
  self.rejected_count = 0
  self.missing_count = 0
}

///|
/// Smooth a whole scalar measurement series and return the state history.
pub fn Kalman1D::filter(
  self : Kalman1D,
  measurements : Array[Double],
  control? : Array[Double],
) -> Array[Double] {
  let result : Array[Double] = []
  for i, measurement in measurements {
    let u = match control {
      None => 0.0
      Some(values) =>
        match values.get(i) {
          None => 0.0
          Some(value) => value
        }
    }
    self.predict(u)
    self.update(measurement)
    result.push(self.state())
  }
  result
}

///|
/// Adapt process noise using the latest innovation while bounding changes.
pub fn Kalman1D::adapt_process_noise(
  self : Kalman1D,
  lower : Double,
  upper : Double,
) -> Unit {
  let low = if lower < 0.0 { 0.0 } else { lower }
  let high = if upper < low { low } else { upper }
  let candidate = self.last_innovation * self.last_innovation - self.p
  self.q = if candidate < low {
    low
  } else if candidate > high {
    high
  } else {
    candidate
  }
}