#pragma once

#include "../core/filters.hpp"

#include <optional>

namespace dv::noise {

/**
 * Memory efficient, spatiotemporal event filter as proposed by Khodamoradi's algorithm;
 * "O(N)-Space Spatiotemporal Filter for Reducing Noise in Neuromorphic Vision Sensors".
 * @tparam EventStoreClass Type of event store to filter.
 */
template<dv::concepts::EventStorage EventStoreClass = dv::EventStore>
class KNoiseFilter : public EventFilterBase<EventStoreClass> {
public:
	/**
	 * Construct a spatiotemporal event filter that discards an event if the number of spatiotemporal neighbors is less
	 * than #numSupportingPixels. The spatial window size is defined by #windowHalfSize, while the temporal window size
	 * is defined by #timeDelta.
	 *
	 * Note that this is similar in functionality to the BackgroundActivityNoiseFilter, where #timeDelta
	 * corresponds to the backgroundActivityTime, and #windowHalfSize and #numSupportingPixels are by default set to 1;
	 * but more memory efficient due to only storing values per row/column.
	 *
	 * @param resolution Sensor resolution.
	 * @param timeDelta Size of the temporal window used to check the number of neighboring events.
	 */
	explicit KNoiseFilter(const cv::Size &resolution, const dv::Duration timeDelta = dv::Duration(2'000)) :
		mResolutionLimits(resolution.width - 1, resolution.height - 1),
		mColumnCells(resolution.width),
		mRowCells(resolution.height),
		mTimeDelta(timeDelta.count()) {
		if (timeDelta.count() < 1) {
			throw dv::exceptions::InvalidArgument<int64_t>("TimeDelta must be greater than 0.", timeDelta.count());
		}
	}

	/**
	 * Accept events using the input stream operator.
	 * @param events 	Input events.
	 * @return
	 */
	KNoiseFilter &operator<<(const EventStoreClass &events) {
		accept(events);
		return *this;
	}

	/**
	 * Test if the given event should be filtered out by checking if the number of neighboring events within the defined
	 * spatiotemporal window is greater than #mNumSupportingPixels.
	 * @param event Event to be checked.
	 * @return True to retain event, false to discard.
	 */
	[[nodiscard]] bool retain(const typename EventStoreClass::value_type &event) noexcept override {
		const bool isSignal = doKNoiseLookup(event);

		// Update values
		KMemCell &xCell     = mColumnCells[event.x()];
		xCell.mTimestamp    = event.timestamp();
		xCell.mPolarity     = event.polarity();
		xCell.mOtherAddress = event.y();

		KMemCell &yCell     = mRowCells[event.y()];
		yCell.mTimestamp    = event.timestamp();
		yCell.mPolarity     = event.polarity();
		yCell.mOtherAddress = event.x();

		return isSignal;
	}

	/**
	 * Get the duration of the temporal window used to check for neighboring events.
	 * @return
	 */
	[[nodiscard]] dv::Duration getTemporalWindowDuration() const {
		return dv::Duration(mTimeDelta);
	}

	/**
	 * Set the duration of the temporal window used to check for neighboring events.
	 * @param timeDelta
	 */
	void setTemporalWindowDuration(const dv::Duration timeDelta) {
		if (timeDelta.count() < 1) {
			throw dv::exceptions::InvalidArgument<int64_t>("TimeDelta must be greater than 0.", timeDelta.count());
		}

		mTimeDelta = timeDelta.count();
	}

private:
	/**
	 * Memory cell used per row/column to store latest event information. Note that, differently from the original
	 * implementation, the number of bits per memory cell is not exactly 64 (unnecessary bookkeeping bit removed).
	 */
	struct KMemCell {
		int64_t mTimestamp;
		int16_t mOtherAddress;
		bool mPolarity;
	};

	/**
	 * Check if for the column given by event.x() + columnShift, the last received event for that column is within a
	 * spatiotemporal window from the current event.
	 * @param event Event to be checked.
	 * @param columnShift Shift in the column position from the current event column.
	 * @return True if the last received event at the given column position is within a spatiotemporal window from the
	 * current event.
	 */
	[[nodiscard]] inline bool checkColumnSupport(const dv::Event event, const int8_t columnShift) {
		const int16_t xCoordinate = event.x() + columnShift;

		const KMemCell &xCell = mColumnCells[xCoordinate];

		const bool validTime     = ((event.timestamp() - xCell.mTimestamp) <= mTimeDelta);
		const bool validPolarity = (event.polarity() == xCell.mPolarity);
		// PERF: this last check has a heavy impact on performance, with 50% to 30% reduction in throughput.
		const bool validOther = (std::abs(xCell.mOtherAddress - event.y()) <= 1);

		return validTime && validPolarity && validOther;
	}

	/**
	 * Check if for the row given by event.y() + rowShift, the last received event for that row is within a
	 * spatiotemporal window from the current event.
	 * @param event Event to be checked.
	 * @param rowShift Shift in the row position from the current event row.
	 * @return True if the last received event at the given row position is within a spatiotemporal window from the
	 * current event.
	 */
	[[nodiscard]] inline bool checkRowSupport(const dv::Event event, const int8_t rowShift) {
		const int16_t yCoordinate = event.y() + rowShift;

		const KMemCell &yCell = mRowCells[yCoordinate];

		const bool validTime     = ((event.timestamp() - yCell.mTimestamp) <= mTimeDelta);
		const bool validPolarity = (event.polarity() == yCell.mPolarity);
		// PERF: this last check has a heavy impact on performance, with 50% to 30% reduction in throughput.
		const bool validOther = (std::abs(yCell.mOtherAddress - event.x()) <= 1);

		return validTime && validPolarity && validOther;
	}

	[[nodiscard]] inline bool doKNoiseLookup_unsafe(const dv::Event event) {
		// 3x column checks, in order.
		if (checkColumnSupport(event, -1)) {
			return true;
		}

		if (checkColumnSupport(event, 0)) {
			return true;
		}

		if (checkColumnSupport(event, 1)) {
			return true;
		}

		// 3x row checks, in order.
		if (checkRowSupport(event, -1)) {
			return true;
		}

		if (checkRowSupport(event, 0)) {
			return true;
		}

		if (checkRowSupport(event, 1)) {
			return true;
		}

		return false;
	}

	[[nodiscard]] inline bool doKNoiseLookup(const dv::Event event) {
		// Compute map limits.
		const bool notBorderLeft  = (event.x() != 0);
		const bool notBorderDown  = (event.y() != mResolutionLimits.height);
		const bool notBorderRight = (event.x() != mResolutionLimits.width);
		const bool notBorderUp    = (event.y() != 0);

		if (notBorderLeft && notBorderDown && notBorderRight && notBorderUp) {
			return doKNoiseLookup_unsafe(event);
		}

		// 3x column checks.
		if (notBorderLeft) {
			if (checkColumnSupport(event, -1)) {
				return true;
			}
		}

		if (checkColumnSupport(event, 0)) {
			return true;
		}

		if (notBorderRight) {
			if (checkColumnSupport(event, 1)) {
				return true;
			}
		}

		// 3x row checks.
		if (notBorderUp) {
			if (checkRowSupport(event, -1)) {
				return true;
			}
		}

		if (checkRowSupport(event, 0)) {
			return true;
		}

		if (notBorderDown) {
			if (checkRowSupport(event, 1)) {
				return true;
			}
		}

		return false;
	}

	/// Sensor resolution
	cv::Size mResolutionLimits;

	/// Temporal size of the spatiotemporal window (in us)
	int64_t mTimeDelta;

	/// Vector containing the memory cells per row/column of the sensor resolution.
	std::vector<KMemCell> mColumnCells;
	std::vector<KMemCell> mRowCells;
};

static_assert(concepts::EventFilter<KNoiseFilter<>, dv::EventStore>);

} // namespace dv::noise
