// Motion player for MoveIt
// Author: Max Schwarz <max.schwarz@uni-bonn.de>

#include <nimbro_keyframe_server/moveit_player.h>

#include <moveit/robot_state/robot_state.h>
#include <moveit/robot_state/conversions.h>
#include <ros/console.h>

#include <Eigen/Geometry>
#include <nimbro_collision_detection/CollisionDetection.h>

namespace nimbro_keyframe_server
{

const std::string UNRESOLVED_COLLISION_DETECTION_SRV = "collision_detection/collisionDetection";

template<class T>
T robustAcos(T x)
{
	return std::acos(std::max<T>(-1.0, std::min<T>(1.0, x)));
}

MoveItPlayer::MoveItPlayer(const robot_model::RobotModelPtr& model,
	const robot_state::RobotStatePtr& state,
	ros::Duration dt)
 : m_model(model)
 , m_state(state)
 , m_dt(dt)

 , m_maxLinearAcceleration("/moveit_player/max_acceleration/linear", 0.0, 0.01, 1.0, 0.3)
 , m_maxAngularAcceleration("/moveit_player/max_acceleration/angular", 0.0, 0.01, 1.0, 0.3)
 , m_maxJointSpaceAcceleration("/moveit_player/max_acceleration/joint_space", 0.0, 0.01, 20.0, 0.9)

 , m_maxLinearVelocity("/moveit_player/max_velocity/linear",0.0,0.01,100.0,30.0)
 , m_maxAngularVelocity("/moveit_player/max_velocity/angular",0.0,0.01,100.0,30.0)
 , m_maxJointSpaceVelocity("/moveit_player/max_velocity/joint_space",0.0,0.01,100.0,30.0)
 , m_maxFinalJointSpaceVelocity("/moveit_player/max_final_joint_velocity_degree", 0.0, 0.1, 100.0, 1.0)

