/* +------------------------------------------------------------------------+
   |                     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 "opengl-precomp.h"	 // Precompiled header
//
#include <mrpt/opengl/CEllipsoid2D.h>
#include <mrpt/opengl/CEllipsoid3D.h>
#include <mrpt/opengl/CPointCloud.h>
#include <mrpt/opengl/CSetOfLines.h>
#include <mrpt/opengl/CSetOfObjects.h>
#include <mrpt/opengl/stock_objects.h>
#include <mrpt/poses/CPointPDFParticles.h>
#include <mrpt/poses/CPointPDFSOG.h>
#include <mrpt/poses/CPose3DPDFParticles.h>
#include <mrpt/poses/CPose3DPDFSOG.h>
#include <mrpt/poses/CPose3DQuatPDFGaussian.h>
#include <mrpt/poses/CPosePDFGaussian.h>
#include <mrpt/poses/CPosePDFParticles.h>
#include <mrpt/poses/CPosePDFSOG.h>

using namespace mrpt;
using namespace mrpt::math;
using namespace mrpt::opengl;
using namespace mrpt::poses;

const double POSE_TAIL_LENGTH = 0.1;
const double POSE_TAIL_WIDTH = 3.0;
const float POSE_POINT_SIZE = 4.0f;
const float POSE_AXIS_SCALE = 0.1f;

#define POSE_COLOR 0, 0, 1
#define POINT_COLOR 1, 0, 0

/** Returns a representation of a the PDF - this is just an auxiliary function,
 * it's more natural to call
 *    mrpt::poses::CPosePDF::getAs3DObject     */
CSetOfObjects::Ptr CSetOfObjects::posePDF2opengl(const CPosePDF& o)
{
	auto outObj = CSetOfObjects::Create();

	if (IS_CLASS(o, CPosePDFSOG))
	{
		const auto* p = dynamic_cast<const CPosePDFSOG*>(&o);
		ASSERT_(p != nullptr);

		opengl::CSetOfLines::Ptr lins = std::make_shared<opengl::CSetOfLines>();
		lins->setColor(0.0f, 0.0f, 1.0f, 0.6f);
		lins->setLineWidth(POSE_TAIL_WIDTH);

		for (const auto& it : *p)
		{
			auto ellip = mrpt::opengl::CEllipsoid2D::Create();

			ellip->setPose(CPose3D(it.mean.x(), it.mean.y(), 0));
			ellip->setCovMatrix(it.cov.blockCopy<2, 2>());
			ellip->setColor(POSE_COLOR, 0.6f);
			ellip->setQuantiles(3);
			ellip->enableDrawSolid3D(false);

			outObj->insert(ellip);

			lins->appendLine(
				it.mean.x(), it.mean.y(), 0,
				it.mean.x() + POSE_TAIL_LENGTH * cos(it.mean.phi()),
				it.mean.y() + POSE_TAIL_LENGTH * sin(it.mean.phi()), 0);
		}
		outObj->insert(lins);
	}
	else if (IS_CLASS(o, CPosePDFGaussian))
	{
		const auto* p = dynamic_cast<const CPosePDFGaussian*>(&o);
		ASSERT_(p != nullptr);

		opengl::CSetOfLines::Ptr lins = std::make_shared<opengl::CSetOfLines>();
		lins->setColor(POSE_COLOR, 0.6f);
		lins->setLineWidth(POSE_TAIL_WIDTH);

		auto ellip = mrpt::opengl::CEllipsoid2D::Create();

		ellip->setPose(CPose3D(p->mean.x(), p->mean.y(), 0));
		ellip->setCovMatrix(p->cov.blockCopy<2, 2>());
		ellip->setColor(POSE_COLOR, 0.6f);

		ellip->setQuantiles(3);
		ellip->enableDrawSolid3D(false);

		outObj->insert(ellip);

		lins->appendLine(
			p->mean.x(), p->mean.y(), 0,
			p->mean.x() + POSE_TAIL_LENGTH * cos(p->mean.phi()),
			p->mean.y() + POSE_TAIL_LENGTH * sin(p->mean.phi()), 0);

		outObj->insert(lins);
	}
	else if (IS_CLASS(o, CPosePDFParticles))
	{
		const auto* p = dynamic_cast<const CPosePDFParticles*>(&o);
		ASSERT_(p != nullptr);

		auto pnts = opengl::CPointCloud::Create();
		pnts->setColor(POSE_COLOR, 0.6f);
		pnts->setPointSize(POSE_POINT_SIZE);

		auto lins = opengl::CSetOfLines::Create();
		lins->setColor(POSE_COLOR, 0.6f);
		lins->setLineWidth(POSE_TAIL_WIDTH);

		for (size_t i = 0; i < p->size(); ++i)
		{
			const auto po = p->m_particles[i].d;
			pnts->insertPoint(d2f(po.x), d2f(po.y), 0);
			lins->appendLine(
				po.x, po.y, 0, po.x + POSE_TAIL_LENGTH * cos(po.phi),
				po.y + POSE_TAIL_LENGTH * sin(po.phi), 0);
		}
		outObj->insert(pnts);
		outObj->insert(lins);
	}

	return outObj;
}

