///|
/// Linear Kalman Filter for 3D state tracking with Constant Velocity (CV) or Constant Acceleration (CA) models.
pub(all) struct KalmanFilter3D {
  state_dim : Int
  state : Array[Double] // x
  mut covariance : DenseMatrix // P
  process_noise : DenseMatrix // Q
  measurement_noise : DenseMatrix // R
} derive(Debug, Eq)

///|
/// Creates a 6-state Constant Velocity 3D Kalman Filter (position + velocity).
pub fn KalmanFilter3D::new_constant_velocity(
  initial_position : Point3D,
  process_noise_std? : Double = 0.1,
  measurement_noise_std? : Double = 0.01,
) -> KalmanFilter3D {
  let state = [
    initial_position.x,
    initial_position.y,
    initial_position.z,
    0.0,
    0.0,
    0.0,
  ]

  let p = DenseMatrix::identity(6)
  let q = DenseMatrix::identity(6).scale(process_noise_std * process_noise_std)
  let r = DenseMatrix::identity(3).scale(
    measurement_noise_std * measurement_noise_std,
  )

  { state_dim: 6, state, covariance: p, process_noise: q, measurement_noise: r }
}

///|
/// Predicts the state forward by time delta dt (seconds).
pub fn KalmanFilter3D::predict(self : KalmanFilter3D, dt : Double) -> Unit {
  // State transition matrix F for CV model: x_{k+1} = x_k + v_k * dt
  let f = DenseMatrix::identity(6)
  f.set(0, 3, dt)
  f.set(1, 4, dt)
  f.set(2, 5, dt)

  // x = F * x
  match f.mul_vec(self.state) {
    Some(new_state) =>
      for i in 0..<6 {
        self.state[i] = new_state[i]
      }
    None => ()
  }

  // P = F * P * F^T + Q
  let ft = f.transpose()
  match f.mul(self.covariance) {
    Some(fp) =>
      match fp.mul(ft) {
        Some(fpft) =>
          match fpft.add(self.process_noise) {
            Some(new_p) => self.covariance = new_p
            None => ()
          }
        None => ()
      }
    None => ()
  }
}

///|
/// Updates the Kalman filter with a new 3D position measurement.
pub fn KalmanFilter3D::update(
  self : KalmanFilter3D,
  measurement : Point3D,
) -> Unit {
  // Measurement matrix H (3x6): maps state to [x, y, z]
  let h = DenseMatrix::zeros(3, 6)
  h.set(0, 0, 1.0)
  h.set(1, 1, 1.0)
  h.set(2, 2, 1.0)

  let z = [measurement.x, measurement.y, measurement.z]

  // Innovation residual y = z - H * x
  let hx = [self.state[0], self.state[1], self.state[2]]
  let y = [z[0] - hx[0], z[1] - hx[1], z[2] - hx[2]]

  // Innovation covariance S = H * P * H^T + R
  let ht = h.transpose()
  let s_opt = match h.mul(self.covariance) {
    Some(hp) =>
      match hp.mul(ht) {
        Some(hpht) => hpht.add(self.measurement_noise)
        None => None
      }
    None => None
  }

  match s_opt {
    Some(s) =>
      match s.inverse() {
        Some(s_inv) =>
          // Kalman gain K = P * H^T * S^(-1) (6x3)
          match self.covariance.mul(ht) {
            Some(pht) =>
              match pht.mul(s_inv) {
                Some(k_gain) => {
                  // State update: x = x + K * y
                  match k_gain.mul_vec(y) {
                    Some(ky) =>
                      for i in 0..<6 {
                        self.state[i] = self.state[i] + ky[i]
                      }
                    None => ()
                  }

                  // Covariance update: P = (I - K * H) * P
                  let i6 = DenseMatrix::identity(6)
                  match k_gain.mul(h) {
                    Some(kh) =>
                      match i6.sub(kh) {
                        Some(i_minus_kh) =>
                          match i_minus_kh.mul(self.covariance) {
                            Some(new_p) => self.covariance = new_p
                            None => ()
                          }
                        None => ()
                      }
                    None => ()
                  }
                }
                None => ()
              }
            None => ()
          }
        None => ()
      }
    None => ()
  }
}

///|
/// Gets the current filtered 3D position estimate.
pub fn KalmanFilter3D::position(self : KalmanFilter3D) -> Point3D {
  Point3D::new(self.state[0], self.state[1], self.state[2])
}

///|
/// Gets the current filtered 3D velocity estimate.
pub fn KalmanFilter3D::velocity(self : KalmanFilter3D) -> Point3D {
  Point3D::new(self.state[3], self.state[4], self.state[5])
}