 , m_maxJerk("/moveit_player/max_jerk", 0.0, 0.01, 200.0, 0.2)
 , m_collisionDetection("/moveit_player/collision_detection/enabled", false)
 , m_skipStep(false)
 , m_interrupt(false)
 , m_limitingJointSpaceVelocity(false)
 , m_inCollision(false)
 , m_useCollisionChecking(true)
{
	// Remember all joint Groups
	const std::vector<moveit::core::JointModelGroup*>& modelGroups = m_model->getJointModelGroups();
	std::vector<moveit::core::JointModelGroup*>::const_iterator group = modelGroups.begin();

	ros::NodeHandle nh;

	for (; group != modelGroups.end(); ++group)
	{
		std::string groupName = (*group)->getName();
		m_groups[groupName] = Group();
		m_groups[groupName].jointGroup = *group;
		m_groups[groupName].interpolationSpace = KeyFrame::IS_NONE;
		std::string rotspace_param = "/robot_description_kinematics/" + groupName + "/moveit_player_rotspace";
		std::string rotspace;
		nh.param(rotspace_param,rotspace, std::string("full"));

		if(rotspace == "without_roll")
			m_groups[groupName].rotationSpace = Group::ROTSPACE_WITHOUT_ROLL;
		else if(rotspace == "without_pitch")
			m_groups[groupName].rotationSpace = Group::ROTSPACE_WITHOUT_PITCH;
		else if(rotspace == "without_yaw")
			m_groups[groupName].rotationSpace = Group::ROTSPACE_WITHOUT_YAW;
		else
			m_groups[groupName].rotationSpace = Group::ROTSPACE_FULL;
	}

	// Create a copy of the state for us to play in with IK
	m_ikState.reset(new robot_state::RobotState(*m_state));

	// Trajectory server for FollowJointTrajectory action
	// TODO: Can we move this somewhere else?
	ros::NodeHandle n("~");

	m_trajectoryServer.reset(
		new TrajectoryServer(n, "moveit_player/followJointTrajectory",
			boost::bind(&MoveItPlayer::trajectoryHandleGoal, this, _1),
			boost::bind(&MoveItPlayer::trajectoryCancelGoal, this, _1),
			false
		)
	);
	m_trajectoryServer->start();

	m_firstExec = true;

	m_unmodifiedRobotState.reset(new robot_state::RobotState(*m_state));

	m_collisionDetectionSrv = n.serviceClient<nimbro_collision_detection::CollisionDetection>(n.resolveName(UNRESOLVED_COLLISION_DETECTION_SRV));
}

MoveItPlayer::~MoveItPlayer()
{
}

bool MoveItPlayer::playMotion(const Motion::Ptr& motion)
{
	m_motion = motion;
	m_keyframeIndex = 0;

	std::string name = motion->name();

	m_skipStep =(name.substr(0, 4)=="skip")?true:false;
	m_interrupt = (name=="stop_motion")?true:false;

	setupReflexxes();

	m_inCollision = false;

	return true;
}

void MoveItPlayer::stopMotion()
{
	m_motion.reset();
}

/**
 * returns: the velocity the system is able to reach if it accelerates at the maximal acceleration from the current postion to the desired position
 */
//FIXME use only linear acceleration limit
static double calcMaxVel(double currentPosition,
	double desiredPosition,
	double currentVelocity,
	double maxAcceleration)
{
	double diff = desiredPosition-currentPosition;

	if(diff == 0.0)
		return 0.0;

	if(diff < 0.0)
	{
			maxAcceleration *= -1.0;
	}

	double head = - currentVelocity/maxAcceleration;
	double tail = std::pow(head, 2.0) + 2 * (diff/maxAcceleration);
	if(tail < 0.0) // should never happen
	{
		ROS_ERROR_STREAM("tail < 0.0");
	}
	tail = std::sqrt(tail);

	double t1 = head - tail;
	double t2 = head + tail;

	// select the correct t
	double t = std::max(t1, t2);

	// should never happen
	if((t1>=0 && t2>=0) || (t1<=0 && t2<=0))
	{
		ROS_ERROR_STREAM("(t1>=0 && t2>=0) || (t1<=0 && t2<=0)");
	}

	double maxVel = currentVelocity + maxAcceleration * t;

	return maxVel;
}

Eigen::Vector3d MoveItPlayer::calculateTargetVelocity(std::string groupName,
	Eigen::Vector3d currentPosition,
	Eigen::Vector3d desiredPosition,
	Eigen::Vector4d currentVelocity,
	double linearVelocity)
{
	// the desired velocity is the maximal possible velocity in the direction which points from the previous position to the next position
	Eigen::Vector3d targetVelocity;
	targetVelocity.setZero();
	if(m_keyframeIndex < (m_motion->size()-1)) // make sure that this keyframe is not the last in the motion
	{
			Eigen::Vector3d lastPosition;
			// if this keyframe is the first of the motion the last position is the current position
			if(m_keyframeIndex == 0)
			{
				lastPosition = currentPosition;
			}
			else
			{
				lastPosition = (&(*m_motion)[m_keyframeIndex-1])->state(groupName).translation();
			}

			Eigen::Vector3d nextPosition = (&(*m_motion)[m_keyframeIndex+1])->state(groupName).translation();
			Eigen::Vector3d velocityDirection = nextPosition - lastPosition;

			// check if this results to zero
			if(!velocityDirection.isZero())
			{
				velocityDirection.normalize();

				targetVelocity = velocityDirection * linearVelocity;
			}
	}

	if(!targetVelocity.isZero())
	{
		// the maximal value for alpha is 1.0, this ensures that we do not try to reach a target velocity which is higher than the desired velocity
		double alpha = 1.0;
		double maxAlpha = std::numeric_limits<double>::lowest();
		Eigen::Vector3d maxVel;
		maxVel.setZero();
		// The proportion of the desired velocity needs to be maintained, otherwise the direction whould not be correct
		for (unsigned int j = 0; j < 3; ++j)
		{
			maxVel(j) = calcMaxVel(currentPosition(j),
								   desiredPosition(j),
								   currentVelocity(j),
								   m_maxLinearAcceleration()
  								);

			double tmpAlpha = maxVel(j)/targetVelocity(j);

			if(tmpAlpha<alpha && tmpAlpha >= 0.0)
			{
					alpha = tmpAlpha;
			}
			if(tmpAlpha>maxAlpha)
			{
					maxAlpha = tmpAlpha;
			}
		}

		if(maxAlpha < 0.0)
		{
			alpha = 0.0;
		}

		// we cannot set the maximal reachable velocity as target velocity as the velocity profile of the q spline prohibits this
		targetVelocity *= alpha * 0.5;
	}

	return targetVelocity;
}

void MoveItPlayer::setupReflexxes()
{
	if (! currentKeyFrame())
		return;
	if(m_keyframeIndex >= m_motion->size())
	{
		return;
	}

	m_state->updateLinkTransforms();

	m_jointSpaceVelocityLimit = 0;

	const KeyFrame* frame = currentKeyFrame();

	unsigned int dofs = 0;
	KeyFrame::GroupMap::const_iterator group_it =frame->jointGroups().begin();
	for(; group_it != frame->jointGroups().end(); ++group_it)
	{
		std::string groupName = group_it->first;
		Group* group = &m_groups.at(groupName);

		switch(group_it->second.interpolationSpace())
		{
			case KeyFrame::IS_CARTESIAN:
			{
				dofs += 4;

				// Get the state of the end effector of this group
				auto srdf = m_model->getSRDF();
				std::string tipFrame;

				for(auto eef : srdf->getEndEffectors())
				{
					if(eef.parent_group_ == groupName)
					{
						tipFrame = eef.parent_link_;
						break;
					}
				}

				if(tipFrame.empty())
				{
					ROS_ERROR_STREAM("No end effector link is specified for the group: "<<groupName);
					throw std::runtime_error("No endeffector link");
				}

				Eigen::Affine3d currentState = m_state->getGlobalLinkTransform(tipFrame);

				group->endEffectorPosition << currentState.translation(), 0;

				if(group->interpolationSpace != KeyFrame::IS_CARTESIAN)
				{
// 					ROS_ERROR("%s: ================== RESETTING VELOCITY =========================, last interpolation space: %d", groupName.c_str(), group->interpolationSpace);
					group->endEffectorVelocity.setZero();
					group->endEffectorAcceleration.setZero();
				}

				// either calculate the target velocity or set it so zero
				if(frame->continuous())
				{
					group->endEffectorTargetVelocity = calculateTargetVelocity(
						groupName,
						currentState.translation(),
						frame->state(groupName).translation(),
						group->endEffectorVelocity,
						frame->linearVelocity()
					);
				}
				else
				{
					group->endEffectorTargetVelocity.setZero();
				}

				group->rotationStart = currentState.rotation();
				group->rotationEnd = group_it->second.state().rotation();

				switch(group->rotationSpace)
				{
					case Group::ROTSPACE_FULL:
						group->rotationDist = group->rotationEnd.angularDistance(group->rotationStart);
						break;
					case Group::ROTSPACE_WITHOUT_ROLL:
					{
						Eigen::Vector3d start = group->rotationStart * Eigen::Vector3d::UnitX();
						Eigen::Vector3d end = group->rotationEnd * Eigen::Vector3d::UnitX();
						group->rotationDist = std::abs(robustAcos(start.dot(end)));
						break;
					}
					case Group::ROTSPACE_WITHOUT_PITCH:
					{
						Eigen::Vector3d start = group->rotationStart * Eigen::Vector3d::UnitY();
						Eigen::Vector3d end = group->rotationEnd * Eigen::Vector3d::UnitY();
						group->rotationDist = std::abs(robustAcos(start.dot(end)));
						break;
					}
					case Group::ROTSPACE_WITHOUT_YAW:
					{
						Eigen::Vector3d start = group->rotationStart * Eigen::Vector3d::UnitZ();
						Eigen::Vector3d end = group->rotationEnd * Eigen::Vector3d::UnitZ();
						group->rotationDist = std::abs(robustAcos(start.dot(end)));
						break;
					}
				}

				// Copy current state into the IK state
				*m_ikState = *m_state;
// 				ROS_INFO_STREAM ("Current State : \n"  << group_it->second.state().matrix());

// 				ROS_INFO_STREAM("Cartesian interpolation starting at " << group->endEffectorPosition.transpose());
				break;
			}
			case KeyFrame::IS_JOINT_SPACE:
			{
				unsigned int groupDofs = group->jointGroup->getVariableCount();
				dofs += groupDofs;

				group->jointSpacePositions.resize(groupDofs);
				group->jointSpaceVelocities.resize(groupDofs);
				group->jointSpaceAccelerations.resize(groupDofs);

				m_state->copyJointGroupPositions(group->jointGroup->getName(), group->jointSpacePositions.data());

				if(group->interpolationSpace != KeyFrame::IS_JOINT_SPACE)
				{
					group->jointSpaceVelocities.setZero();
					group->jointSpaceAccelerations.setZero();
				}

				break;
			}
			case KeyFrame::IS_NONE:
				continue;
		}
	}

	for(std::map<std::string, Group>::iterator it = m_groups.begin(); it != m_groups.end(); ++it)
		it->second.interpolationSpace = KeyFrame::IS_NONE;

	for(group_it = frame->jointGroups().begin(); group_it != frame->jointGroups().end(); ++group_it)
	{
		std::string groupName = group_it->first;
		Group* group = &m_groups.at(groupName);

		ROS_INFO("%s: Setting interpolation space: %d", groupName.c_str(), group_it->second.interpolationSpace());
		group->interpolationSpace = group_it->second.interpolationSpace();
	}

	m_reflexxes.reset(new ReflexxesAPI(dofs, m_dt.toSec()));
	m_reflexxesInput.reset(new RMLPositionInputParameters(dofs));
	m_reflexxesOutput.reset(new RMLPositionOutputParameters(dofs));

	if(m_startCallback)
		m_startCallback(*currentKeyFrame());
}

void MoveItPlayer::trajectoryHandleGoal(const TrajectoryGoalHandle& goal)
{
	if(m_trajectoryGoal.isValid())
	{
		ROS_INFO("Preempting active trajectory goal");
		m_trajectoryGoal.setCanceled();
	}

	m_firstExec = true;
	m_trajectoryPoint = 0;
	m_trajectoryGoal = goal;
	m_trajectoryGoal.setAccepted();
}

void MoveItPlayer::trajectoryCancelGoal(TrajectoryGoalHandle goal)
{
	if(goal == m_trajectoryGoal)
		m_trajectoryGoal = TrajectoryGoalHandle();

	goal.setCanceled();
}

void MoveItPlayer::useInternalCollisionChecking(bool use)
{
	m_useCollisionChecking = use;
}

bool MoveItPlayer::step(nimbro_keyframe_server::ExecutionFeedback* feedback)
{
	// backup the current state
	*m_unmodifiedRobotState = *m_state;

	bool result = executeReflexxes(feedback);
	bool revert = false;

	if(m_collisionDetection() && m_useCollisionChecking)
	{
		if(!m_collisionDetectionSrv.exists())
		{
			ROS_ERROR("Could not check for collisions.");
			revert = true;
		}
		else
		{
			nimbro_collision_detection::CollisionDetection srv;
			moveit::core::robotStateToJointStateMsg(*m_state, srv.request.state);

			if(!m_collisionDetectionSrv.call(srv))
			{
				ROS_ERROR("Could not check for collisions.");
				revert = true;
			}

			bool inCollision = (srv.response.distance <= 0.0);
			if(inCollision)
			{
				for(auto& pair: srv.response.pairs)
				{
					ROS_WARN("collision between: '%s' and '%s'.",pair.first.c_str(), pair.second.c_str());
				}
				revert = true;
				m_inCollision = true;
			}
		}

		if(revert)
		{
			ROS_WARN("Stopping execution. Robot is nearly in collision.");

			// revert the changes
			*m_state = *m_unmodifiedRobotState;

			// stop the current motion
			m_motion.reset();
			result = false;
		}
	}

	return result;
}

bool MoveItPlayer::executeReflexxes(nimbro_keyframe_server::ExecutionFeedback* feedback)
{
	if(m_trajectoryGoal.isValid())
	{
		auto trajectory = m_trajectoryGoal.getGoal()->trajectory;

		if(m_trajectoryPoint < trajectory.points.size())
		{
			auto point = trajectory.points[m_trajectoryPoint];

			if(m_firstExec)
			{
				m_firstExec = false;
				m_execTime = ros::Time::now();
			}

			if((ros::Time::now()-m_execTime).toNSec() >= point.time_from_start.toNSec())
			{
				for(unsigned int j=0; j < trajectory.joint_names.size(); ++j)
				{
					m_state->setVariableAcceleration(trajectory.joint_names[j], 0.0 /*m_trajectory.points[0].accelerations[j]*/);
					m_state->setVariableVelocity(trajectory.joint_names[j], point.velocities[j]);
					m_state->setVariablePosition(trajectory.joint_names[j], point.positions[j]);
				}

				m_trajectoryPoint++;
			}
		}
		else
		{
			control_msgs::FollowJointTrajectoryResult rs;
			rs.error_code = rs.SUCCESSFUL;
			m_trajectoryGoal.setSucceeded(rs);
			m_trajectoryGoal = TrajectoryGoalHandle();
		}

		if(feedback)
			feedback->playing = false;

		return false;
	}

	if(!m_motion)
	{
		if(feedback)
			feedback->playing = false;

		return false;
	}

	if(m_keyframeIndex >= m_motion->size() || m_interrupt)
	{
		m_motion.reset();

		if(feedback)
			feedback->playing = false;

		return false;
	}

	if(feedback)
	{
		feedback->playing = true;
		feedback->keyframe_index = m_keyframeIndex;
	}

	const KeyFrame* frame = currentKeyFrame();

	// Accelerate with m_maxJointSpaceAcceleration until we reach the velocity
	// limit
	double currentMaxJointSpaceVelocity = std::isfinite(frame->jointSpaceVelocity()) ? frame->jointSpaceVelocity() : m_maxJointSpaceVelocity();
	double currentMaxJointSpaceAcceleration = std::isfinite(frame->jointSpaceAcceleration()) ? frame->jointSpaceAcceleration() : m_maxJointSpaceAcceleration();

	if(m_jointSpaceVelocityLimit < currentMaxJointSpaceVelocity)
	{
		m_jointSpaceVelocityLimit = std::min<double>(
			currentMaxJointSpaceVelocity,
			m_jointSpaceVelocityLimit + m_dt.toSec() * currentMaxJointSpaceAcceleration
		);
	}

	KeyFrame::GroupMap::const_iterator group_it = frame->jointGroups().begin();
	unsigned int dof_offset = 0;
	for (; group_it != frame->jointGroups().end(); ++group_it)
	{
		const std::string& groupName = group_it->first;
		Group* group = &m_groups.at(groupName);

		switch(group_it->second.interpolationSpace())
		{
			case KeyFrame::IS_CARTESIAN:
			{
// 				ROS_INFO("dof_offset: %u", dof_offset);
				for(unsigned int j = 0; j < 4; ++j)
				{
// 					ROS_INFO("pos input %d: %f", j, group->endEffectorPosition[j]);
					m_reflexxesInput->CurrentPositionVector->VecData[dof_offset + j] = group->endEffectorPosition[j];
					m_reflexxesInput->CurrentVelocityVector->VecData[dof_offset + j] = group->endEffectorVelocity[j];
					m_reflexxesInput->CurrentAccelerationVector->VecData[dof_offset + j] = group->endEffectorAcceleration[j];
				}

				for (unsigned int j = 0; j < 3; ++j)
				{
					m_reflexxesInput->TargetPositionVector->VecData[dof_offset + j] = frame->state(groupName).translation()[j];
					m_reflexxesInput->TargetVelocityVector->VecData[dof_offset + j] = group->endEffectorTargetVelocity(j);

					m_reflexxesInput->MaxVelocityVector->VecData[dof_offset + j] = std::isfinite(frame->linearVelocity()) ?
						frame->linearVelocity() : m_maxLinearVelocity();
					m_reflexxesInput->MaxAccelerationVector->VecData[dof_offset + j] = std::isfinite(frame->linearAcceleration()) ?
						frame->linearAcceleration() : m_maxLinearAcceleration();
					m_reflexxesInput->MaxJerkVector->VecData[dof_offset + j] = m_maxJerk();
					m_reflexxesInput->SelectionVector->VecData[dof_offset + j] = true;
				}

				m_reflexxesInput->TargetPositionVector->VecData[dof_offset + 3] = group->rotationDist;
				m_reflexxesInput->MaxVelocityVector->VecData[dof_offset + 3] = std::isfinite(frame->angularVelocity()) ?
					frame->angularVelocity() : m_maxAngularVelocity();
				m_reflexxesInput->MaxAccelerationVector->VecData[dof_offset + 3] = std::isfinite(frame->angularAcceleration()) ?
					frame->angularAcceleration() : m_maxAngularAcceleration();
				m_reflexxesInput->MaxJerkVector->VecData[dof_offset + 3] = m_maxJerk();
				m_reflexxesInput->SelectionVector->VecData[dof_offset + 3] = true;
				dof_offset += 4;
				break;
			}
			case KeyFrame::IS_JOINT_SPACE:
			{
				std::vector<std::string> gJoints = group->jointGroup->getVariableNames();
				unsigned int varCount = group->jointGroup->getVariableCount();
				for(unsigned int j = 0; j < varCount; ++j)
				{
					m_reflexxesInput->CurrentPositionVector->VecData[dof_offset + j] = group->jointSpacePositions[j];
					m_reflexxesInput->CurrentVelocityVector->VecData[dof_offset + j] = group->jointSpaceVelocities[j];
					m_reflexxesInput->CurrentAccelerationVector->VecData[dof_offset + j] = group->jointSpaceAccelerations[j];

					// Target position

					std::map<std::string, double>::const_iterator it = frame->jointPositions().find(gJoints[j]);
					if(it == frame->jointPositions().end())
						m_reflexxesInput->TargetPositionVector->VecData[dof_offset + j] = group->jointSpacePositions[j];
					else
						m_reflexxesInput->TargetPositionVector->VecData[dof_offset + j] = it->second;

					m_reflexxesInput->MaxVelocityVector->VecData[dof_offset + j] = std::isfinite(frame->jointSpaceVelocity()) ?
						frame->jointSpaceVelocity() : m_maxJointSpaceVelocity();

					m_reflexxesInput->MaxAccelerationVector->VecData[dof_offset + j] = std::isfinite(frame->jointSpaceAcceleration()) ?
							frame->jointSpaceAcceleration() : m_maxJointSpaceAcceleration();

					m_reflexxesInput->MaxJerkVector->VecData[dof_offset + j] = m_maxJerk();
					m_reflexxesInput->SelectionVector->VecData[dof_offset + j] = true;
				}
				dof_offset += varCount;
				break;
			}
			case KeyFrame::IS_NONE:
				continue;
		}
	}

	// Reflexxes call
	int ret = m_reflexxes->RMLPosition(*m_reflexxesInput, m_reflexxesOutput.get(), m_reflexxesFlags);

	m_state->updateLinkTransforms();

// 	ROS_INFO("Reflexxes: %d", ret);

	dof_offset = 0;

	double maxJointSpaceVelocity = 0.0f;

	if(feedback)
		feedback->active_limits = 0;

	// Reflexxes output
	group_it = frame->jointGroups().begin();
	dof_offset = 0;
	for (; group_it != frame->jointGroups().end(); ++group_it)
	{
		const std::string& groupName = group_it->first;
		Group* group = &m_groups[groupName];

		switch(group_it->second.interpolationSpace())
		{
			case KeyFrame::IS_CARTESIAN:
			{
				// We need to use the result of the Reflexxes call in this iteration
				Eigen::Vector4d nextEndEffectorPosition = Eigen::Map<Eigen::Vector4d>(m_reflexxesOutput->NewPositionVector->VecData + dof_offset);
				Eigen::Vector4d nextEndEffectorVelocity = Eigen::Map<Eigen::Vector4d>(m_reflexxesOutput->NewVelocityVector->VecData + dof_offset);
				Eigen::Vector4d nextEndEffectorAcceleration = Eigen::Map<Eigen::Vector4d>(m_reflexxesOutput->NewAccelerationVector->VecData + dof_offset);
				group->endEffectorPosition = nextEndEffectorPosition;
				group->endEffectorVelocity = nextEndEffectorVelocity;
				group->endEffectorAcceleration = nextEndEffectorAcceleration;

				// Check if we are close to vel/acc limits for operator feedback
				if(feedback)
				{
					double linVel = group->endEffectorVelocity.head<3>().array().abs().maxCoeff();
					double linAcc = group->endEffectorAcceleration.head<3>().array().abs().maxCoeff();
					double angVel = std::abs(group->endEffectorVelocity[3]);
					double angAcc = std::abs(group->endEffectorAcceleration[3]);

					if(std::abs(linVel - m_reflexxesInput->MaxVelocityVector->VecData[dof_offset]) < 1e-2)
						feedback->active_limits |= ExecutionFeedback::LIN_VEL_LIMIT;
					if(std::abs(linAcc - m_reflexxesInput->MaxAccelerationVector->VecData[dof_offset]) < 1e-2)
						feedback->active_limits |= ExecutionFeedback::LIN_ACC_LIMIT;
					if(std::abs(angVel - m_reflexxesInput->MaxVelocityVector->VecData[dof_offset + 3]) < 1e-2)
						feedback->active_limits |= ExecutionFeedback::ANG_VEL_LIMIT;
					if(std::abs(angAcc - m_reflexxesInput->MaxAccelerationVector->VecData[dof_offset + 3]) < 1e-2)
						feedback->active_limits |= ExecutionFeedback::ANG_ACC_LIMIT;
				}

				// IK
				Eigen::Affine3d state;
				state = Eigen::Translation3d(group->endEffectorPosition.head<3>());

				if(fabs(group->rotationDist) > 1e-4)
				{
					state = state * group->rotationStart.slerp(group->endEffectorPosition[3] / group->rotationDist, group->rotationEnd);
				}
				else
				{
					state = state * group->rotationEnd;
				}

				*m_ikState = *m_state;

// 				ROS_INFO_STREAM("interpolated state: " << state.translation().transpose() << ", euler: " << state.rotation().eulerAngles(0,1,2).transpose());

				if(!m_ikState->setFromIK(group->jointGroup, state) )
				{
					if(!m_skipStep)
					{
						ROS_ERROR("IK failed at keyframe %d, stopping motion...", m_keyframeIndex);
						ROS_ERROR_STREAM("state was: " << state.translation().transpose() << ", euler: " << state.rotation().eulerAngles(0,1,2).transpose());
						m_motion.reset();
						m_inCollision = true;
						return true;
					}
					// if skip is possible try with the next keyframe
					ROS_WARN("********************    Skipping Step     ******************** ");
					m_keyframeIndex++;
					return false;
				}

				Eigen::VectorXd oldPos(group->jointGroup->getVariableCount());
				m_state->copyJointGroupPositions(group->jointGroup, oldPos.data());

				Eigen::VectorXd newPos(group->jointGroup->getVariableCount());
				m_ikState->copyJointGroupPositions(group->jointGroup, newPos.data());

				Eigen::VectorXd jointDelta = newPos - oldPos;

				double maxDelta = jointDelta.array().abs().maxCoeff();
				double maxVel = maxDelta / m_dt.toSec();

				if(maxVel > m_jointSpaceVelocityLimit)
				{
					// Scale jointDelta back so that we are within our limits
					double alpha = m_jointSpaceVelocityLimit / maxVel;

					jointDelta = alpha * jointDelta;
					maxVel = alpha * maxVel;

					// Taylor approximation
					nextEndEffectorPosition = group->endEffectorPosition + alpha * (nextEndEffectorPosition - group->endEffectorPosition);
					nextEndEffectorVelocity = alpha * nextEndEffectorVelocity;
					nextEndEffectorAcceleration = alpha * nextEndEffectorAcceleration;

					// Write joint values
					Eigen::VectorXd nextJointAngles = oldPos + jointDelta;
					m_state->setJointGroupPositions(group->jointGroup, nextJointAngles.data());

					m_limitingJointSpaceVelocity = true;

					if(feedback)
						feedback->active_limits |= ExecutionFeedback::CART_JS_VEL_LIMIT;
				}
				else
				{
					m_state->setJointGroupPositions(group->jointGroup, newPos.data());
					m_limitingJointSpaceVelocity = false;
				}

				maxJointSpaceVelocity = std::max(maxVel, maxJointSpaceVelocity);

				Eigen::VectorXd jointVelocity = jointDelta / m_dt.toSec();
				{
					const std::vector<int> &il = group->jointGroup->getVariableIndexList();
					for (std::size_t i = 0 ; i < il.size() ; ++i)
					{
						m_state->setVariableVelocity(il[i], jointVelocity[i]);
					}
				}

				// Update reflexxes state so the motion continues.
				group->endEffectorPosition = nextEndEffectorPosition;
				group->endEffectorVelocity = nextEndEffectorVelocity;
				group->endEffectorAcceleration = nextEndEffectorAcceleration;

				dof_offset += 4;
				break;
			}
			case KeyFrame::IS_JOINT_SPACE:
			{
				std::vector<std::string> gJoints = group->jointGroup->getVariableNames();
				int varCount  = group->jointGroup->getVariableCount();
				for (int j = 0; j < varCount; ++j)
				{
					group->jointSpacePositions[j] = m_reflexxesOutput->NewPositionVector->VecData[j + dof_offset];
					group->jointSpaceVelocities[j] = m_reflexxesOutput->NewVelocityVector->VecData[j + dof_offset];
					group->jointSpaceAccelerations[j] = m_reflexxesOutput->NewAccelerationVector->VecData[j + dof_offset];

					maxJointSpaceVelocity = std::max(std::abs(group->jointSpaceVelocities[j]), maxJointSpaceVelocity);

					// Check if we are close to vel/acc limits for operator feedback
					if(feedback)
					{
						if(std::abs(group->jointSpaceVelocities[j] - m_reflexxesInput->MaxVelocityVector->VecData[j + dof_offset]) < 1e-2)
							feedback->active_limits |= ExecutionFeedback::JS_VEL_LIMIT;
						if(std::abs(group->jointSpaceAccelerations[j] - m_reflexxesInput->MaxAccelerationVector->VecData[j + dof_offset]) < 1e-2)
							feedback->active_limits |= ExecutionFeedback::JS_ACC_LIMIT;
					}
				}
				m_state->setJointGroupPositions(group->jointGroup, group->jointSpacePositions.data());

				{
					const std::vector<int> &il = group->jointGroup->getVariableIndexList();
					for (std::size_t i = 0 ; i < il.size() ; ++i)
						m_state->setVariableVelocity(il[i],group->jointSpaceVelocities[i]);
				}

				dof_offset += varCount;
				break;
			}
			case KeyFrame::IS_NONE:
				continue;
		}
	}

	if(m_keyframeIndex + 1 >= m_motion->size())
	{
		// We are in the last keyframe of the motion. Make sure that we are at
		// complete rest before ending the trajectory. In particular, the
		// IK might still be performing nullspace optimization.

		if(maxJointSpaceVelocity > m_maxFinalJointSpaceVelocity() * M_PI / 180.0)
			return false; // not finished
	}

	if(ret == ReflexxesAPI::RML_FINAL_STATE_REACHED)
	{
// 		ROS_INFO_STREAM("ReflexxesAPI::RML_FINAL_STATE_REACHED");

		if(m_endCallback)
			m_endCallback(*currentKeyFrame());

		m_keyframeIndex++;
		setupReflexxes();

		if(m_keyframeIndex >= m_motion->size())
		{
			// the final state of the last keyframe was reached
			return true;
		}
	}

	return false;
}

void MoveItPlayer::setStartCallback(const Callback& cb)
{
	m_startCallback = cb;
}

void MoveItPlayer::setEndCallback(const Callback& cb)
{
	m_endCallback = cb;
}

Eigen::Vector4d MoveItPlayer::getEndEffectorPosition(std::string groupName)
{
	auto it=m_groups.find(groupName);
	if(it==m_groups.end())
	{
		ROS_ERROR("Can't find group '%s'.",groupName.c_str());
		return Eigen::Vector4d().setZero();
	}
		return m_groups[groupName].endEffectorPosition;
}

Eigen::Vector4d MoveItPlayer::getEndEffectorVelocity(std::string groupName)
{
	auto it=m_groups.find(groupName);
	if(it==m_groups.end())
	{
		ROS_ERROR("Can't find group '%s'.",groupName.c_str());
		return Eigen::Vector4d().setZero();
	}
	return m_groups[groupName].endEffectorVelocity;
}

Eigen::Vector4d MoveItPlayer::getEndEffectorAcceleration(std::string groupName)
{
	auto it=m_groups.find(groupName);
	if(it==m_groups.end())
	{
		ROS_ERROR("Can't find group '%s'.",groupName.c_str());
		return Eigen::Vector4d().setZero();
	}
	return m_groups[groupName].endEffectorAcceleration;
}

void MoveItPlayer::debugDump()
{
	ROS_INFO("Current MoveItPlayer state:");
	for(auto& group : m_groups)
	{
		ROS_INFO(" - Group %s:", group.first.c_str());
		ROS_INFO_STREAM(" endEffectorPosition:\n" << group.second.endEffectorPosition.matrix());
		ROS_INFO_STREAM(" jointSpacePositions: " << group.second.jointSpacePositions.transpose());
		ROS_INFO_STREAM(" rotationStart: " << group.second.rotationStart.coeffs().transpose());
		ROS_INFO_STREAM(" rotationEnd: " << group.second.rotationEnd.coeffs().transpose());
		ROS_INFO_STREAM(" rotationDist: " << group.second.rotationDist);
	}

	if(m_keyframeIndex < m_motion->size())
	{
		ROS_INFO("Current MoveItPlayer keyframe:");
		ROS_INFO("%s", m_motion->serializeToYAML().c_str());
	}
}

}


