/* +------------------------------------------------------------------------+
   |                     Mobile Robot Programming Toolkit (MRPT)            |
   |                          https://www.mrpt.org/                         |
   |                                                                        |
   | Copyright (c) 2005-2023, Individual contributors, see AUTHORS file     |
   | See: https://www.mrpt.org/Authors - All rights reserved.               |
   | Released under BSD License. See: https://www.mrpt.org/License          |
   +------------------------------------------------------------------------+ */

#ifdef __EMSCRIPTEN__  // are we actually in a Wasm JS instance?
#include <emscripten.h>
#include <emscripten/html5.h>

#define GL_GLEXT_PROTOTYPES
#define EGL_EGLEXT_PROTOTYPES
#include <GLFW/glfw3.h>

#include <functional>

std::function<void()> loop;
void main_loop() { loop(); }
#endif

#include <mrpt/core/exceptions.h>
#include <mrpt/core/lock_helper.h>
#include <mrpt/gui/CDisplayWindowGUI.h>
#include <mrpt/io/lazy_load_path.h>
#include <mrpt/obs/CObservation3DRangeScan.h>
#include <mrpt/obs/CRawlog.h>
#include <mrpt/opengl/CAxis.h>
#include <mrpt/opengl/CGridPlaneXY.h>
#include <mrpt/opengl/CPointCloudColoured.h>
#include <mrpt/opengl/stock_objects.h>

#include <iostream>

#if MRPT_HAS_NANOGUI

constexpr int SMALL_FONT_SIZE = 15;
constexpr int LARGE_FONT_SIZE = 20;

// We need to define all variables here to avoid emscripten memory errors
// if they are defined as *locals* in the function.
struct AppData
{
	AppData() = default;

	mrpt::gui::CDisplayWindowGUI::Ptr win;
	mrpt::opengl::CAxis::Ptr gl_corner_reference;
	mrpt::obs::CObservation3DRangeScan::Ptr obs;

	mrpt::img::TCamera originalCalib;
	mrpt::poses::CPose3D originalSensorPose;

	mrpt::opengl::CPointCloudColoured::Ptr gl_pts;
	mrpt::opengl::CSetOfObjects::Ptr gl_sensorPoseCorner;

	std::vector<nanogui::TextBox*> edBoxes /* cx,cy,... dist[7]*/
		,
		edBoxesExtrinsics /*x,y,... pitch, roll*/;

	std::vector<nanogui::Slider*> sliders /* cx,cy,... dist[7] */,
		slidersExtrinsics /*x,y,... pitch, roll*/;

	nanogui::ComboBox* cbDistType;

	void obs_params_to_gui();
	void gui_params_to_obs();
};

static AppData app;

// The main function: update all graphs
static void recalcAll()
{
	auto lck = mrpt::lockHelper(app.win->background_scene_mtx);

	if (!app.obs || !app.obs->hasRangeImage)
	{
		app.gl_pts->clear();
		return;
	}

	mrpt::obs::T3DPointsProjectionParams pp;
	pp.takeIntoAccountSensorPoseOnRobot = true;

	app.obs->unprojectInto(*app.gl_pts, pp);

	const auto bbox = app.gl_pts->getBoundingBox();

	app.gl_pts->recolorizeByCoordinate(
		bbox.min.z, bbox.max.z, 2 /*0,1,2=x,y,z*/);

	app.gl_pts->setPointSize(2.0f);

	app.gl_sensorPoseCorner->setPose(app.obs->sensorPose);
};

