#include "dvsnoisefilter.hpp"

#include <boost/tokenizer.hpp>

void DVSNoiseFilter::initInputs(dv::InputDefinitionList &in) {
	in.addEventInput("events");
}

void DVSNoiseFilter::initOutputs(dv::OutputDefinitionList &out) {
	out.addEventOutput("events");
}

const char *DVSNoiseFilter::initDescription() {
	return ("Filters out noise from DVS change (polarity) events.");
}

void DVSNoiseFilter::initConfigOptions(dv::RuntimeConfig &config) {
	config.add("hotPixelEnable", dv::ConfigOption::boolOption("Enable the hot pixel filter.", false));
	config.add("hotPixelLearn",
		dv::ConfigOption::buttonOption(
			"Learn the position of current hot (abnormally active) pixels, so they can be filtered out.", "Run"));
	config.add("hotPixelList",
		dv::ConfigOption::stringOption(
			"List of hot pixels currently being filtered (comma separated 'x,y' addresses with '|' separated pixels).",
			""));
	config.add(
		"hotPixelTime", dv::ConfigOption::intOption(
							"Time in µs to accumulate events for learning new hot pixels.", 1000000, 0, 30000000));
	config.add("hotPixelCount",
		dv::ConfigOption::intOption(
			"Number of events needed in a learning time period for a pixel to be considered hot.", 10000, 0, 10000000));

	config.add("statistics/hotPixelFilteredOn",
		dv::ConfigOption::statisticOption("Number of ON events filtered out by the hot pixel filter."));
	config.add("statistics/hotPixelFilteredOff",
		dv::ConfigOption::statisticOption("Number of OFF events filtered out by the hot pixel filter."));

	config.add(
		"backgroundActivityEnable", dv::ConfigOption::boolOption("Enable the background activity filter.", true));
	config.add(
		"backgroundActivityTwoLevels", dv::ConfigOption::boolOption("Use two-level background activity filtering."));
	config.add("backgroundActivityCheckPolarity",
		dv::ConfigOption::boolOption("Consider polarity when filtering background activity."));
	config.add("backgroundActivitySupportMin",
		dv::ConfigOption::intOption(
			"Minimum number of direct neighbor pixels that must support this pixel for it to be valid.", 1, 1, 8));
	config.add("backgroundActivitySupportMax",
		dv::ConfigOption::intOption(
			"Maximum number of direct neighbor pixels that can support this pixel for it to be valid.", 8, 1, 8));
	config.add("backgroundActivityTime",
		dv::ConfigOption::intOption(
			"Maximum time difference in µs for events to be considered correlated and not be filtered out.", 2000, 0,
			10000000));

	config.add("statistics/backgroundActivityFilteredOn",
		dv::ConfigOption::statisticOption("Number of ON events filtered out by the background activity filter."));
	config.add("statistics/backgroundActivityFilteredOff",
		dv::ConfigOption::statisticOption("Number of OFF events filtered out by the background activity filter."));

	config.add("refractoryPeriodEnable", dv::ConfigOption::boolOption("Enable the refractory period filter.", true));
	config.add("refractoryPeriodTime",
		dv::ConfigOption::intOption("Minimum time between events to not be filtered out.", 100, 0, 10000000));

	config.add("statistics/refractoryPeriodFilteredOn",
		dv::ConfigOption::statisticOption("Number of ON events filtered out by the refractory period filter."));
	config.add("statistics/refractoryPeriodFilteredOff",
		dv::ConfigOption::statisticOption("Number of OFF events filtered out by the refractory period filter."));

	config.setPriorityOptions({"backgroundActivityEnable", "backgroundActivityTime"});
}

DVSNoiseFilter::DVSNoiseFilter() :
	hotPixelLearningStarted(false),
	hotPixelLearningStartTime(0),
	hotPixelStatOn(0),
	hotPixelStatOff(0),
	backgroundActivityStatOn(0),
	backgroundActivityStatOff(0),
	refractoryPeriodStatOn(0),
	refractoryPeriodStatOff(0) {
	auto eventInput = inputs.getEventInput("events");

	sizeX = static_cast<int16_t>(eventInput.sizeX());
	sizeY = static_cast<int16_t>(eventInput.sizeY());

	timestampsMap.resize(static_cast<size_t>(sizeX * sizeY));

	// Populate event output info node, keep same as input info node.
	outputs.getEventOutput("events").setup(inputs.getEventInput("events"));
}