//==================================================================================================
/*!
  @file

  @copyright 2015 NumScale SAS
  @copyright 2015 J.T. Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_ROR_HPP_INCLUDED
#define BOOST_SIMD_ARCH_COMMON_GENERIC_FUNCTION_ROR_HPP_INCLUDED

#include <boost/simd/meta/cardinal_of.hpp>
#include <boost/simd/detail/assert_utils.hpp>
#include <boost/simd/function/bitwise_cast.hpp>
#include <boost/simd/function/bitwise_or.hpp>
#include <boost/simd/function/minus.hpp>
#include <boost/simd/function/shift_left.hpp>
#include <boost/simd/function/shr.hpp>
#include <boost/simd/function/splat.hpp>
#include <boost/dispatch/function/overload.hpp>
#include <boost/dispatch/meta/as_integer.hpp>
#include <boost/dispatch/meta/scalar_of.hpp>
#include <boost/config.hpp>

namespace boost { namespace simd { namespace ext
{
  namespace bd = boost::dispatch;
  namespace bs = boost::simd;

  BOOST_DISPATCH_OVERLOAD ( ror_
                          , (typename A0, typename A1)
                          , bd::cpu_
                          , bd::generic_< bd::integer_<A0> >
                          , bd::generic_< bd::integer_<A1> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 const& a0, A0 const& a1
//  TODO                                   , typename std::enable_if<bs::cardinal_of<A1>::value
//                                      == bs::cardinal_of<A0>::value>::type* = 0
                                    ) const
    {
      using s_t = bd::scalar_of_t<A0>;
      BOOST_ASSERT_MSG(assert_good_shift<A0>(a1), "ror : rotation is out of range");

      s_t const width = sizeof(s_t)*CHAR_BIT;
      return shr(a0, a1) | shift_left(a0, (width-splat<A0>(a1)) & (width-1));
    }
  };

  BOOST_DISPATCH_OVERLOAD ( ror_
                          , (typename A0, typename A1)
                          , bd::cpu_
                          , bd::generic_< bd::floating_<A0> >
                          , bd::generic_< bd::integer_<A1> >
                          )
  {
    BOOST_FORCEINLINE A0 operator() ( A0 const& a0, A1 const& a1
//   TODO                                  , typename std::enable_if<bs::cardinal_of<A1>::value
//                                      == bs::cardinal_of<A0>::value>::type* = 0
                                    ) const
    {
      using i_t = bd::as_integer_t<A0, unsigned>;
      return bitwise_cast<A0>( ror ( bitwise_cast<i_t>(a0)
                                            , splat<i_t>(a1)
                                            )
                                      );
    }
  };
} } }

#endif


#include <Spirit_Defines.h>
#include <engine/Method_MC.hpp>
#include <engine/Vectormath.hpp>
#include <data/Spin_System.hpp>
#include <data/Spin_System_Chain.hpp>
#include <io/IO.hpp>
#include <utility/Logging.hpp>
#include <utility/Constants.hpp>

#include <Eigen/Dense>

#include <iostream>
#include <ctime>
#include <math.h>

using namespace Utility;

namespace Engine
{
    Method_MC::Method_MC(std::shared_ptr<Data::Spin_System> system, int idx_img, int idx_chain) :
        Method(system->mc_parameters, idx_img, idx_chain)
    {
        // Currently we only support a single image being iterated at once:
        this->systems = std::vector<std::shared_ptr<Data::Spin_System>>(1, system);
        this->SenderName = Log_Sender::MC;

        this->noi = this->systems.size();
        this->nos = this->systems[0]->geometry->nos;
        this->nos_nonvacant = this->systems[0]->geometry->nos_nonvacant;

        this->xi = vectorfield(this->nos, {0,0,0});

        // We assume it is not converged before the first iteration
        // this->max_torque = system->mc_parameters->force_convergence + 1.0;

        // History
        this->history = std::map<std::string, std::vector<scalar>>{
            {"max_torque", {this->max_torque}},
            {"E", {this->max_torque}},
            {"M_z", {this->max_torque}} };

        this->parameters_mc = system->mc_parameters;

        // Starting cone angle
        this->cone_angle = Constants::Pi * this->parameters_mc->metropolis_cone_angle / 180.0;
        this->n_rejected = 0;
        this->acceptance_ratio_current = this->parameters_mc->acceptance_ratio_target;
    }

    // This implementation is mostly serial as parallelization is nontrivial
    //      if the range of neighbours for each atom is not pre-defined.
    void Method_MC::Iteration()
    {
        // Temporaries
        auto& spins_old       = *this->systems[0]->spins;
        auto  spins_new       = spins_old;

        // Generate randomly displaced spin configuration according to cone radius
        // Vectormath::get_random_vectorfield_unitsphere(this->parameters_mc->prng, random_unit_vectors);

        // TODO: add switch between Metropolis and heat bath
        // One Metropolis step
        Metropolis(spins_old, spins_new);
        Vectormath::set_c_a(1, spins_new, spins_old);
    }

    // Simple metropolis step
    void Method_MC::Metropolis(const vectorfield & spins_old, vectorfield & spins_new)
    {
        auto distribution = std::uniform_real_distribution<scalar>(0, 1);
        auto distribution_idx = std::uniform_int_distribution<>(0, this->nos-1);
        scalar kB_T = Constants::k_B * this->parameters_mc->temperature;

        scalar diff = 0.01;

        // Cone angle feedback algorithm
        if( this->parameters_mc->metropolis_step_cone && this->parameters_mc->metropolis_cone_adaptive )
        {
            this->acceptance_ratio_current = 1 - (scalar)this->n_rejected / (scalar)this->nos_nonvacant;

            if( (this->acceptance_ratio_current < this->parameters_mc->acceptance_ratio_target) && (this->cone_angle > diff) )
                this->cone_angle -= diff;

            if( (this->acceptance_ratio_current > this->parameters_mc->acceptance_ratio_target) && (this->cone_angle < Constants::Pi-diff) )
                this->cone_angle += diff;

            this->parameters_mc->metropolis_cone_angle = this->cone_angle * 180.0 / Constants::Pi;
        }
        this->n_rejected = 0;

        // One Metropolis step for each spin
        Vector3 e_z{0, 0, 1};
        scalar costheta, sintheta, phi;
        Matrix3 local_basis;
        scalar cos_cone_angle = std::cos(cone_angle);

        // Loop over NOS samples (on average every spin should be hit once per Metropolis step)
        for (int idx=0; idx < this->nos; ++idx)
        {
            int ispin;
            if( this->parameters_mc->metropolis_random_sample )
                // Better statistics, but additional calculation of random number
                ispin = distribution_idx(this->parameters_mc->prng);
            else
                // Faster, but worse statistics
                ispin = idx;

            if( Vectormath::check_atom_type(this->systems[0]->geometry->atom_types[ispin]) )
            {
                // Sample a cone
                if( this->parameters_mc->metropolis_step_cone )
                {
                    // Calculate local basis for the spin
                    if( spins_old[ispin].z() < 1-1e-10 )
                    {
                        local_basis.col(2) = spins_old[ispin];
                        local_basis.col(0) = (local_basis.col(2).cross(e_z)).normalized();
                        local_basis.col(1) = local_basis.col(2).cross(local_basis.col(0));
                    }
                    else
                    {
                        local_basis = Matrix3::Identity();
                    }

                    // Rotation angle between 0 and cone_angle degrees
                    costheta = 1 - (1 - cos_cone_angle) * distribution(this->parameters_mc->prng);

                    sintheta = std::sqrt(1 - costheta*costheta);

                    // Random distribution of phi between 0 and 360 degrees
                    phi = 2*Constants::Pi * distribution(this->parameters_mc->prng);

                    // New spin orientation in local basis
                    Vector3 local_spin_new{ sintheta * std::cos(phi),
                                            sintheta * std::sin(phi),
                                            costheta };

                    // New spin orientation in regular basis
                    spins_new[ispin] = local_basis * local_spin_new;
                }
                // Sample the entire unit sphere
                else
                {
                    // Rotation angle between 0 and 180 degrees
                    costheta = distribution(this->parameters_mc->prng);

                    sintheta = std::sqrt(1 - costheta*costheta);

                    // Random distribution of phi between 0 and 360 degrees
                    phi = 2*Constants::Pi * distribution(this->parameters_mc->prng);

                    // New spin orientation in local basis
                    spins_new[ispin] = Vector3{ sintheta * std::cos(phi),
                                                sintheta * std::sin(phi),
                                                costheta };
                }

                // Energy difference of configurations with and without displacement
                scalar Eold  = this->systems[0]->hamiltonian->Energy_Single_Spin(ispin, spins_old);
                scalar Enew  = this->systems[0]->hamiltonian->Energy_Single_Spin(ispin, spins_new);
                scalar Ediff = Enew-Eold;

                // Metropolis criterion: reject the step if energy rose
                if( Ediff > 1e-14 )
                {
                    if( this->parameters_mc->temperature < 1e-12 )
                    {
                        // Restore the spin
                        spins_new[ispin] = spins_old[ispin];
                        // Counter for the number of rejections
                        ++this->n_rejected;
                    }
                    else
                    {
                        // Exponential factor
                        scalar exp_ediff    = std::exp( -Ediff/kB_T );
                        // Metropolis random number
                        scalar x_metropolis = distribution(this->parameters_mc->prng);

                        // Only reject if random number is larger than exponential
                        if( exp_ediff < x_metropolis )
                        {
                            // Restore the spin
                            spins_new[ispin] = spins_old[ispin];
                            // Counter for the number of rejections
                            ++this->n_rejected;
                        }
                    }
                }
            }
        }
    }

    // TODO:
    // Implement heat bath algorithm, see Y. Miyatake et al, J Phys C: Solid State Phys 19, 2539 (1986)
    // void Method_MC::HeatBath(const vectorfield & spins_old, vectorfield & spins_new)
    // {
    // }

    void Method_MC::Hook_Pre_Iteration()
    {
    }

    void Method_MC::Hook_Post_Iteration()
    {
    }

    void Method_MC::Initialize()
    {
    }

    void Method_MC::Finalize()
    {
        this->systems[0]->iteration_allowed = false;
    }

    void Method_MC::Message_Start()
    {
        //---- Log messages
        std::vector<std::string> block(0);
        block.push_back(fmt::format("------------  Started  {} Calculation  ------------", this->Name()));
        block.push_back(fmt::format("    Going to iterate {} step(s)", this->n_log));
        block.push_back(fmt::format("                with {} iterations per step", this->n_iterations_log));
        if( this->parameters_mc->metropolis_step_cone )
        {
            if( this->parameters_mc->metropolis_cone_adaptive )
            {
                block.push_back(fmt::format("   Target acceptance {:>6.3f}", this->parameters_mc->acceptance_ratio_target));
                block.push_back(fmt::format("   Cone angle (deg): {:>6.3f} (adaptive)", this->cone_angle*180/Constants::Pi));
            }
            else
            {
                block.push_back(fmt::format("   Target acceptance {:>6.3f}", this->parameters_mc->acceptance_ratio_target));
                block.push_back(fmt::format("   Cone angle (deg): {:>6.3f} (non-adaptive)", this->cone_angle*180/Constants::Pi));
            }
        }
        block.push_back("-----------------------------------------------------");
        Log.SendBlock(Log_Level::All, this->SenderName, block, this->idx_image, this->idx_chain);
    }

    void Method_MC::Message_Step()
    {
        // Update time of current step
        auto t_current = system_clock::now();

        // Update the system's energy
        this->systems[0]->UpdateEnergy();

        // Send log message
        std::vector<std::string> block(0);
        block.push_back(fmt::format("----- {} Calculation: {}", this->Name(), Timing::DateTimePassed(t_current - this->t_start)));
        block.push_back(fmt::format("    Completed                 {} / {} step(s) (step size {})", this->step, this->n_log, this->n_iterations_log));
        block.push_back(fmt::format("    Iteration                 {} / {}", this->iteration, this->n_iterations));
        block.push_back(fmt::format("    Time since last step:     {}", Timing::DateTimePassed(t_current - this->t_last)));
        block.push_back(fmt::format("    Iterations / sec:         {}", this->n_iterations_log / Timing::SecondsPassed(t_current - this->t_last)));
        if( this->parameters_mc->metropolis_step_cone )
        {
            if( this->parameters_mc->metropolis_cone_adaptive )
            {
                block.push_back(fmt::format("    Current acceptance ratio: {:>6.3f} (target {})", this->acceptance_ratio_current, this->parameters_mc->acceptance_ratio_target));
                block.push_back(fmt::format("    Current cone angle (deg): {:>6.3f} (adaptive)", this->cone_angle*180/Constants::Pi));
            }
            else
            {
                block.push_back(fmt::format("    Current acceptance ratio: {:>6.3f}", this->acceptance_ratio_current));
                block.push_back(fmt::format("    Current cone angle (deg): {:>6.3f} (non-adaptive)", this->cone_angle*180/Constants::Pi));
            }
        }
        block.push_back(fmt::format("    Total energy:             {:20.10f}", this->systems[0]->E));
        Log.SendBlock(Log_Level::All, this->SenderName, block, this->idx_image, this->idx_chain);

        // Update time of last step
        this->t_last = t_current;
    }

    void Method_MC::Message_End()
    {
        //---- End timings
        auto t_end = system_clock::now();

        //---- Termination reason
        std::string reason = "";
        if( this->StopFile_Present() )
            reason = "A STOP file has been found";
        else if( this->Walltime_Expired(t_end - this->t_start) )
            reason = "The maximum walltime has been reached";

        // Update the system's energy
        this->systems[0]->UpdateEnergy();

        //---- Log messages
        std::vector<std::string> block;
        block.push_back(fmt::format("------------ Terminated {} Calculation ------------", this->Name()));
        if( reason.length() > 0 )
            block.push_back(fmt::format("----- Reason:   {}", reason));
        block.push_back(fmt::format("----- Duration:       {}", Timing::DateTimePassed(t_end - this->t_start)));
        block.push_back(fmt::format("    Completed         {} / {} step(s)", this->step, this->n_log));
        block.push_back(fmt::format("    Iteration         {} / {}", this->iteration, this->n_iterations));
        block.push_back(fmt::format("    Iterations / sec: {}", this->iteration / Timing::SecondsPassed(t_end - this->t_start)));
        if( this->parameters_mc->metropolis_step_cone )
        {
            if( this->parameters_mc->metropolis_cone_adaptive )
            {
                block.push_back(fmt::format("    Acceptance ratio: {:>6.3f} (target {})", this->acceptance_ratio_current, this->parameters_mc->acceptance_ratio_target));
                block.push_back(fmt::format("    Cone angle (deg): {:>6.3f} (adaptive)", this->cone_angle*180/Constants::Pi));
            }
            else
            {
                block.push_back(fmt::format("    Acceptance ratio: {:>6.3f}", this->acceptance_ratio_current));
                block.push_back(fmt::format("    Cone angle (deg): {:>6.3f} (non-adaptive)", this->cone_angle*180/Constants::Pi));
            }
        }
        block.push_back(fmt::format("    Total energy:     {:20.10f}", this->systems[0]->E));
        block.push_back("-----------------------------------------------------");
        Log.SendBlock(Log_Level::All, this->SenderName, block, this->idx_image, this->idx_chain);
    }

    void Method_MC::Save_Current(std::string starttime, int iteration, bool initial, bool final)
    {
    }

    // Method name as string
    std::string Method_MC::Name() { return "MC"; }
}


/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkBirdEyeViewSnap.cxx
  Author: Pierre Guilbert

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

  This software is distributed WITHOUT ANY WARRANTY; without even
  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
  PURPOSE.  See the above copyright notice for more information.

=========================================================================*/

