#include "../../include/dv-processing/data/generate.hpp"
#include "../../include/dv-processing/optimization/contrast_maximization_translation_and_depth.hpp"
#include "../../include/dv-processing/optimization/contrast_maximization_wrapper.hpp"

#include <benchmark/benchmark.h>

#include <random>

static const size_t numberOfTranslations = 5;
static const cv::Size resolution(640, 480);
static const size_t numberOfLearningRates = 10;
static const float initialLearningRate    = 1e-1f;
static const float learningRateStep       = 5e-2f;

[[nodiscard]] static dv::EventStore generateTestEvents() {
	dv::EventStore rectangle = dv::data::generate::eventRectangle(10000LL, cv::Point(50, 50), cv::Point(100, 100));

	dv::EventStore events;
	// We translate the line in x direction
	dv::EventStore translatedEvents;
	cv::Point2i translation(5, 0);
	for (int i = 0; i < numberOfTranslations; i++) {
		translatedEvents  = dv::EventStore();
		int64_t timestamp = rectangle.getLowestTime() + (1000 * i);
		for (const auto &event : rectangle) {
			translatedEvents.emplace_back(
				timestamp, event.x() + translation.x * i, event.y() + translation.y * i, event.polarity());
		}
		events.add(translatedEvents);
	}

	return events;
}

static void bmContrastMaximization(benchmark::State &state) {
	using TranslationLoss = dv::optimization::TranslationLossFunctor;

	// Initialize test data
	const auto store = generateTestEvents();

	// Initialize parameters used for contrast maximization
	auto cameraGeometry
		= std::make_shared<dv::camera::CameraGeometry>(resolution.width, resolution.height, 50., 50., resolution);
	static constexpr float contribution              = 0.02f;
	static constexpr int numberOfVariablesToOptimize = 4;
	static constexpr int numberOfMeasurementsForGradient
		= numberOfVariablesToOptimize; // Must be >= numberOfVariablesToOptimize
	const auto learningRate = static_cast<float>(state.range(0) * learningRateStep + initialLearningRate);
	const Eigen::Vector4f initTranslationDepthEig{0.f, 0.f, 0.f, 3.f};

	// Run benchmark and report number of processed events
	size_t counter = 0;
	for (auto _ : state) {
		dv::optimization::ContrastMaximizationWrapper<TranslationLoss> contrastMaximization(
			std::make_unique<TranslationLoss>(
				cameraGeometry, store, contribution, numberOfVariablesToOptimize, numberOfMeasurementsForGradient),
			learningRate);
		auto optimizationOutput = contrastMaximization.optimize(initTranslationDepthEig);
		if (!optimizationOutput.optimizationSuccessful) {
			throw std::runtime_error("Optimization did not converge");
		}
		counter += store.size();
	}
	state.SetItemsProcessed(static_cast<int64_t>(counter));
	state.SetLabel("events");
}

int main(int argc, char **argv) {
	BENCHMARK(bmContrastMaximization)->DenseRange(0, numberOfLearningRates);

	::benchmark::Initialize(&argc, argv);
	if (::benchmark::ReportUnrecognizedArguments(argc, argv)) {
		return 1;
	}
	::benchmark::RunSpecifiedBenchmarks();

	return EXIT_SUCCESS;
}
