/* +------------------------------------------------------------------------+
   |                     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/core/round.h>
#include <mrpt/img/TPixelCoord.h>
#include <mrpt/math/KDTreeCapable.h>
#include <mrpt/vision/types.h>

#include <functional>

namespace mrpt::vision
{
/** \addtogroup  mrptvision_features
	@{ */

/** Simple structure for image key points. Descriptors are stored separately in
 * CFeatureList. This is a template to allow using both integers
 * (TKeyPoint) or floats (TKeyPointf) as pixel coordinates. \sa
 * TKeyPoint, TKeyPointf
 */
template <typename PIXEL_COORD_TYPE>
struct TKeyPoint_templ
{
	/** The type of \a pt */
	using pixel_coords_t = PIXEL_COORD_TYPE;
	/** The type of pt.x and pt.y */
	using pixel_coord_t = typename PIXEL_COORD_TYPE::pixel_coord_t;

	/** Coordinates in the image */
	pixel_coords_t pt;

	/** ID of the feature */
	TFeatureID ID{static_cast<TFeatureID>(-1)};

	/** Status of the feature tracking process */
	TFeatureTrackStatus track_status{status_IDLE};

	/** A measure of the "goodness" of the feature (typically, the KLT_response
	 * value) */
	float response{0};
	/** The image octave the image was found in: 0=original image, 1=1/2 image,
	 * 2=1/4 image, etc. */
	uint8_t octave{0};
	/** A field for any other flags needed by the user (this has not a
	 * predefined meaning) */
	uint8_t user_flags{0};

	/** Constructor that only sets the pt.{x,y} values, leaving all other values
	 * to *undefined values*. */
	template <typename COORD_TYPE>
	inline TKeyPoint_templ(const COORD_TYPE x, const COORD_TYPE y) : pt(x, y)
	{
	}

	/** Default constructor, leaves all fields uninitialized */
	inline TKeyPoint_templ() = default;
	template <typename OTHER_TKeyPoint>
	explicit TKeyPoint_templ(const OTHER_TKeyPoint& o)
		: pt(o.pt.x, o.pt.y),
		  ID(o.ID),
		  track_status(o.track_status),
		  response(o.response),
		  octave(o.octave),
		  user_flags(o.user_flags)
	{
	}
};

/** Simple structure for image key points.
 *  \sa TKeyPointf, CFeature, TKeyPointList
 */
using TKeyPoint = TKeyPoint_templ<mrpt::img::TPixelCoord>;

/** A version of TKeyPoint with subpixel precision */
using TKeyPointf = TKeyPoint_templ<mrpt::img::TPixelCoordf>;

template <typename FEATURE>
struct TKeyPointTraits;

template <>
struct TKeyPointTraits<TKeyPoint>
{
	using coord_t = int;

	static inline coord_t f2coord(float f) { return mrpt::round(f); }
};

template <>
struct TKeyPointTraits<TKeyPointf>
{
	using coord_t = float;

	static inline coord_t f2coord(float f) { return f; }
};

/** A list of image features using the structure TKeyPoint for each feature
 * Users normally will use directly: TKeyPointList, TKeyPointfList
 */
template <typename FEATURE>
struct TKeyPointList_templ
{
   public:
	using TFeatureVector = std::vector<FEATURE>;
	using feature_t = FEATURE;

	/** @name Utilities
		@{ */

	/** Returns a const ref to the actual std::vector<> container */
	const TFeatureVector& getVector() const { return m_feats; }
	/** Returns the maximum ID of all features in the list, or 0 if it's empty
	 */
	TFeatureID getMaxID() const
	{
		if (this->empty()) return 0;
		TFeatureID maxID = m_feats[0].ID;
		size_t N = m_feats.size() - 1;
		for (; N; --N)
			mrpt::keep_max(maxID, m_feats[N].ID);
		return maxID;
	}

