///|
/// A weighted state candidate used to compare independent estimators.
pub struct StateCandidate {
  name : String
  state : Array[Double]
  covariance : Matrix
  weight : Double
  result : UpdateResult
} derive(Debug)

///|
pub fn StateCandidate::new(
  name : String,
  state : Array[Double],
  covariance : Matrix,
  weight : Double,
  result : UpdateResult,
) -> StateCandidate {
  {
    name,
    state: state.copy(),
    covariance: covariance.copy(),
    weight: if weight < 0.0 {
      0.0
    } else {
      weight
    },
    result,
  }
}

///|
pub fn StateCandidate::name(self : StateCandidate) -> String {
  self.name
}

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

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

///|
pub fn StateCandidate::weight(self : StateCandidate) -> Double {
  self.weight
}

///|
pub fn StateCandidate::result(self : StateCandidate) -> UpdateResult {
  self.result
}

///|
pub fn StateCandidate::is_usable(self : StateCandidate) -> Bool {
  self.weight > 0.0 &&
  vector_is_finite(self.state) &&
  covariance_is_psd(self.covariance, 0.001) &&
  self.result is Accepted
}

///|
/// Weighted consensus over a set of independent state estimates.
pub struct EstimatorEnsemble {
  dimension : Int
  candidates : Array[StateCandidate]
  mut rejected : Int
}

///|
pub fn EstimatorEnsemble::new(dimension : Int) -> EstimatorEnsemble {
  {
    dimension: if dimension < 0 {
      0
    } else {
      dimension
    },
    candidates: [],
    rejected: 0,
  }
}

///|
pub fn EstimatorEnsemble::add(
  self : EstimatorEnsemble,
  candidate : StateCandidate,
) -> Bool {
  if candidate.state().length() != self.dimension || !candidate.is_usable() {
    self.rejected = self.rejected + 1
    return false
  }
  self.candidates.push(candidate)
  true
}

///|
pub fn EstimatorEnsemble::dimension(self : EstimatorEnsemble) -> Int {
  self.dimension
}

///|
pub fn EstimatorEnsemble::length(self : EstimatorEnsemble) -> Int {
  self.candidates.length()
}

///|
pub fn EstimatorEnsemble::rejected(self : EstimatorEnsemble) -> Int {
  self.rejected
}

///|
pub fn EstimatorEnsemble::candidates(
  self : EstimatorEnsemble,
) -> Array[StateCandidate] {
  self.candidates.copy()
}

///|
pub fn EstimatorEnsemble::consensus(self : EstimatorEnsemble) -> Array[Double] {
  let mut result = Array::make(self.dimension, 0.0)
  let mut total_weight = 0.0
  for candidate in self.candidates {
    let weight = candidate.weight()
    result = vector_axpy(weight, candidate.state(), result)
    total_weight = total_weight + weight
  }
  if total_weight <= 0.0 {
    Array::make(self.dimension, 0.0)
  } else {
    vector_scale(result, 1.0 / total_weight)
  }
}

///|
pub fn EstimatorEnsemble::consensus_covariance(
  self : EstimatorEnsemble,
) -> Matrix {
  if self.candidates.length() == 0 {
    return Matrix::zeros(self.dimension, self.dimension)
  }
  let center = self.consensus()
  let mut result = Matrix::zeros(self.dimension, self.dimension)
  let mut total_weight = 0.0
  for candidate in self.candidates {
    let weight = candidate.weight()
    let delta = vector_sub(candidate.state(), center)
    result = result.add(
      candidate.covariance().add(Matrix::outer(delta, delta)).scale(weight),
    )
    total_weight = total_weight + weight
  }
  if total_weight <= 0.0 {
    Matrix::zeros(self.dimension, self.dimension)
  } else {
    result.scale(1.0 / total_weight).symmetric_part()
  }
}

///|
pub fn EstimatorEnsemble::best(self : EstimatorEnsemble) -> StateCandidate? {
  if self.candidates.length() == 0 {
    return None
  }
  let mut index = 0
  for i in 1.. self.candidates[index].weight() {
      index = i
    }
  }
  Some(self.candidates[index])
}

///|
pub fn EstimatorEnsemble::disagreement(self : EstimatorEnsemble) -> Double {
  if self.candidates.length() < 2 {
    return 0.0
  }
  let center = self.consensus()
  let mut maximum = 0.0
  for candidate in self.candidates {
    let distance = vector_distance(candidate.state(), center)
    if distance > maximum {
      maximum = distance
    }
  }
  maximum
}

///|
pub fn EstimatorEnsemble::clear(self : EstimatorEnsemble) -> Unit {
  self.candidates.clear()
  self.rejected = 0
}

///|
/// Return normalized weights for a set of positive confidence scores.
pub fn normalize_confidences(confidences : Array[Double]) -> Array[Double] {
  let mut total = 0.0
  for confidence in confidences {
    if confidence > 0.0 && !confidence.is_nan() {
      total = total + confidence
    }
  }
  if total <= 0.0 {
    return Array::make(confidences.length(), 0.0)
  }
  confidences.map(value => {
    if value <= 0.0 || value.is_nan() {
      0.0
    } else {
      value / total
    }
  })
}

///|
pub fn candidate_from_filter(
  name : String,
  filter : KalmanND,
  result : UpdateResult,
  confidence : Double,
) -> StateCandidate {
  StateCandidate::new(
    name,
    filter.state(),
    filter.covariance(),
    confidence,
    result,
  )
}

///|
pub fn candidate_from_scalar(
  name : String,
  filter : Kalman1D,
  result : UpdateResult,
  confidence : Double,
) -> StateCandidate {
  StateCandidate::new(
    name,
    [filter.state()],
    Matrix::from_rows([[filter.uncertainty()]]),
    confidence,
    result,
  )
}

///|
pub fn ensemble_quality(ensemble : EstimatorEnsemble) -> Double {
  if ensemble.length() == 0 {
    return 0.0
  }
  let disagreement = ensemble.disagreement()
  let base = ensemble.length().to_double() /
    (ensemble.length().to_double() + 1.0)
  base / (1.0 + disagreement)
}