#include <OGRE/OgreSceneNode.h>
#include <OGRE/OgreSceneManager.h>

#include <tf/transform_listener.h>

#include <rviz/visualization_manager.h>
#include <rviz/properties/color_property.h>
#include <rviz/properties/float_property.h>
#include <rviz/default_plugin/covariance_property.h>
#include <rviz/properties/int_property.h>
#include <rviz/ogre_helpers/shape.h>
#include <rviz/ogre_helpers/arrow.h>
#include <rviz/ogre_helpers/movable_text.h>
#include <rviz/ogre_helpers/billboard_line.h>
#include <rviz/frame_manager.h>

#include "state_visual.h"
#include "perceptions_display.h"

namespace rviz_multiception
{

PerceptionsDisplay::PerceptionsDisplay()
{
	color_property_ = new rviz::ColorProperty( "Color", QColor( 10, 10, 10 ),
	                                           "Color to draw the arrows.",
	                                           this, SLOT( updateColorAndAlpha() ));

	alpha_property_ = new rviz::FloatProperty( "Alpha", 1.0,
	                                           "State transparency, 0 is fully transparent, 1.0 is fully opaque.",
	                                           this, SLOT( updateColorAndAlpha() ));

	shape_alpha_property_ = new rviz::FloatProperty( "Shape Alpha", 0.4,
	                                           "Shape transparency, 0 is fully transparent, 1.0 is fully opaque.",
	                                           this, SLOT( updateColorAndAlpha() ));

	trace_duration_property_ = new rviz::FloatProperty( "Trace Duration (s)", 1,
	                                           "Duration of state traces, in seconds.",
	                                           this, SLOT( filterTraceBuffer() ));

	text_color_property_ = new rviz::ColorProperty( "Debug Text Color", QColor( 10, 10, 10 ),
	                                           "Color of the debug text.",
	                                           this, SLOT( updateColorAndAlpha() ));

	covariance_property_ = new rviz::CovarianceProperty( "Covariance", true, "Whether or not the covariances of the messages should be shown.",
	                                                     this, SLOT( queueRender() ));
}

void PerceptionsDisplay::onInitialize()
{
	MFDClass::onInitialize();
	frame_node_ = scene_node_->createChildSceneNode();
}

PerceptionsDisplay::~PerceptionsDisplay()
{
	scene_manager_->destroySceneNode(frame_node_);
}

// Clear the visuals by deleting their objects.
void PerceptionsDisplay::reset()
{
	MFDClass::reset();
	resetVectors();
}

void PerceptionsDisplay::resetVectors(size_t N)
{
	pose_visuals_.clear();  pose_visuals_.reserve(N);
	shape_visuals_.clear(); shape_visuals_.reserve(N);
	cov_visuals_.clear();   cov_visuals_.reserve(N);
	text_visuals_.clear();  text_visuals_.reserve(N);
	for(Ogre::SceneNode* textNode: text_nodes_) scene_manager_->destroySceneNode(textNode);
	text_nodes_.clear();    text_nodes_.reserve(N);

	frame_node_->detachAllObjects();
	frame_node_->removeAndDestroyAllChildren();
}

// Set the current color and alpha values for each visual.
void PerceptionsDisplay::updateColorAndAlpha()
{
	const float alpha = alpha_property_->getFloat();
	const float shape_alpha = shape_alpha_property_->getFloat();
	const Ogre::ColourValue color = color_property_->getOgreColor();
	const Ogre::ColourValue textColor = text_color_property_->getOgreColor();

	for(size_t i=0; i<pose_visuals_.size(); i++)
	{
		pose_visuals_[i]->setColor(color.r, color.g, color.b, alpha);
		shape_visuals_[i]->setColor(color.r, color.g, color.b, alpha*shape_alpha);
// 		cov_visuals_[i]->setColor(color.r, color.g, color.b, alpha*shape_alpha);
	}

	for(size_t i=0; i<text_visuals_.size(); i++)
	{
		text_visuals_[i]->setColor(textColor);
	}
}

void PerceptionsDisplay::filterTraceBuffer()
{
	float duration = trace_duration_property_->getFloat();
	double oldestTimePossible = ros::Time::now().toSec() - duration;

	auto newEnd = std::remove_if(trace_buffer_.begin(), trace_buffer_.end(), [oldestTimePossible](const multiception::Perceptions::ConstPtr& s)
	{
		return !s || s->header.stamp.toSec()<oldestTimePossible;
	});
	trace_buffer_.erase(newEnd, trace_buffer_.end());

	std::sort(trace_buffer_.begin(), trace_buffer_.end(), [](const multiception::Perceptions::ConstPtr& a, const multiception::Perceptions::ConstPtr& b)
	{
		return a->header.stamp<b->header.stamp;
	});
}

void PerceptionsDisplay::processMessage(const multiception::Perceptions::ConstPtr& msg)
{
	if(msg->perceptions.size()>0)
	{
		resetVectors(msg->perceptions.size());
	}
	else
	{
		return;
	}

	Ogre::Vector3 position;
	Ogre::Quaternion orientation;
	if(context_->getFrameManager()->getTransform(msg->header.frame_id, msg->header.stamp, position, orientation))
	{
		frame_node_->setPosition(position);
		frame_node_->setOrientation(orientation);
	}
	else
	{
		ROS_DEBUG("Error transforming from frame '%s' to frame '%s'", msg->header.frame_id.c_str(), qPrintable(fixed_frame_));
		return;
	}

	// Process the perceptions in the message
	for(const multiception::Perception& perception : msg->perceptions)
	{
		boost::shared_ptr<rviz::Arrow> arrowVis { new rviz::Arrow(scene_manager_, frame_node_) };
		boost::shared_ptr<rviz::Shape> shapeVis { new rviz::Shape(rviz::Shape::Cube, scene_manager_, frame_node_) };
		boost::shared_ptr<rviz::CovarianceVisual> covVis = covariance_property_->createAndPushBackVisual(scene_manager_, frame_node_);

		computeStateObjects(perception.state, *arrowVis, *covVis);

		// Process the shape
		const float z = (perception.height>0.1) ? perception.height/2+.1 : 1;
		shapeVis->setScale( Ogre::Vector3(perception.length, perception.width, std::max(0.1f, perception.height)) );
		shapeVis->setPosition( Ogre::Vector3(perception.state.x, perception.state.y, z) );
		shapeVis->setOrientation( Ogre::Quaternion(Ogre::Radian(perception.state.yaw), Ogre::Vector3::UNIT_Z) );

		// Process the debug text
		if(!perception.debug.empty())
		{
			Ogre::SceneNode* childNode = frame_node_->createChildSceneNode();
			boost::shared_ptr<rviz::MovableText> debugText { new rviz::MovableText(perception.debug, "Liberation Sans", .7) };
			debugText->setTextAlignment(rviz::MovableText::H_CENTER, rviz::MovableText::V_CENTER);
			childNode->setPosition(Ogre::Vector3(perception.state.x, perception.state.y, z+perception.height/2+.1+1) );
			childNode->attachObject(debugText.get());
			text_visuals_.push_back(debugText);
			text_nodes_.push_back(childNode);
		}

	// Process the buffered trace
	trace_buffer_.push_back(msg); filterTraceBuffer(); trace_visuals_.clear();

	for(const auto& perceptions : trace_buffer_)
	{
		for(const auto& p : perceptions->perceptions)
		{
			// ID 0 represent untracked perception, thus a perception trace is not relevant
			if(p.object_id!=0)
			{
				// Create the visual for this object if non-existent
				if(trace_visuals_.count(p.object_id)==0) 
				{
					trace_visuals_[p.object_id] = boost::make_shared<rviz::BillboardLine>(scene_manager_, scene_node_);
				}

				trace_visuals_[p.object_id]->addPoint(Ogre::Vector3(p.state.x, p.state.y, 0), color);
			}
		}

		pose_visuals_.push_back(arrowVis);
		shape_visuals_.push_back(shapeVis);
		cov_visuals_.push_back(covVis);
	}

	updateColorAndAlpha();
	context_->queueRender();
}

} // end namespace rviz_multiception

#include <pluginlib/class_list_macros.h>
PLUGINLIB_EXPORT_CLASS(rviz_multiception::PerceptionsDisplay,rviz::Display)