// SPDX-FileCopyrightText: Fondazione Istituto Italiano di Tecnologia (IIT)
// SPDX-License-Identifier: BSD-3-Clause

#ifndef ATTITUDE_QUATERNION_EKF_H
#define ATTITUDE_QUATERNION_EKF_H

#include <iDynTree/AttitudeEstimator.h>
#include <iDynTree/Direction.h>
#include <iDynTree/ExtendedKalmanFilter.h>

namespace iDynTree
{
const unsigned int output_dimensions_with_magnetometer = 4; ///< dimension of \f$ \mathbb{R}^3
                                                            ///< \times \mathbb{R} \f$ accelerometer
                                                            ///< measurements and magnetometer yaw
                                                            ///< measurement
const unsigned int output_dimensions_without_magnetometer = 3; ///< dimension of \f$ \mathbb{R}^3
                                                               ///< \f$ accelerometer measurements
const unsigned int input_dimensions = 3; ///< dimension of \f$ \mathbb{R}^3 \f$ gyroscope
                                         ///< measurements

/**
 * @struct AttitudeQuaternionEKFParameters Parameters to set up the quaternion EKF
 * @var AttitudeQuaternionEKFParameters::time_step_in_seconds
 * discretization time step in seconds, default value: \f$ 0.01 s \f$
 * @var AttitudeQuaternionEKFParameters::bias_correlation_time_factor
 * time factor modeling how fast the bias can vary, default value: \f$ 0.01 \f$
 * @var AttitudeQuaternionEKFParameters::accelerometer_noise_variance
 * measurement noise covariance \f$ \sigma_{acc}^{2} \f$ of accelerometer measurement, default
 * value: \f$ 0.001 \f$
 * @var AttitudeQuaternionEKFParameters::magnetometer_noise_variance
 * measurement noise covariance \f$ \sigma_{mag}^{2} \f$ of magnetometer measurement, default value:
 * \f$ 0.001 \f$
 * @var AttitudeQuaternionEKFParameters::gyroscope_noise_variance
 * system noise covariance \f$ \sigma_{gyro}^{2} \f$ of gyroscope measurement, since it is the input
 * to the system, default value: \f$ 0.001 \f$
 * @var AttitudeQuaternionEKFParameters::gyro_bias_noise_variance
 * system noise covariance \f$ \sigma_{gyrobias}^{2} \f$ of gyroscope bias estimate, default value:
 * \f$ 0.0001 \f$
 * @var AttitudeQuaternionEKFParameters::initial_orientation_error_variance
 * initial state covariance \f$ \sigma_{q_0}^{2} \f$ of orientation, default value: \f$ 10 \f$
 * @var AttitudeQuaternionEKFParameters::initial_ang_vel_error_variance
 * initial state covariance \f$ \sigma_{\omega_0}^{2} \f$ of angular velocity, default value: \f$ 10
 * \f$
 * @var AttitudeQuaternionEKFParameters::initial_gyro_bias_error_variance
 * measurement noise covariance \f$ \sigma_{acc}^{2} \f$ of gyroscope bias, default value: \f$ 10
 * \f$
 * @var AttitudeQuaternionEKFParameters::use_magnetometer_measurements
 * flag to enable the use of magnetometer measurement for yaw correction, default value: false
 */
struct AttitudeQuaternionEKFParameters
{
    double time_step_in_seconds{0.01};
    double bias_correlation_time_factor{0.01};

    // measurement noise - zero mean, and a given variance
    double accelerometer_noise_variance{0.001};
    double magnetometer_noise_variance{0.001};

    // process noise - zero mean, and a given covariance
    double gyroscope_noise_variance{0.001};
    double gyro_bias_noise_variance{0.0001};

    double initial_orientation_error_variance{10};
    double initial_ang_vel_error_variance{10};
    double initial_gyro_bias_error_variance{10};

