// The filters that take any plant: the extended one linearises the model at the
// current estimate, the unscented one carries a small set of sigma points
// through it. Both implement the same `Estimator` contract, so a caller swaps
// one for the other without touching the loop.

///|
/// An extended Kalman filter: the linear recursion with the Jacobians the model
/// supplies, whether analytically or by the contract's default differences.
pub(all) struct ExtendedKalman {
  model : &@contract.Model
  process : @linear.Mat
  measurement : @linear.Mat
  mut state : @linear.Vec
  mut covariance : @linear.Mat
}

///|
/// An extended Kalman filter for any plant that implements the model contract.
pub fn ExtendedKalman::new(
  model~ : &@contract.Model,
  process~ : @linear.Mat,
  measurement~ : @linear.Mat,
  state~ : @linear.Vec,
  covariance~ : @linear.Mat,
) -> ExtendedKalman raise @linear.Refused {
  let n = state.len()
  if covariance.rows != n || covariance.cols != n {
    raise @linear.Refused::Dimension(got=covariance.rows, want=n)
  }
  if process.rows != n || process.cols != n {
    raise @linear.Refused::Dimension(got=process.rows, want=n)
  }
  let m = measurement.rows
  if measurement.cols != m {
    raise @linear.Refused::Dimension(got=measurement.cols, want=m)
  }
  { model, process, measurement, state: state.clone(), covariance, }
}

///|
pub impl @contract.Estimator for ExtendedKalman with fn predict(self, u, dt) {
  let transition = self.model.step_jacobian(self.state, u, dt)
  self.state = self.model.step(self.state, u, dt)
  let propagated = transition.mul(self.covariance).mul(transition.transpose())
  self.covariance = propagated.add(self.process)
}

///|
pub impl @contract.Estimator for ExtendedKalman with fn correct(self, z) {
  let observation = self.model.observe_jacobian(self.state)
  let predicted = self.model.observe(self.state)
  let innovation = z.sub(predicted)
  let s = observation
    .mul(self.covariance)
    .mul(observation.transpose())
    .add(self.measurement)
  let gain = @linear.solve(s, observation.mul(self.covariance)).transpose()
  let n = self.state.len()
  self.state = self.state.add(gain.apply(innovation))
  let identity = @linear.Mat::identity(n)
  let factor = identity.sub(gain.mul(observation))
  self.covariance = factor
    .mul(self.covariance)
    .mul(factor.transpose())
    .add(gain.mul(self.measurement).mul(gain.transpose()))
}

///|
pub impl @contract.Estimator for ExtendedKalman with fn state(self) {
  self.state.clone()
}

///|
pub impl @contract.Estimator for ExtendedKalman with fn covariance(self) {
  self.covariance
}

///|
pub extend ExtendedKalman with @contract.Estimator::{
  predict,
  correct,
  state,
  covariance,
}

///|
impl @debug.Debug for ExtendedKalman with fn to_repr(self) {
  @debug.Repr::literal(
    "ExtendedKalman(" + @linear.Vec::to_string(self.state) + ")",
  )
}

///|
pub extend ExtendedKalman with @debug.Debug::{to_repr}

///|
/// The spread of the sigma points in the unscented transform.
///
/// `alpha` sets how far the points sit from the mean, `beta` folds in the
/// knowledge that the distribution is Gaussian, and `kappa` is the usual
/// secondary scaling. The preset is the one the reference implementations use.
#derive(Debug)
pub(all) struct UnscentedConfig {
  alpha : Double
  beta : Double
  kappa : Double
}

///|
/// The preset spread: a narrow alpha, the Gaussian beta and a zero kappa.
pub let unscented_config : UnscentedConfig = {
  alpha: 1.0e-3,
  beta: 2.0,
  kappa: 0.0,
}

///|
/// An unscented Kalman filter.
pub(all) struct UnscentedKalman {
  model : &@contract.Model
  process : @linear.Mat
  measurement : @linear.Mat
  mut state : @linear.Vec
  mut covariance : @linear.Mat
  config : UnscentedConfig
}

///|
/// An unscented Kalman filter for any plant that implements the model contract.
pub fn UnscentedKalman::new(
  model~ : &@contract.Model,
  process~ : @linear.Mat,
  measurement~ : @linear.Mat,
  state~ : @linear.Vec,
  covariance~ : @linear.Mat,
  config~ : UnscentedConfig = unscented_config,
) -> UnscentedKalman raise @linear.Refused {
  let n = state.len()
  if covariance.rows != n || covariance.cols != n {
    raise @linear.Refused::Dimension(got=covariance.rows, want=n)
  }
  if process.rows != n || process.cols != n {
    raise @linear.Refused::Dimension(got=process.rows, want=n)
  }
  { model, process, measurement, state: state.clone(), covariance, config, }
}

///|
/// A sigma point set with the weights that go with it.
priv struct Sigma {
  points : Array[@linear.Vec]
  means : Array[Double]
  covariances : Array[Double]
}

///|
/// The sigma points of `state` and `covariance`.
fn sigma_points(
  state : @linear.Vec,
  covariance : @linear.Mat,
  config : UnscentedConfig,
) -> Sigma {
  let n = state.len()
  let lambda = config.alpha * config.alpha * (n.to_double() + config.kappa) -
    n.to_double()
  let raw = n.to_double() + lambda
  // A spread that collapses would give every sigma point the same weight; the
  // preset never does, but a caller's alpha can.
  let scale = if raw > 1.0e-6 { raw } else { 1.0e-6 }
  let factor = matrix_sqrt(covariance)
  let spread = scale.sqrt()
  let points = Array::make(2 * n + 1, state)
  let means = Array::make(2 * n + 1, 0.0)
  let covariances = Array::make(2 * n + 1, 0.0)
  points[0] = state.clone()
  means[0] = lambda / scale
  covariances[0] = lambda / scale +
    (1.0 - config.alpha * config.alpha + config.beta)
  for i in 0.. @linear.Vec {
  let mut mean = @linear.Vec::zeros(points[0].len())
  for i in 0.. @linear.Mat {
  let n = mean.len()
  let mut out = noise
  for i in 0..