/* +------------------------------------------------------------------------+
   |                     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          |
   +------------------------------------------------------------------------+ */

// This seems to be assumed by OpenNI.h and undefined for some reason in
// GCC/Ubuntu
#if !defined(MRPT_OS_WINDOWS)
#define linux 1
#endif
#include <OpenNI.h>
#include <PS1080.h>
#include <mrpt/img/CImage.h>
#include <mrpt/io/CFileGZOutputStream.h>
#include <mrpt/obs/CObservation3DRangeScan.h>
#include <mrpt/opengl.h>
#include <mrpt/serialization/CArchive.h>
#include <mrpt/system/CTicTac.h>
#include <mrpt/system/filesystem.h>
#include <mrpt/system/os.h>

#include <iostream>

using namespace mrpt;
using namespace mrpt::serialization;
using namespace std;

// This simple demo records form an OpenNI2 device into a rawlog as 3D
// observations. It's meant as a temporary workaround before OpenNI2 is
// integrated as a generic sensor so that it works with rawlog-grabber.

bool setONI2StreamMode(
	openni::VideoStream& stream, int w, int h, int fps,
	openni::PixelFormat format);
bool initONI2Stream(
	openni::Device& device, openni::SensorType sensorType,
	openni::VideoStream& stream, int w, int h, int fps,
	openni::PixelFormat format);
bool synchONI2MirrorMode(
	openni::VideoStream& rgb, openni::VideoStream& depth, bool& isMirrorMode);
bool startONI2Streams(openni::VideoStream& rgb, openni::VideoStream& depth);

