/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, 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/CMatrixFixed.h>
#include <mrpt/math/CVectorDynamic.h>
#include <mrpt/math/ops_matrices.h>
#include <mrpt/math/wrap2pi.h>

namespace mrpt::math
{
/** \addtogroup stats_grp
 * @{
 */

/** @name Probability density distributions (pdf) distance metrics
@{ */

/** Computes the squared mahalanobis distance of a vector X given the mean MU
 * and the covariance *inverse* COV_inv
 *  \f[ d^2 =  (X-MU)^\top \Sigma^{-1} (X-MU)  \f]
 */
template <class VECTORLIKE1, class VECTORLIKE2, class MAT>
typename MAT::Scalar mahalanobisDistance2(
	const VECTORLIKE1& X, const VECTORLIKE2& MU, const MAT& COV)
{
	MRPT_START
#if defined(_DEBUG) || (MRPT_ALWAYS_CHECKS_DEBUG_MATRICES)
	ASSERT_(!X.empty());
	ASSERT_(X.size() == MU.size());
	ASSERT_(X.size() == COV.rows() && COV.isSquare());
#endif
	const size_t N = X.size();
	CVectorDynamic<typename MAT::Scalar> X_MU(N);
	for (size_t i = 0; i < N; i++)
		X_MU[i] = X[i] - MU[i];
	auto z = COV.llt_solve(X_MU);
	return z.dot(z);
	MRPT_END
}

/** Computes the mahalanobis distance of a vector X given the mean MU and the
 * covariance *inverse* COV_inv
 *  \f[ d = \sqrt{ (X-MU)^\top \Sigma^{-1} (X-MU) }  \f]
 */
template <class VECTORLIKE1, class VECTORLIKE2, class MAT>
inline typename VECTORLIKE1::Scalar mahalanobisDistance(
	const VECTORLIKE1& X, const VECTORLIKE2& MU, const MAT& COV)
{
	return std::sqrt(mahalanobisDistance2(X, MU, COV));
}

/** Computes the squared mahalanobis distance between two *non-independent*
 * Gaussians, given the two covariance matrices and the vector with the
 * difference of their means.
 *  \f[ d^2 = \Delta_\mu^\top (\Sigma_1 + \Sigma_2 - 2 \Sigma_12 )^{-1}
 * \Delta_\mu  \f]
 */
template <class VECTORLIKE, class MAT1, class MAT2, class MAT3>
typename MAT1::Scalar mahalanobisDistance2(
	const VECTORLIKE& mean_diffs, const MAT1& COV1, const MAT2& COV2,
	const MAT3& CROSS_COV12)
{
	MRPT_START
#if defined(_DEBUG) || (MRPT_ALWAYS_CHECKS_DEBUG_MATRICES)
	ASSERT_(!mean_diffs.empty());
	ASSERT_(mean_diffs.size() == COV1.rows());
	ASSERT_(COV1.isSquare() && COV2.isSquare());
	ASSERT_(COV1.rows() == COV2.rows());
#endif
	MAT1 COV = COV1;
	COV += COV2;
	COV.substract_An(CROSS_COV12, 2);
	MAT1 COV_inv;
	COV.inverse_LLt(COV_inv);
	return multiply_HCHt_scalar(mean_diffs, COV_inv);
	MRPT_END
}

/** Computes the mahalanobis distance between two *non-independent* Gaussians
 * (or independent if CROSS_COV12=nullptr), given the two covariance matrices
 * and the vector with the difference of their means.
 *  \f[ d = \sqrt{ \Delta_\mu^\top (\Sigma_1 + \Sigma_2 - 2 \Sigma_12 )^{-1}
 * \Delta_\mu } \f]
 */
template <class VECTORLIKE, class MAT1, class MAT2, class MAT3>
inline typename VECTORLIKE::Scalar mahalanobisDistance(
	const VECTORLIKE& mean_diffs, const MAT1& COV1, const MAT2& COV2,
	const MAT3& CROSS_COV12)
{
	return std::sqrt(mahalanobisDistance(mean_diffs, COV1, COV2, CROSS_COV12));
}

/** Computes the squared mahalanobis distance between a point and a Gaussian,
 * given the covariance matrix and the vector with the difference between the
 * mean and the point.
 *  \f[ d^2 = \Delta_\mu^\top \Sigma^{-1} \Delta_\mu  \f]
 */
template <class VECTORLIKE, class MATRIXLIKE>
inline typename MATRIXLIKE::Scalar mahalanobisDistance2(
	const VECTORLIKE& delta_mu, const MATRIXLIKE& cov)
{
	ASSERTDEB_(cov.isSquare());
	ASSERTDEB_(size_t(cov.cols()) == size_t(delta_mu.size()));
	return multiply_HtCH_scalar(delta_mu, cov.inverse());
}

/** Computes the mahalanobis distance between a point and a Gaussian, given the
 * covariance matrix and the vector with the difference between the mean and the
 * point.
 *  \f[ d^2 = \sqrt( \Delta_\mu^\top \Sigma^{-1} \Delta_\mu ) \f]
 */
template <class VECTORLIKE, class MATRIXLIKE>
inline typename MATRIXLIKE::Scalar mahalanobisDistance(
	const VECTORLIKE& delta_mu, const MATRIXLIKE& cov)
{
	return std::sqrt(mahalanobisDistance2(delta_mu, cov));
}

/** Computes the integral of the product of two Gaussians, with means separated
 * by "mean_diffs" and covariances "COV1" and "COV2".
 *  \f[ D = \frac{1}{(2 \pi)^{0.5 N} \sqrt{}  }  \exp( \Delta_\mu^\top
 * (\Sigma_1 + \Sigma_2 - 2 \Sigma_12)^{-1} \Delta_\mu)  \f]
 */
template <typename T>
T productIntegralTwoGaussians(
	const std::vector<T>& mean_diffs, const CMatrixDynamic<T>& COV1,
	const CMatrixDynamic<T>& COV2)
{
	const size_t vector_dim = mean_diffs.size();
	ASSERT_(vector_dim >= 1);

	CMatrixDynamic<T> C = COV1;
	C += COV2;	// Sum of covs:
	const T cov_det = C.det();
	CMatrixDynamic<T> C_inv;
	C.inverse_LLt(C_inv);

	return std::pow(M_2PI, -0.5 * vector_dim) * (1.0 / std::sqrt(cov_det)) *
		exp(-0.5 * mean_diffs.multiply_HCHt_scalar(C_inv));
}

/** Computes the integral of the product of two Gaussians, with means separated
 * by "mean_diffs" and covariances "COV1" and "COV2".
 *  \f[ D = \frac{1}{(2 \pi)^{0.5 N} \sqrt{}  }  \exp( \Delta_\mu^\top
 * (\Sigma_1 + \Sigma_2)^{-1} \Delta_\mu)  \f]
 */
template <typename T, size_t DIM>
T productIntegralTwoGaussians(
	const std::vector<T>& mean_diffs, const CMatrixFixed<T, DIM, DIM>& COV1,
	const CMatrixFixed<T, DIM, DIM>& COV2)
{
	ASSERT_(mean_diffs.size() == DIM);

	CMatrixFixed<T, DIM, DIM> C = COV1;
	C += COV2;	// Sum of covs:
	const T cov_det = C.det();
	CMatrixFixed<T, DIM, DIM> C_inv(mrpt::math::UNINITIALIZED_MATRIX);
	C.inverse_LLt(C_inv);

	return std::pow(M_2PI, -0.5 * DIM) * (1.0 / std::sqrt(cov_det)) *
		exp(-0.5 * mean_diffs.multiply_HCHt_scalar(C_inv));
}

/** Computes both, the integral of the product of two Gaussians and their square
 * Mahalanobis distance.
 * \sa productIntegralTwoGaussians, mahalanobisDistance2
 */
template <typename T, class VECLIKE, class MATLIKE1, class MATLIKE2>
void productIntegralAndMahalanobisTwoGaussians(
	const VECLIKE& mean_diffs, const MATLIKE1& COV1, const MATLIKE2& COV2,
	T& maha2_out, T& intprod_out, const MATLIKE1* CROSS_COV12 = nullptr)
{
	const size_t vector_dim = mean_diffs.size();
	ASSERT_(vector_dim >= 1);

	MATLIKE1 C = COV1;
	C += COV2;	// Sum of covs:
	if (CROSS_COV12)
	{
		C -= *CROSS_COV12;
		C -= *CROSS_COV12;
	}
	const T cov_det = C.det();
	MATLIKE1 C_inv;
	C.inverse_LLt(C_inv);

	maha2_out = mrpt::math::multiply_HCHt_scalar(mean_diffs, C_inv);
	intprod_out = std::pow(M_2PI, -0.5 * vector_dim) *
		(1.0 / std::sqrt(cov_det)) * exp(-0.5 * maha2_out);
}

/** Computes both, the logarithm of the PDF and the square Mahalanobis distance
 * between a point (given by its difference wrt the mean) and a Gaussian.
 * \sa productIntegralTwoGaussians, mahalanobisDistance2, normalPDF,
 * mahalanobisDistance2AndPDF
 */
template <typename T, class VECLIKE, class MATRIXLIKE>
void mahalanobisDistance2AndLogPDF(
	const VECLIKE& diff_mean, const MATRIXLIKE& cov, T& maha2_out,
	T& log_pdf_out)
{
	MRPT_START
	ASSERTDEB_(cov.isSquare());
	ASSERTDEB_(size_t(cov.cols()) == size_t(diff_mean.size()));
	const MATRIXLIKE C_inv = cov.inverse_LLt();
	maha2_out = multiply_HtCH_scalar(diff_mean, C_inv);
	log_pdf_out = static_cast<typename MATRIXLIKE::Scalar>(-0.5) *
		(maha2_out +
		 static_cast<typename MATRIXLIKE::Scalar>(cov.cols()) *
			 ::log(static_cast<typename MATRIXLIKE::Scalar>(M_2PI)) +
		 ::log(cov.det()));
	MRPT_END
}

/** Computes both, the PDF and the square Mahalanobis distance between a point
 * (given by its difference wrt the mean) and a Gaussian.
 * \sa productIntegralTwoGaussians, mahalanobisDistance2, normalPDF
 */
template <typename T, class VECLIKE, class MATRIXLIKE>
inline void mahalanobisDistance2AndPDF(
	const VECLIKE& diff_mean, const MATRIXLIKE& cov, T& maha2_out, T& pdf_out)
{
	mahalanobisDistance2AndLogPDF(diff_mean, cov, maha2_out, pdf_out);
	pdf_out = std::exp(pdf_out);  // log to linear
}

/** Computes covariances and mean of any vector of containers, given optional
 * weights for the different samples.
 * \param elements Any kind of vector of vectors/arrays, eg.
 * std::vector<mrpt::math::CVectorDouble>, with all the input samples, each
 * sample in a "row".
 * \param covariances Output estimated covariance; it can be a fixed/dynamic
 * matrix or a matrixview.
 * \param means Output estimated mean; it can be
 * CVectorDouble/CVectorFixedDouble, etc... \param weights_mean If !=nullptr, it
 * must point to a vector of size()==number of elements, with normalized weights
 * to take into account for the mean. \param weights_cov If !=nullptr, it must
 * point to a vector of size()==number of elements, with normalized weights to
 * take into account for the covariance. \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 This method is used in mrpt::math::unscented_transform_gaussian
 * \ingroup stats_grp
 */
template <
	class VECTOR_OF_VECTORS, class MATRIXLIKE, class VECTORLIKE,
	class VECTORLIKE2, class VECTORLIKE3>
inline void
	covariancesAndMeanWeighted(	 // Done inline to speed-up the special case
								 // expanded in covariancesAndMean() below.
		const VECTOR_OF_VECTORS& elements, MATRIXLIKE& covariances,
		VECTORLIKE& means, const VECTORLIKE2* weights_mean,
		const VECTORLIKE3* weights_cov, const bool* elem_do_wrap2pi = nullptr)
{
	ASSERTMSG_(
		elements.size() != 0,
		"No samples provided, so there is no way to deduce the output size.");
	using T = typename MATRIXLIKE::Scalar;
	const size_t DIM = elements[0].size();
	means.resize(DIM);
	covariances.resize(DIM, DIM);
	const size_t nElms = elements.size();
	const T NORM = 1.0 / nElms;
	if (weights_mean)
	{
		ASSERTDEB_(size_t(weights_mean->size()) == size_t(nElms));
	}
	// The mean goes first:
	for (size_t i = 0; i < DIM; i++)
	{
		T accum = 0;
		if (!elem_do_wrap2pi || !elem_do_wrap2pi[i])
		{  // i'th dimension is a "normal", real number:
			if (weights_mean)
			{
				for (size_t j = 0; j < nElms; j++)
					accum += (*weights_mean)[j] * elements[j][i];
			}
			else
			{
				for (size_t j = 0; j < nElms; j++)
					accum += elements[j][i];
				accum *= NORM;
			}
		}
		else
		{  // i'th dimension is a circle in [-pi,pi]: we need a little trick
			// here:
			double accum_L = 0, accum_R = 0;
			double Waccum_L = 0, Waccum_R = 0;
			for (size_t j = 0; j < nElms; j++)
			{
				double ang = elements[j][i];
				const double w =
					weights_mean != nullptr ? (*weights_mean)[j] : NORM;
				if (fabs(ang) > 0.5 * M_PI)
				{  // LEFT HALF: 0,2pi
					if (ang < 0) ang = (M_2PI + ang);
					accum_L += ang * w;
					Waccum_L += w;
				}
				else
				{  // RIGHT HALF: -pi,pi
					accum_R += ang * w;
					Waccum_R += w;
				}
			}
			if (Waccum_L > 0) accum_L /= Waccum_L;	// [0,2pi]
			if (Waccum_R > 0) accum_R /= Waccum_R;	// [-pi,pi]
			if (accum_L > M_PI)
				accum_L -= M_2PI;  // Left side to [-pi,pi] again:
			accum =
				(accum_L * Waccum_L +
				 accum_R * Waccum_R);  // The overall result:
		}
		means[i] = accum;
	}
	// Now the covariance:
	for (size_t i = 0; i < DIM; i++)
		for (size_t j = 0; j <= i; j++)	 // Only 1/2 of the matrix
		{
			typename MATRIXLIKE::Scalar elem = 0;
			if (weights_cov)
			{
				ASSERTDEB_(size_t(weights_cov->size()) == size_t(nElms));
				for (size_t k = 0; k < nElms; k++)
				{
					const T Ai = (elements[k][i] - means[i]);
					const T Aj = (elements[k][j] - means[j]);
					if (!elem_do_wrap2pi || !elem_do_wrap2pi[i])
						elem += (*weights_cov)[k] * Ai * Aj;
					else
						elem += (*weights_cov)[k] * mrpt::math::wrapToPi(Ai) *
							mrpt::math::wrapToPi(Aj);
				}
			}
			else
			{
				for (size_t k = 0; k < nElms; k++)
				{
					const T Ai = (elements[k][i] - means[i]);
					const T Aj = (elements[k][j] - means[j]);
					if (!elem_do_wrap2pi || !elem_do_wrap2pi[i])
						elem += Ai * Aj;
					else
						elem +=
							mrpt::math::wrapToPi(Ai) * mrpt::math::wrapToPi(Aj);
				}
				elem *= NORM;
			}
			covariances(i, j) = elem;
			if (i != j) covariances(j, i) = elem;
		}
}

/** Computes covariances and mean of any vector of containers.
 * \param elements Any kind of vector of vectors/arrays, eg.
 * std::vector<mrpt::math::CVectorDouble>, with all the input samples, each
 * sample in a "row".
 * \param covariances Output estimated covariance; it can be a fixed/dynamic
 * matrix or a matrixview.
 * \param means Output estimated mean; it can be
 * CVectorDouble/CVectorFixedDouble, etc... \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.
 * \ingroup stats_grp
 */
template <class VECTOR_OF_VECTORS, class MATRIXLIKE, class VECTORLIKE>
void covariancesAndMean(
	const VECTOR_OF_VECTORS& elements, MATRIXLIKE& covariances,
	VECTORLIKE& means, const bool* elem_do_wrap2pi = nullptr)
{  // The function below is inline-expanded here:
	covariancesAndMeanWeighted<
		VECTOR_OF_VECTORS, MATRIXLIKE, VECTORLIKE, CVectorDouble,
		CVectorDouble>(
		elements, covariances, means, nullptr, nullptr, elem_do_wrap2pi);
}

/** Computes the weighted histogram for a vector of values and their
 * corresponding weights.
 *  \param values [IN] The N values
 *  \param weights [IN] The weights for the corresponding N values (don't need
 * to be normalized)
 *  \param binWidth [IN] The desired width of the bins
 *  \param out_binCenters [OUT] The centers of the M bins generated to cover
 * from the minimum to the maximum value of "values" with the given "binWidth"
 *  \param out_binValues [OUT] The ratio of values at each given bin, such as
 * the whole vector sums up the unity.
 *  \sa weightedHistogramLog
 */
template <class VECTORLIKE1, class VECTORLIKE2>
void weightedHistogram(
	const VECTORLIKE1& values, const VECTORLIKE1& weights, float binWidth,
	VECTORLIKE2& out_binCenters, VECTORLIKE2& out_binValues)
{
	MRPT_START
	using TNum = typename mrpt::math::ContainerType<VECTORLIKE1>::element_t;

	ASSERT_(values.size() == weights.size());
	ASSERT_(binWidth > 0);
	TNum minBin = minimum(values);
	auto nBins =
		static_cast<unsigned>(ceil((maximum(values) - minBin) / binWidth));

	// Generate bin center and border values:
	out_binCenters.resize(nBins);
	out_binValues.clear();
	out_binValues.resize(nBins, 0);
	TNum halfBin = TNum(0.5) * binWidth;
	;
	VECTORLIKE2 binBorders(nBins + 1, minBin - halfBin);
	for (unsigned int i = 0; i < nBins; i++)
	{
		binBorders[i + 1] = binBorders[i] + binWidth;
		out_binCenters[i] = binBorders[i] + halfBin;
	}

	// Compute the histogram:
	TNum totalSum = 0;
	typename VECTORLIKE1::const_iterator itVal, itW;
	for (itVal = values.begin(), itW = weights.begin(); itVal != values.end();
		 ++itVal, ++itW)
	{
		int idx = round(((*itVal) - minBin) / binWidth);
		if (idx >= int(nBins)) idx = nBins - 1;
		ASSERTDEB_(idx >= 0);
		out_binValues[idx] += *itW;
		totalSum += *itW;
	}

	if (totalSum) out_binValues /= totalSum;

	MRPT_END
}

/** Computes the weighted histogram for a vector of values and their
 * corresponding log-weights.
 *  \param values [IN] The N values
 *  \param weights [IN] The log-weights for the corresponding N values (don't
 * need to be normalized)
 *  \param binWidth [IN] The desired width of the bins
 *  \param out_binCenters [OUT] The centers of the M bins generated to cover
 * from the minimum to the maximum value of "values" with the given "binWidth"
 *  \param out_binValues [OUT] The ratio of values at each given bin, such as
 * the whole vector sums up the unity.
 *  \sa weightedHistogram
 */
template <class VECTORLIKE1, class VECTORLIKE2>
void weightedHistogramLog(
	const VECTORLIKE1& values, const VECTORLIKE1& log_weights, float binWidth,
	VECTORLIKE2& out_binCenters, VECTORLIKE2& out_binValues)
{
	MRPT_START
	using TNum = typename mrpt::math::ContainerType<VECTORLIKE1>::element_t;

	ASSERT_(values.size() == log_weights.size());
	ASSERT_(binWidth > 0);
	TNum minBin = minimum(values);
	auto nBins =
		static_cast<unsigned>(ceil((maximum(values) - minBin) / binWidth));

	// Generate bin center and border values:
	out_binCenters.resize(nBins);
	out_binValues.clear();
	out_binValues.resize(nBins, 0);
	TNum halfBin = TNum(0.5) * binWidth;
	;
	VECTORLIKE2 binBorders(nBins + 1, minBin - halfBin);
	for (unsigned int i = 0; i < nBins; i++)
	{
		binBorders[i + 1] = binBorders[i] + binWidth;
		out_binCenters[i] = binBorders[i] + halfBin;
	}

	// Compute the histogram:
	const TNum max_log_weight = maximum(log_weights);
	TNum totalSum = 0;
	typename VECTORLIKE1::const_iterator itVal, itW;
	for (itVal = values.begin(), itW = log_weights.begin();
		 itVal != values.end(); ++itVal, ++itW)
	{
		int idx = round(((*itVal) - minBin) / binWidth);
		if (idx >= int(nBins)) idx = nBins - 1;
		ASSERTDEB_(idx >= 0);
		const TNum w = exp(*itW - max_log_weight);
		out_binValues[idx] += w;
		totalSum += w;
	}

	if (totalSum) out_binValues /= totalSum;

	MRPT_END
}

/** A numerically-stable method to compute average likelihood values with
 * strongly different ranges (unweighted likelihoods: compute the arithmetic
 * mean).
 *  This method implements this equation:
 *
 *  \f[ return = - \log N + \log  \sum_{i=1}^N e^{ll_i-ll_{max}} + ll_{max} \f]
 *
 * See also the <a
 * href="http://www.mrpt.org/Averaging_Log-Likelihood_Values:Numerical_Stability">tutorial
 * page</a>.
 * \ingroup stats_grp
 */
double averageLogLikelihood(const CVectorDouble& logLikelihoods);

/** Computes the average of a sequence of angles in radians taking into account
 * the correct wrapping in the range \f$ ]-\pi,\pi [ \f$, for example, the mean
 * of (2,-2) is \f$ \pi \f$, not 0.
 * \ingroup stats_grp
 */
double averageWrap2Pi(const CVectorDouble& angles);

/** A numerically-stable method to average likelihood values with strongly
 * different ranges (weighted likelihoods).
 *  This method implements this equation:
 *
 *  \f[ return = \log \left( \frac{1}{\sum_i e^{lw_i}} \sum_i  e^{lw_i}
 * e^{ll_i}  \right) \f]
 *
 * See also the <a
 * href="http://www.mrpt.org/Averaging_Log-Likelihood_Values:Numerical_Stability">tutorial
 * page</a>.
 * \ingroup stats_grp
 */
double averageLogLikelihood(
	const CVectorDouble& logWeights, const CVectorDouble& logLikelihoods);

/**  @} */	// end of grouping container_ops_grp

}  // namespace mrpt::math
