/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2024, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */
#pragma once

#include <mrpt/math/CMatrixDynamic.h>
#include <mrpt/math/CMatrixFixed.h>
#include <mrpt/math/data_utils.h>
#include <mrpt/math/math_frwds.h>
#include <mrpt/math/num_jacobian.h>
#include <mrpt/math/ops_matrices.h>
#include <mrpt/random.h>

#include <functional>
#include <vector>

namespace mrpt::math
{
/** @addtogroup  gausspdf_transform_grp Gaussian PDF transformation functions
 *  \ingroup mrpt_math_grp
 * @{ */

/** Scaled unscented transformation (SUT) for estimating the Gaussian
 * distribution of a variable Y=f(X) for an arbitrary function f() provided by
 * the user.
 *  The user must supply the function in "functor" which takes points in the X
 * space and returns the mapped point in Y, optionally using an extra, constant
 * parameter ("fixed_param") which remains constant.
 *
 *  The parameters alpha, K and beta are the common names of the SUT method,
 * and the default values are those recommended in most papers.
 *
 * \param elem_do_wrap2pi If !=nullptr; it must point to an array of "bool" of
 * size()==dimension of each element, stating if it's needed to do a wrap to
 * [-pi,pi] to each dimension.
 * \sa The example in MRPT/samples/unscented_transform_test
 * \sa transform_gaussian_montecarlo, transform_gaussian_linear
 */
template <
    class VECTORLIKE1,
    class MATLIKE1,
    class USERPARAM,
    class VECTORLIKE2,
    class VECTORLIKE3,
    class MATLIKE2>
void transform_gaussian_unscented(
    const VECTORLIKE1& x_mean,
    const MATLIKE1& x_cov,
    void (*functor)(const VECTORLIKE1& x, const USERPARAM& fixed_param, VECTORLIKE3& y),
    const USERPARAM& fixed_param,
    VECTORLIKE2& y_mean,
    MATLIKE2& y_cov,
    const bool* elem_do_wrap2pi = nullptr,
    const double alpha = 1e-3,
    const double K = 0,
    const double beta = 2.0)
{
  MRPT_START
  const size_t Nx = x_mean.size();  // Dimensionality of inputs X
  const double lambda = alpha * alpha * (Nx + K) - Nx;
  const double c = Nx + lambda;

  // Generate weights:
  const double Wi = 0.5 / c;
  std::vector<double> W_mean(1 + 2 * Nx, Wi), W_cov(1 + 2 * Nx, Wi);
  W_mean[0] = lambda / c;
  W_cov[0] = W_mean[0] + (1 - alpha * alpha + beta);

  // Generate X_i samples:
  MATLIKE1 L;
  const bool valid = x_cov.chol(L);
  if (!valid)
    throw std::runtime_error(
        "transform_gaussian_unscented: Singular covariance matrix in "
        "Cholesky.");
  L *= sqrt(c);

  // Propagate the samples X_i -> Y_i:
  // We don't need to store the X sigma points: just use one vector to compute
  // all the Y sigma points:
  std::vector<VECTORLIKE3> Y(1 + 2 * Nx);  // 2Nx+1 sigma points
  VECTORLIKE1 X = x_mean;
  functor(X, fixed_param, Y[0]);
  VECTORLIKE1 delta;  // i'th row of L:
  delta.resize(Nx);
  size_t row = 1;
  for (size_t i = 0; i < Nx; i++)
  {
    for (size_t k = 0; k < Nx; k++) delta[k] = L(i, k);
    X = x_mean;
    X -= delta;
    functor(X, fixed_param, Y[row++]);
    X = x_mean;
    X += delta;
    functor(X, fixed_param, Y[row++]);
  }

  // Estimate weighted cov & mean:
  mrpt::math::covariancesAndMeanWeighted(Y, y_cov, y_mean, &W_mean, &W_cov, elem_do_wrap2pi);
  MRPT_END
}

/** Simple Montecarlo-base estimation of the Gaussian distribution of a variable
 * Y=f(X) for an arbitrary function f() provided by the user.
 *  The user must supply the function in "functor" which takes points in the X
 * space and returns the mapped point in Y, optionally using an extra, constant
 * parameter ("fixed_param") which remains constant.
 * \param out_samples_y If !=nullptr, this vector will contain, upon return,
 * the sequence of random samples generated and propagated through the
 * functor().
 * \sa The example in MRPT/samples/unscented_transform_test
 * \sa transform_gaussian_unscented, transform_gaussian_linear
 */
template <
    class VECTORLIKE1,
    class MATLIKE1,
    class USERPARAM,
    class VECTORLIKE2,
    class VECTORLIKE3,
    class MATLIKE2>
void transform_gaussian_montecarlo(
    const VECTORLIKE1& x_mean,
    const MATLIKE1& x_cov,
    void (*functor)(const VECTORLIKE1& x, const USERPARAM& fixed_param, VECTORLIKE3& y),
    const USERPARAM& fixed_param,
    VECTORLIKE2& y_mean,
    MATLIKE2& y_cov,
    const size_t num_samples = 1000,
    std::vector<VECTORLIKE3>* out_samples_y = nullptr)
{
  MRPT_START
  std::vector<VECTORLIKE1> samples_x;
  mrpt::random::getRandomGenerator().drawGaussianMultivariateMany(
      samples_x, num_samples, x_cov, &x_mean);
  std::vector<VECTORLIKE3> samples_y(num_samples);
  for (size_t i = 0; i < num_samples; i++) functor(samples_x[i], fixed_param, samples_y[i]);
  mrpt::math::covariancesAndMean(samples_y, y_cov, y_mean);
  if (out_samples_y)
  {
    out_samples_y->clear();
    samples_y.swap(*out_samples_y);
  }
  MRPT_END
}

/** First order uncertainty propagation estimator of the Gaussian distribution
 * of a variable Y=f(X) for an arbitrary function f() provided by the user.
 *  The user must supply the function in "functor" which takes points in the X
 * space and returns the mapped point in Y, optionally using an extra, constant
 * parameter ("fixed_param") which remains constant.
 *  The Jacobians are estimated numerically using the vector of small
 * increments "x_increments".
 * \sa The example in MRPT/samples/unscented_transform_test
 * \sa transform_gaussian_unscented, transform_gaussian_montecarlo
 * \note This function requires `#include <Eigen/Dense>`
 */
template <
    class VECTORLIKE1,
    class MATLIKE1,
    class USERPARAM,
    class VECTORLIKE2,
    class VECTORLIKE3,
    class MATLIKE2>
void transform_gaussian_linear(
    const VECTORLIKE1& x_mean,
    const MATLIKE1& x_cov,
    void (*functor)(const VECTORLIKE1& x, const USERPARAM& fixed_param, VECTORLIKE3& y),
    const USERPARAM& fixed_param,
    VECTORLIKE2& y_mean,
    MATLIKE2& y_cov,
    const VECTORLIKE1& x_increments)
{
  MRPT_START
  // Mean: simple propagation:
  functor(x_mean, fixed_param, y_mean);
  // Cov: COV = H C Ht
  CMatrixFixed<double, VECTORLIKE3::RowsAtCompileTime, VECTORLIKE1::RowsAtCompileTime> H;
  mrpt::math::estimateJacobian(
      x_mean,
      std::function<void(const VECTORLIKE1& x, const USERPARAM& fixed_param, VECTORLIKE3& y)>(
          functor),
      x_increments, fixed_param, H);
  mrpt::math::multiply_HCHt(H, x_cov, y_cov);
  MRPT_END
}

/** @} */

}  // namespace mrpt::math