int main(int argc, char** argv)
{
	bool manual_mode = false;
	//						Properties
	//========================================================================================
	const unsigned width = 640, height = 480, fps = 30;

	try
	{
		if (argc < 2)
		{
			cerr << "Usage: " << argv[0] << " OUTPUT_PATH [-m]\n";
			cerr << "Example: " << argv[0]
				 << " /home/jdoe/myrawlog will "
					"produce file myrawlog.rawlog in "
					"folder /home/jdoe\n";
			cerr << "-m will activate manual mode\n";
			return 1;
		}
		else if ((argc == 3) && (string(argv[2]) == "-m"))
			manual_mode = true;

		openni::Status rc = openni::STATUS_OK;

		openni::Device device;
		openni::VideoMode options;
		openni::VideoStream depth, rgb;

		//									Device is opened
		//=======================================================================================
		const char* deviceURI = openni::ANY_DEVICE;

		rc = openni::OpenNI::initialize();

		if (rc != openni::STATUS_OK)
		{
			printf(
				"After initialization:\n %s\n",
				openni::OpenNI::getExtendedError());
		}

		rc = device.open(deviceURI);

		// cout << endl << "Do we have IR sensor? " <<
		// device.hasSensor(openni::SENSOR_IR);
		// cout << endl << "Do we have RGB sensor? " <<
		// device.hasSensor(openni::SENSOR_COLOR);
		// cout << endl << "Do we have Depth sensor? " <<
		// device.hasSensor(openni::SENSOR_DEPTH);

		if (rc != openni::STATUS_OK)
		{
			printf(
				"Device open failed:\n%s\n",
				openni::OpenNI::getExtendedError());
			openni::OpenNI::shutdown();
			return 1;
		}

		//								Create RGB and Depth channels
		//========================================================================================
		if (initONI2Stream(
				device, openni::SENSOR_COLOR, rgb, width, height, fps,
				openni::PIXEL_FORMAT_RGB888) == false)
		{ return 1; }
		if (initONI2Stream(
				device, openni::SENSOR_DEPTH, depth, width, height, fps,
				openni::PIXEL_FORMAT_DEPTH_1_MM) == false)
		{ return 1; }
		bool isMirror;
		if (synchONI2MirrorMode(rgb, depth, isMirror) == false) { return 1; }
		if (startONI2Streams(rgb, depth) == false) { return 1; }
		if (device.isImageRegistrationModeSupported(
				openni::IMAGE_REGISTRATION_DEPTH_TO_COLOR))
		{
			rc = device.setImageRegistrationMode(
				openni::IMAGE_REGISTRATION_DEPTH_TO_COLOR);
		}
		else
		{
			cout << "Device doesn't do image registration!" << endl;
		}

		if (!depth.isValid() || !rgb.isValid())
		{
			printf("No valid streams. Exiting\n");
			openni::OpenNI::shutdown();
			return 2;
		}

		if (rc != openni::STATUS_OK)
		{
			openni::OpenNI::shutdown();
			return 3;
		}

		//	if (device.setDepthColorSyncEnabled(true) == openni::STATUS_OK)
		//		cout << "setDepthColorSyncEnabled" << endl;
		//	else
		//		cout << "setDepthColorSyncEnabled failed!" << endl;

		// rc = device.setImageRegistrationMode(openni::IMAGE_REGISTRATION_OFF);

		// Allow detection of closer points (although they will flicker)
		// bool CloseRange;
		// depth.setProperty(XN_STREAM_PROPERTY_CLOSE_RANGE, 1);
		// depth.getProperty(XN_STREAM_PROPERTY_CLOSE_RANGE, &CloseRange);
		// printf("\nClose range: %s", CloseRange?"On":"Off");

		//							Output files setup
		//=================depth=======================================================================
		// slam::CColouredPointsMap points;
		openni::VideoFrameRef framed, framergb;

		const string out_name = string(argv[1]);

		// Create output directory for images ------------------------------
		const string out_img_dir = out_name + string("_Images");

		cout << "Creating images directory: " << out_img_dir << endl;
		mrpt::system::createDirectory(out_img_dir);

		// Create rawlog file ----------------------------------------------
		const string out_rawlog_fil = out_name + string(".rawlog");
		cout << "Creating rawlog: " << out_rawlog_fil << endl;
		mrpt::io::CFileGZOutputStream f_out(out_rawlog_fil);

		// Fill out the common field to all entries:
		mrpt::obs::CObservation3DRangeScan obs;
		obs.sensorLabel = "OpenNI2";
		obs.hasConfidenceImage = false;

		// Kinect style: ranges are actually depth values, not Euclidean
		// distances.
		obs.range_is_depth = true;

		// Calculate depth camera parameters (from pcl WIP)
		float hFov = depth.getHorizontalFieldOfView();
		float fx = width / (2.0f * tan(hFov / 2.0f));
		float vFov = depth.getVerticalFieldOfView();
		float fy = height / (2.0f * tan(vFov / 2.0f));

		printf(
			"Got Fov from the device, calculated fx and fy:\n"
			"hFov:%f\n"
			"vFov:%f\n"
			"fx:%f\n"
			"fy:%f\n",
			hFov, vFov, fx, fy);

		obs.cameraParams.nrows = width;
		obs.cameraParams.ncols = height;
		obs.cameraParams.fx(fx);
		obs.cameraParams.fy(fy);
		obs.cameraParams.cx((width - 1) * 0.5);
		obs.cameraParams.cy((height - 1) * 0.5);
		obs.cameraParamsIntensity = obs.cameraParams;

		// Images are already registered from OpenNI2
		obs.relativePoseIntensityWRTDepth =
			mrpt::poses::CPose3D(0, 0, 0, 0, 0, 0);
		obs.hasPoints3D = false;

		//							Grab frames and record
		//========================================================================================

		mrpt::system::CTicTac tictac;
		unsigned int frame = 0;

		bool depthfirst = true;
		bool sampleinsync = false;

		if (manual_mode)
			cout << "Press any key to grab a frame. 'ESC' to stop\n";
		else
			cout << "Press any key to stop\n";

		while ((!system::os::kbhit() && !manual_mode) || manual_mode)
		{
			if (manual_mode)
			{
				if (system::os::getch() == 27) break;
			}

			frame++;

			// This loop makes sure our samples of depth and RGB are in sync.
			// It should only iterate once if samples are being recorded out of
			// sync.
			sampleinsync = false;
			while (!sampleinsync)
			{
				tictac.Tic();
				if (depthfirst) depth.readFrame(&framed);
				else
					rgb.readFrame(&framergb);
				double t1 = tictac.Tac() * 1000;

				tictac.Tic();
				if (depthfirst) rgb.readFrame(&framergb);
				else
					depth.readFrame(&framed);

				double t2 = tictac.Tac() * 1000;
				cout << "frame " << frame << ". t1: " << t1 << ", t2: " << t2;

				if (t2 > t1)
				{
					cout << " --- out of sync";
					depthfirst = !depthfirst;
				}
				else
				{
					cout << " --- in sync    ";
					sampleinsync = true;
				}

				cout << "\r";
			}

			if ((framed.getWidth() != framergb.getWidth()) ||
				(framed.getHeight() != framergb.getHeight()))
			{ cout << "\nBoth frames don't have the same size."; }
			else
			{
				// Read one frame
				const char* pDepthRow = (const char*)framed.getData();
				const char* pRgbRow = (const char*)framergb.getData();
				int d_step = framed.getStrideInBytes();
				int rgb_step = framergb.getStrideInBytes();

				obs.timestamp = mrpt::system::getCurrentTime();

				obs.hasRangeImage = true;
				obs.rangeImage_forceResetExternalStorage();
				obs.rangeImage_setSize(height, width);

				obs.hasIntensityImage = true;
				obs.intensityImageChannel =
					mrpt::obs::CObservation3DRangeScan::CH_VISIBLE;

				const string sDepthfile = mrpt::format("%i_depth.bin", frame);
				const string sRGBfile = mrpt::format("%i_rgb.png", frame);

				img::CImage iimage(width, height, CH_RGB);
				for (unsigned int yc = 0; yc < height; ++yc)
				{
					const openni::DepthPixel* pDepth =
						(const openni::DepthPixel*)pDepthRow;
					const openni::RGB888Pixel* pRgb =
						(const openni::RGB888Pixel*)pRgbRow;
					for (unsigned int xc = 0; xc < width;
						 ++xc, ++pDepth, ++pRgb)
					{
						int _x = xc;
						if (isMirror) { _x = width - _x - 1; }
						obs.rangeImage(yc, _x) = (*pDepth) * 1.0 / 1000;
						iimage.setPixel(
							_x, yc, (pRgb->r << 16) + (pRgb->g << 8) + pRgb->b);

						// obs.intensityImage.setPixel(xc,yc,(*pRgb));
					}
					pDepthRow += d_step;
					pRgbRow += rgb_step;
				}
				obs.rangeImage_convertToExternalStorage(
					sDepthfile, out_img_dir + string("/"));
				obs.intensityImage = iimage;
				obs.intensityImage.saveToFile(
					out_img_dir + string("/") + sRGBfile);
				obs.intensityImage.setExternalStorage(
					out_img_dir + string("/") + sRGBfile);

				archiveFrom(f_out) << obs;
			}
		}

		cout << "\nStopping...\n";

		depth.destroy();
		rgb.destroy();
		openni::OpenNI::shutdown();

		return 0;
	}
	catch (const std::exception& e)
	{
		std::cout << "MRPT exception caught: " << e.what() << std::endl;
		return -1;
	}
}