    bool use_magnetometer_measurements{false};
};

/**
 * @class AttitudeQuaternionEKF implements a Quaternion based Discrete Extended Kalman Filter fusing
 * IMU measurements, to give estimates of orientation, angular velocity and gyroscope bias
 *
 * It follows the implementation detailed in
 * <a
 * href="https://wuecampus2.uni-wuerzburg.de/moodle/pluginfile.php/1109745/mod_resource/content/1/QEKF_Floatsat_WS16.pdf">Quaternion
 * Based Extended Kalman Filter, slides by Michael Stohmeier</a> The filter is used to estimate the
 * states \f$ X = \begin{bmatrix} {^A}q_B \\ {^B}\Omega_{A,B} \\ {^B}b \end{bmatrix}^T \f$ where \f$
 * {^A}q_B \in \mathbb{R}^4 \f$  is the quaternion representing the orientation of a body(IMU) frame
 * with respect to an inertial frame ,
 *       \f$ {^B}\Omega_{A,B} \in \mathbb{R}^3 \f$  is the angular velocity of a body(IMU) frame
 * with respect to an inertial frame, expressed in the body frame and
 *       \f$ {^B}b \in \mathbb{R}^3 \f$  is the gyroscope bias expressed in the body frame.
 * @note: we will drop the subscripts and superscripts in the rest of the documentation for
 * convenience
 *
 * Discretized dynamics during the prediction step,
 * \f[ \hat{{x}}_{k+1} = \begin{bmatrix} q_{k} \otimes \text{exp}(\omega \Delta T) \\ y_{gyro_{k}} -
 * b_{k} \\ (1 - \lambda_{b} \Delta t)b_k \end{bmatrix} \f]
 *
 * Measurement model for accelerometer is given as,
 * \f[ h_{acc}(\hat{x}_{k+1}) = \begin{bmatrix} 2(q_1q_3 - q_0q_2) \\ 2(q_2q_3 - q_0q_1) \\ q_0^2 -
 * q_1^2 - q_2^2 + q_3^2 \end{bmatrix} \f] obtained from \f$ {^w}R_b^T e_3 \f$ of the assumed
 * gravity direction.
 *
 * Measurement model for magnetometer measurement is given as,
 * \f[ h_{mag}(\hat{x}_{k+1}) = atan2( 2(q_0q_3 + q_1q_2),1 - 2(q_2^2 + q_3^2) ) \f]
 *
 * The linearized system propogation and measurement model is obtained by computing Jacobins F and H
 * with respect to the state.
 *
 * The zero mean, additive Gaussian noise can be set using the covariance matrices which will be
 * used during predict and update steps.
 *
 * The propagateStates() method is called to set the input vector for the EKF, then ekfPredict() is
 * called to propagate the state through the propagation function f() and propate the state
 * covariance using the Jacobian F.
 *
 * The updateFilterWithMeasurements() is called to set the measurement vector for the EKF, and then
 * ekfUpdate is used to correct the state estimate and its covariance using the measurement model
 * function h() and the measurement Jacobian H.
 *
 * The usage of the QEKF should follow the decribed procedure below,
 * - instantiate the filter
 * - set parameters
 * - call initializeFilter() (this is necessary for resizing the buffers, the user should call this
 * method after setting parameters)
 * - use setInternalState() to set initial state (The filter will throw an error, if this is not
 * called atleast once, this enforces the user to set intial state)
 * - Once initialized, the following filter methods can be run in a loop to get the orientation
 * estimates,
 *     - propagateStates() method to propagate the states and covariance
 *     - updateFilterWithMeasurements() method to correct the predicted states and covariance
 *     - getInternalState() or getOrientationEstimate*() methods to get the entire state estimate or
 * only the attitude estimated in desired representation
 *
 * @warning calling the method useMagnetometerMeasurements() while the estimator is running, will
 * reset the filter, reinitialize the filter to resize buffers and sets the previous estiamted state
 * as the inital state.
 * @note calling other set parameter methods does not reset the filter, since they are not
 * associated with changing the output dimensions
 *
 */
class AttitudeQuaternionEKF : public IAttitudeEstimator, public DiscreteExtendedKalmanFilterHelper
{
public:
    AttitudeQuaternionEKF();

    /**
     * @brief Get filter parameters as a struct.
     * @param[out] params object of AttitudeQuaternionEKFParameters passed as reference
     */
    void getParameters(AttitudeQuaternionEKFParameters& params)
    {
        params = m_params_qekf;
    }