void AppData::obs_params_to_gui()
{
	ASSERT_(obs);
	ASSERT_EQUAL_(edBoxes.size(), sliders.size());
	ASSERT_EQUAL_(edBoxes.size(), 12U);

	const auto& p = obs->cameraParams;

	edBoxes[0]->setValue(mrpt::format("%.04f", p.cx()));
	edBoxes[1]->setValue(mrpt::format("%.04f", p.cy()));
	edBoxes[2]->setValue(mrpt::format("%.04f", p.fx()));
	edBoxes[3]->setValue(mrpt::format("%.04f", p.fy()));

	cbDistType->setSelectedIndex(static_cast<int>(p.distortion));

	for (int i = 0; i < 8; i++)
		edBoxes[4 + i]->setValue(mrpt::format("%.04e", p.dist[i]));

	// slider ranges:
	const double maxPixelRanges =
		std::max(std::max(p.cx(), p.cy()), std::max(p.fx(), p.fy()));
	const double minPixelRanges =
		std::min(std::min(p.cx(), p.cy()), std::min(p.fx(), p.fy()));
	for (int i = 0; i < 4; i++)
	{
		sliders[i]->setRange({0.75 * minPixelRanges, 1.5 * maxPixelRanges});
	}
	for (int i = 4; i < 12; i++)
	{
		const double absVal =
			std::max(1e-3, std::abs(std::stod(edBoxes[i]->value())));
		sliders[i]->setRange({-2.0 * absVal, 2.0 * absVal});
	}

	// slider values:
	for (size_t i = 0; i < edBoxes.size(); i++)
		sliders[i]->setValue(std::stod(edBoxes[i]->value()));

	// Extrinsics ranges:
	for (int i = 0; i < 3; i++)
		slidersExtrinsics[i]->setRange({-2.0, 2.0});

	slidersExtrinsics[3]->setRange({-180.0, 180.0});
	slidersExtrinsics[4]->setRange({-90.0, 90.0});
	slidersExtrinsics[5]->setRange({-180.0, 180.0});

	// Extrinsics values:
	for (int i = 0; i < 6; i++)
	{
		double val = obs->sensorPose[i];
		if (i > 3) val = mrpt::RAD2DEG(val);

		edBoxesExtrinsics[i]->setValue(mrpt::format("%.04f", val));
		slidersExtrinsics[i]->setValue(val);
	}
}

void AppData::gui_params_to_obs()
{
	ASSERT_(obs);
	ASSERT_EQUAL_(edBoxes.size(), sliders.size());
	ASSERT_EQUAL_(edBoxes.size(), 12U);

	auto& p = obs->cameraParams;
	p.distortion = mrpt::img::DistortionModel::plumb_bob;

	p.cx(std::stod(edBoxes[0]->value()));
	p.cy(std::stod(edBoxes[1]->value()));
	p.fx(std::stod(edBoxes[2]->value()));
	p.fy(std::stod(edBoxes[3]->value()));

	for (int i = 0; i < 8; i++)
		p.dist.at(i) = std::stod(edBoxes[4 + i]->value());

	p.distortion = mrpt::typemeta::str2enum<mrpt::img::DistortionModel>(
		cbDistType->items().at((cbDistType->selectedIndex())));

	// Extrinsics:
	app.obs->sensorPose = mrpt::poses::CPose3D::FromXYZYawPitchRoll(
		std::stod(edBoxesExtrinsics[0]->value()),  // x
		std::stod(edBoxesExtrinsics[1]->value()),  // y
		std::stod(edBoxesExtrinsics[2]->value()),  // z
		mrpt::DEG2RAD(std::stod(edBoxesExtrinsics[3]->value())),  // yaw
		mrpt::DEG2RAD(std::stod(edBoxesExtrinsics[4]->value())),  // pitch
		mrpt::DEG2RAD(std::stod(edBoxesExtrinsics[5]->value()))	 // roll
	);
}

