#pragma once

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

#include <optional>

namespace dv::noise {

enum class FrequencyFilterType {
	PASS,
	CUT,
};

inline std::ostream &operator<<(std::ostream &os, const FrequencyFilterType &var) {
	switch (var) {
		case FrequencyFilterType::PASS:
			return os << "PASS";

		case FrequencyFilterType::CUT:
			return os << "CUT";
	}

	return os;
}

/**
 * A base class for basic frequency filters. Different frequency filters
 * (low-pass, high-pass, band-pass, etc.) are derived as special cases of this base filter. Handles data input and
 * output, as well as checking if an event should be retained. Derived classes only have to define the required behavior
 * by passing correct arguments to the constructor.
 * @tparam EventStoreClass Type of event store.
 */
template<dv::concepts::EventStorage EventStoreClass = dv::EventStore>
class BaseFrequencyFilter : public EventFilterBase<EventStoreClass> {
protected:
	dv::TimeSurface mTimeSurface;
	std::optional<int64_t> mLowCutOffPeriod;
	std::optional<int64_t> mHighCutOffPeriod;
	FrequencyFilterType mFilterType;

public:
	/**
	 * Construct a base event frequency filter for filtering events based on defined thresholds for the upper/lower
	 * cutoff frequencies. Different filter behaviors can be implemented by passing different
	 * arguments to the constructor (low-pass, high-pass, band-pass, etc.)
	 * @param resolution Sensor resolution.
	 * @param lowerCutOffFrequency Lower cutoff frequency for filter. Ignored if std::nullopt
	 * @param upperCutOffFrequency Upper cutoff frequency for filter. Ignored if std::nullopt
	 * @param filterType Filter type, whether the filter behaves as a cut or a pass for the given cutoff frequencies
	 */
	explicit BaseFrequencyFilter(const cv::Size &resolution, const std::optional<float> lowerCutOffFrequency,
		const std::optional<float> upperCutOffFrequency, const FrequencyFilterType filterType) :
		mTimeSurface(resolution),
		mFilterType(filterType) {
		if (lowerCutOffFrequency.has_value()) {
			dv::runtime_assert(
				[lowerCutOffFrequency] {
					return lowerCutOffFrequency.value() > 0;
				},
				[] {
					return "In frequency filter, invalid cutoff frequency value. Cutoff frequency must be greater than "
						   "zero.";
				});

			mHighCutOffPeriod = static_cast<std::optional<int64_t>>(1e6f / lowerCutOffFrequency.value());
		}
		else {
			mHighCutOffPeriod = std::nullopt;
		}

		if (upperCutOffFrequency.has_value()) {
			dv::runtime_assert(
				[upperCutOffFrequency] {
					return upperCutOffFrequency.value() > 0;
				},
				[] {
					return "In frequency filter, invalid cutoff frequency value. Cutoff frequency must be greater than "
						   "zero.";
				});

			mLowCutOffPeriod = static_cast<std::optional<int64_t>>(1e6f / upperCutOffFrequency.value());
		}
		else {
			mLowCutOffPeriod = std::nullopt;
		}
	}

