#pragma once

#include "../../../include/dv-processing/features/mean_shift_tracker.hpp"

#include <pybind11/pybind11.h>

namespace py = pybind11;

/**
 * @brief Bind features::FeatureTracks to python
 * @param m pybind The module to bind to
 */
void bind_mean_shift_tracker(py::module &m) {
	using pybind11::operator""_a;

	py::class_<dv::features::MeanShiftTracker>(
		m, "MeanShiftTracker", "Track event blobs using mean shift algorithm on time surface event data.")
		.def(py::init([](const cv::Size &resolution, const int bandwidth, const dv::Duration timeWindow,
						  const float stepSize, const float weightMultiplier, float convergenceNorm, int maxIters) {
			return dv::features::MeanShiftTracker(resolution, bandwidth, timeWindow, nullptr, nullptr, stepSize,
				weightMultiplier, convergenceNorm, maxIters);
		}),
			"resolution"_a, "bandwidth"_a, "timeWindow"_a, "stepSize"_a = 0.5, "weightMultiplier"_a = 1,
			"convergenceNorm"_a = 0.01, "maxIters"_a = 2000,
			"Constructor for mean shift tracker using Epanechnikov kernel as weights for the time surface of events "
			"used to update track location. The kernel weights have highest value on the previous track location. This "
			"assumption is based on the idea that the new track location is \"close\" to last track location. The "
			"consecutive track updates are performed until the maximum number of iteration is reached or the shift "
			"between consecutive updates is below a threshold.")
		.def("accept", &dv::features::MeanShiftTracker::accept, py::arg("store"),
			"Add events to time surface and update latest timestamp.")
		.def(
			"setRedetectionStrategy",
			[](dv::features::MeanShiftTracker &self, const dv::features::NoRedetection &strategy) {
				self.setRedetectionStrategy(std::make_unique<dv::features::NoRedetection>(strategy));
			},
			py::arg("strategy"), "Set a new redetection strategy.")
		.def(
			"setRedetectionStrategy",
			[](dv::features::MeanShiftTracker &self, const dv::features::FeatureCountRedetection &strategy) {
				self.setRedetectionStrategy(std::make_unique<dv::features::FeatureCountRedetection>(strategy));
			},
			py::arg("strategy"), "Set a new redetection strategy.")
		.def(
			"setRedetectionStrategy",
			[](dv::features::MeanShiftTracker &self, const dv::features::UpdateIntervalRedetection &strategy) {
				self.setRedetectionStrategy(std::make_unique<dv::features::UpdateIntervalRedetection>(strategy));
			},
			py::arg("strategy"), "Set a new redetection strategy.")
		.def(
			"setRedetectionStrategy",
			[](dv::features::MeanShiftTracker &self,
				const dv::features::UpdateIntervalOrFeatureCountRedetection &strategy) {
				self.setRedetectionStrategy(
					std::make_unique<dv::features::UpdateIntervalOrFeatureCountRedetection>(strategy));
			},
			py::arg("strategy"), "Set a new redetection strategy.")
		.def(
			"setDetector",
			[](dv::features::MeanShiftTracker &self, const dv::features::EventFeatureBlobDetector &detector) {
				self.setDetector(std::make_unique<dv::features::EventFeatureBlobDetector>(detector));
			},
			py::arg("detector"), "Set a new feature (blob) detector from accumulated images.")
		.def("getMaxTracks", &dv::features::MeanShiftTracker::getMaxTracks, "Get the maximum number of tracks.")
		.def("setMaxTracks", &dv::features::MeanShiftTracker::setMaxTracks, py::arg("maxTracks"),
			"Set the maximum number of tracks.")
		.def("runTracking", &dv::features::MeanShiftTracker::runTracking, "Perform the tracking and cache the results.")
		.def("removeTracks", &dv::features::MeanShiftTracker::removeTracks, py::arg("trackIds"),
			"Remove tracks from cached results, so the wouldn't be tracked anymore. TrackIds are the `class_id` value "
			"of the keypoint structure.")
		.def("getBandwidth", &dv::features::MeanShiftTracker::getBandwidth,
			"Getter for bandwidth value that defines the search area for a new track.")
		.def("setBandwidth", &dv::features::MeanShiftTracker::setBandwidth, py::arg("bandwidth"),
			"Setter for bandwidth value.")
		.def("getTimeWindow", &dv::features::MeanShiftTracker::getTimeWindow,
			"Get time duration used to compute time surface.")
		.def("setTimeWindow", &dv::features::MeanShiftTracker::setTimeWindow, py::arg("timeWindow"),
			"Setter for window size for normalized time surface.")
		.def("getStepSize", &dv::features::MeanShiftTracker::getStepSize,
			"Get learning rate value used for track location update.")
		.def("setStepSize", &dv::features::MeanShiftTracker::setStepSize, py::arg("stepSize"),
			"Setter for learning rate for motion towards new center during one mean shift iteration.")
		.def("getWeightMultiplier", &dv::features::MeanShiftTracker::getWeightMultiplier,
			"Get scaling factor for Epanechnikov weights used in the computation of the mean shift cost update.")
		.def("setWeightMultiplier", &dv::features::MeanShiftTracker::setWeightMultiplier, py::arg("multiplier"),
			"Setter for scaling factor used in the computation of the mean shift cost update.");
}