    /**
     * @brief Set filter parameters with the struct members.
     * This resets filter since it also calls useMagnetometerMeasurements(flag)
     * (if the use_magnetometer_measurements flag has been toggled).
     * @param[in] params object of AttitudeQuaternionEKFParameters passed as a const reference
     * @return true/false if successful/not
     */
    void setParameters(const AttitudeQuaternionEKFParameters& params)
    {
        m_params_qekf = params;
        useMagnetometerMeasurements(params.use_magnetometer_measurements);
    }

    /**
     * @brief Set the gravity direction assumed by the filter
     * This affects the measurement model function h() and Jacobian H
     * @param[in] gravity_dir gravity direction
     */
    void setGravityDirection(const iDynTree::Direction& gravity_dir);

    /**
     * @brief set discretization time step in seconds
     * @param[in] time_step_in_seconds time step
     */
    void setTimeStepInSeconds(double time_step_in_seconds)
    {
        m_params_qekf.time_step_in_seconds = time_step_in_seconds;
    }

    /**
     * @brief set bias correlation time factor
     * @param[in] bias_correlation_time_factor time factor for bias evolution
     */
    void setBiasCorrelationTimeFactor(double bias_correlation_time_factor)
    {
        m_params_qekf.bias_correlation_time_factor = bias_correlation_time_factor;
    }

    /**
     * @brief set flag to use magnetometer measurements
     * @param[in] use_magnetometer_measurements enable/disable magnetometer measurements
     * @note calling this method with the flag same as current flag value will not change anything,
     *  meanwhile a new flag setting will reset the filter, reinitialize the filter and
     *  set the previous state as filter's initial state and previous state covariance as filter's
     * intial state covariance
     */
    bool useMagnetometerMeasurements(bool use_magnetometer_measurements);

    /**
     * @brief prepares the measurement noise covariance matrix and calls
     * ekfSetMeasurementNoiseMeanAndCovariance() measurement noise depends only on accelerometer xyz
     * (and magnetometer z)
     * @note the noise has zero mean (basically passes a zero vector with covariance matrix)
     * @param[in] acc variance for accelerometer measurements
     * @param[in] mag variance for magnetometer measurements
     * @return true/false if successful/not
     */
    bool setMeasurementNoiseVariance(double acc, double mag);

    /**
     * @brief prepares the system noise covariance matrix and calls
     * ekfSetSystemNoiseMeanAndCovariance() process noise depends on gyro measurement and gyro bias
     * estimate - since gyro measurement is passed as input
     * @note the noise has zero mean (basically passes a zero vector with covariance matrix)
     * measurement noise depends only on accelerometer xyz (and magnetometer z)
     * @param[in] gyro variance for gyroscope measurements
     * @param[in] gyro_bias variance for gyroscope bias estimates
     * @return true/false if successful/not
     */
    bool setSystemNoiseVariance(double gyro, double gyro_bias);

    /**
     * @brief prepares the state covariance matrix and calls ekfSetStateCovariance()
     * @param[in] orientation_var variance for intial orientation state estimate
     * @param[in] ang_vel_var variance for initial angular velocity state estimate
     * @param[in] gyro_bias_var variance for intial gyro bias state estimate
     * @return true/false if successful/not
     */
    bool
    setInitialStateCovariance(double orientation_var, double ang_vel_var, double gyro_bias_var);

    /**
     * @brief intializes the filter by resizing buffers and setting parameters
     * - sets state, output and input dimensions for the ekf
     * - resizes internal buffers
     * - calls ekfInit()
     * - sets system noise, measurement noise and initial state covariance
     * - if successful sets initialized flag to true
     * @return true/false if successful/not
     */
    bool initializeFilter();