/** Returns a representation of a the PDF - this is just an auxiliary function,
 * it's more natural to call
 *    mrpt::poses::CPointPDF::getAs3DObject     */
CSetOfObjects::Ptr CSetOfObjects::posePDF2opengl(const CPointPDF& o)
{
	auto outObj = std::make_shared<CSetOfObjects>();

	if (IS_CLASS(o, CPointPDFSOG))
	{
		const auto* p = dynamic_cast<const CPointPDFSOG*>(&o);
		ASSERT_(p != nullptr);

		// For each gaussian node
		for (const auto& it : *p)
		{
			mrpt::opengl::CRenderizable::Ptr obj;
			if (it.val.cov(2, 2) == 0)
			{
				auto ellip = mrpt::opengl::CEllipsoid2D::Create();
				ellip->setCovMatrix(it.val.cov.blockCopy<2, 2>());
				ellip->setQuantiles(3);
				ellip->enableDrawSolid3D(false);
				obj = ellip;
			}
			else
			{
				auto ellip = mrpt::opengl::CEllipsoid3D::Create();
				ellip->setCovMatrix(it.val.cov);
				ellip->setQuantiles(3);
				ellip->enableDrawSolid3D(false);
				obj = ellip;
			}
			obj->setPose(it.val.mean);
			obj->setColor(POINT_COLOR);

			outObj->insert(obj);
		}  // end for each gaussian node
	}
	else if (IS_CLASS(o, CPointPDFGaussian))
	{
		const auto* p = dynamic_cast<const CPointPDFGaussian*>(&o);
		ASSERT_(p != nullptr);

		mrpt::opengl::CRenderizable::Ptr obj;
		if (p->cov(2, 2) == 0)
		{
			auto ellip = mrpt::opengl::CEllipsoid2D::Create();
			ellip->setCovMatrix(p->cov.blockCopy<2, 2>());
			ellip->setQuantiles(3);
			ellip->enableDrawSolid3D(false);
			obj = ellip;
		}
		else
		{
			auto ellip = mrpt::opengl::CEllipsoid3D::Create();
			ellip->setCovMatrix(p->cov);
			ellip->setQuantiles(3);
			ellip->enableDrawSolid3D(false);
			obj = ellip;
		}
		obj->setLocation(p->mean.x(), p->mean.y(), p->mean.z());
		obj->setColor(POINT_COLOR);
		outObj->insert(obj);
	}
	else if (IS_CLASS(o, CPointPDFParticles))
	{
		const auto* p = dynamic_cast<const CPointPDFParticles*>(&o);
		ASSERT_(p != nullptr);

		mrpt::opengl::CPointCloud::Ptr obj =
			mrpt::opengl::CPointCloud::Create();
		const size_t N = p->size();

		obj->resize(N);
		obj->setColor(POINT_COLOR);
		for (size_t i = 0; i < N; i++)
			obj->setPoint(
				i, p->m_particles[i].d->x, p->m_particles[i].d->y,
				p->m_particles[i].d->z);
		outObj->insert(obj);
	}

	return outObj;
}

