/* +------------------------------------------------------------------------+
   |                     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 <mrpt/gui/CDisplayWindow.h>
#include <mrpt/gui/CDisplayWindow3D.h>
#include <mrpt/maps/CLandmarksMap.h>
#include <mrpt/opengl/CGridPlaneXY.h>
#include <mrpt/opengl/Scene.h>
#include <mrpt/vision/CFeatureExtraction.h>

#include <iostream>

using mrpt::format;
using namespace mrpt::gui;
using namespace mrpt::vision;
using namespace mrpt::img;
using namespace mrpt::opengl;
using namespace mrpt::system;
using namespace std;

#include <mrpt/examples_config.h>
string myDataDir(
	MRPT_EXAMPLES_BASE_DIRECTORY +
	string("vision_keypoint_matching_example/imgs/"));

// ------------------------------------------------------
//	TestExtractMatchProjectAndPaint
// ------------------------------------------------------
void TestExtractMatchProjectAndPaint()
{
	CFeatureExtraction fExt;
	CFeatureList featsHarris_L, featsHarris_R;
	CMatchedFeatureList mHarris, mSIFT, mSURF;
	CImage imL, imR;

	string imgL = myDataDir + string("imL_p01.jpg");  // Left image
	string imgR = myDataDir + string("imR_p01.jpg");  // Right image

	// Load and check images
	if (!imL.loadFromFile(imgL))
	{
		cerr << "Cannot load " << imgL << endl;
		return;
	}
	cout << "Loaded test image: " << imgL << endl;

	if (!imR.loadFromFile(imgR))
	{
		cerr << "Cannot load " << imgR << endl;
		return;
	}
	cout << "Loaded test image: " << imgR << endl;

	cout << "***************************************************" << endl;
	cout << "***************************************************" << endl;

	// Extract features:
	// HARRIS
	cout << "Detecting HARRIS features in LEFT image" << endl;
	fExt.options.featsType = featKLT;
	fExt.detectFeatures(imL, featsHarris_L);
	cout << "Detected " << featsHarris_L.size() << endl;

	cout << "Detecting HARRIS features in RIGHT image" << endl;
	fExt.detectFeatures(imR, featsHarris_R);
	cout << "Detected " << featsHarris_R.size() << endl;

	cout << "***************************************************" << endl;
	cout << "***************************************************" << endl;

	// Match features:
	TMatchingOptions opt;
	cout << "Matching HARRIS features by CORRELATION" << endl;
	matchFeatures(featsHarris_L, featsHarris_R, mHarris);
	cout << "Matches found: " << mHarris.size() << endl;

	cout << "***************************************************" << endl;

	// Project features:
	mrpt::maps::CLandmarksMap outMap;
	TStereoSystemParams stereoOptions;	// Default options: Bumblebee + 640x480
	cout << "Projecting matched features" << endl;
	mrpt::vision::projectMatchedFeatures(mHarris, stereoOptions, outMap);

	CDisplayWindow3D win3D("3D Map");
	Scene::Ptr& scene3D = win3D.get3DSceneAndLock();
	CGridPlaneXY::Ptr gridXY = CGridPlaneXY::Create(-10, 10, -10, 10, 0, 1);
	scene3D->insert(gridXY);

	scene3D->insert(outMap.getVisualization());

	win3D.unlockAccess3DScene();
	win3D.repaint();

	mrpt::system::pause();

}  // end TestExtractMatchProjectAndPaint

// ------------------------------------------------------
//		TestMatchFeatures
// ------------------------------------------------------
void TestMatchFeatures(bool showMatches)
{
	CFeatureExtraction fExt;
	CFeatureList featsHarris_L, featsHarris_R, featsSIFT_L, featsSIFT_R,
		featsSURF_L, featsSURF_R, featsFAST_L, featsFAST_R;
	CMatchedFeatureList mHarris, mSIFT, mSURF, mHarris_SAD, mFAST_CC, mFAST_SAD;
	CImage imL, imR;

	string imgL = myDataDir + string("imL_p01.jpg");  // Left image
	string imgR = myDataDir + string("imR_p01.jpg");  // Right image

	// Load and check images
	if (!imL.loadFromFile(imgL))
	{
		cerr << "Cannot load " << imgL << endl;
		return;
	}
	cout << "Loaded LEFT image: " << endl << imgL << endl;

	if (!imR.loadFromFile(imgR))
	{
		cerr << "Cannot load " << imgR << endl;
		return;
	}
	cout << "Loaded RIGHT image: " << endl << imgR << endl;

	cout << "***************************************************" << endl;
	cout << "***************************************************" << endl;

	// Extract features:
	// HARRIS
	cout << "Detecting HARRIS features in LEFT image" << endl;
	fExt.options.featsType = featHarris;
	fExt.detectFeatures(imL, featsHarris_L, 250);
	cout << "Detected " << featsHarris_L.size() << endl;

	cout << "Detecting HARRIS features in RIGHT image" << endl;
	fExt.detectFeatures(imR, featsHarris_R, 250);
	cout << "Detected " << featsHarris_R.size() << endl;
	cout << "***************************************************" << endl;

	// SIFT
	cout << "Detecting SIFT features in LEFT image" << endl;
	fExt.options.featsType = featSIFT;
	fExt.detectFeatures(imL, featsSIFT_L);
	cout << "Detected " << featsSIFT_L.size() << endl;

	cout << "Detecting SIFT features in RIGHT image" << endl;
	fExt.options.featsType = featSIFT;
	fExt.detectFeatures(imR, featsSIFT_R);
	cout << "Detected " << featsSIFT_R.size() << endl;
	cout << "***************************************************" << endl;

	// SURF
	cout << "Detecting SURF features in LEFT image" << endl;
	fExt.options.featsType = featSURF;
	fExt.detectFeatures(imL, featsSURF_L);
	cout << "Detected " << featsSURF_L.size() << endl;

	cout << "Detecting SURF features in RIGHT image" << endl;
	fExt.detectFeatures(imR, featsSURF_R);
	cout << "Detected " << featsSURF_R.size() << endl;
	cout << "***************************************************" << endl;

	// FAST
	cout << "Detecting FAST features in LEFT image" << endl;
	fExt.options.featsType = featFAST;
	fExt.detectFeatures(imL, featsFAST_L, 0, 250);
	cout << "Detected " << featsFAST_L.size() << endl;

	cout << "Detecting FAST features in RIGHT image" << endl;
	fExt.detectFeatures(imR, featsFAST_R, 0, 250);
	cout << "Detected " << featsFAST_R.size() << endl;
	cout << "***************************************************" << endl;
	cout << "***************************************************" << endl;

	// Match features:
	TMatchingOptions opt;

	//	// HARRIS
	CTicTac tictac;
	double T = 0.0;
	cout << "Matching HARRIS features" << endl;
	tictac.Tic();
	matchFeatures(featsHarris_L, featsHarris_R, mHarris);
	T = tictac.Tac();
	cout << "[NCC] Matches found: " << mHarris.size() << " in " << T * 1000.0f
		 << " ms " << endl;

	opt.matching_method = TMatchingOptions::mmSAD;
	tictac.Tic();
	matchFeatures(featsHarris_L, featsHarris_R, mHarris_SAD, opt);
	T = tictac.Tac();
	cout << "[SAD] Matches found: " << mHarris_SAD.size() << " in "
		 << T * 1000.0f << " ms " << endl;
	cout << "***************************************************" << endl;

	// SIFT
	cout << "Matching SIFT features by DESCRIPTOR" << endl;
	opt.matching_method = TMatchingOptions::mmDescriptorSIFT;
	tictac.Tic();
	matchFeatures(featsSIFT_L, featsSIFT_R, mSIFT, opt);
	T = tictac.Tac();
	cout << "Matches found: " << mSIFT.size() << " in " << T * 1000.0f << " ms "
		 << endl;
	cout << "***************************************************" << endl;

	// SURF
	cout << "Matching SURF features by DESCRIPTOR" << endl;
	opt.matching_method = TMatchingOptions::mmDescriptorSURF;
	tictac.Tic();
	matchFeatures(featsSURF_L, featsSURF_R, mSURF, opt);
	T = tictac.Tac();
	cout << "Matches found: " << mSURF.size() << " in " << T * 1000.0f << " ms "
		 << endl;
	cout << "***************************************************" << endl;

	// FAST
	cout << "Matching FAST features" << endl;
	tictac.Tic();
	matchFeatures(featsFAST_L, featsFAST_R, mFAST_CC);
	T = tictac.Tac();
	cout << "[NCC] Matches found: " << mFAST_CC.size() << " in " << T * 1000.0f
		 << " ms " << endl;

	opt.matching_method = TMatchingOptions::mmSAD;
	tictac.Tic();
	matchFeatures(featsFAST_L, featsFAST_R, mFAST_SAD, opt);
	T = tictac.Tac();
	cout << "[SAD] Matches found: " << mFAST_SAD.size() << " in " << T * 1000.0f
		 << " ms " << endl;
	cout << "***************************************************" << endl;

	if (showMatches)
	{
		CDisplayWindow winHarrisSAD, winHarrisNCC, winFASTSAD, winFASTNCC,
			winSIFT, winSURF;

		winHarrisSAD.setWindowTitle("Matches with Harris + SAD");
		winHarrisNCC.setWindowTitle("Matches with Harris + NCC");
		winFASTSAD.setWindowTitle("Matches with FAST + SAD");
		winFASTNCC.setWindowTitle("Matches with FAST + NCC");
		winSIFT.setWindowTitle("Matches with SIFT");
		winSURF.setWindowTitle("Matches with SURF");

		winHarrisNCC.showImagesAndMatchedPoints(
			imL, imR, mHarris, TColor(0, 0, 255));
		winHarrisSAD.showImagesAndMatchedPoints(
			imL, imR, mHarris_SAD, TColor(0, 0, 255));
		winSIFT.showImagesAndMatchedPoints(imL, imR, mSIFT, TColor(0, 255, 0));
		winSURF.showImagesAndMatchedPoints(imL, imR, mSURF, TColor(0, 255, 0));
		winFASTSAD.showImagesAndMatchedPoints(
			imL, imR, mFAST_SAD, TColor(0, 255, 0));
		winFASTNCC.showImagesAndMatchedPoints(
			imL, imR, mFAST_CC, TColor(0, 255, 0));

		mrpt::system::pause();
	}

}  // end TestMatchFeatures

// ------------------------------------------------------
//		TestMatchingComparative
// ------------------------------------------------------
void TestMatchingComparative()
{
	// Take two images
	string imgL = myDataDir + string("imL_p01.jpg");  // Left image
	string imgR = myDataDir + string("imR_p01.jpg");  // Right image

	CImage im1, im2;
	im1.loadFromFile(imgL);
	im2.loadFromFile(imgR);

	size_t imW = im1.getWidth();
	size_t imH = im1.getHeight();

	CFeatureExtraction fExt;
	fExt.options.featsType = featFAST;
	fExt.options.patchSize = 21;

	// Find FAST features
	CFeatureList list1, list2;
	fExt.detectFeatures(im1, list1, 150);
	// Compute SIFT & SURF descriptors
	fExt.computeDescriptors(im1, list1, descSIFT);
	fExt.computeDescriptors(im1, list1, descSURF);

	fExt.detectFeatures(im2, list2, 150);
	// Compute SIFT & SURF descriptors
	fExt.computeDescriptors(im2, list2, descSIFT);
	fExt.computeDescriptors(im2, list2, descSURF);

	im1.drawFeatures(list1);
	im2.drawFeatures(list2);

	CDisplayWindow win, win2;
	win.setPos(0, 0);
	win2.setPos(0, imH * 1.5);
	CImage joinimage, copyjoinimage, copyInfoImage;
	size_t imW2 = 1280;
	size_t imH2 = 150;

	CImage infoimage(imW2, imH2, CH_RGB);

	joinimage.joinImagesHorz(im1, im2);
	infoimage.filledRectangle(0, 0, imW2, imH2, TColor(150, 150, 150));
	infoimage.textOut(20, imH2 - 53, "SAD", TColor::blue());
	infoimage.textOut(20, imH2 - 41, "NCC", TColor::blue());
	infoimage.textOut(20, imH2 - 29, "SIFT", TColor::blue());
	infoimage.textOut(20, imH2 - 17, "SURF", TColor::blue());
	for (auto it1 = list1.begin(); it1 != list1.end(); ++it1)
	{
		const auto& pt1 = it1->keypoint.pt;

		copyInfoImage = infoimage;
		copyjoinimage = joinimage;
		copyjoinimage.line(pt1.x, 0, pt1.x, imH, TColor::green());	// Horiz
		copyjoinimage.line(
			pt1.x + imW, 0, pt1.x + imW, imH,
			TColor::green());  // Horiz
		copyjoinimage.line(
			0, pt1.y, imW + imW, pt1.y, TColor::green());  // Epipolar
		copyjoinimage.drawCircle(
			pt1.x, pt1.y, 4, TColor::green(), 2);  // Keypoint

		copyInfoImage.update_patch(*it1->patch, 0, 0);
		bool firstMatch = true;
		int cnt = 0;
		int px = 80;
		double minsad = 1.0, maxncc = 0.0;
		float minsiftd = 1.0f, minsurfd = 1.0f;
		int idxsad = 0, idxncc = 0, idxsiftd = 0, idxsurfd = 0;

		for (auto it2 = list2.begin(); it2 != list2.end(); ++it2)
		{
			const auto& pt2 = it2->keypoint.pt;

			if (fabs(pt1.y - pt2.y) <= 1.0 && pt1.x > pt2.x)
			{
				// Compute matching with SAD and Correlation and SIFT/SURF?
				// Use epipolar constraints
				// Compute SAD
				double sad = mrpt::vision::computeSAD(*it1->patch, *it2->patch);
				if (sad < minsad)
				{
					minsad = sad;
					idxsad = cnt;
				}
				// Compute Correlation
				double ncc;
				size_t u, v;
				mrpt::vision::openCV_cross_correlation(
					*it1->patch, *it2->patch, u, v, ncc);
				if (ncc > maxncc)
				{
					maxncc = ncc;
					idxncc = cnt;
				}

				// Compute distance between descriptors SIFT
				float siftd = it1->descriptorSIFTDistanceTo(*it2);
				if (siftd < minsiftd)
				{
					minsiftd = siftd;
					idxsiftd = cnt;
				}

				// Compute distance between descriptors SIFT
				float surfd = it1->descriptorSURFDistanceTo(*it2);
				if (surfd < minsurfd)
				{
					minsurfd = surfd;
					idxsurfd = cnt;
				}

				// Plot images + features + each candidate + difference score
				if (firstMatch)
				{
					copyjoinimage.line(
						pt1.x + imW, 0, pt1.x + imW, imH,
						TColor::green());  // Limit line (only the first time)
					firstMatch = false;
				}  // end-if

				copyjoinimage.drawCircle(
					pt2.x + imW, pt2.y, 4, TColor::blue(),
					2);	 // Keypoint
				double rx0, rx1, ry0, ry1, tx, ty;
				rx0 = pt2.x + imW - 15;
				rx1 = pt2.x + imW;
				tx = pt2.x + imW - 13;
				if (cnt % 2)
				{
					ry0 = pt2.y - 20;
					ry1 = pt2.y - 10;
					ty = pt2.y - 22;
				}
				else
				{
					ry0 = pt2.y + 10;
					ry1 = pt2.y + 20;
					ty = pt2.y + 8;
				}
				copyjoinimage.filledRectangle(
					rx0, ry0, rx1, ry1, TColor(150, 150, 150));
				copyjoinimage.textOut(
					tx, ty, format("%d", cnt), TColor::blue());

				px = 80 + cnt * 50;
				if (px + fExt.options.patchSize > imW2) continue;

				copyInfoImage.update_patch(*it2->patch, px, 30);

				copyInfoImage.textOut(
					px, imH2 - 70, format("%d", cnt), TColor::blue());
				copyInfoImage.textOut(
					px, imH2 - 53, format("%.2f", sad), TColor::blue());
				copyInfoImage.textOut(
					px, imH2 - 41, format("%.2f", ncc), TColor::blue());
				copyInfoImage.textOut(
					px, imH2 - 29, format("%.2f", siftd), TColor::blue());
				copyInfoImage.textOut(
					px, imH2 - 17, format("%.2f", surfd), TColor::blue());

				cnt++;
			}  // end if
		}  // end for it2
		copyInfoImage.textOut(
			80 + idxsad * 50, imH2 - 53, format("%.2f", minsad),
			TColor::green());
		copyInfoImage.textOut(
			80 + idxncc * 50, imH2 - 41, format("%.2f", maxncc),
			TColor::green());
		copyInfoImage.textOut(
			80 + idxsiftd * 50, imH2 - 29, format("%.2f", minsiftd),
			TColor::green());
		copyInfoImage.textOut(
			80 + idxsurfd * 50, imH2 - 17, format("%.2f", minsurfd),
			TColor::green());

		win.showImage(copyjoinimage);
		win2.showImage(copyInfoImage);
		mrpt::system::pause();
	}  // end for it1

	// Save to file
	// Check number of good features

}  // end TestMatchingComparative

int main(int argc, char** argv)
{
	try
	{
		if (argc == 1)
		{
			cerr << "Usage: " << argv[0] << endl;
			cerr << "Options:" << endl;
			cerr << " -match [-s]: TestMatchFeatures (if -s is set, final "
					"matches in images will be shown)."
				 << endl;
			cerr << " -comp: TestMatchingComparative." << endl;
			cerr << " -proj: TestExtractMatchProjectAndPaint." << endl;
		}
		else
		{
			if (!strcmp(argv[1], "-match"))
			{
				if (argc == 3 && !strcmp(argv[2], "-s"))
					TestMatchFeatures(true);
				else
					TestMatchFeatures(false);
			}
			else if (!strcmp(argv[1], "-proj"))
				TestExtractMatchProjectAndPaint();
			else if (!strcmp(argv[1], "-comp"))
			{
				cout << "Press ^C to finish program." << endl;
				TestMatchingComparative();
			}
			else
			{
				cerr << "Usage: " << argv[0] << endl;
				cerr << "Options:" << endl;
				cerr << " -match [-s]: TestMatchFeatures (if -s is set, final "
						"matches in images will be shown)."
					 << endl;
				cerr << " -comp: TestMatchingComparative." << endl;
				cerr << " -proj: TestExtractMatchProjectAndPaint." << endl;
			}
		}
		return 0;
	}
	catch (const std::exception& e)
	{
		std::cerr << "MRPT error: " << mrpt::exception_to_str(e) << std::endl;
		return -1;
	}
	catch (...)
	{
		printf("Another exception!!");
		return -1;
	}
}