bool setONI2StreamMode(
	openni::VideoStream& stream, int w, int h, int fps,
	openni::PixelFormat format)
{
	/*
  void openni::VideoMode::setResolution()
	Setter function for the resolution of this VideoMode. Application use of
  this function is not recommended.
	Instead, use SensorInfo::getSupportedVideoModes() to obtain a list of valid
  video modes

  -- cited from OpenNI2 help. setResolution() is not recommended.
  */
	bool found = false;
	const openni::Array<openni::VideoMode>& modes =
		stream.getSensorInfo().getSupportedVideoModes();
	for (int i = 0, i_end = modes.getSize(); i < i_end; ++i)
	{
		if (modes[i].getResolutionX() != w) { continue; }
		if (modes[i].getResolutionY() != h) { continue; }
		if (modes[i].getPixelFormat() != format) { continue; }
		openni::Status rc = stream.setVideoMode(modes[i]);
		if (rc != openni::STATUS_OK)
		{
			printf(
				"%s:Couldn't find RGB stream:\n%s\n", __FUNCTION__,
				openni::OpenNI::getExtendedError());
			return false;
		}
		return true;
	}
	return false;
}

#if 0
bool initONI2RGBStream(openni::Device& device, openni::VideoStream& rgb, int w, int h, int fps, openni::PixelFormat format){
	openni::Status rc = openni::STATUS_OK;
	rc = rgb.create(device, openni::SENSOR_COLOR);
	if(rc != openni::STATUS_OK){
		printf("%s:Couldn't find RGB stream:\n%s\n", __FUNCTION__, openni::OpenNI::getExtendedError());
		return false;
	}
	rc = rgb.setMirroringEnabled(false);
	if (rc != openni::STATUS_OK){
		printf("%s:setMirroringEnabled(false) failed:\n%s\n", __FUNCTION__, openni::OpenNI::getExtendedError());
		return false;
	}
	openni::VideoMode options = rgb.getVideoMode();
	printf("Initial resolution RGB (%d, %d) FPS %d Format %d\n", options.getResolutionX(), options.getResolutionY(), options.getFps(), options.getPixelFormat());
	if(setONI2StreamMode(rgb, w, h, fps, format) == false){
		printf("%s:Can't find desired rgb mode\n", __FUNCTION__ );
		return false;
	}
	options = rgb.getVideoMode();
	printf("  -> (%d, %d) FPS %d Format %d\n", options.getResolutionX(), options.getResolutionY(), options.getFps(), options.getPixelFormat());
	rc = rgb.start();
	if (rc != openni::STATUS_OK){
		printf("%s:Couldn't start RGB stream:\n%s\n", __FUNCTION__, openni::OpenNI::getExtendedError());
		rgb.destroy();
		return false;
	}
	return true;
}

