text: """This module implements EKF localization using GNSS and IMU."""
  
  from collections import deque
  from functools import partial
  
  import erdos
  from erdos import Message, ReadStream, Timestamp, WriteStream
  
  import numpy as np
  
  from pylot.utils import Location, Pose, Quaternion, Rotation, Transform, \
          Vector3D
  
  
  class LocalizationOperator(erdos.Operator):
      """Localizes vehicle using GPS and IMU.
  
      The operator implements localization using a Kalman filter.
  
      Args:
          imu_stream (:py:class:`erdos.ReadStream`): Stream on which IMU
              sensors readings are received.
          gnss_stream (:py:class:`erdos.ReadStream`): Stream on which
              GNSS readings are received.
          ground_pose_stream (:py:class:`erdos.ReadStream`): Stream on which
              the operator receives noisy poses. It refines these poses using
              GPS and IMU.
          pose_stream (:py:class:`erdos.ReadStream`): Stream on which the
              operator outputs pose messages.
          flags (absl.flags): Object to be used to access absl flags.
      """
      def __init__(self, imu_stream: ReadStream, gnss_stream: ReadStream,
                   ground_pose_stream: ReadStream, pose_stream: WriteStream,
                   flags):
          # Register callbacks on read streams.
          self._imu_updates = deque()
          imu_stream.add_callback(
              partial(self.buffer_msg, msg_type="IMU", queue=self._imu_updates))
  
          self._gnss_updates = deque()
          gnss_stream.add_callback(
              partial(self.buffer_msg, msg_type="GNSS",
                      queue=self._gnss_updates))
  
          self._ground_pose_updates = deque()
          ground_pose_stream.add_callback(
              partial(self.buffer_msg,
                      msg_type="pose",
                      queue=self._ground_pose_updates))
          erdos.add_watermark_callback(
              [imu_stream, gnss_stream, ground_pose_stream], [],
              self.on_watermark)
  
          # Save the output stream.
          self._pose_stream = pose_stream
  
          # Initialize a logger.
          self._flags = flags
          self._logger = erdos.utils.setup_logging(self.config.name,
                                                   self.config.log_file_name)
  
          # Gravity vector.
          self._g = np.array([0, 0, -9.81])
  
          # Previous timestamp values.
          self._last_pose_estimate = None
          self._last_timestamp = None
  
          # NOTE: At the start of the simulation, the vehicle drops down from
          # the sky, during which the IMU values screw up the calculations.
          # This boolean flag takes care to start the prediction only when the
          # values have stabilized.
          self._is_started = False
  
          # Constants required for the Kalman filtering.
          var_imu_f, var_imu_w, var_gnss = 0.5, 0.5, 0.1
          self.__Q = np.identity(6)
          self.__Q[0:3, 0:3] = self.__Q[0:3, 0:3] * var_imu_f
          self.__Q[3:6, 3:6] = self.__Q[3:6, 3:6] * var_imu_w
  
          self.__F = np.identity(9)
  
          self.__L = np.zeros([9, 6])
          self.__L[3:9, :] = np.identity(6)
  
          self.__R_GNSS = np.identity(3) * var_gnss
  
          self._last_covariance = np.zeros((9, 9))
  
      @staticmethod
      def connect(imu_stream: ReadStream, gnss_stream: ReadStream,
                  ground_pose_stream: ReadStream):
          pose_stream = erdos.WriteStream()
          return [pose_stream]
  
      def destroy(self):
          self._logger.warn('destroying {}'.format(self.config.name))
  
      def buffer_msg(self, msg: Message, msg_type: str, queue: deque):
          """Callback which buffers the received message."""
          self._logger.debug("@{}: received {} message.".format(
              msg.timestamp, msg_type))
          queue.append(msg)
  
      def get_message(self, queue: deque, timestamp: Timestamp, name: str):
          msg = None
          if queue:
              while len(queue) > 0:
                  retrieved_msg = queue.popleft()
                  if retrieved_msg.timestamp == timestamp:
                      msg = retrieved_msg
                      break
              if not msg:
                  raise ValueError(
                      "The message for {} with timestamp {} was not found".
                      format(name, timestamp))
          return msg
  
      def __skew_symmetric(self, v):
          """Skew symmetric form of a 3x1 vector."""
          return np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]],
                          dtype=np.float64)
  
      def __update_using_gnss(self, location_estimate, velocity_estimate,
                              rotation_estimate, gnss_reading, delta_t):
          # Construct H_k = [I, 0, 0] (shape=(3, 9))
          H_k = np.zeros((3, 9))
          H_k[:, :3] = np.identity(3)
  
          # Propogate uncertainty.
          Q = self.__Q * delta_t * delta_t
          self._last_covariance = (self.__F.dot(self._last_covariance).dot(
              self.__F.T)) + (self.__L.dot(Q).dot(self.__L.T))
  
          # Compute Kalman gain. (shape=(9, 3))
          K_k = self._last_covariance.dot(
              H_k.T.dot(
                  np.linalg.inv(
                      H_k.dot(self._last_covariance.dot(H_k.T)) +
                      self.__R_GNSS)))
  
          # Compute error state. (9x3) x ((3x1) - (3x1)) = shape(9, 1)
          delta_x_k = K_k.dot(gnss_reading - location_estimate)
  
          # Correct predicted state.
          corrected_location_estimate = location_estimate + delta_x_k[0:3]
          corrected_velocity_estimate = velocity_estimate + delta_x_k[3:6]
          roll, pitch, yaw = delta_x_k[6:]
          corrected_rotation_estimate = Quaternion.from_rotation(
              Rotation(roll=roll, pitch=pitch, yaw=yaw)) * rotation_estimate
  
          # Fix the covariance.
          self._last_covariance = (np.identity(9) - K_k.dot(H_k)).dot(
              self._last_covariance)
  
          return (
              corrected_location_estimate,
              corrected_velocity_estimate,
              corrected_rotation_estimate,
          )
  
      @erdos.profile_method()
      def on_watermark(self, timestamp: Timestamp):
          self._logger.debug("@{}: received watermark.".format(timestamp))
          if timestamp.is_top:
              self._pose_stream.send(erdos.WatermarkMessage(timestamp))
              return
  
          # Retrieve the messages for this timestamp.
          pose_msg = self.get_message(self._ground_pose_updates, timestamp,
                                      "pose")
          gnss_msg = self.get_message(self._gnss_updates, timestamp, "GNSS")
          imu_msg = self.get_message(self._imu_updates, timestamp, "IMU")
  
          if self._last_pose_estimate is None or \
             (abs(imu_msg.acceleration.y) > 100 and not self._is_started):
              self._logger.debug(
                  "@{}: The initial pose estimate is not initialized.".format(
                      timestamp))
              # If this is the first update or values have not stabilized,
              # save the estimates.
              if pose_msg:
                  self._last_pose_estimate = pose_msg.data
                  self._last_timestamp = timestamp.coordinates[0]
                  self._pose_stream.send(pose_msg)
                  self._pose_stream.send(erdos.WatermarkMessage(timestamp))
              else:
                  raise NotImplementedError(
                      "Need pose message to initialize the estimates.")
          else:
              self._is_started = True
  
              # Initialize the delta_t
              current_ts = timestamp.coordinates[0]
              delta_t = (current_ts - self._last_timestamp) / 1000.0
  
              # Estimate the rotation.
              last_rotation_estimate = Quaternion.from_rotation(
                  self._last_pose_estimate.transform.rotation)
              rotation_estimate = (
                  last_rotation_estimate *
                  Quaternion.from_angular_velocity(imu_msg.gyro, delta_t))
  
              # Transform the IMU accelerometer data from the body frame to the
              # world frame, and retrieve location and velocity estimates.
              accelerometer_data = last_rotation_estimate.matrix.dot(
                  imu_msg.acceleration.as_numpy_array()) + self._g
              last_location_estimate = \
                  self._last_pose_estimate.transform.location.as_numpy_array()
              last_velocity_estimate = \
                  self._last_pose_estimate.velocity_vector.as_numpy_array()
  
              # Estimate the location.
              location_estimate = last_location_estimate + (
                  delta_t * last_velocity_estimate) + ((
                      (delta_t**2) / 2.0) * accelerometer_data)
  
              # Estimate the velocity.
              velocity_estimate = last_velocity_estimate + (delta_t *
                                                            accelerometer_data)
  
              # Fuse the GNSS values using an EKF to fix drifts and noise in
              # the estimates.
  
              # Linearize the motion model and compute Jacobians.
              self.__F[0:3, 3:6] = np.identity(3) * delta_t
              self.__F[3:6, 6:9] = last_rotation_estimate.matrix.dot(
                  -self.__skew_symmetric(accelerometer_data.reshape(
                      (3, 1)))) * delta_t
  
              # Fix estimates using GNSS
              gnss_reading = Location.from_gps(
                  gnss_msg.latitude, gnss_msg.longitude,
                  gnss_msg.altitude).as_numpy_array()
              (
                  location_estimate,
                  velocity_estimate,
                  rotation_estimate,
              ) = self.__update_using_gnss(location_estimate, velocity_estimate,
                                           rotation_estimate, gnss_reading,
                                           delta_t)
  
              # Create the PoseMessage and send it downstream.
              current_pose = Pose(
                  transform=Transform(location=Location(*location_estimate),
                                      rotation=rotation_estimate.as_rotation()),
                  forward_speed=Vector3D(*velocity_estimate).magnitude(),
                  velocity_vector=Vector3D(*velocity_estimate),
                  localization_time=current_ts,
              )
              self._logger.debug("@{}: Predicted pose: {}".format(
                  timestamp, current_pose))
              self._pose_stream.send(erdos.Message(timestamp, current_pose))
              self._pose_stream.send(erdos.WatermarkMessage(timestamp))
  
              # Set the estimates for the next iteration.
              self._last_timestamp = current_ts
              self._last_pose_estimate = current_pose
  