static void AppDepthCamDemo()
{
	nanogui::init();

	// Create main window:
	mrpt::gui::CDisplayWindowGUI_Params cp;
	// cp.fullscreen = true;

	// Init values:
	app.gl_corner_reference = mrpt::opengl::CAxis::Create(
		-5.0f, -5.0f, .0f, 5.0f, 5.0f, 2.0f, 0.2f, 1.0f, true);
	app.gl_corner_reference->setTextScale(0.04);

	app.gl_pts = mrpt::opengl::CPointCloudColoured::Create();
	app.gl_sensorPoseCorner =
		mrpt::opengl::stock_objects::CornerXYZSimple(0.10f);

	// Create GUI:
	app.win = mrpt::gui::CDisplayWindowGUI::Create(
		"Depth camera distortion demo", 1000, 700, cp);

	// Add main window:
	// -----------------------------
	nanogui::Window* w = new nanogui::Window(&(*app.win), "Camera parameters");
	w->setPosition(nanogui::Vector2i(10, 50));
	w->setLayout(new nanogui::GridLayout(
		nanogui::Orientation::Horizontal, 1, nanogui::Alignment::Fill, 5, 0));
	w->setFixedWidth(350);

	// Load button:
	w->add<nanogui::Label>("Operations")->setFontSize(LARGE_FONT_SIZE);

	w->add<nanogui::Button>("Load depth image from rawlog...")
		->setCallback([]() {
			try
			{
				const std::string loadFile = nanogui::file_dialog(
					{{"rawlog", "MRPT rawlog dataset files"},
					 {"*", "Any file"}},
					false);
				if (loadFile.empty()) return;

				// Enable MRPT to find externalized rawlogs:
				mrpt::io::setLazyLoadPathBase(
					mrpt::obs::CRawlog::detectImagesDirectory(loadFile));

				mrpt::obs::CRawlog rawlog;
				if (!rawlog.loadFromRawLogFile(loadFile)) return;

				ASSERT_(!rawlog.empty());

				mrpt::obs::CObservation3DRangeScan::Ptr obs;

				// Obs-only format:
				if (rawlog.getType(0) == mrpt::obs::CRawlog::etObservation)
				{
					obs = std::dynamic_pointer_cast<
						mrpt::obs::CObservation3DRangeScan>(
						rawlog.getAsObservation(0));
					ASSERTMSG_(
						obs,
						"First observation in rawlog is not "
						"CObservation3DRangeScan");
				}
				else
				{
					// SF format:
					ASSERT_(
						rawlog.getType(0) ==
						mrpt::obs::CRawlog::etSensoryFrame);

					obs = rawlog.getAsObservations(0)
							  ->getObservationByClass<
								  mrpt::obs::CObservation3DRangeScan>();
					ASSERT_(obs);
				}

				app.obs = obs;
				app.originalCalib = obs->cameraParams;
				app.originalSensorPose = obs->sensorPose;
				app.obs_params_to_gui();
				recalcAll();
			}
			catch (const std::exception& e)
			{
				std::cerr << e.what() << std::endl;
				// auto dlg =
				new nanogui::MessageDialog(
					app.win->screen(), nanogui::MessageDialog::Type::Warning,
					"Exception", e.what());
			}
		});

	w->add<nanogui::Button>("Revert to original calibration")
		->setCallback([]() {
			if (!app.obs) return;
			app.obs->cameraParams = app.originalCalib;
			app.obs->sensorPose = app.originalSensorPose;
			app.obs_params_to_gui();
			recalcAll();
		});

	// Sliders and value controls:
	w->add<nanogui::Label>("Intrinsics")->setFontSize(LARGE_FONT_SIZE);
	;
	{
		auto pn = w->add<nanogui::Widget>();
		pn->setLayout(new nanogui::GridLayout(
			nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5,
			0));

		pn->add<nanogui::Label>("Distortion model");
		app.cbDistType = pn->add<nanogui::ComboBox>(
			std::vector<std::string>({"none", "plumb_bob", "kannala_brandt"}));
		pn->add<nanogui::Label>(" ");

		app.cbDistType->setCallback([](int) {
			app.gui_params_to_obs();
			recalcAll();
		});

		std::vector<std::string> lb = {
			"cx=",			 "cy=",			  "fx=",		   "fy=",
			"dist[0] (k1)=", "dist[1] (k2)=", "dist[2] (t1)=", "dist[3] (t2)=",
			"dist[4] (k3)=", "dist[5] (k4)=", "dist[6] (k5)=", "dist[7] (k6)=",
		};

		for (size_t i = 0; i < lb.size(); i++)
		{
			pn->add<nanogui::Label>(lb[i])->setFontSize(SMALL_FONT_SIZE);

			nanogui::TextBox* ed = pn->add<nanogui::TextBox>(" ");
			app.edBoxes.push_back(ed);

			ed->setEditable(true);
			ed->setFormat("[-+]?[0-9.e+-]*");
			ed->setCallback([](const std::string&) {
				app.gui_params_to_obs();
				recalcAll();
				return true;
			});
			nanogui::Slider* sl = pn->add<nanogui::Slider>();
			app.sliders.push_back(sl);

			sl->setRange({-1.0f, 1.0f});
			sl->setCallback([&, ed](float val) {
				ed->setValue(mrpt::format("%.03g", val));
				ed->callback()(ed->value());
			});

			ed->setFontSize(SMALL_FONT_SIZE);
		}
	}

	// Extrinsics:
	w->add<nanogui::Label>("Extrinsics")->setFontSize(LARGE_FONT_SIZE);
	{
		auto pn = w->add<nanogui::Widget>();
		pn->setLayout(new nanogui::GridLayout(
			nanogui::Orientation::Horizontal, 3, nanogui::Alignment::Fill, 5,
			0));

		std::vector<std::string> lb = {
			"x=", "y=", "z=", "yaw=", "pitch=", "roll="};

		for (size_t i = 0; i < lb.size(); i++)
		{
			pn->add<nanogui::Label>(lb[i])->setFontSize(SMALL_FONT_SIZE);

			nanogui::TextBox* ed = pn->add<nanogui::TextBox>(" ");
			app.edBoxesExtrinsics.push_back(ed);

			ed->setEditable(true);
			ed->setFormat("[-+]?[0-9.e+-]*");
			ed->setCallback([](const std::string&) {
				app.gui_params_to_obs();
				recalcAll();
				return true;
			});
			ed->setFontSize(SMALL_FONT_SIZE);

			nanogui::Slider* sl = pn->add<nanogui::Slider>();
			app.slidersExtrinsics.push_back(sl);

			sl->setRange({-1.0f, 1.0f});
			sl->setCallback([&, ed](float val) {
				ed->setValue(mrpt::format("%.03g", val));
				ed->callback()(ed->value());
			});
		}
	}

	// Add top menu subwindow:
	// -----------------------------
	{
		nanogui::ref<nanogui::Window> winMenu =
			new nanogui::Window(&(*app.win), "");
		winMenu->setPosition(nanogui::Vector2i(0, 0));
		winMenu->setLayout(new nanogui::BoxLayout(
			nanogui::Orientation::Horizontal, nanogui::Alignment::Middle, 5));
		nanogui::Theme* modTheme =
			new nanogui::Theme(app.win->screen()->nvgContext());
		modTheme->mWindowHeaderHeight = 1;
		winMenu->setTheme(modTheme);

#ifndef __EMSCRIPTEN__	// are we actually in a Wasm JS instance?
		winMenu->add<nanogui::Button>("Quit", ENTYPO_ICON_ARROW_BOLD_LEFT)
			->setCallback([]() { app.win->setVisible(false); });
#endif

		winMenu->add<nanogui::Label>("      ");	 // separator

		winMenu
			->add<nanogui::CheckBox>(
				"Show reference frame",
				[&](bool b) { app.gl_corner_reference->setVisibility(b); })
			->setChecked(true);

		winMenu
			->add<nanogui::CheckBox>(
				"Ortho. view",
				[&](bool b) { app.win->camera().setCameraProjective(!b); })
			->setChecked(false);
	}

	// Add a background scene:
	// -----------------------------
	{
		auto scene = mrpt::opengl::Scene::Create();
		scene->insert(mrpt::opengl::CGridPlaneXY::Create());

		scene->insert(app.gl_corner_reference);
		scene->insert(app.gl_pts);
		scene->insert(app.gl_sensorPoseCorner);

		auto lck = mrpt::lockHelper(app.win->background_scene_mtx);
		app.win->background_scene = std::move(scene);
	}

	app.win->performLayout();

	app.win->camera().setZoomDistance(7.0f);

	// Update view and process events:
	app.win->drawAll();
	app.win->setVisible(true);

#if !defined(__EMSCRIPTEN__)  // are we in a Wasm JS instance?
	// No: regular procedure:

	nanogui::mainloop();
	nanogui::shutdown();
#else

	// Yes, we are in a web browser running on JS:
	loop = [] {
		// Check if any events have been activated (key pressed, mouse moved
		// etc.) and call corresponding response functions
		glfwPollEvents();

		glClearColor(0.2f, 0.25f, 0.3f, 1.0f);
		glClear(GL_COLOR_BUFFER_BIT);

		// Draw nanogui
		app.win->drawContents();
		app.win->drawWidgets();

		glfwSwapBuffers(app.win->glfwWindow());
	};
	emscripten_set_main_loop(main_loop, 0, true);

	glfwDestroyWindow(app.win->nanogui_screen()->glfwWindow());

	// Terminate GLFW, clearing any resources allocated by GLFW.
	glfwTerminate();

#endif
}

#endif	// MRPT_HAS_NANOGUI

int main()
{
	try
	{
#if MRPT_HAS_NANOGUI
		AppDepthCamDemo();
#else
		std::cerr << "This example requires MRPT built with NANOGUI.\n";
#endif
		return 0;
	}
	catch (const std::exception& e)
	{
		std::cerr << mrpt::exception_to_str(e) << std::endl;
		return -1;
	}
}