    bool updateFilterWithMeasurements(const iDynTree::LinearAccelerometerMeasurements& linAccMeas,
                                      const iDynTree::GyroscopeMeasurements& gyroMeas) override;
    bool updateFilterWithMeasurements(const iDynTree::LinearAccelerometerMeasurements& linAccMeas,
                                      const iDynTree::GyroscopeMeasurements& gyroMeas,
                                      const iDynTree::MagnetometerMeasurements& magMeas) override;
    bool propagateStates() override;
    bool getOrientationEstimateAsRotationMatrix(iDynTree::Rotation& rot) override;
    bool getOrientationEstimateAsQuaternion(iDynTree::UnitQuaternion& q) override;
    bool getOrientationEstimateAsRPY(iDynTree::RPY& rpy) override;
    size_t getInternalStateSize() const override;
    bool getInternalState(const iDynTree::Span<double>& stateBuffer) const override;
    bool getDefaultInternalInitialState(const iDynTree::Span<double>& stateBuffer) const override;
    bool setInternalState(const iDynTree::Span<double>& stateBuffer) override;
    bool
    setInternalStateInitialOrientation(const iDynTree::Span<double>& orientationBuffer) override;

protected:
    AttitudeEstimatorState m_state_qekf, m_initial_state_qekf;
    AttitudeQuaternionEKFParameters m_params_qekf; ///< struct holding the QEKF parameters

private:
    /**
     * discrete system propagation \f$ f(X, u) = f(X, y_gyro) \f$
     * where \f$ X = \begin{bmatrix} q_0 &  q_1 & q_2 & q_3 & \omega_x & \omega_y & \omega_z & b_x &
     * b_y & b_z \end{bmatrix}^T \f$
     * \f$ u = \begin{bmatrix} {y_{gyro}}_x & {y_{gyro}}_y & {y_{gyro}}_z \end{bmatrix}^T \f$
     * \f$ f(X, u) = \begin{bmatrix} q_{k} \otimes \text{exp}(\omega \Delta T) \\ y_{gyro} - b \\ (1
     * - \lambda_{b} \Delta t)b \end{bmatrix}\f$
     */
    bool ekf_f(const iDynTree::VectorDynSize& x_k,
               const iDynTree::VectorDynSize& u_k,
               iDynTree::VectorDynSize& xhat_k_plus_one) override;

    /**
     * discrete measurement prediction
     * where \f$ h(X) = \begin{bmatrix} h_{acc}(X) & h_{mag}(X) \end{bmatrix}^T \f$
     * \f$ h_{acc}(X) = R^T \begin{bmatrix} 0 \\  0 \\ -1 \end{bmatrix} \f$
     * \f$ h_{mag}(X) = atan2(tan(yaw))\f$
     */
    bool ekf_h(const iDynTree::VectorDynSize& xhat_k_plus_one,
               iDynTree::VectorDynSize& zhat_k_plus_one) override;

    /**
     * @brief Describes the system Jacobian necessary for the propagation of predicted state
     * covariance The analytical Jacobian describing the partial derivative of the system
     * propagation with respect to the state
     * @param[in] x system state
     * @param[out] F system Jacobian
     * @return bool true/false if successful or not
     */
    bool ekfComputeJacobianF(iDynTree::VectorDynSize& x, iDynTree::MatrixDynSize& F) override;
    bool ekfComputeJacobianF(iDynTree::VectorDynSize& x,
                             iDynTree::VectorDynSize& u,
                             iDynTree::MatrixDynSize& F) override;

    /**
     * @brief Describes the measurement Jacobian necessary for computing Kalman gain and updating
     * the predicted state and its covariance The analytical Jacobian describing the partial
     * derivative of the measurement model with respect to the state
     * @param[in] x system state
     * @param[out] H measurement Jacobian
     * @return bool true/false if successful or not
     */
    bool ekfComputeJacobianH(iDynTree::VectorDynSize& x, iDynTree::MatrixDynSize& H) override;