void DVSNoiseFilter::run() {
	auto evt_in  = inputs.getEventInput("events").events();
	auto evt_out = outputs.getEventOutput("events").events();

	bool hotPixelEnabled           = config.get<dv::CfgType::BOOL>("hotPixelEnable");
	bool refractoryPeriodEnabled   = config.get<dv::CfgType::BOOL>("refractoryPeriodEnable");
	bool backgroundActivityEnabled = config.get<dv::CfgType::BOOL>("backgroundActivityEnable");

	int32_t hotPixelTime                 = config.get<dv::CfgType::INT>("hotPixelTime");
	int32_t refractoryPeriodTime         = config.get<dv::CfgType::INT>("refractoryPeriodTime");
	int32_t backgroundActivitySupportMin = config.get<dv::CfgType::INT>("backgroundActivitySupportMin");
	int32_t backgroundActivitySupportMax = config.get<dv::CfgType::INT>("backgroundActivitySupportMax");
	int32_t backgroundActivityTime       = config.get<dv::CfgType::INT>("backgroundActivityTime");
	bool backgroundActivityCheckPolarity = config.get<dv::CfgType::BOOL>("backgroundActivityCheckPolarity");
	bool backgroundActivityTwoLevels     = config.get<dv::CfgType::BOOL>("backgroundActivityTwoLevels");

	// Hot Pixel learning: initialize and store packet-level timestamp.
	if (config.get<dv::CfgType::BOOL>("hotPixelLearn") && !hotPixelLearningStarted) {
		// Initialize hot pixel learning.
		hotPixelLearningMap.resize(static_cast<size_t>(sizeX * sizeY));

		hotPixelLearningStarted = true;

		// Store start timestamp.
		hotPixelLearningStartTime = evt_in[0].timestamp();

		log.debug << "HotPixel Learning: started on ts=" << hotPixelLearningStartTime << "." << std::endl;
	}

	for (const auto &evt : evt_in) {
		size_t pixelIndex = static_cast<size_t>((evt.y() * sizeX) + evt.x()); // Target pixel.

		// Hot Pixel learning: determine which pixels are abnormally active,
		// by counting how many times they spike in a given time period. The
		// ones above a given threshold will be considered "hot".
		// This is done first, so that other filters (including the Hot Pixel
		// filter itself) don't influence the learning operation.
		if (hotPixelLearningStarted) {
			hotPixelLearningMap[pixelIndex]++;

			if (evt.timestamp() > (hotPixelLearningStartTime + hotPixelTime)) {
				// Enough time has passed, we can proceed with data evaluation.
				hotPixelGenerateArray();

				// Done, reset and notify end of learning.
				hotPixelLearningMap.clear();
				hotPixelLearningMap.shrink_to_fit();

				hotPixelLearningStarted = false;

				config.set<dv::CfgType::BOOL>("hotPixelLearn", false);

				log.debug << "HotPixel Learning: completed on ts=" << evt.timestamp() << "." << std::endl;
			}
		}

		// Hot Pixel filter: filter out abnormally active pixels by their address.
		if (hotPixelEnabled) {
			bool filteredOut = std::any_of(hotPixelArray.cbegin(), hotPixelArray.cend(), [&evt](const auto &px) {
				return ((evt.x() == px.x) && (evt.y() == px.y));
			});

			if (filteredOut) {
				if (evt.polarity()) {
					hotPixelStatOn++;
				}
				else {
					hotPixelStatOff++;
				}

				// Go to next event, don't execute other filters and don't
				// update timestamps map. Hot pixels don't provide any useful
				// timing information, as they are repeating noise.
				continue;
			}
		}

		// Refractory Period filter.
		// Execute before BAFilter, as this is a much simpler check, so if we
		// can we try to eliminate the event early in a less costly manner.
		if (refractoryPeriodEnabled) {
			if ((evt.timestamp() - GET_TS(timestampsMap[pixelIndex])) < refractoryPeriodTime) {
				if (evt.polarity()) {
					refractoryPeriodStatOn++;
				}
				else {
					refractoryPeriodStatOff++;
				}

				// Invalid, jump.
				goto WriteTimestamp;
			}
		}

		if (backgroundActivityEnabled) {
			size_t supportPixelNum = doBackgroundActivityLookup(evt.x(), evt.y(), pixelIndex, evt.timestamp(),
				evt.polarity(), supportPixelIndexes.data(), backgroundActivityTime, backgroundActivityCheckPolarity);

			bool filteredOut = true;

			if ((supportPixelNum >= static_cast<size_t>(backgroundActivitySupportMin))
				&& (supportPixelNum <= static_cast<size_t>(backgroundActivitySupportMax))) {
				if (backgroundActivityTwoLevels) {
					// Do the check again for all previously discovered supporting pixels.
					for (size_t i = 0; i < supportPixelNum; i++) {
						size_t supportPixelIndex = supportPixelIndexes[i];
						int16_t supportPixelX    = static_cast<int16_t>(supportPixelIndex % static_cast<size_t>(sizeX));
						int16_t supportPixelY    = static_cast<int16_t>(supportPixelIndex / static_cast<size_t>(sizeX));

						if (doBackgroundActivityLookup(supportPixelX, supportPixelY, supportPixelIndex, evt.timestamp(),
								evt.polarity(), nullptr, backgroundActivityTime, backgroundActivityCheckPolarity)
							> 0) {
							filteredOut = false;
							break;
						}
					}
				}
				else {
					filteredOut = false;
				}
			}

			if (filteredOut) {
				if (evt.polarity()) {
					backgroundActivityStatOn++;
				}
				else {
					backgroundActivityStatOff++;
				}

				// Invalid, jump.
				goto WriteTimestamp;
			}
		}

		// Valid event.
		evt_out.emplace_back(evt);

	WriteTimestamp:
		// Update pixel timestamp (one write). Always update so filters are
		// ready at enable-time right away.
		timestampsMap[pixelIndex] = SET_TSPOL(evt.timestamp(), evt.polarity());
	}

	evt_out.commit();

	// Update statistics.
	config.set<dv::CfgType::LONG>("statistics/hotPixelFilteredOn", hotPixelStatOn);
	config.set<dv::CfgType::LONG>("statistics/hotPixelFilteredOff", hotPixelStatOff);
	config.set<dv::CfgType::LONG>("statistics/backgroundActivityFilteredOn", backgroundActivityStatOn);
	config.set<dv::CfgType::LONG>("statistics/backgroundActivityFilteredOff", backgroundActivityStatOff);
	config.set<dv::CfgType::LONG>("statistics/refractoryPeriodFilteredOn", refractoryPeriodStatOn);
	config.set<dv::CfgType::LONG>("statistics/refractoryPeriodFilteredOff", refractoryPeriodStatOff);
}

