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

#include "apps-precomp.h"  // Precompiled headers
//
#include <mrpt/vision/CStereoRectifyMap.h>

#include "rawlog-edit-declarations.h"

using namespace mrpt;
using namespace mrpt::obs;
using namespace mrpt::system;
using namespace mrpt::apps;
using namespace std;
using namespace mrpt::io;

// ======================================================================
//		op_stereo_rectify
// ======================================================================
DECLARE_OP_FUNCTION(op_stereo_rectify)
{
	// A class to do this operation:
	class CRawlogProcessor_StereoRectify
		: public CRawlogProcessorOnEachObservation
	{
	   protected:
		TOutputRawlogCreator outrawlog;

		string target_label;
		string outDir;
		string imgFileExtension;
		double rectify_alpha;  // [0,1] see cvStereoRectify()

		bool m_this_obs_is_ok;

		mrpt::vision::CStereoRectifyMap rectify_map;

		size_t m_num_external_files_failures;

	   public:
		size_t m_changedCams;

		CRawlogProcessor_StereoRectify(
			CFileGZInputStream& in_rawlog, TCLAP::CmdLine& cmdline,
			bool Verbose)
			: CRawlogProcessorOnEachObservation(in_rawlog, cmdline, Verbose)
		{
			m_changedCams = 0;
			m_num_external_files_failures = 0;

			// Load .ini file with poses:
			string str;
			getArgValue<string>(cmdline, "stereo-rectify", str);

			vector<string> lstTokens;
			tokenize(str, ",", lstTokens);
			if (lstTokens.size() != 2)
				throw std::runtime_error(
					"--stereo-rectify op: argument must be in the format: "
					"--stereo-rectify LABEL,ALPHA_VALUE");

			target_label = lstTokens[0];

			const string& sAlpha = lstTokens[1];
			rectify_alpha = atof(sAlpha.c_str());
			if (rectify_alpha != -1 &&
				!(rectify_alpha >= 0 && rectify_alpha <= 1))
				throw std::runtime_error(
					"--stereo-rectify op: Invalid ALPHA value. Use '-1' for "
					"auto guess.");

			getArgValue<string>(cmdline, "image-format", imgFileExtension);

			// Create a "/Images_Rectified" directory.
			const string out_rawlog_basedir =
				extractFileDirectory(outrawlog.out_rawlog_filename);

			outDir = (out_rawlog_basedir.empty()
						  ? string()
						  : (out_rawlog_basedir + string("/"))) +
				extractFileName(outrawlog.out_rawlog_filename) +
				string("_Images");
			if (directoryExists(outDir))
				throw runtime_error(
					string("*ABORTING*: Output directory for rectified images "
						   "already exists: ") +
					outDir +
					string("\n. Select a different output path or "
						   "remove the directory."));

			VERBOSE_COUT << "Creating directory: " << outDir << endl;

			mrpt::system::createDirectory(outDir);
			if (!fileExists(outDir))
				throw runtime_error(
					string("*ABORTING*: Couldn't create directory: ") + outDir);

			// Add the final /
			outDir += "/";

			// Optional argument:  "--image-size=640x480"
			string strResize;
			if (getArgValue<string>(cmdline, "image-size", str))
			{
				tokenize(str, "x", lstTokens);
				if (lstTokens.size() != 2)
					throw std::runtime_error(
						"--stereo-rectify op: Expected format: --image-size "
						"NCOLSxNROWS");

				const int nCols = atoi(lstTokens[0].c_str());
				const int nRows = atoi(lstTokens[1].c_str());
				VERBOSE_COUT << "Will rectify and resize to " << nCols << "x"
							 << nRows << " simultaneously.\n";

				rectify_map.enableResizeOutput(true, nCols, nRows);
			}

			// Optional argument: "--rectify-centers-coincide"
			if (isFlagSet(cmdline, "rectify-centers-coincide"))
			{
				VERBOSE_COUT
					<< "Will rectify such that both image centers coincide.\n";
				rectify_map.enableBothCentersCoincide(true);
			}
		}

		bool processOneObservation(CObservation::Ptr& obs) override
		{
			m_this_obs_is_ok = true;

			if (strCmpI(obs->sensorLabel, target_label))
			{
				if (IS_CLASS(*obs, CObservationStereoImages))
				{
					CObservationStereoImages::Ptr o =
						std::dynamic_pointer_cast<CObservationStereoImages>(
							obs);

					try
					{
						// Already initialized the rectification map?
						if (!rectify_map.isSet())
						{
							// On the first ocassion, initialize map:
							rectify_map.setAlpha(rectify_alpha);
							rectify_map.setFromCamParams(*o);
						}

						// This is needed to raise an exception of the correct
						// type that reveal any missing external file:
						o->imageLeft.getWidth();
						o->imageRight.getWidth();

						// This call rectifies the images in-place and also
						// updates
						// all the camera parameters as needed:
						rectify_map.rectify(*o);

						const string label_time = format(
							"%s_%f", o->sensorLabel.c_str(),
							timestampTotime_t(o->timestamp));
						{
							const string fileName = string("img_") +
								label_time + string("_left.") +
								imgFileExtension;
							o->imageLeft.saveToFile(outDir + fileName);
							o->imageLeft.setExternalStorage(fileName);
						}
						{
							const string fileName = string("img_") +
								label_time + string("_right.") +
								imgFileExtension;
							o->imageRight.saveToFile(outDir + fileName);
							o->imageRight.setExternalStorage(fileName);
						}
						m_changedCams++;
					}
					catch (mrpt::img::CExceptionExternalImageNotFound&)
					{
						const size_t MAX_FAILURES = 1000;
						m_num_external_files_failures++;

						if (m_num_external_files_failures < MAX_FAILURES)
						{
							m_this_obs_is_ok = false;
							cerr << "\n *WARNING*: Dropping one observation "
									"due to missing external image file at "
									"rawlog entry "
								 << m_rawlogEntry << endl;
						}
						else
						{
							throw std::runtime_error(
								"*ERROR* Too many external images missing, "
								"this doesn't seem spureous missings!");
						}
					}
				}
			}
			return true;
		}

		// This method can be reimplemented to save the modified object to an
		// output stream.
		void OnPostProcess(
			mrpt::obs::CActionCollection::Ptr& actions,
			mrpt::obs::CSensoryFrame::Ptr& SF,
			mrpt::obs::CObservation::Ptr& obs) override
		{
			if (!m_this_obs_is_ok) return;

			ASSERT_((actions && SF) || obs);
			if (actions) (*outrawlog.out_rawlog) << actions << SF;
			else
				(*outrawlog.out_rawlog) << obs;
		}
	};

	// Process
	// ---------------------------------
	CRawlogProcessor_StereoRectify proc(in_rawlog, cmdline, verbose);
	proc.doProcessRawlog();

	// Dump statistics:
	// ---------------------------------
	VERBOSE_COUT << "Time to process file (sec)        : " << proc.m_timToParse
				 << "\n";
	VERBOSE_COUT << "Number of modified entries        : " << proc.m_changedCams
				 << "\n";
}
