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

#include <ros/console.h>
#include <rviz/ogre_helpers/arrow.h>
#include <rviz/default_plugin/covariance_visual.h>
#include <rviz/default_plugin/covariance_property.h>
#include <geometry_msgs/Vector3.h>

#include <tf2/LinearMath/Quaternion.h>
#include <tf2_geometry_msgs/tf2_geometry_msgs.h>

#include "state_visual.h"

namespace rviz_multiception
{

StateVisual::StateVisual(Ogre::SceneManager* scene_manager, Ogre::SceneNode* parent_node, rviz::CovarianceProperty* covariance_property)
{
	scene_manager_ = scene_manager;

	// Ogre::SceneNode s form a tree, with each node storing the
	// transform (position and orientation) of itself relative to its
	// parent. Ogre does the math of combining those transforms when it
	// is time to render.
	//
	// Here we create a node to store the pose of the header frame
	// relative to the RViz fixed frame.
	frame_node_ = parent_node->createChildSceneNode();

	// We create the arrow object within the frame node so that we can
	// set its position and direction relative to its header frame.
	pose_arrow_.reset(new rviz::Arrow(scene_manager_, frame_node_));
	covariance_ = covariance_property->createAndPushBackVisual(scene_manager, frame_node_);
}

StateVisual::~StateVisual()
{
	// Destroy the frame node since we don't need it anymore.
	scene_manager_->destroySceneNode(frame_node_);
}

void StateVisual::setContent(const multiception::State& state)
{
	// Convert the geometry_msgs::Vector3 to an Ogre::Vector3.
	Ogre::Vector3 pose(state.x, state.y, 1.5);
	pose_arrow_->setPosition(pose);

	// Scale the arrow by the bounded velocity
	float length = std::min(4.0, std::max(1.0, state.v/5.0));
	Ogre::Vector3 scale(length, 1, 1);
	pose_arrow_->setScale(scale);

	// Set the orientation of the arrow according to the heading.
	Ogre::Vector3 direction(std::cos(state.yaw), std::sin(state.yaw), 0);
	pose_arrow_->setDirection(direction);

	// Set the covariance ellipsis
	geometry_msgs::PoseWithCovariance cov;
	//cov.pose.position.x = state.x; cov.pose.position.y = state.y;
	tf2::Quaternion q; 	q.setRPY(0.0, 0.0, state.yaw);
	cov.pose.orientation = tf2::toMsg(q);
	cov.covariance[0*6+0] = state.cov[0*6+0]; // xx
	cov.covariance[0*6+1] = state.cov[0*6+1]; // xy
	cov.covariance[1*6+0] = state.cov[1*6+0]; // yx
	cov.covariance[1*6+1] = state.cov[1*6+1]; // yy

	// Special case, if yaw is 0, we consider the state to be static
	// without orientation (e.g: pole) and displaying an orientation makes no sense
	if(state.v==0) pose_arrow_->setScale(Ogre::Vector3(0.1, 0.1, 0.1));
	else if(state.cov[2*6+2]<M_PI)
	{
		cov.covariance[0*6+5] = state.cov[0*6+2]; // xyaw
		cov.covariance[1*6+5] = state.cov[1*6+2]; // yyaw
		cov.covariance[5*6+0] = state.cov[2*6+0]; // yawx
		cov.covariance[5*6+1] = state.cov[2*6+1]; // yawy
		cov.covariance[5*6+5] = std::max(.0, state.cov[2*6+2]); // yawyaw
	}

	covariance_->setCovariance(cov);
	covariance_->setPosition(pose);
	Ogre::Quaternion covq(Ogre::Radian(state.yaw), Ogre::Vector3::UNIT_Z);
	covariance_->setOrientation(covq);
}

// Position and orientation are passed through to the SceneNode.
void StateVisual::setFramePosition(const Ogre::Vector3& position)
{
	frame_node_->setPosition(position);
}

void StateVisual::setFrameOrientation(const Ogre::Quaternion& orientation)
{
	frame_node_->setOrientation(orientation);
}

// Color is passed through to the Arrow object.
void StateVisual::setColor(float r, float g, float b, float a)
{
	pose_arrow_->setColor(r, g, b, a);
}

} // namespace rviz_multiception