inline size_t DVSNoiseFilter::doBackgroundActivityLookup(int16_t x, int16_t y, size_t pixelIndex, int64_t timestamp,
	bool polarity, size_t *supportIndexes, int32_t backgroundActivityTime, bool backgroundActivityCheckPolarity) {
	// Compute map limits.
	bool notBorderLeft  = (x != 0);
	bool notBorderDown  = (y != (sizeY - 1));
	bool notBorderRight = (x != (sizeX - 1));
	bool notBorderUp    = (y != 0);

	// Background Activity filter: if difference between current timestamp
	// and stored neighbor timestamp is smaller than given time limit, it
	// means the event is supported by a neighbor and thus valid. If it is
	// bigger, then the event is not supported, and we need to check the
	// next neighbor. If all are bigger, the event is invalid.
	size_t result = 0;

	{
		if (notBorderLeft) {
			pixelIndex--;

			if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
				if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
					if (supportIndexes != nullptr) {
						supportIndexes[result] = pixelIndex;
					}
					result++;
				}
			}

			pixelIndex++;
		}

		if (notBorderRight) {
			pixelIndex++;

			if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
				if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
					if (supportIndexes != nullptr) {
						supportIndexes[result] = pixelIndex;
					}
					result++;
				}
			}

			pixelIndex--;
		}
	}

	if (notBorderUp) {
		pixelIndex -= static_cast<size_t>(sizeX); // Previous row.

		if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
			if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
				if (supportIndexes != nullptr) {
					supportIndexes[result] = pixelIndex;
				}
				result++;
			}
		}

		if (notBorderLeft) {
			pixelIndex--;

			if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
				if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
					if (supportIndexes != nullptr) {
						supportIndexes[result] = pixelIndex;
					}
					result++;
				}
			}

			pixelIndex++;
		}

		if (notBorderRight) {
			pixelIndex++;

			if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
				if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
					if (supportIndexes != nullptr) {
						supportIndexes[result] = pixelIndex;
					}
					result++;
				}
			}

			pixelIndex--;
		}

		pixelIndex += static_cast<size_t>(sizeX); // Back to middle row.
	}

	if (notBorderDown) {
		pixelIndex += static_cast<size_t>(sizeX); // Next row.

		if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
			if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
				if (supportIndexes != nullptr) {
					supportIndexes[result] = pixelIndex;
				}
				result++;
			}
		}

		if (notBorderLeft) {
			pixelIndex--;

			if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
				if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
					if (supportIndexes != nullptr) {
						supportIndexes[result] = pixelIndex;
					}
					result++;
				}
			}

			pixelIndex++;
		}

		if (notBorderRight) {
			pixelIndex++;

			if ((timestamp - GET_TS(timestampsMap[pixelIndex])) < backgroundActivityTime) {
				if (!backgroundActivityCheckPolarity || polarity == GET_POL(timestampsMap[pixelIndex])) {
					if (supportIndexes != nullptr) {
						supportIndexes[result] = pixelIndex;
					}
					result++;
				}
			}
		}
	}

	return (result);
}