// LOCAL
#include "vtkBirdEyeViewSnap.h"

// STD
#include <iostream>
#include <fstream>
#include <sstream>
#include <cmath>

// VTK
#include <vtkCellArray.h>
#include <vtkCellData.h>
#include <vtkDataArray.h>
#include <vtkDoubleArray.h>
#include <vtkFloatArray.h>
#include <vtkImageData.h>
#include <vtkInformation.h>
#include <vtkInformationVector.h>
#include <vtkMath.h>
#include <vtkNew.h>
#include <vtkObjectFactory.h>
#include <vtkPointData.h>
#include <vtkPoints.h>
#include <vtkPolyData.h>
#include <vtkPolyLine.h>
#include <vtkSmartPointer.h>
#include <vtkStreamingDemandDrivenPipeline.h>
#include <vtkTransform.h>
#include <vtkTupleInterpolator.h>
#include <vtkUnsignedCharArray.h>
#include <vtkUnsignedIntArray.h>
#include <vtkUnsignedShortArray.h>
#include <vtkPNGWriter.h>

// BOOST
#include <boost/algorithm/string.hpp>

// Eigen
#include <Eigen/Dense>

// Implementation of the New function
vtkStandardNewMacro(vtkBirdEyeViewSnap)

//----------------------------------------------------------------------------
vtkBirdEyeViewSnap::vtkBirdEyeViewSnap()
{
  this->pixelResX = 0.20; // 20cm
  this->pixelResY = 0.20; // 20cm
  this->Count = 0;
  this->RadicalFileName = "NoRadical";
  this->ExtensionFileName = "NoExtension";

  // identity => projection removing z-index
  this->Orientation << 1, 0, 0,
                       0, 1, 0,
                       0, 0, 1;
}

