///|
/// Extended Kalman Filter (EKF) for 6-DOF tracking of 3D Position, Velocity, Orientation Quaternion, and Angular Velocity.
/// State vector (13D): [x, y, z, vx, vy, vz, qw, qx, qy, qz, wx, wy, wz]
pub(all) struct ExtendedKalmanFilter6DOF {
  state : Array[Double]
  mut covariance : DenseMatrix // 13x13
  process_noise : DenseMatrix // 13x13
  measurement_noise : DenseMatrix // 6x6 for [x, y, z, rvec_x, rvec_y, rvec_z]
} derive(Debug, Eq)

///|
/// Creates a new 6DOF Extended Kalman Filter initialized with a given 3D pose.
pub fn ExtendedKalmanFilter6DOF::new(
  initial_pose : Pose3D,
  pos_process_noise? : Double = 0.05,
  rot_process_noise? : Double = 0.02,
  pos_meas_noise? : Double = 0.01,
  rot_meas_noise? : Double = 0.02,
) -> ExtendedKalmanFilter6DOF {
  let q = Quaternion::from_rvec(initial_pose.rvec).normalized()
  let state = [
    initial_pose.tvec.x,
    initial_pose.tvec.y,
    initial_pose.tvec.z,
    0.0,
    0.0,
    0.0,
    q.w,
    q.x,
    q.y,
    q.z,
    0.0,
    0.0,
    0.0,
  ]

  let cov = DenseMatrix::identity(13).scale(0.1)
  let q_mat = DenseMatrix::identity(13)
  for i in 0..<3 {
    q_mat.set(i, i, pos_process_noise * pos_process_noise)
  }
  for i in 3..<6 {
    q_mat.set(i, i, pos_process_noise * pos_process_noise * 4.0)
  }
  for i in 6..<10 {
    q_mat.set(i, i, rot_process_noise * rot_process_noise)
  }
  for i in 10..<13 {
    q_mat.set(i, i, rot_process_noise * rot_process_noise * 4.0)
  }

  let r_mat = DenseMatrix::identity(6)
  for i in 0..<3 {
    r_mat.set(i, i, pos_meas_noise * pos_meas_noise)
  }
  for i in 3..<6 {
    r_mat.set(i, i, rot_meas_noise * rot_meas_noise)
  }

  { state, covariance: cov, process_noise: q_mat, measurement_noise: r_mat }
}

///|
/// Predicts the 6DOF state forward by dt seconds.
pub fn ExtendedKalmanFilter6DOF::predict(
  self : ExtendedKalmanFilter6DOF,
  dt : Double,
) -> Unit {
  // Integrate position: p = p + v * dt
  self.state[0] += self.state[3] * dt
  self.state[1] += self.state[4] * dt
  self.state[2] += self.state[5] * dt

  // Integrate quaternion with angular velocity w = [wx, wy, wz]
  let qw = self.state[6]
  let qx = self.state[7]
  let qy = self.state[8]
  let qz = self.state[9]
  let wx = self.state[10]
  let wy = self.state[11]
  let wz = self.state[12]

  let delta_qw = 0.5 * (-qx * wx - qy * wy - qz * wz) * dt
  let delta_qx = 0.5 * (qw * wx + qy * wz - qz * wy) * dt
  let delta_qy = 0.5 * (qw * wy - qx * wz + qz * wx) * dt
  let delta_qz = 0.5 * (qw * wz + qx * wy - qy * wx) * dt

  let new_q = Quaternion::new(
    qw + delta_qw,
    qx + delta_qx,
    qy + delta_qy,
    qz + delta_qz,
  ).normalized()

  self.state[6] = new_q.w
  self.state[7] = new_q.x
  self.state[8] = new_q.y
  self.state[9] = new_q.z

  // Propagate Jacobian F (13x13)
  let f = DenseMatrix::identity(13)
  f.set(0, 3, dt)
  f.set(1, 4, dt)
  f.set(2, 5, dt)

  // Covariance 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 EKF with a new measured 3D Pose (rvec, tvec).
pub fn ExtendedKalmanFilter6DOF::update(
  self : ExtendedKalmanFilter6DOF,
  measured_pose : Pose3D,
) -> Unit {
  // Measurement function maps state to [x, y, z, rvec_x, rvec_y, rvec_z]
  let q_est = Quaternion::new(
    self.state[6],
    self.state[7],
    self.state[8],
    self.state[9],
  )
  let rvec_est = q_est.to_rvec()

  let z = [
    measured_pose.tvec.x,
    measured_pose.tvec.y,
    measured_pose.tvec.z,
    measured_pose.rvec.x,
    measured_pose.rvec.y,
    measured_pose.rvec.z,
  ]

  let h_x = [
    self.state[0],
    self.state[1],
    self.state[2],
    rvec_est.x,
    rvec_est.y,
    rvec_est.z,
  ]

  let y = Array::make(6, 0.0)
  for i in 0..<6 {
    y[i] = z[i] - h_x[i]
  }

  // Measurement Jacobian H (6x13)
  let h_mat = DenseMatrix::zeros(6, 13)
  h_mat.set(0, 0, 1.0)
  h_mat.set(1, 1, 1.0)
  h_mat.set(2, 2, 1.0)
  h_mat.set(3, 7, 1.0)
  h_mat.set(4, 8, 1.0)
  h_mat.set(5, 9, 1.0)

  let ht = h_mat.transpose()
  let s_opt = match h_mat.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) =>
          match self.covariance.mul(ht) {
            Some(pht) =>
              match pht.mul(s_inv) {
                Some(k_gain) => {
                  match k_gain.mul_vec(y) {
                    Some(ky) => {
                      for i in 0..<13 {
                        self.state[i] = self.state[i] + ky[i]
                      }
                      // Re-normalize quaternion
                      let q_norm = Quaternion::new(
                        self.state[6],
                        self.state[7],
                        self.state[8],
                        self.state[9],
                      ).normalized()
                      self.state[6] = q_norm.w
                      self.state[7] = q_norm.x
                      self.state[8] = q_norm.y
                      self.state[9] = q_norm.z
                    }
                    None => ()
                  }

                  // Covariance update: P = (I - K * H) * P
                  let i13 = DenseMatrix::identity(13)
                  match k_gain.mul(h_mat) {
                    Some(kh) =>
                      match i13.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 6DOF Pose.
pub fn ExtendedKalmanFilter6DOF::pose(
  self : ExtendedKalmanFilter6DOF,
) -> Pose3D {
  let tvec = Point3D::new(self.state[0], self.state[1], self.state[2])
  let q = Quaternion::new(
    self.state[6],
    self.state[7],
    self.state[8],
    self.state[9],
  )
  let rvec = q.to_rvec()
  Pose3D::new(rvec, tvec)
}