/*
 * Copyright 2016 The Cartographer Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef CARTOGRAPHER_MAPPING_INTERNAL_3D_IMU_INTEGRATION_H_
#define CARTOGRAPHER_MAPPING_INTERNAL_3D_IMU_INTEGRATION_H_

#include <algorithm>
#include <deque>

#include "Eigen/Core"
#include "Eigen/Geometry"
#include "cartographer/common/time.h"
#include "cartographer/sensor/imu_data.h"
#include "cartographer/transform/rigid_transform.h"
#include "cartographer/transform/transform.h"
#include "glog/logging.h"

namespace cartographer {
namespace mapping {

template <typename T>
struct IntegrateImuResult {
  Eigen::Matrix<T, 3, 1> delta_velocity;
  Eigen::Matrix<T, 3, 1> delta_translation;
  Eigen::Quaternion<T> delta_rotation;
};

template <typename T, typename RangeType, typename IteratorType>
IntegrateImuResult<T> IntegrateImu(
    const RangeType& imu_data,
    const Eigen::Transform<T, 3, Eigen::Affine>&
        linear_acceleration_calibration,
    const Eigen::Transform<T, 3, Eigen::Affine>& angular_velocity_calibration,
    const common::Time start_time, const common::Time end_time,
    IteratorType* const it) {
  CHECK_LE(start_time, end_time);
  CHECK(*it != imu_data.end());
  CHECK_LE((*it)->time, start_time);
  if (std::next(*it) != imu_data.end()) {
    CHECK_GT(std::next(*it)->time, start_time);
  }

  common::Time current_time = start_time;

  IntegrateImuResult<T> result = {Eigen::Matrix<T, 3, 1>::Zero(),
                                  Eigen::Matrix<T, 3, 1>::Zero(),
                                  Eigen::Quaterniond::Identity().cast<T>()};
  while (current_time < end_time) {
    common::Time next_imu_data = common::Time::max();
    if (std::next(*it) != imu_data.end()) {
      next_imu_data = std::next(*it)->time;
    }
    common::Time next_time = std::min(next_imu_data, end_time);
    const T delta_t(common::ToSeconds(next_time - current_time));

    const Eigen::Matrix<T, 3, 1> delta_angle =
        (angular_velocity_calibration *
         (*it)->angular_velocity.template cast<T>()) *
        delta_t;
    result.delta_rotation *=
        transform::AngleAxisVectorToRotationQuaternion(delta_angle);
    result.delta_velocity += result.delta_rotation *
                             ((linear_acceleration_calibration *
                               (*it)->linear_acceleration.template cast<T>()) *
                              delta_t);
    result.delta_translation += delta_t * result.delta_velocity;
    current_time = next_time;
    if (current_time == next_imu_data) {
      ++(*it);
    }
  }
  return result;
}

// Returns velocity delta in map frame.
template <typename RangeType, typename IteratorType>
IntegrateImuResult<double> IntegrateImu(const RangeType& imu_data,
                                        const common::Time start_time,
                                        const common::Time end_time,
                                        IteratorType* const it) {
  return IntegrateImu<double, RangeType, IteratorType>(
      imu_data, Eigen::Affine3d::Identity(), Eigen::Affine3d::Identity(),
      start_time, end_time, it);
}

template <typename T>
struct ExtrapolatePoseResult {
  transform::Rigid3<T> pose;
  Eigen::Matrix<T, 3, 1> velocity;
};

// Returns pose and linear velocity at 'time' which is equal to
// 'prev_from_tracking' extrapolated using IMU data.
template <typename T, typename RangeType, typename IteratorType>
ExtrapolatePoseResult<T> ExtrapolatePoseWithImu(
    const transform::Rigid3<T>& prev_from_tracking,
    const Eigen::Matrix<T, 3, 1>& prev_velocity_in_tracking,
    const common::Time prev_time, const Eigen::Matrix<T, 3, 1>& gravity,
    const common::Time time, const RangeType& imu_data,
    IteratorType* const imu_it) {
  const IntegrateImuResult<T> result =
      IntegrateImu(imu_data, Eigen::Transform<T, 3, Eigen::Affine>::Identity(),
                   Eigen::Transform<T, 3, Eigen::Affine>::Identity(), prev_time,
                   time, imu_it);

  const T delta_t = static_cast<T>(common::ToSeconds(time - prev_time));
  const Eigen::Matrix<T, 3, 1> translation =
      prev_from_tracking.translation() +
      prev_from_tracking.rotation() *
          (delta_t * prev_velocity_in_tracking + result.delta_translation) -
      static_cast<T>(.5) * delta_t * delta_t * gravity;
  const Eigen::Quaternion<T> rotation =
      prev_from_tracking.rotation() * result.delta_rotation;

  const Eigen::Matrix<T, 3, 1> velocity =
      prev_from_tracking.rotation() *
          (prev_velocity_in_tracking + result.delta_velocity) -
      delta_t * gravity;
  return ExtrapolatePoseResult<T>{transform::Rigid3<T>(translation, rotation),
                                  velocity};
}

// Same as above but given the last two poses.
template <typename T, typename RangeType, typename IteratorType>
ExtrapolatePoseResult<T> ExtrapolatePoseWithImu(
    const transform::Rigid3<T>& prev_from_tracking,
    const common::Time prev_time,
    const transform::Rigid3<T>& prev_prev_from_tracking,
    const common::Time prev_prev_time, const Eigen::Matrix<T, 3, 1>& gravity,
    const common::Time time, const RangeType& imu_data,
    IteratorType* const imu_it) {
  // TODO(danielsievers): Really we should integrate velocities starting from
  // the midpoint in between two poses, since this is how we fit them to poses
  // in the optimization.
  const T prev_delta_t =
      static_cast<T>(common::ToSeconds(prev_time - prev_prev_time));
  const Eigen::Matrix<T, 3, 1> prev_velocity_in_tracking =
      prev_from_tracking.inverse().rotation() *
      (prev_from_tracking.translation() -
       prev_prev_from_tracking.translation()) /
      prev_delta_t;

  return ExtrapolatePoseWithImu(prev_from_tracking, prev_velocity_in_tracking,
                                prev_time, gravity, time, imu_data, imu_it);
}

}  // namespace mapping
}  // namespace cartographer

#endif  // CARTOGRAPHER_MAPPING_INTERNAL_3D_IMU_INTEGRATION_H_