//----------------------------------------------------------------------------
vtkBirdEyeViewSnap::~vtkBirdEyeViewSnap()
{
}

//-----------------------------------------------------------------------------
void vtkBirdEyeViewSnap::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
}

//-----------------------------------------------------------------------------
int vtkBirdEyeViewSnap::RequestData(vtkInformation *vtkNotUsed(request),
  vtkInformationVector **inputVector, vtkInformationVector *outputVector)
{
  // Get the input
  vtkPolyData * input = vtkPolyData::GetData(inputVector[0]->GetInformationObject(0));

  // Get the output
  vtkPolyData *output = vtkPolyData::GetData(outputVector->GetInformationObject(0));
  output->ShallowCopy(input);

  // Check that the provided filename is valid
  if (this->RadicalFileName == "NoRadical" ||
      this->ExtensionFileName == "NoExtension")
  {
    vtkGenericWarningMacro("Filename has not been settled or is invalid");
    return 0;
  }

  // transform the input
  Eigen::Matrix<double, 3, 1> point;
  double vtkpoint[3];
  for (unsigned int k = 0; k < input->GetNumberOfPoints(); ++k)
  {
    input->GetPoint(k, vtkpoint);
    point << vtkpoint[0], vtkpoint[1], vtkpoint[2];
    point = this->Orientation * point;
    vtkpoint[0] = point(0);
    vtkpoint[1] = point(1);
    vtkpoint[2] = point(2);
    output->GetPoints()->SetPoint(k, vtkpoint);
  }

  // Create the bird eye view image
  double bounds[6];
  output->GetBounds(bounds);
  unsigned int x, y;

  unsigned int H = std::ceil((bounds[1] - bounds[0]) / this->pixelResX);
  unsigned int W = std::ceil((bounds[3] - bounds[2]) / this->pixelResY);

  // Image
  vtkSmartPointer<vtkImageData> image = vtkSmartPointer<vtkImageData>::New();
  image->SetDimensions(H, W, 1);
  image->AllocateScalars(VTK_UNSIGNED_CHAR, 1);
  unsigned char* dataPointer = static_cast<unsigned char*>(image->GetScalarPointer());
  std::fill(dataPointer, dataPointer + H * W, 0);

  // Get the reflectivity array
  vtkDataArray* intensity = output->GetPointData()->GetArray("intensity");

  for (unsigned int k = 0; k < output->GetNumberOfPoints(); ++k)
  {
    // Compute pixel coordinate
    output->GetPoint(k, vtkpoint);
    x = std::floor((vtkpoint[0] - bounds[0]) / (bounds[1] - bounds[0]) * (H - 1));
    y = std::floor((vtkpoint[1] - bounds[2]) / (bounds[3] - bounds[2]) * (W - 1));

    // set value
    unsigned char value = intensity->GetVariantValue(k).ToUnsignedChar();
    image->SetScalarComponentFromDouble(x, y, 0, 0, value);
  }

  // Save the image
  std::stringstream ss;
  ss << this->RadicalFileName << this->Count << "." << ExtensionFileName;
  vtkSmartPointer<vtkPNGWriter> writer = vtkSmartPointer<vtkPNGWriter>::New();
  writer->SetFileName(ss.str().c_str());
  writer->SetInputData(image);
  writer->Write();

  this->Count++;

  return 1;
}

//-----------------------------------------------------------------------------
void vtkBirdEyeViewSnap::SetPlaneParam(double params[4])
{
  // Here we will construct a new base of R3 using
  // the normal of the plane as the Z-axis. The new
  // x axis will be the projection of the older x-axis
  // on the plane and the new y-axis will be completed
  // so that we obtain a orthonormal direct base
  Eigen::Matrix<double, 3, 1> ex, ey, ez, ex2, ey2, ez2;
  ex << 1, 0, 0;
  ey << 0, 1, 0;
  ez << 0, 0, 1;

  ez2 << params[0], params[1], params[2];
  ez2.normalize();
  ez2 = ez.dot(ez2) / std::abs(ez.dot(ez2)) * ez2;

  ex2 = ex - ex.dot(ez2) * ez2;
  ex2.normalize();

  ey2 = ez2.cross(ex2);
  ey2.normalize();

  this->Orientation.row(0) = ex2;
  this->Orientation.row(1) = ey2;
  this->Orientation.row(2) = ez2;
}

//-----------------------------------------------------------------------------
void vtkBirdEyeViewSnap::SetFolderName(std::string filename)
{
  std::vector<std::string> strs;
  boost::split(strs, filename, boost::is_any_of("."));

  // Check that there is an extension
  if (strs.size() <= 1)
  {
    vtkGenericWarningMacro("filename is not valid: " << filename);
    return;
  }

  // Check that the extension is supported
  std::string extension = strs[strs.size() - 1];
  if ((extension != "png") && (extension != "jpg"))
  {
    vtkGenericWarningMacro("file format not supported: " << extension
                           << " supported file format are: jpg, png");
    return;
  }

  std::stringstream ss;
  for (unsigned int k = 0; k < strs.size() - 1; ++k)
  {
    ss << strs[k];
  }

  this->ExtensionFileName = extension;
  this->RadicalFileName = ss.str();
}