	/**
	 * Test whether the frequency of an event is less than the cutoff frequency.
	 * @param event 	Event to be tested.
	 * @return 			True - the frequency of the tested event at that pixel location is less than the cutoff
	 * 					frequency, false otherwise.
	 */
	[[nodiscard]] bool retain(const typename EventStoreClass::value_type &event) noexcept override {
		// By default, if no cutoff frequencies are defined, retain event
		if (!mLowCutOffPeriod.has_value() && !mHighCutOffPeriod.has_value()) {
			return true;
		}

		// We check if event should be retained if the filter is a pass filter, and invert the logic if it is a cut
		// filter
		int64_t &tsTime   = mTimeSurface(event.y(), event.x());
		const auto period = event.timestamp() - tsTime;
		bool isSignal     = true;
		tsTime            = event.timestamp();

		// Filter events based on lower cutoff period (high cutoff frequency)
		if (mLowCutOffPeriod.has_value()) {
			isSignal &= (period >= mLowCutOffPeriod);
		}

		// Filter events based on higher cutoff period (lower cutoff frequency)
		if (mHighCutOffPeriod.has_value()) {
			isSignal &= (period <= mHighCutOffPeriod);
		}

		// Invert logic if filter is a cut filter
		if (mFilterType == FrequencyFilterType::CUT) {
			isSignal = !isSignal;
		}

		// Add event timestamp to time surface and return check
		return isSignal;
	}
};

/**
 * A low-pass event frequency filter. Discards events at a pixel location with a frequency above a given cutoff
 * frequency. This is also commonly referred to as refractory period filtering, where the refractory period is the
 * inverse of the cutoff frequency.
 * @tparam EventStoreClass Type of event store.
 */
template<dv::concepts::EventStorage EventStoreClass = dv::EventStore>
class LowPassFilter : public BaseFrequencyFilter<EventStoreClass> {
public:
	/**
	 * A low-pass event frequency filter. Discards events at a pixel location with a frequency above a given cutoff
	 * frequency. This is also commonly referred to as refractory period filtering, where the refractory period is the
	 * inverse of #cutOffFrequency.
	 * @param resolution Sensor resolution.
	 * @param cutOffFrequency Filter cutoff frequency. All events with a frequency above this given cutoff are
	 * discarded.
	 */
	explicit LowPassFilter(const cv::Size &resolution, const float cutOffFrequency) :
		BaseFrequencyFilter<EventStoreClass>(resolution, std::nullopt, cutOffFrequency, FrequencyFilterType::PASS) {
	}

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

	/**
	 * Get the cutoff frequency for the frequency filter.
	 * @return Currently configured cutoff frequency.
	 */
	[[nodiscard]] float getCutOffFrequency() const {
		return 1e6f / static_cast<float>(this->mLowCutOffPeriod.value());
	}

	/**
	 * Set a new cutoff frequency for the frequency filter.
	 * @param frequency New cutoff frequency value.
	 */
	void setCutOffFrequency(const float frequency) {
		this->mLowCutOffPeriod = static_cast<int64_t>(1e6f / frequency);
	}
};

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

/**
 * A high-pass event frequency filter. Discards events at a pixel location with a frequency below a given cutoff
 * frequency.
 * @tparam EventStoreClass Type of event store.
 */
template<dv::concepts::EventStorage EventStoreClass = dv::EventStore>
class HighPassFilter : public BaseFrequencyFilter<EventStoreClass> {
public:
	/**
	 * A high-pass event frequency filter. Discards events at a pixel location with a frequency below a given cutoff
	 * frequency.
	 * @param resolution Sensor resolution.
	 * @param cutOffFrequency Filter cutoff frequency. All events with a frequency below this given cutoff are
	 * discarded.
	 */
	explicit HighPassFilter(const cv::Size &resolution, const float cutOffFrequency) :
		BaseFrequencyFilter<EventStoreClass>(resolution, cutOffFrequency, std::nullopt, FrequencyFilterType::PASS) {
	}

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

	/**
	 * Get the cutoff frequency for the frequency filter.
	 * @return Currently configured cutoff frequency.
	 */
	[[nodiscard]] float getCutOffFrequency() const {
		return 1e6f / static_cast<float>(this->mHighCutOffPeriod.value());
	}

	/**
	 * Set a new cutoff frequency for the frequency filter.
	 * @param frequency New cutoff frequency value.
	 */
	void setCutOffFrequency(const float frequency) {
		this->mHighCutOffPeriod = static_cast<int64_t>(1e6f / frequency);
	}
};

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

/**
 * A band-pass event frequency filter. Discards events at a pixel location with a frequency outside a given frequency
 * band defined by a lower cutoff frequency and an upper cutoff frequency.
 * @tparam EventStoreClass Type of event store.
 */
template<dv::concepts::EventStorage EventStoreClass = dv::EventStore>
class BandPassFilter : public BaseFrequencyFilter<EventStoreClass> {
public:
	/**
	 * A band-pass event frequency filter. Discards events at a pixel location with a frequency outside a given
	 * frequency band defined by a lower cutoff frequency and an upper cutoff frequency.
	 * @param resolution Sensor resolution.
	 * @param lowerCutOffFrequency Lower filter cutoff frequency. Together with #highCutOffFrequency, defines the
	 * frequency band for the band-pass filter. All events with a frequency outside of this band are discarded.
	 * @param upperCutOffFrequency Upper filter cutoff frequency. Together with #lowCutOffFrequency, defines the
	 * frequency band for the band-pass filter. All events with a frequency outside of this band are discarded.
	 */
	explicit BandPassFilter(
		const cv::Size &resolution, const float lowerCutOffFrequency, const float upperCutOffFrequency) :
		BaseFrequencyFilter<EventStoreClass>(
			resolution, lowerCutOffFrequency, upperCutOffFrequency, FrequencyFilterType::PASS) {
	}

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

