/* +------------------------------------------------------------------------+
   |                     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/CMatrixDynamic.h>
#include <mrpt/math/CMatrixFixed.h>

namespace mrpt
{
namespace math
{
namespace detail
{
// Auxiliary method: templatized for working with float/double's.
template <typename SCALAR>
double internal_kmeans(
	const bool use_kmeansplusplus_method, const size_t nPoints, const size_t k,
	const size_t dims, const SCALAR* points, const size_t attempts,
	SCALAR* out_center, int* out_assignments);

// Auxiliary method, the actual code of the two front-end functions offered to
// the user below.
template <class LIST_OF_VECTORS1, class LIST_OF_VECTORS2>
double stub_kmeans(
	[[maybe_unused]] const bool use_kmeansplusplus_method, const size_t k,
	const LIST_OF_VECTORS1& points, std::vector<int>& assignments,
	LIST_OF_VECTORS2* out_centers, const size_t attempts)
{
	MRPT_START
	ASSERT_(k >= 1);
	const size_t N = points.size();
	assignments.resize(N);
	if (out_centers) out_centers->clear();
	if (!N) return 0;  // No points, we're done.
	// Parse to required format:
	size_t dims = 0;
	const auto it_first = points.begin();
	const auto it_end = points.end();
	using TInnerVector = typename LIST_OF_VECTORS1::value_type;
	using TInnerVectorCenters = typename LIST_OF_VECTORS2::value_type;
	std::vector<typename TInnerVector::value_type> raw_vals;
	typename TInnerVector::value_type* trg_ptr = nullptr;
	for (typename LIST_OF_VECTORS1::const_iterator it = it_first; it != it_end;
		 ++it)
	{
		if (it == it_first)
		{
			dims = it->size();
			ASSERTMSG_(dims > 0, "Dimensionality of points can't be zero.");
			raw_vals.resize(N * dims);
			trg_ptr = &raw_vals[0];
		}
		else
		{
			ASSERTMSG_(
				size_t(dims) == size_t(it->size()),
				"All points must have the same dimensionality.");
		}

		std::memcpy(
			trg_ptr, &(*it)[0],
			dims * sizeof(typename TInnerVector::value_type));
		trg_ptr += dims;
	}
	// Call the internal implementation:
	std::vector<typename TInnerVectorCenters::value_type> centers(dims * k);
	const double ret = detail::internal_kmeans(
		false, N, k, points.begin()->size(), &raw_vals[0], attempts,
		&centers[0], &assignments[0]);
	// Centers:
	if (out_centers)
	{
		const typename TInnerVectorCenters::value_type* center_ptr =
			&centers[0];
		for (size_t i = 0; i < k; i++)
		{
			TInnerVectorCenters c;
			c.resize(dims);
			for (size_t j = 0; j < dims; j++)
				c[j] = *center_ptr++;
			out_centers->push_back(c);
		}
	}
	return ret;
	MRPT_END
}
}  // namespace detail

/** @name k-means algorithms
@{ */

/** k-means algorithm to cluster a list of N points of arbitrary dimensionality
 *into exactly K clusters.
 *   The list of input points can be any template CONTAINER<POINT> with:
 *		- CONTAINER can be: Any STL container:
 *std::vector,std::list,std::deque,...
 *		- POINT can be:
 *			- std::vector<double/float>
 *			- CVectorFixedDouble<N> / CVectorFixedFloat<N>
 *
 *  \param k [IN] Number of cluster to look for.
 *  \param points [IN] The list of N input points. It can be any STL-like
 *containers of std::vector<float/double>, for example a
 *std::vector<mrpt::math::CVectorDouble>, a std::list<CVectorFloat>, etc...
 *  \param assignments [OUT] At output it will have a number [0,k-1] for each
 *of the N input points.
 *  \param out_centers [OUT] If not nullptr, at output will have the centers of
 *each group. Can be of any of the supported types of "points", but the basic
 *coordinates should be float or double exactly as in "points".
 *  \param attempts [IN] Number of attempts.
 *
 * \sa A more advanced algorithm, see: kmeanspp
 * \note Uses the kmeans++ implementation by David Arthur (2009,
 *http://www.stanford.edu/~darthur/kmpp.zip).
 */
template <class LIST_OF_VECTORS1, class LIST_OF_VECTORS2>
inline double kmeans(
	const size_t k, const LIST_OF_VECTORS1& points,
	std::vector<int>& assignments, LIST_OF_VECTORS2* out_centers = nullptr,
	const size_t attempts = 3)
{
	return detail::stub_kmeans(
		false /* standard method */, k, points, assignments, out_centers,
		attempts);
}

/** k-means++ algorithm to cluster a list of N points of arbitrary
 *dimensionality into exactly K clusters.
 *   The list of input points can be any template CONTAINER<POINT> with:
 *		- CONTAINER can be: Any STL container:
 *std::vector,std::list,std::deque,...
 *		- POINT can be:
 *			- std::vector<double/float>
 *			- CVectorFixedDouble<N> / CVectorFixedFloat<N>
 *
 *  \param k [IN] Number of cluster to look for.
 *  \param points [IN] The list of N input points. It can be any STL-like
 *containers of std::vector<float/double>, for example a
 *std::vector<mrpt::math::CVectorDouble>, a std::list<CVectorFloat>, etc...
 *  \param assignments [OUT] At output it will have a number [0,k-1] for each
 *of the N input points.
 *  \param out_centers [OUT] If not nullptr, at output will have the centers of
 *each group. Can be of any of the supported types of "points", but the basic
 *coordinates should be float or double exactly as in "points".
 *  \param attempts [IN] Number of attempts.
 *
 * \sa The standard kmeans algorithm, see: kmeans
 * \note Uses the kmeans++ implementation by David Arthur (2009,
 *http://www.stanford.edu/~darthur/kmpp.zip).
 */
template <class LIST_OF_VECTORS1, class LIST_OF_VECTORS2 = LIST_OF_VECTORS1>
inline double kmeanspp(
	const size_t k, const LIST_OF_VECTORS1& points,
	std::vector<int>& assignments, LIST_OF_VECTORS2* out_centers = nullptr,
	const size_t attempts = 3)
{
	return detail::stub_kmeans(
		true /* kmeans++ algorithm*/, k, points, assignments, out_centers,
		attempts);
}

/** @} */

}  // namespace math
}  // namespace mrpt