bool initONI2DepthStream(openni::Device& device, openni::VideoStream& depth, int w, int h, int fps, openni::PixelFormat format){
	openni::Status rc = depth.create(device, openni::SENSOR_DEPTH);
	if (rc != openni::STATUS_OK){
		printf("%s:Couldn't find depth stream:\n%s\n", __FUNCTION__, openni::OpenNI::getExtendedError());
		return false;
	}
	openni::VideoMode options = depth.getVideoMode();
	printf("Initial resolution Depth(%d, %d) FPS %d Format %d\n", options.getResolutionX(), options.getResolutionY(), options.getFps(), options.getPixelFormat());
	if(setONI2StreamMode(depth, w, h, fps, format) == false){
		printf("%s:Can't find desired depth mode\n", __FUNCTION__ );
		return false;
	}
	options = depth.getVideoMode();
	printf("  -> (%d, %d) FPS %d Format %d\n", options.getResolutionX(), options.getResolutionY(), options.getFps(), options.getPixelFormat());
	depth.setMirroringEnabled(false);
	rc = depth.start();
	if (rc != openni::STATUS_OK){
		printf("Couldn't start depth stream:\n%s\n", openni::OpenNI::getExtendedError());
		depth.destroy();
		return false;
	}
	return true;
}
#endif

bool initONI2Stream(
	openni::Device& device, openni::SensorType sensorType,
	openni::VideoStream& stream, int w, int h, int fps,
	openni::PixelFormat format)
{
	openni::Status rc = openni::STATUS_OK;
	const char* strSensor;
	if (sensorType == openni::SENSOR_COLOR)
	{ strSensor = "openni::SENSOR_COLOR"; }
	else if (sensorType == openni::SENSOR_DEPTH)
	{
		strSensor = "openni::SENSOR_DEPTH";
	}
	else
	{
		printf("%s:Unknown SensorType -> %d\n", __FUNCTION__, sensorType);
		return false;
	}
	rc = stream.create(device, sensorType);
	if (rc != openni::STATUS_OK)
	{
		printf(
			"%s:Couldn't find sensor %s: %s\n", __FUNCTION__, strSensor,
			openni::OpenNI::getExtendedError());
		return false;
	}
	openni::VideoMode options = stream.getVideoMode();
	printf(
		"%s:Initial resolution %s (%d, %d) FPS %d Format %d\n", __FUNCTION__,
		strSensor, options.getResolutionX(), options.getResolutionY(),
		options.getFps(), options.getPixelFormat());
	if (setONI2StreamMode(stream, w, h, fps, format) == false)
	{
		printf(
			"%s:Can't find desired mode in the %s\n", __FUNCTION__, strSensor);
		return false;
	}
	options = stream.getVideoMode();
	printf(
		"  -> (%d, %d) FPS %d Format %d\n", options.getResolutionX(),
		options.getResolutionY(), options.getFps(), options.getPixelFormat());
	return true;
}

bool synchONI2MirrorMode(
	openni::VideoStream& rgb, openni::VideoStream& depth, bool& isMirrorMode)
{
	static const int SIZE = 2;
	openni::VideoStream* streams[SIZE] = {&rgb, &depth};
	static const char* strNames[SIZE] = {
		"openni::SENSOR_COLOR", "openni::SENSOR_DEPTH"};
	isMirrorMode = false;
	// Check whether both stream support mirroring.
	for (int i = 0; i < SIZE; ++i)
	{
		if (streams[i]->isPropertySupported(
				openni::STREAM_PROPERTY_MIRRORING) == false)
		{
			printf(
				"%s:openni::STREAM_PROPERTY_MIRRORING is not supported on %s.\n"
				"  We assume this is MS Kinect and taken images are inverted "
				"to right and left.\n",
				__FUNCTION__, strNames[i]);
			// In this case, getMirroringEnabled() method always returns false.
			// So we cannot confirm whether the images are inverted or not.
			isMirrorMode = true;
			break;
		}
	}
	// Set both stream to same mirror mode.
	for (int i = 0; i < SIZE; ++i)
	{
		if (streams[i]->isPropertySupported(
				openni::STREAM_PROPERTY_MIRRORING) == false)
		{ break; }
		if (streams[i]->setMirroringEnabled(isMirrorMode) != openni::STATUS_OK)
		{
			printf(
				"%s:setMirroringEnabled() failed: %s\n", __FUNCTION__,
				openni::OpenNI::getExtendedError());
			return false;
		}
	}
	return true;
}

bool startONI2Streams(openni::VideoStream& rgb, openni::VideoStream& depth)
{
	static const int SIZE = 2;
	openni::VideoStream* streams[SIZE] = {&rgb, &depth};
	static const char* strNames[SIZE] = {
		"openni::SENSOR_COLOR", "openni::SENSOR_DEPTH"};
	for (int i = 0; i < SIZE; ++i)
	{
		openni::Status rc = streams[i]->start();
		if (rc != openni::STATUS_OK)
		{
			printf(
				"%s:Couldn't start RGB stream:\n%s\n", __FUNCTION__,
				openni::OpenNI::getExtendedError());
			streams[i]->destroy();
			return false;
		}
	}
	return true;
}