    /** @brief prepares the system noise covariance matrix using internal struct params
     * system  model is as good as gyroscope measurement and bias estimate
     * system noise covariance can be descibed as \f$ Q = F_u  U  {F_u}^T \f$
     * where \f$ F = \begin{bmatrix} \frac{\partial f}{\partial y_gyro} & \frac{\partial f}{
     * \partial x_gyrobias} \end{bmatrix} \f$
     * \f$ = \begin{bmatrix} 0_{4 \times 3} & 0_{4 \times 3} \\ I_{3 \times 3} & 0_{3 \times 3}
     * \\ 0_{3 \times 3} & I_{3 \times 3} \end{bmatrix}\f$
     * \f$ U = diag(\begin{bmatrix} \sigma_{gyro}^{2} I_{3 \times 3} & \sigma_{gyrobias}^{2} I_{3
     * \times 3} \end{bmatrix}) \f$
     * @param[in] Q matrix container as reference
     */
    void prepareSystemNoiseCovarianceMatrix(iDynTree::MatrixDynSize& Q);

    /** @brief prepares the measurement noise covariance matrix using internal struct parameters
     * measurement noise depends only on accelerometer measurement along x-,y- and z- directions
     * along with magnetometer z-direction if included
     * measurement noise covariance can be descibed as,
     * \f$ R = \begin{bmatrix} \sigma_{acc}^{2} I_{3 \times 3} & 0_{3 \times 1} \\ 0_{1 \times 3} &
     * \sigma_{mag}^{2}\f$ if magnetometer measurements is also considered. In case of magnetometer
     * measurements not being considered, it is reduced to the \f$ 3 \times 3 \f$ matrix
     * @param[in] R matrix container as reference
     */
    void prepareMeasurementNoiseCovarianceMatrix(iDynTree::MatrixDynSize& R);

    /**
     * @brief serializes the state struct to state x of VectorDynSize
     */
    void serializeStateVector();

    /**
     * @brief deserializes state x of VectorDynSize to the state struct
     */
    void deserializeStateVector();

    /**
     * @brief serializes the accelerometer and magenetometer measurements into y vector
     * since DiscreteExtendedKalmanFilter expects a VectorDynSize including all necessary
     * measurements
     */
    void serializeMeasurementVector();

    /**
     * @brief serializes measurements, calls ekfUpdate step and
     * gets state estimate corrected by measurements
     *
     * @return true/false, if successful/not
     */
    bool callEkfUpdate();

    iDynTree::Rotation m_orientationInSO3; ///< orientation estimate as rotation matrix \f$ {^A}R_B
                                           ///< \f$ where \f$ A \f$ is inertial frame and \f$ B \f$
                                           ///< is the frame attached to the body
    iDynTree::RPY m_orientationInRPY; ///< orientation estimate as a 3D vector in RPY
                                      ///< representation, where \f$ {^A}R_B =
                                      ///< Rot_z(yaw)Rot_y(pitch)Rot_x(roll) \f$

    iDynTree::GyroscopeMeasurements m_Omega_y; ///< 3d gyroscope measurement giving angular velocity
                                               ///< of body wrt inertial frame, expressed in body
                                               ///< frame
    iDynTree::LinearAccelerometerMeasurements m_Acc_y; ///< 3d accelerometer measurement giving
                                                       ///< proper classical acceleration expressed
                                                       ///< in body frame
    double m_Mag_y; ///< magnetometer yaw measurement expressed in body frame

    iDynTree::VectorDynSize m_x; ///< state vector for the EKF - orientation, angular velocity, gyro
                                 ///< bias
    iDynTree::VectorDynSize m_y; ///< measurement vector for the EKF - accelerometer (and
                                 ///< magnetometer yaw)
    iDynTree::VectorDynSize m_u; ///< input vector for the EKF - gyroscope measurement

    size_t m_state_size; ///< state dimensions
    size_t m_output_size; ///< output dimensions
    size_t m_input_size; ///< input dimensions
    bool m_initialized{false}; ///< flag to check if QEKF is initialized

    iDynTree::Matrix4x4 m_Id4; ///< \f$ 4 \times 4 \f$  identity matrix
    iDynTree::Matrix3x3 m_Id3; ///< \f$ 3 \times 3 \f$  identity matrix
    iDynTree::Direction m_gravity_direction; ///< direction of the gravity vector expressed in the
                                             ///< inertial frame denoted by \f$ A \f$, default set
                                             ///< to \f$ e_3 = \begin{bmatrix} 0 & 0 & 1.0
                                             ///< \end{bmatrix}^T \f$
};

} // namespace iDynTree

#endif