	/** Returns a vector with a LUT of the first feature index per row, to
	 * efficiently look for neighbors, etc.
	 *  By default this vector is empty, so if a feature detector is used that
	 * doesn't fill this out, it will remain empty and useless.
	 *  \note FASTER detectors do fill this out. In general, a feature list
	 * that dynamically changes will not use this LUT.
	 */
	const std::vector<size_t>& getFirstIndexPerRowLUT() const
	{
		return m_first_index_per_row;
	}
	/// \overload
	std::vector<size_t>& getFirstIndexPerRowLUT()
	{
		return m_first_index_per_row;
	}

	/** @} */

	/** @name Method and datatypes to emulate a STL container
		@{ */
	using iterator = typename TFeatureVector::iterator;
	using const_iterator = typename TFeatureVector::const_iterator;

	using reverse_iterator = typename TFeatureVector::reverse_iterator;
	using const_reverse_iterator =
		typename TFeatureVector::const_reverse_iterator;

	inline iterator begin() { return m_feats.begin(); }
	inline iterator end() { return m_feats.end(); }
	inline const_iterator begin() const { return m_feats.begin(); }
	inline const_iterator end() const { return m_feats.end(); }
	inline reverse_iterator rbegin() { return m_feats.rbegin(); }
	inline reverse_iterator rend() { return m_feats.rend(); }
	inline const_reverse_iterator rbegin() const { return m_feats.rbegin(); }
	inline const_reverse_iterator rend() const { return m_feats.rend(); }
	inline iterator erase(const iterator& it) { return m_feats.erase(it); }
	inline bool empty() const { return m_feats.empty(); }
	inline size_t size() const { return m_feats.size(); }
	inline void clear()
	{
		m_feats.clear();
		m_first_index_per_row.clear();
	}
	inline void resize(size_t N) { m_feats.resize(N); }
	inline void reserve(size_t N) { m_feats.reserve(N); }
	inline void push_back(const FEATURE& f) { m_feats.push_back(f); }
	inline void emplace_back(const int x, const int y)
	{
		m_feats.emplace_back(x, y);
	}

	inline FEATURE& operator[](const std::size_t index)
	{
		return m_feats[index];
	}
	inline const FEATURE& operator[](const std::size_t index) const
	{
		return m_feats[index];
	}

	inline FEATURE& back() { return m_feats.back(); }
	inline const FEATURE& back() const { return m_feats.back(); }
	inline FEATURE& front() { return m_feats.front(); }
	inline const FEATURE& front() const { return m_feats.front(); }
	/** @} */

	/** @name getFeature*() methods for template-based access to feature list
		@{ */
	inline typename TKeyPointTraits<FEATURE>::coord_t getFeatureX(
		size_t i) const
	{
		return m_feats[i].pt.x;
	}
	inline typename TKeyPointTraits<FEATURE>::coord_t getFeatureY(
		size_t i) const
	{
		return m_feats[i].pt.y;
	}
	inline TFeatureID getFeatureID(size_t i) const { return m_feats[i].ID; }
	inline float getFeatureResponse(size_t i) const
	{
		return m_feats[i].response;
	}
	inline bool isPointFeature([[maybe_unused]] size_t i) const { return true; }
	inline float getScale(size_t i) const
	{
		return d2f(1 << m_feats[i].octave);
	}
	inline TFeatureTrackStatus getTrackStatus(size_t i)
	{
		return m_feats[i].track_status;
	}

	inline void setFeatureX(
		size_t i, typename TKeyPointTraits<FEATURE>::coord_t x)
	{
		m_feats[i].pt.x = x;
	}
	inline void setFeatureY(
		size_t i, typename TKeyPointTraits<FEATURE>::coord_t y)
	{
		m_feats[i].pt.y = y;
	}

	inline void setFeatureXf(size_t i, float x)
	{
		m_feats[i].pt.x = TKeyPointTraits<FEATURE>::f2coord(x);
	}
	inline void setFeatureYf(size_t i, float y)
	{
		m_feats[i].pt.y = TKeyPointTraits<FEATURE>::f2coord(y);
	}

