#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/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
{

// The constructor must have no arguments, so we can't give the
// constructor the parameters it needs to fully initialize.
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() ));
}

// After the top-level rviz::Display::initialize() does its own setup,
// it calls the subclass's onInitialize() function. This is where we
// instantiate all the workings of the class. We make sure to also
// call our immediate super-class's onInitialize() function, since it
// does important stuff setting up the message filter.
//
//  Note that "MFDClass" is a typedef of
// ``MessageFilterDisplay<message type>``, to save typing that long
// templated class name every time you need to refer to the
// superclass.
void PerceptionsDisplay::onInitialize()
{
	MFDClass::onInitialize();
}

PerceptionsDisplay::~PerceptionsDisplay()
{
}

// Clear the visuals by deleting their objects.
void PerceptionsDisplay::reset()
{
	MFDClass::reset();
	state_visuals_.clear();
	shape_visuals_.clear();
	text_visuals_.clear();
	trace_buffer_.clear();
	trace_visuals_.clear();
	for(Ogre::SceneNode* textNode: text_nodes_) context_->getSceneManager()->destroySceneNode(textNode);
	text_nodes_.clear();
	text_visuals_.clear();
}

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

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

	for(size_t i=0; i<state_visuals_.size(); i++)
	{
		trace_visuals_[i]->setColor(color.r, color.g, color.b, alpha);
	}
}

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;
	});
}

// This is our callback to handle an incoming message.
void PerceptionsDisplay::processMessage(const multiception::Perceptions::ConstPtr& msg)
{
	// Here we call the rviz::FrameManager to get the transform from the
	// fixed frame to the frame in the header of this message. If
	// it fails, we can't do anything else so we return.
	const size_t N = msg->perceptions.size();
	if(N==0)
	{
		return;
	}

	// We are keeping a vector of state visual pointers
	state_visuals_.clear(); state_visuals_.reserve(N);
	shape_visuals_.clear(); shape_visuals_.reserve(N);
	for(Ogre::SceneNode* textNode: text_nodes_) context_->getSceneManager()->destroySceneNode(textNode);
	text_nodes_.clear(); text_nodes_.reserve(N);
	text_visuals_.clear(); text_visuals_.reserve(N);
	
	const float alpha = alpha_property_->getFloat();
	const float shape_alpha = shape_alpha_property_->getFloat();
	const Ogre::ColourValue color = color_property_->getOgreColor();

	// Process the perceptions in the message
	for(const multiception::Perception& perception : msg->perceptions)
	{
		Ogre::Quaternion orientation;
		Ogre::Vector3 position;

		if(!context_->getFrameManager()->getTransform(perception.header.frame_id, perception.header.stamp, position, orientation))
		{
			ROS_DEBUG("Error transforming from frame '%s' to frame '%s'", perception.header.frame_id.c_str(), qPrintable(fixed_frame_));
			continue;
		}

		boost::shared_ptr<StateVisual> state(new StateVisual(context_->getSceneManager(), scene_node_, covariance_property_));
		boost::shared_ptr<rviz::Shape> shape(new rviz::Shape(rviz::Shape::Cube, context_->getSceneManager(), scene_node_));

		// Set the state visual's pose and content
		state->setContent(perception.state);
		state->setFramePosition(position);
		state->setFrameOrientation(orientation);

		// Set the shape visual's pose and shape
		Ogre::Vector3 shapePose(perception.state.x, perception.state.y, 0.1);
		Ogre::Quaternion shapeOrientation(Ogre::Radian(perception.state.yaw), Ogre::Vector3::UNIT_Z);
		Ogre::Vector3 shapeScale (perception.length, perception.width, std::max(0.1f, perception.height));
		shape->setPosition(shapePose);
		shape->setOrientation(shapeOrientation);
		shape->setScale(shapeScale);

		state->setColor(color.r, color.g, color.b, alpha);
		shape->setColor(color.r, color.g, color.b, shape_alpha);

		// Create a hovering debug text if one is provided
		if(!perception.debug.empty())
		{
			Ogre::SceneNode* childNode = scene_node_->createChildSceneNode();
			boost::shared_ptr<rviz::MovableText> debugText { new rviz::MovableText(perception.debug, "Liberation Sans", .7, text_color_property_->getOgreColor()) };
			debugText->setTextAlignment(rviz::MovableText::H_CENTER, rviz::MovableText::V_CENTER);
			shapePose.z = z+perception.height/2+.1+1;
			childNode->setPosition(shapePose);
			childNode->attachObject(debugText.get());
			text_visuals_.push_back(debugText);
			text_nodes_.push_back(childNode);
		}

		// And add it to the vector
		state_visuals_.push_back(state);
		shape_visuals_.push_back(shape);
	}

	// 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);
			}
		}
	}

	context_->queueRender();
}

} // end namespace rviz_multiception

// Tell pluginlib about this class. It is important to do this in
// global scope, outside our package's namespace.
#include <pluginlib/class_list_macros.h>
PLUGINLIB_EXPORT_CLASS(rviz_multiception::PerceptionsDisplay,rviz::Display)