//-----------------------------------------------------------------------------
void vtkBirdEyeViewSnap::SetResolution(double sX, double sY)
{
  this->pixelResX = sX;
  this->pixelResY = sY;
}

//-----------------------------------------------------------------------------
void vtkBirdEyeViewSnap::SetCount(unsigned int count)
{
  this->Count = count;
}


//==================================================================================================
/*!
  @file

  Copyright 2016 NumScale SAS
  Copyright 2016 NumScale J.T.Lapreste

  Distributed under the Boost Software License, Version 1.0.
  (See accompanying file LICENSE.md or copy at http://boost.org/LICENSE_1_0.txt)
*/
//==================================================================================================
#ifndef BOOST_SIMD_FUNCTION_DEFINITION_SATURATED_ABS_HPP_INCLUDED
#define BOOST_SIMD_FUNCTION_DEFINITION_SATURATED_ABS_HPP_INCLUDED

#include <boost/simd/config.hpp>
#include <boost/simd/function/definition/abss.hpp>
#include <boost/simd/detail/dispatch.hpp>

namespace boost { namespace simd
{
#if defined(DOXYGEN_ONLY)
  /*!
  @ingroup group-arithmetic

    Computes the saturated absolute value of its parameter.

    This is an alias of function @c abss

    @see abss
  **/
  template<typename T> auto saturated_abs(T const& x);
#endif
  BOOST_DISPATCH_CALLABLE_DEFINITION(tag::abss_,saturated_abs);

} }

#endif


#include <stdlib.h>
#include <stdio.h>
#include <omp.h>
#include <chrono>
#include <vector>
#include <string>
#include <sstream>
#include <iostream>
#include <stdexcept>
#include <Eigen/Geometry>
#include <visualization_msgs/Marker.h>
#include <arc_utilities/eigen_helpers.hpp>
#include <arc_utilities/voxel_grid.hpp>
#include <arc_utilities/pretty_print.hpp>
#include <sdf_tools/sdf.hpp>

#ifndef SDF_GENERATION_HPP
#define SDF_GENERATION_HPP

namespace sdf_generation
{
struct bucket_cell
{
  double distance_square;
  int32_t update_direction;
  uint32_t location[3];
  uint32_t closest_point[3];
};

typedef VoxelGrid::VoxelGrid<bucket_cell> DistanceField;

inline int32_t GetDirectionNumber(
    const int32_t dx, const int32_t dy, const int32_t dz)
{
  return ((dx + 1) * 9) + ((dy + 1) * 3) + (dz + 1);
}

inline std::vector<std::vector<std::vector<std::vector<int32_t>>>>
MakeNeighborhoods()
{
  // First vector<>: 2 - the first bucket queue, the points we know are zero
  // distance, start with a complete set of neighbors to check. Every other
  // bucket queue checks fewer neighbors.
  // Second vector<>: 27 (# of source directions in fully-connected 3d grid).
  // Third vector<>:
  std::vector<std::vector<std::vector<std::vector<int32_t>>>> neighborhoods;
  // I don't know why there are 2 initial neighborhoods.
  neighborhoods.resize(2);
  for (size_t n = 0; n < neighborhoods.size(); n++)
  {
    neighborhoods[n].resize(27);
    // Loop through the source directions.
    for (int32_t dx = -1; dx <= 1; dx++)
    {
      for (int32_t dy = -1; dy <= 1; dy++)
      {
        for (int32_t dz = -1; dz <= 1; dz++)
        {
          const int32_t direction_number = GetDirectionNumber(dx, dy, dz);
          // Loop through the target directions.
          for (int32_t tdx = -1; tdx <= 1; tdx++)
          {
            for (int32_t tdy = -1; tdy <= 1; tdy++)
            {
              for (int32_t tdz = -1; tdz <= 1; tdz++)
              {
                // Ignore the case of ourself.
                if (tdx == 0 && tdy == 0 && tdz == 0)
                {
                  continue;
                }
                // Why is one set of neighborhoods larger than the other?
                if (n >= 1)
                {
                  if ((abs(tdx) + abs(tdy) + abs(tdz)) != 1)
                  {
                    continue;
                  }
                  if ((dx * tdx) < 0 || (dy * tdy) < 0 || (dz * tdz) < 0)
                  {
                    continue;
                  }
                }
                std::vector<int32_t> new_point;
                new_point.resize(3);
                new_point[0] = tdx;
                new_point[1] = tdy;
                new_point[2] = tdz;
                neighborhoods[n][direction_number].push_back(new_point);
              }
            }
          }
        }
      }
    }
  }
  return neighborhoods;
}

inline double ComputeDistanceSquared(
    const int32_t x1, const int32_t y1, const int32_t z1,
    const int32_t x2, const int32_t y2, const int32_t z2)
{
  const int32_t dx = x1 - x2;
  const int32_t dy = y1 - y2;
  const int32_t dz = z1 - z2;
  return double((dx * dx) + (dy * dy) + (dz * dz));
}

class MultipleThreadIndexQueueWrapper
{
public:

  explicit MultipleThreadIndexQueueWrapper(const size_t max_queues)
  {
    per_thread_queues_.resize(
          GetNumOMPThreads(), ThreadIndexQueues(max_queues));
  }

  const VoxelGrid::GRID_INDEX& Query(
      const int32_t distance_squared, const size_t idx) const
  {
    size_t working_index = idx;
    for (size_t thread = 0; thread < per_thread_queues_.size(); thread++)
    {
      const auto& current_thread_queue =
          per_thread_queues_.at(thread).at(distance_squared);
      const size_t current_thread_queue_size = current_thread_queue.size();
      if (working_index < current_thread_queue_size)
      {
        return current_thread_queue.at(working_index);
      }
      else
      {
        working_index -= current_thread_queue_size;
      }
    }
    throw std::runtime_error("Failed to find item");
  }

  size_t NumQueues() const
  {
    return per_thread_queues_.at(0).size();
  }

  size_t Size(const int32_t distance_squared) const
  {
    size_t total_size = 0;
    for (size_t thread = 0; thread < per_thread_queues_.size(); thread++)
    {
      total_size += per_thread_queues_.at(thread).at(distance_squared).size();
    }
    return total_size;
  }

  void Enqueue(
      const int32_t distance_squared, const VoxelGrid::GRID_INDEX& index)
  {
#if defined(_OPENMP)
    const size_t thread_num = (size_t)omp_get_thread_num();
#else
    const size_t thread_num = 0;
#endif
    per_thread_queues_.at(thread_num).at(distance_squared).push_back(index);
  }

  void ClearCompletedQueues(const int32_t distance_squared)
  {
    for (size_t thread = 0; thread < per_thread_queues_.size(); thread++)
    {
      per_thread_queues_.at(thread).at(distance_squared).clear();
    }
  }

private:

  inline static size_t GetNumOMPThreads()
  {
#if defined(_OPENMP)
    size_t num_threads = 0;
#pragma omp parallel
    {
      num_threads = (size_t)omp_get_num_threads();
    }
    return num_threads;
#else
    return 1;
#endif
  }