/** Returns a representation of a the PDF - this is just an auxiliary function,
 * it's more natural to call
 *    mrpt::poses::CPose3DPDF::getAs3DObject     */
CSetOfObjects::Ptr CSetOfObjects::posePDF2opengl(const CPose3DPDF& o)
{
	CSetOfObjects::Ptr outObj = std::make_shared<CSetOfObjects>();

	if (IS_CLASS(o, CPose3DPDFSOG))
	{
		const auto* p = dynamic_cast<const CPose3DPDFSOG*>(&o);
		ASSERT_(p != nullptr);

		// For each gaussian node
		for (const auto& it : *p)
		{
			opengl::CEllipsoid3D::Ptr obj =
				std::make_shared<opengl::CEllipsoid3D>();

			obj->setPose(it.val.mean);
			obj->setCovMatrix(it.val.cov.blockCopy<3, 3>());

			obj->setQuantiles(3);
			obj->enableDrawSolid3D(false);
			obj->setColor(POSE_COLOR);

			outObj->insert(obj);

			opengl::CSetOfObjects::Ptr axes =
				opengl::stock_objects::CornerXYZ();
			axes->setPose(it.val.mean);
			axes->setScale(POSE_AXIS_SCALE);
			outObj->insert(axes);
		}  // end for each gaussian node
	}
	else if (IS_CLASS(o, CPose3DPDFGaussian))
	{
		const auto* p = dynamic_cast<const CPose3DPDFGaussian*>(&o);
		ASSERT_(p != nullptr);

		auto obj = std::make_shared<opengl::CEllipsoid3D>();

		obj->setPose(p->mean);
		obj->setCovMatrix(p->cov.blockCopy<3, 3>());

		obj->setQuantiles(3);
		obj->enableDrawSolid3D(false);
		obj->setColor(POSE_COLOR);

		outObj->insert(obj);

		opengl::CSetOfObjects::Ptr axes = opengl::stock_objects::CornerXYZ();
		axes->setPose(p->mean);
		axes->setScale(POSE_AXIS_SCALE);
		outObj->insert(axes);
	}
	else if (IS_CLASS(o, CPose3DPDFParticles))
	{
		const auto* p = dynamic_cast<const CPose3DPDFParticles*>(&o);
		ASSERT_(p != nullptr);

		for (size_t i = 0; i < p->size(); i++)
		{
			opengl::CSetOfObjects::Ptr axes =
				opengl::stock_objects::CornerXYZSimple(POSE_AXIS_SCALE);
			axes->setPose(p->m_particles[i].d);
			outObj->insert(axes);
		}
	}

	return outObj;
}

/** Returns a representation of a the PDF - this is just an auxiliary function,
 * it's more natural to call
 *    mrpt::poses::CPose3DQuatPDF::getAs3DObject     */
CSetOfObjects::Ptr CSetOfObjects::posePDF2opengl(const CPose3DQuatPDF& o)
{
	CSetOfObjects::Ptr outObj = std::make_shared<CSetOfObjects>();

	if (IS_CLASS(o, CPose3DQuatPDFGaussian))
	{
		const auto* p = dynamic_cast<const CPose3DQuatPDFGaussian*>(&o);
		ASSERT_(p != nullptr);

		auto obj = mrpt::opengl::CEllipsoid3D::Create();

		obj->setPose(CPose3D(p->mean));
		obj->setCovMatrix(p->cov.blockCopy<3, 3>());

		obj->setQuantiles(3);
		obj->enableDrawSolid3D(false);
		obj->setColor(POSE_COLOR);

		outObj->insert(obj);

		opengl::CSetOfObjects::Ptr axes = opengl::stock_objects::CornerXYZ();
		axes->setPose(CPose3D(p->mean));
		axes->setScale(POSE_AXIS_SCALE);
		outObj->insert(axes);
	}

	return outObj;
}
