/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2022, 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/bayes/CProbabilityParticle.h>
#include <mrpt/random.h>

namespace mrpt
{
/// \ingroup mrpt_bayes_grp
namespace bayes
{
/** A base class for implementing rejection sampling in a generic state space.
 *   See the main method CRejectionSamplingCapable::rejectionSampling
 *  To use this class, create your own class as a child of this one and
 * implement the desired
 *   virtual methods, and add any required internal data.
 * \ingroup mrpt_bayes_grp
 */
template <
	class TStateSpace,
	mrpt::bayes::particle_storage_mode STORAGE =
		mrpt::bayes::particle_storage_mode::POINTER>
class CRejectionSamplingCapable
{
   public:
	using TParticle = CProbabilityParticle<TStateSpace, STORAGE>;

	/** Virtual destructor
	 */
	virtual ~CRejectionSamplingCapable() = default;
	/** Generates a set of N independent samples via rejection sampling.
	 * \param desiredSamples The number of desired samples to generate
	 * \param outSamples The output samples.
	 * \param timeoutTrials The maximum number of rejection trials for each
	 * generated sample (i.e. the maximum number of iterations). This can be
	 * used to set a limit to the time complexity of the algorithm for difficult
	 * probability densities.
	 *  All will have equal importance weights (a property of rejection
	 * sampling), although those samples
	 *   generated at timeout will have a different importance weights.
	 */
	void rejectionSampling(
		size_t desiredSamples, std::vector<TParticle>& outSamples,
		size_t timeoutTrials = 1000)
	{
		MRPT_START

		TStateSpace x;
		typename std::vector<TParticle>::iterator it;

		// Set output size:
		if (outSamples.size() != desiredSamples)
		{
			// Free old memory:
			outSamples.clear();

			// Reserve new memory:
			outSamples.resize(desiredSamples);
			for (it = outSamples.begin(); it != outSamples.end(); ++it)
				it->d.reset(new TStateSpace);
		}

		// Rejection sampling loop:
		double acceptanceProb;
		for (it = outSamples.begin(); it != outSamples.end(); ++it)
		{
			size_t timeoutCount = 0;
			double bestLik = -1e250;
			TStateSpace bestVal;
			do
			{
				RS_drawFromProposal(*it->d);
				acceptanceProb = RS_observationLikelihood(*it->d);
				ASSERT_(acceptanceProb >= 0 && acceptanceProb <= 1);
				if (acceptanceProb > bestLik)
				{
					bestLik = acceptanceProb;
					bestVal = *it->d;
				}
			} while (acceptanceProb <
						 mrpt::random::getRandomGenerator().drawUniform(
							 0.0, 0.999) &&
					 (++timeoutCount) < timeoutTrials);

			// Save weights:
			if (timeoutCount >= timeoutTrials)
			{
				it->log_w = log(bestLik);
				*it->d = bestVal;
			}
			else
			{
				it->log_w = 0;	// log(1.0);
			}
		}  // end for it

		MRPT_END
	}

   protected:
	/** Generates one sample, drawing from some proposal distribution.
	 */
	virtual void RS_drawFromProposal(TStateSpace& outSample) = 0;

	/** Returns the NORMALIZED observation likelihood (linear, not
	 * exponential!!!) at a given point of the state space (values in the range
	 * [0,1]).
	 */
	virtual double RS_observationLikelihood(const TStateSpace& x) = 0;

};	// End of class def.

}  // namespace bayes
}  // namespace mrpt