  typedef std::vector<std::vector<VoxelGrid::GRID_INDEX>> ThreadIndexQueues;
  std::vector<ThreadIndexQueues> per_thread_queues_;

};

inline DistanceField BuildDistanceFieldSerial(
    const Eigen::Isometry3d& grid_origin_transform,
    const double grid_resolution,
    const int64_t grid_num_x_cells,
    const int64_t grid_num_y_cells,
    const int64_t grid_num_z_cells,
    const std::vector<VoxelGrid::GRID_INDEX>& points)
{
  const std::chrono::time_point<std::chrono::steady_clock> start_time
      = std::chrono::steady_clock::now();
  // Make the DistanceField container
  bucket_cell default_cell;
  default_cell.distance_square = std::numeric_limits<double>::infinity();
  DistanceField distance_field(grid_origin_transform, grid_resolution,
                               grid_num_x_cells, grid_num_y_cells,
                               grid_num_z_cells, default_cell);
  // Compute maximum distance square
  const int64_t max_distance_square =
      (distance_field.GetNumXCells() * distance_field.GetNumXCells())
      + (distance_field.GetNumYCells() * distance_field.GetNumYCells())
      + (distance_field.GetNumZCells() * distance_field.GetNumZCells());
  // Make bucket queue
  std::vector<std::vector<bucket_cell>> bucket_queue(max_distance_square + 1);
  bucket_queue[0].reserve(points.size());
  // Set initial update direction
  int32_t initial_update_direction = GetDirectionNumber(0, 0, 0);
  // Mark all provided points with distance zero and add to the bucket queue
  for (size_t index = 0; index < points.size(); index++)
  {
    const VoxelGrid::GRID_INDEX& current_index = points[index];
    auto query = distance_field.GetMutable(current_index);
    if (query)
    {
      query.Value().location[0] = (uint32_t)current_index.x;
      query.Value().location[1] = (uint32_t)current_index.y;
      query.Value().location[2] = (uint32_t)current_index.z;
      query.Value().closest_point[0] = (uint32_t)current_index.x;
      query.Value().closest_point[1] = (uint32_t)current_index.y;
      query.Value().closest_point[2] = (uint32_t)current_index.z;
      query.Value().distance_square = 0.0;
      query.Value().update_direction = initial_update_direction;
      bucket_queue[0].push_back(query.Value());
    }
    // If the point is outside the bounds of the SDF, skip
    else
    {
      throw std::runtime_error("Point for BuildDistanceField out of bounds");
    }
  }
  // HERE BE DRAGONS
  // Process the bucket queue
  const std::vector<std::vector<std::vector<std::vector<int>>>> neighborhoods =
      MakeNeighborhoods();
  for (size_t bq_idx = 0; bq_idx < bucket_queue.size(); bq_idx++)
  {
    for (const auto& cur_cell : bucket_queue[bq_idx])
    {
      // Get the current location
      const double x = cur_cell.location[0];
      const double y = cur_cell.location[1];
      const double z = cur_cell.location[2];
      // Pick the update direction
      // Only the first bucket queue gets the larger set of neighborhoods?
      // Don't really userstand why.
      const size_t direction_switch = (bq_idx > 0) ? 1 : 0;
      // Make sure the update direction is valid
      if (cur_cell.update_direction < 0 || cur_cell.update_direction > 26)
      {
        continue;
      }
      // Get the current neighborhood list
      const std::vector<std::vector<int>>& neighborhood =
          neighborhoods[direction_switch][cur_cell.update_direction];
      // Update the distance from the neighboring cells
      for (size_t nh_idx = 0; nh_idx < neighborhood.size(); nh_idx++)
      {
        // Get the direction to check
        const int dx = neighborhood[nh_idx][0];
        const int dy = neighborhood[nh_idx][1];
        const int dz = neighborhood[nh_idx][2];
        const int nx = x + dx;
        const int ny = y + dy;
        const int nz = z + dz;
        auto neighbor_query =
            distance_field.GetMutable((int64_t)nx, (int64_t)ny, (int64_t)nz);
        if (!neighbor_query)
        {
          // "Neighbor" is outside the bounds of the SDF
          continue;
        }
        // Update the neighbor's distance based on the current
        const int32_t new_distance_square =
            (int32_t)ComputeDistanceSquared(nx, ny, nz,
                                            cur_cell.closest_point[0],
                                            cur_cell.closest_point[1],
                                            cur_cell.closest_point[2]);
        if (new_distance_square > max_distance_square)
        {
          // Skip these cases
          continue;
        }
        if (new_distance_square < neighbor_query.Value().distance_square)
        {
          // If the distance is better, time to update the neighbor
          neighbor_query.Value().distance_square = new_distance_square;
          neighbor_query.Value().closest_point[0] = cur_cell.closest_point[0];
          neighbor_query.Value().closest_point[1] = cur_cell.closest_point[1];
          neighbor_query.Value().closest_point[2] = cur_cell.closest_point[2];
          neighbor_query.Value().location[0] = nx;
          neighbor_query.Value().location[1] = ny;
          neighbor_query.Value().location[2] = nz;
          neighbor_query.Value().update_direction =
              GetDirectionNumber(dx, dy, dz);
          // Add the neighbor into the bucket queue
          bucket_queue[new_distance_square].push_back(neighbor_query.Value());
        }
      }
    }
    // Clear the current queue now that we're done with it
    bucket_queue[bq_idx].clear();
  }
  const std::chrono::time_point<std::chrono::steady_clock> end_time
      = std::chrono::steady_clock::now();
  const std::chrono::duration<double> elapsed = end_time - start_time;
  std::cout << "Computed DistanceField for grid size (" << grid_num_x_cells
            << ", " << grid_num_y_cells << ", " << grid_num_z_cells << ") in "
            << elapsed.count() << " seconds" << std::endl;
  return distance_field;
}

inline DistanceField BuildDistanceFieldParallel(
    const Eigen::Isometry3d& grid_origin_transform,
    const double grid_resolution,
    const int64_t grid_num_x_cells,
    const int64_t grid_num_y_cells,
    const int64_t grid_num_z_cells,
    const std::vector<VoxelGrid::GRID_INDEX>& points)
{
  const std::chrono::time_point<std::chrono::steady_clock> start_time
      = std::chrono::steady_clock::now();
  // Make the DistanceField container
  bucket_cell default_cell;
  default_cell.distance_square = std::numeric_limits<double>::infinity();
  DistanceField distance_field(grid_origin_transform, grid_resolution,
                               grid_num_x_cells, grid_num_y_cells,
                               grid_num_z_cells, default_cell);
  // Compute maximum distance square
  const int64_t max_distance_square =
      (distance_field.GetNumXCells() * distance_field.GetNumXCells())
      + (distance_field.GetNumYCells() * distance_field.GetNumYCells())
      + (distance_field.GetNumZCells() * distance_field.GetNumZCells());
  // Make bucket queue
  std::vector<std::vector<bucket_cell>> bucket_queue(max_distance_square + 1);
  bucket_queue[0].reserve(points.size());
  MultipleThreadIndexQueueWrapper bucket_queues(max_distance_square + 1);
  // Set initial update direction
  int32_t initial_update_direction = GetDirectionNumber(0, 0, 0);
  // Mark all provided points with distance zero and add to the bucket queues
  // points MUST NOT CONTAIN DUPLICATE ENTRIES!
#if defined(_OPENMP)
#pragma omp parallel for
#endif
  for (size_t index = 0; index < points.size(); index++)
  {
    const VoxelGrid::GRID_INDEX& current_index = points[index];
    auto query = distance_field.GetMutable(current_index);
    if (query)
    {
      query.Value().location[0] = (uint32_t)current_index.x;
      query.Value().location[1] = (uint32_t)current_index.y;
      query.Value().location[2] = (uint32_t)current_index.z;
      query.Value().closest_point[0] = (uint32_t)current_index.x;
      query.Value().closest_point[1] = (uint32_t)current_index.y;
      query.Value().closest_point[2] = (uint32_t)current_index.z;
      query.Value().distance_square = 0.0;
      query.Value().update_direction = initial_update_direction;
      bucket_queues.Enqueue(0, current_index);
    }
    // If the point is outside the bounds of the SDF, skip
    else
    {
      throw std::runtime_error("Point for BuildDistanceField out of bounds");
    }
  }
  // HERE BE DRAGONS
  // Process the bucket queue
  const std::vector<std::vector<std::vector<std::vector<int>>>> neighborhoods =
      MakeNeighborhoods();
  for (int32_t current_distance_square = 0;
       current_distance_square < (int32_t)bucket_queues.NumQueues();
       current_distance_square++)
  {
#if defined(_OPENMP)
#pragma omp parallel for
#endif
    for (size_t idx = 0; idx < bucket_queues.Size(current_distance_square);
         idx++)
    {
      const VoxelGrid::GRID_INDEX& current_index =
          bucket_queues.Query(current_distance_square, idx);
      // Get the current location
      const bucket_cell& cur_cell =
          distance_field.GetImmutable(current_index).Value();
      const double x = cur_cell.location[0];
      const double y = cur_cell.location[1];
      const double z = cur_cell.location[2];
      // Pick the update direction
      // Only the first bucket queue gets the larger set of neighborhoods?
      // Don't really userstand why.
      const size_t direction_switch = (current_distance_square > 0) ? 1 : 0;
      // Make sure the update direction is valid
      if (cur_cell.update_direction < 0 || cur_cell.update_direction > 26)
      {
        continue;
      }
      // Get the current neighborhood list
      const std::vector<std::vector<int32_t>>& neighborhood =
          neighborhoods[direction_switch][cur_cell.update_direction];
      // Update the distance from the neighboring cells
      for (size_t nh_idx = 0; nh_idx < neighborhood.size(); nh_idx++)
      {
        // Get the direction to check
        const int32_t dx = neighborhood[nh_idx][0];
        const int32_t dy = neighborhood[nh_idx][1];
        const int32_t dz = neighborhood[nh_idx][2];
        const int32_t nx = (int32_t)(x + dx);
        const int32_t ny = (int32_t)(y + dy);
        const int32_t nz = (int32_t)(z + dz);
        const VoxelGrid::GRID_INDEX neighbor_index(
              (int64_t)nx, (int64_t)ny, (int64_t)nz);
        auto neighbor_query = distance_field.GetMutable(neighbor_index);
        if (!neighbor_query)
        {
          // "Neighbor" is outside the bounds of the SDF
          continue;
        }
        // Update the neighbor's distance based on the current
        const int32_t new_distance_square =
            (int32_t)ComputeDistanceSquared(nx, ny, nz,
                                            cur_cell.closest_point[0],
                                            cur_cell.closest_point[1],
                                            cur_cell.closest_point[2]);
        if (new_distance_square > max_distance_square)
        {
          // Skip these cases
          continue;
        }
        if (new_distance_square < neighbor_query.Value().distance_square)
        {
          // If the distance is better, time to update the neighbor
          neighbor_query.Value().distance_square = new_distance_square;
          neighbor_query.Value().closest_point[0] = cur_cell.closest_point[0];
          neighbor_query.Value().closest_point[1] = cur_cell.closest_point[1];
          neighbor_query.Value().closest_point[2] = cur_cell.closest_point[2];
          neighbor_query.Value().location[0] = nx;
          neighbor_query.Value().location[1] = ny;
          neighbor_query.Value().location[2] = nz;
          neighbor_query.Value().update_direction =
              GetDirectionNumber(dx, dy, dz);
          // Add the neighbor into the bucket queue
          bucket_queues.Enqueue(new_distance_square, neighbor_index);
        }
      }
    }
    // Clear the current queues now that we're done with it
    bucket_queues.ClearCompletedQueues(current_distance_square);
  }
  const std::chrono::time_point<std::chrono::steady_clock> end_time
      = std::chrono::steady_clock::now();
  const std::chrono::duration<double> elapsed = end_time - start_time;
  std::cout << "Computed DistanceField for grid size (" << grid_num_x_cells
            << ", " << grid_num_y_cells << ", " << grid_num_z_cells << ") in "
            << elapsed.count() << " seconds" << std::endl;
  return distance_field;
}

#define COMPARE_DISTANCE_FIELD_GENERATION

inline DistanceField BuildDistanceField(
    const Eigen::Isometry3d& grid_origin_transform,
    const double grid_resolution,
    const int64_t grid_num_x_cells,
    const int64_t grid_num_y_cells,
    const int64_t grid_num_z_cells,
    const std::vector<VoxelGrid::GRID_INDEX>& points)
{
#ifdef COMPARE_DISTANCE_FIELD_GENERATION
  const DistanceField legacy_field =
      BuildDistanceFieldSerial(grid_origin_transform, grid_resolution,
                               grid_num_x_cells, grid_num_y_cells,
                               grid_num_z_cells, points);
  const DistanceField new_field =
      BuildDistanceFieldParallel(grid_origin_transform, grid_resolution,
                                 grid_num_x_cells, grid_num_y_cells,
                                 grid_num_z_cells, points);
  for (int64_t x_index = 0; x_index < grid_num_x_cells; x_index++)
  {
    for (int64_t y_index = 0; y_index < grid_num_y_cells; y_index++)
    {
      for (int64_t z_index = 0; z_index < grid_num_z_cells; z_index++)
      {
        const bucket_cell& legacy_cell =
            legacy_field.GetImmutable(x_index, y_index, z_index).Value();
        const bucket_cell& new_cell =
            new_field.GetImmutable(x_index, y_index, z_index).Value();
        const double legacy_distance = legacy_cell.distance_square;
        const double new_distance = new_cell.distance_square;
        const double new_delta = std::abs(new_distance - legacy_distance);
        assert(new_delta < 1e-6);
      }
    }
  }
  return new_field;
#else
  // Should we use the parallelizable new variant or the serial legacy variant?
  if (points.size() > 1000)
  {
    return BuildDistanceFieldParallel(grid_origin_transform, grid_resolution,
                                      grid_num_x_cells, grid_num_y_cells,
                                      grid_num_z_cells, points);
  }
  else
  {
    return BuildDistanceFieldSerial(grid_origin_transform, grid_resolution,
                                    grid_num_x_cells, grid_num_y_cells,
                                    grid_num_z_cells, points);
  }
#endif
}

template<typename T>
inline std::pair<sdf_tools::SignedDistanceField, std::pair<double, double>>
ExtractSignedDistanceField(
    const Eigen::Isometry3d& grid_origin_tranform,
    const double grid_resolution,
    const int64_t grid_num_x_cells,
    const int64_t grid_num_y_cells,
    const int64_t grid_num_z_cells,
    const std::function<bool(const VoxelGrid::GRID_INDEX&)>& is_filled_fn,
    const float oob_value,
    const std::string& frame)
{
  const std::chrono::time_point<std::chrono::steady_clock> start_time
      = std::chrono::steady_clock::now();
  std::vector<VoxelGrid::GRID_INDEX> filled;
  std::vector<VoxelGrid::GRID_INDEX> free;
  for (int64_t x_index = 0; x_index < grid_num_x_cells; x_index++)
  {
    for (int64_t y_index = 0; y_index < grid_num_y_cells; y_index++)
    {
      for (int64_t z_index = 0; z_index < grid_num_z_cells; z_index++)
      {
        const VoxelGrid::GRID_INDEX current_index(x_index, y_index, z_index);
        if (is_filled_fn(current_index))
        {
          // Mark as filled
          filled.push_back(current_index);
        }
        else
        {
          // Mark as free space
          free.push_back(current_index);
        }
      }
    }
  }
  // Make two distance fields, one for distance to filled voxels, one for
  // distance to free voxels.
  const DistanceField filled_distance_field =
      BuildDistanceField(grid_origin_tranform, grid_resolution,
                         grid_num_x_cells, grid_num_y_cells, grid_num_z_cells,
                         filled);
  const DistanceField free_distance_field =
      BuildDistanceField(grid_origin_tranform, grid_resolution,
                         grid_num_x_cells, grid_num_y_cells, grid_num_z_cells,
                         free);
  // Generate the SDF
  sdf_tools::SignedDistanceField new_sdf(
        grid_origin_tranform, frame, grid_resolution, grid_num_x_cells,
        grid_num_y_cells, grid_num_z_cells, oob_value);
  double max_distance = -std::numeric_limits<double>::infinity();
  double min_distance = std::numeric_limits<double>::infinity();
  for (int64_t x_index = 0; x_index < new_sdf.GetNumXCells(); x_index++)
  {
    for (int64_t y_index = 0; y_index < new_sdf.GetNumYCells(); y_index++)
    {
      // Parallelize across the Z-axis, since VoxelGrid ensures Z-axis cells are
      // contiguous in memory.
#if defined(_OPENMP)
#pragma omp parallel for
#endif
      for (int64_t z_index = 0; z_index < new_sdf.GetNumZCells(); z_index++)
      {
        const double distance1 =
            std::sqrt(
              filled_distance_field.GetImmutable(x_index, y_index, z_index)
                .Value().distance_square)
            * new_sdf.GetResolution();
        const double distance2 =
            std::sqrt(
              free_distance_field.GetImmutable(x_index, y_index, z_index)
                .Value().distance_square)
            * new_sdf.GetResolution();
        const double distance = distance1 - distance2;
        if (distance > max_distance)
        {
          max_distance = distance;
        }
        if (distance < min_distance)
        {
          min_distance = distance;
        }
        new_sdf.SetValue(x_index, y_index, z_index, (float)distance);
      }
    }
  }
  const std::chrono::time_point<std::chrono::steady_clock> end_time
      = std::chrono::steady_clock::now();
  const std::chrono::duration<double> elapsed = end_time - start_time;
  std::cout << "Computed SDF for grid size (" << grid_num_x_cells << ", "
            << grid_num_y_cells << ", " << grid_num_z_cells << ") in "
            << elapsed.count() << " seconds" << std::endl;
  std::pair<double, double> extrema(max_distance, min_distance);
  return std::pair<sdf_tools::SignedDistanceField,
                   std::pair<double, double>>(new_sdf, extrema);
}

template<typename T, typename BackingStore=std::vector<T>>
inline std::pair<sdf_tools::SignedDistanceField, std::pair<double, double>>
ExtractSignedDistanceField(
    const VoxelGrid::VoxelGrid<T, BackingStore>& grid,
    const std::function<bool(const VoxelGrid::GRID_INDEX&)>& is_filled_fn,
    const float oob_value, const std::string& frame,
    const bool add_virtual_border)
{
  (void)(add_virtual_border);
  const Eigen::Vector3d cell_sizes = grid.GetCellSizes();
  if ((cell_sizes.x() != cell_sizes.y()) || (cell_sizes.x() != cell_sizes.z()))
  {
    throw std::invalid_argument("Grid must have uniform resolution");
  }
  if (add_virtual_border == false)
  {
    // This is the conventional single-pass result
    return ExtractSignedDistanceField<T>(
          grid.GetOriginTransform(), cell_sizes.x(), grid.GetNumXCells(),
          grid.GetNumYCells(), grid.GetNumZCells(), is_filled_fn, oob_value,
          frame);
  }
  else
  {
    const int64_t x_axis_size_offset =
        (grid.GetNumXCells() > 1) ? (int64_t)2 : (int64_t)0;
    const int64_t x_axis_query_offset =
        (grid.GetNumXCells() > 1) ? (int64_t)1 : (int64_t)0;
    const int64_t y_axis_size_offset =
        (grid.GetNumYCells() > 1) ? (int64_t)2 : (int64_t)0;
    const int64_t y_axis_query_offset =
        (grid.GetNumYCells() > 1) ? (int64_t)1 : (int64_t)0;
    const int64_t z_axis_size_offset =
        (grid.GetNumZCells() > 1) ? (int64_t)2 : (int64_t)0;
    const int64_t z_axis_query_offset =
        (grid.GetNumZCells() > 1) ? (int64_t)1 : (int64_t)0;
    // We need to lie about the size of the grid to add a virtual border
    const int64_t num_x_cells = grid.GetNumXCells() + x_axis_size_offset;
    const int64_t num_y_cells = grid.GetNumYCells() + y_axis_size_offset;
    const int64_t num_z_cells = grid.GetNumZCells() + z_axis_size_offset;
    // Make some deceitful helper functions that hide our lies about size
    // For the free space SDF, we lie and say the virtual border is filled
    const std::function<bool(const VoxelGrid::GRID_INDEX&)> free_is_filled_fn
        = [&] (const VoxelGrid::GRID_INDEX& virtual_border_grid_index)
    {
      // Is there a virtual border on our axis?
      if (x_axis_size_offset > 0)
      {
        // Are we a virtual border cell?
        if ((virtual_border_grid_index.x == 0)
            || (virtual_border_grid_index.x == (num_x_cells - 1)))
        {
          return true;
        }
      }
      // Is there a virtual border on our axis?
      if (y_axis_size_offset > 0)
      {
        // Are we a virtual border cell?
        if ((virtual_border_grid_index.y == 0)
            || (virtual_border_grid_index.y == (num_y_cells - 1)))
        {
          return true;
        }
      }
      // Is there a virtual border on our axis?
      if (z_axis_size_offset > 0)
      {
        // Are we a virtual border cell?
        if ((virtual_border_grid_index.z == 0)
            || (virtual_border_grid_index.z == (num_z_cells - 1)))
        {
          return true;
        }
      }
      const VoxelGrid::GRID_INDEX real_grid_index(
            virtual_border_grid_index.x - x_axis_query_offset,
            virtual_border_grid_index.y - y_axis_query_offset,
            virtual_border_grid_index.z - z_axis_query_offset);
      return is_filled_fn(real_grid_index);
    };
    // For the filled space SDF, we lie and say the virtual border is empty
    const std::function<bool(const VoxelGrid::GRID_INDEX&)> filled_is_filled_fn
        = [&] (const VoxelGrid::GRID_INDEX& virtual_border_grid_index)
    {
      // Is there a virtual border on our axis?
      if (x_axis_size_offset > 0)
      {
        // Are we a virtual border cell?
        if ((virtual_border_grid_index.x == 0)
            || (virtual_border_grid_index.x == (num_x_cells - 1)))
        {
          return false;
        }
      }
      // Is there a virtual border on our axis?
      if (y_axis_size_offset > 0)
      {
        // Are we a virtual border cell?
        if ((virtual_border_grid_index.y == 0)
            || (virtual_border_grid_index.y == (num_y_cells - 1)))
        {
          return false;
        }
      }
      // Is there a virtual border on our axis?
      if (z_axis_size_offset > 0)
      {
        // Are we a virtual border cell?
        if ((virtual_border_grid_index.z == 0)
            || (virtual_border_grid_index.z == (num_z_cells - 1)))
        {
          return false;
        }
      }
      const VoxelGrid::GRID_INDEX real_grid_index(
            virtual_border_grid_index.x - x_axis_query_offset,
            virtual_border_grid_index.y - y_axis_query_offset,
            virtual_border_grid_index.z - z_axis_query_offset);
      return is_filled_fn(real_grid_index);
    };
    // Make both SDFs
    auto free_sdf_result = ExtractSignedDistanceField<T>(
                             grid.GetOriginTransform(), cell_sizes.x(),
                             num_x_cells, num_y_cells, num_z_cells,
                             free_is_filled_fn, oob_value, frame);
    auto filled_sdf_result = ExtractSignedDistanceField<T>(
                               grid.GetOriginTransform(), cell_sizes.x(),
                               num_x_cells, num_y_cells, num_z_cells,
                               filled_is_filled_fn, oob_value, frame);
    // Combine to make a single SDF
    sdf_tools::SignedDistanceField combined_sdf(
          grid.GetOriginTransform(), frame, cell_sizes.x(), grid.GetNumXCells(),
          grid.GetNumYCells(), grid.GetNumZCells(), oob_value);
    for (int64_t x_idx = 0; x_idx < combined_sdf.GetNumXCells(); x_idx++)
    {
      for (int64_t y_idx = 0; y_idx < combined_sdf.GetNumYCells(); y_idx++)
      {
        // Parallelize across the Z-axis, since VoxelGrid ensures Z-axis cells
        // are contiguous in memory.
#if defined(_OPENMP)
#pragma omp parallel for
#endif
        for (int64_t z_idx = 0; z_idx < combined_sdf.GetNumZCells(); z_idx++)
        {
          const int64_t query_x_idx = x_idx + x_axis_query_offset;
          const int64_t query_y_idx = y_idx + y_axis_query_offset;
          const int64_t query_z_idx = z_idx + z_axis_query_offset;
          const float free_sdf_value
              = free_sdf_result.first.GetImmutable(
                  query_x_idx, query_y_idx, query_z_idx).Value();
          const float filled_sdf_value
              = filled_sdf_result.first.GetImmutable(
                  query_x_idx, query_y_idx, query_z_idx).Value();
          if (free_sdf_value >= 0.0)
          {
            combined_sdf.SetValue(x_idx, y_idx, z_idx, free_sdf_value);
          }
          else if (filled_sdf_value <= -0.0)
          {
            combined_sdf.SetValue(x_idx, y_idx, z_idx, filled_sdf_value);
          }
          else
          {
            combined_sdf.SetValue(x_idx, y_idx, z_idx, 0.0f);
          }
        }
      }
    }
    // Get the combined max/min values
    const std::pair<double, double> combined_extrema(
          free_sdf_result.second.first, filled_sdf_result.second.second);
    return std::make_pair(combined_sdf, combined_extrema);
  }
}

template<typename T, typename BackingStore=std::vector<T>>
inline std::pair<sdf_tools::SignedDistanceField, std::pair<double, double>>
ExtractSignedDistanceField(const VoxelGrid::VoxelGrid<T, BackingStore>& grid,
                           const std::function<bool(const T&)>& is_filled_fn,
                           const float oob_value, const std::string& frame)
{
  const std::function<bool(const VoxelGrid::GRID_INDEX&)> real_is_filled_fn =
      [&] (const VoxelGrid::GRID_INDEX& index)
  {
    const T& stored = grid.GetImmutable(index).Value();
    // If it matches an object to use OR there are no objects supplied
    if (is_filled_fn(stored))
    {
      // Mark as filled
      return true;
    }
    else
    {
      // Mark as free space
      return false;
    }
  };
  return ExtractSignedDistanceField(
        grid, real_is_filled_fn, oob_value, frame, false);
}
}

#endif // SDF_GENERATION_HPP