void DVSNoiseFilter::hotPixelGenerateArray() {
	std::vector<pixel_with_count> hotPixels;

	int32_t hotPixelCount = config.get<dv::CfgType::INT>("hotPixelCount");

	// Find abnormally active pixels.
	for (size_t i = 0; i < hotPixelLearningMap.size(); i++) {
		if (hotPixelLearningMap[i] >= hotPixelCount) {
			hotPixels.emplace_back(static_cast<int16_t>(i % static_cast<size_t>(sizeX)),
				static_cast<int16_t>(i / static_cast<size_t>(sizeX)), hotPixelLearningMap[i]);
		}
	}

	// Sort in descending order by count.
	std::sort(hotPixels.begin(), hotPixels.end(), [](const pixel_with_count &a, const pixel_with_count &b) {
		return (a.count > b.count);
	});

	// Update manual hot pixel list.
	std::string pixelsList;

	for (const auto &px : hotPixels) {
		pixelsList += fmt::format("{:d},{:d}|", px.x, px.y);
	}

	// Prefix with origin description.
	if (!pixelsList.empty()) {
		pixelsList = inputs.getEventInput("events").getOriginDescription() + ':' + pixelsList;
	}

	config.setString("hotPixelList", pixelsList);

	// Print list of hot pixels for debugging.
	for (size_t i = 0; i < hotPixels.size(); i++) {
		log.debug.format(
			"HotPixel found {:d}: X={:d}, Y={:d}, count={:d}.", i, hotPixels[i].x, hotPixels[i].y, hotPixels[i].count);
	}
}

void DVSNoiseFilter::hotPixelLoadConfig() {
	// Clear out current values.
	hotPixelArray.clear();

	// Parse hot-pixels list and reset search array.
	auto hotPixelList = config.getString("hotPixelList");

	// Search for origin description: if present, only apply list if origin
	// is the same, if not present always apply list.
	const auto colonPos = hotPixelList.find(":");

	if (colonPos != std::string::npos) {
		const auto origin = hotPixelList.substr(0, colonPos);

		if (inputs.getEventInput("events").getOriginDescription() != origin) {
			log.info.format("Hot-pixels list origin '{}' differs from current origin '{}', not loading hot-pixels "
							"list. Please learn new hot-pixels.",
				origin, inputs.getEventInput("events").getOriginDescription());
			return;
		}

		hotPixelList = hotPixelList.substr(colonPos + 1);
	}

	boost::tokenizer<boost::char_separator<char>> hotPixels(hotPixelList, boost::char_separator<char>("|", nullptr));

	for (const auto &px : hotPixels) {
		if (px.empty()) {
			continue;
		}

		// Split on comma.
		const auto commaPos = px.find(",");

		if (commaPos == std::string::npos) {
			// Invalid token.
			continue;
		}

		const auto x = std::stol(px.substr(0, commaPos));
		const auto y = std::stol(px.substr(commaPos + 1));

		hotPixelArray.emplace_back(x, y);
	}

	hotPixelArray.shrink_to_fit();

	// Print list of hot pixels for debugging.
	for (size_t i = 0; i < hotPixelArray.size(); i++) {
		log.debug.format("HotPixel loaded {:d}: X={:d}, Y={:d}.", i, hotPixelArray[i].x, hotPixelArray[i].y);
	}
}

void DVSNoiseFilter::configUpdate() {
	// Configuration changed, reload hot-pixels from manual list.
	hotPixelLoadConfig();
}