	inline void setFeatureID(size_t i, TFeatureID id) { m_feats[i]->ID = id; }
	inline void setFeatureResponse(size_t i, float r)
	{
		m_feats[i]->response = r;
	}
	inline void setScale(size_t i, float s)
	{
		m_feats[i]->octave = mrpt::round(std::log(s) / std::log(2));
	}
	inline void setTrackStatus(size_t i, TFeatureTrackStatus s)
	{
		m_feats[i].track_status = s;
	}

	inline void mark_as_outdated() const {}
	/** @} */

   private:
	/** The actual container with the list of features */
	TFeatureVector m_feats;
	/** A LUT of the first feature index per row, to efficiently look for
	 * neighbors, etc. */
	std::vector<size_t> m_first_index_per_row;
	mrpt::math::CMatrixBool m_occupied_sections;

};	// end of class

/** A list of image features using the structure TKeyPoint for each feature
 * - capable of KD-tree computations */
using TKeyPointList = TKeyPointList_templ<TKeyPoint>;

/** A list of image features using the structure TKeyPointf for each
 * feature - capable of KD-tree computations */
using TKeyPointfList = TKeyPointList_templ<TKeyPointf>;

/** A helper struct to sort keypoints by their response: It can be used with
 *these types:
 *	  - std::vector<cv::KeyPoint>
 *	  - mrpt::vision::TKeyPointList
 */
template <typename FEATURE_LIST>
struct KeypointResponseSorter : public std::function<bool(size_t, size_t)>
{
	const FEATURE_LIST& m_data;
	KeypointResponseSorter(const FEATURE_LIST& data) : m_data(data) {}
	bool operator()(size_t k1, size_t k2) const
	{
		return (m_data[k1].response > m_data[k2].response);
	}
};

/** Helper class: KD-tree search class for vector<KeyPoint>:
 *  Call mark_as_outdated() to force rebuilding the kd-tree after modifying the
 * linked feature list.
 *  \tparam FEAT Can be cv::KeyPoint or mrpt::vision::TKeyPoint
 */
template <typename FEAT>
class CFeatureListKDTree
	: public mrpt::math::KDTreeCapable<CFeatureListKDTree<FEAT>>
{
   public:
	inline void mark_as_outdated()
	{
		mrpt::math::KDTreeCapable<
			CFeatureListKDTree<FEAT>>::kdtree_mark_as_outdated();
	}

	const std::vector<FEAT>& m_data;
	CFeatureListKDTree(const std::vector<FEAT>& data) : m_data(data) {}
	/** @name Methods that MUST be implemented by children classes of
	   KDTreeCapable
		@{ */

	/// Must return the number of data points
	inline size_t kdtree_get_point_count() const { return m_data.size(); }
	/// Returns the dim'th component of the idx'th point in the class:
	inline float kdtree_get_pt(size_t idx, int dim) const
	{
		ASSERTDEB_(dim == 0 || dim == 1);
		if (dim == 0) return m_data[idx].pt.x;
		else
			return m_data[idx].pt.y;
	}

	/// Returns the distance between the vector "p1[0:size-1]" and the data
	/// point with index "idx_p2" stored in the class:
	inline float kdtree_distance(
		const float* p1, size_t idx_p2, [[maybe_unused]] size_t size) const
	{
		ASSERTDEB_(size == 2);

		const float d0 = p1[0] - m_data[idx_p2].pt.x;
		const float d1 = p1[1] - m_data[idx_p2].pt.y;
		return d0 * d0 + d1 * d1;
	}

	// Optional bounding-box computation: return false to default to a standard
	// bbox computation loop.
	//   Return true if the BBOX was already computed by the class and returned
	//   in "bb" so it can be avoided to redo it again.
	//   Look at bb.size() to find out the expected dimensionality (e.g. 2 or 3
	//   for point clouds)
	template <typename BBOX>
	bool kdtree_get_bbox([[maybe_unused]] BBOX& bb) const
	{
		return false;
	}

	/** @} */

};	// end CFeatureListKDTree

/** @} */  // End of add to module: mrptvision_features

}  // namespace mrpt::vision