	/**
	 * Get the lower cutoff frequency of the frequency band.
	 * @return  Currently configured low cutoff frequency of the frequency band.
	 */
	[[nodiscard]] float getLowerCutOffFrequency() const {
		return 1e6f / static_cast<float>(this->mHighCutOffPeriod.value());
	}

	/**
	 * Get the upper cutoff frequency of the frequency band.
	 * @return  Currently configured upper cutoff frequency of the frequency band.
	 */
	[[nodiscard]] float getUpperCutOffFrequency() const {
		return 1e6f / static_cast<float>(this->mLowCutOffPeriod.value());
	}

	/**
	 * Set a new value for the lower cutoff frequency of the frequency band.
	 * @param frequency New lower cutoff frequency value.
	 */
	void setLowerCutOffFrequency(const float frequency) {
		this->mHighCutOffPeriod = static_cast<int64_t>(1e6f / frequency);
	}

	/**
	 * Set a new value for the upper cutoff frequency of the frequency band.
	 * @param frequency New upper cutoff frequency value.
	 */
	void setUpperCutOffFrequency(const float frequency) {
		this->mLowCutOffPeriod = static_cast<int64_t>(1e6f / frequency);
	}
};

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

template<dv::concepts::EventStorage EventStoreClass = dv::EventStore>
class BandCutFilter : public BaseFrequencyFilter<EventStoreClass> {
public:
	/**
	 * A band-cut event frequency filter. Discards events at a pixel location with a frequency inside a given frequency
	 * band defined by a lower cutoff frequency and an upper cutoff frequency.
	 * @param resolution Sensor resolution.
	 * @param lowerCutOffFrequency Lower filter cutoff frequency. Together with #highCutOffFrequency, defines the
	 * frequency band for the band-pass filter. All events with a frequency outside of this band are discarded.
	 * @param upperCutOffFrequency Upper filter cutoff frequency. Together with #lowCutOffFrequency, defines the
	 * frequency band for the band-pass filter. All events with a frequency outside of this band are discarded.
	 */
	explicit BandCutFilter(
		const cv::Size &resolution, const float lowerCutOffFrequency, const float upperCutOffFrequency) :
		BaseFrequencyFilter<EventStoreClass>(
			resolution, lowerCutOffFrequency, upperCutOffFrequency, FrequencyFilterType::CUT) {
	}

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

	/**
	 * Get the lower cutoff frequency of the frequency band.
	 * @return  Currently configured low cutoff frequency of the frequency band.
	 */
	[[nodiscard]] float getLowerCutOffFrequency() const {
		return 1e6f / static_cast<float>(this->mHighCutOffPeriod.value());
	}

	/**
	 * Get the upper cutoff frequency of the frequency band.
	 * @return  Currently configured upper cutoff frequency of the frequency band.
	 */
	[[nodiscard]] float getUpperCutOffFrequency() const {
		return 1e6f / static_cast<float>(this->mLowCutOffPeriod.value());
	}

	/**
	 * Set a new value for the lower cutoff frequency of the frequency band.
	 * @param frequency New lower cutoff frequency value.
	 */
	void setLowerCutOffFrequency(const float frequency) {
		this->mHighCutOffPeriod = static_cast<int64_t>(1e6f / frequency);
	}

	/**
	 * Set a new value for the upper cutoff frequency of the frequency band.
	 * @param frequency New upper cutoff frequency value.
	 */
	void setUpperCutOffFrequency(const float frequency) {
		this->mLowCutOffPeriod = static_cast<int64_t>(1e6f / frequency);
	}
};

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

} // namespace dv::noise

// fmt compatibility for enum class printing.
template<>
struct fmt::formatter<dv::noise::FrequencyFilterType> : fmt::ostream_formatter {};
