#include <iostream>
#include <boost/bind.hpp>
#include <boost/algorithm/string.hpp>
#include "../include/grasp_sampler/RRTFunction.h"
#include <string>
#include <cstring>

using namespace std;
using namespace Eigen;

#define REACHED 0
#define ADVANCED 1
#define TRAPPED -1

bool RRT::solveRRT(Robotmodel& model, std::ostream& sout) {

	_model = model;

	nt_start._nodes.clear(); // NodeTree
	nt_goal._nodes.clear(); // NodeTree

	initConfig.clear();// std::vector<double>
	DOF_weights.clear();// std::vector<double>
	int ActiveDoFs = dofSize;
	step_size = 0.3 / M_PI *180.0;
	std::string start_string[dofSize];

	for (int i = 0; i < ActiveDoFs; i++) {
		initConfig.push_back(qinit(i) * 180.0 / M_PI);
		DOF_weights.push_back(1.0);
	}

	auto start_node = std::make_shared<RRTNode>(initConfig);

	nt_start.addNode(start_node);
	nt_pointer = &nt_start;

	// mulitple seed
	for (int i = 0; i < qgoals.size(); i++)
	{
		std::vector<double> goalConfig;
		for (int j = 0; j < ActiveDoFs; j++)
		{
			goalConfig.push_back(qgoals[i][j] * 180.0 / M_PI);
		}

		auto goal_node = std::make_shared<RRTNode>(goalConfig);

		nt_goal.addNode(goal_node);
	}

	bool finished = false;
	int iter = 1;

    std::vector<double> q_rand;
    RRTNodePtr q_a_near;
    std::vector<double> q_a_reached ;
    RRTNodePtr q_b_near;
    std::vector<double> q_b_reached ;
	int root_tree_index, goal_tree_index;

	while (iter < 50) {

		nt_pointer = &nt_start;								  //		// start tree
		q_rand = generateRandomConfig();
		q_a_near = nt_pointer->getNearest(q_rand, DOF_weights);

		q_a_reached = this->extendNode(q_rand, q_a_near);

		nt_pointer = &nt_goal;
		q_b_near = nt_pointer->getNearest(q_a_reached, DOF_weights); // random config�� ���� ����� c_tree�� ��� ã��
		q_b_reached = this->extendNode(q_a_reached, q_b_near); // near��� ���� random config�� Extend

		if (getEuclideanNorm(q_a_reached, q_b_reached, dofSize) < step_size)
		{
			nt_pointer = &nt_start; //		// start tree
			root_tree_index = nt_pointer->getNearestIndex(q_a_reached, DOF_weights);
			nt_pointer = &nt_goal;
			goal_tree_index = nt_pointer->getNearestIndex(q_b_reached, DOF_weights);

			// q_a_reached가 tree내에 있는 노드가 아니여서 불연속이 생김..
			//cout << "t_turn \t" << t_turn << endl;
			//cout << q_a_reached[0] << "\t" << q_a_reached[1] << "\t" << q_a_reached[2] << "\t" << q_a_reached[3] << "\t" << q_a_reached[4] << "\t" << q_a_reached[5] << "\t" << q_a_reached[6] << endl;
			//cout << q_b_reached[0] << "\t" << q_b_reached[1] << "\t" << q_b_reached[2] << "\t" << q_b_reached[3] << "\t" << q_b_reached[4] << "\t" << q_b_reached[5] << "\t" << q_b_reached[6] << endl;

			//cout << getEuclideanNorm(q_a_reached, q_b_reached, dofSize) << endl;
			//cout << "reached" << endl;
			finished = true;

			break;
		}
		iter++;
	}
	if (finished) {

		path = nt_start.getPathWithIndex(root_tree_index);
		vector<vector<double>> p2;
		p2 = nt_goal.getPathWithIndex(goal_tree_index);

		reverse(p2.begin(), p2.end());
		path.insert(path.end(), p2.begin(), p2.end());

		for (int i = 0; i < path.size(); i++) {
			std::vector<double> node = path[i];
			for (int j = 0; j<node.size() - 1; j++)
				sout << node[j] << ",";

			sout << node[node.size() - 1] << "\n";
		}
	}
	else {
		return false;
	}

	return true;
}
bool RRT::solveCRRT(Robotmodel& model, std::ostream& sout) {
	_model = model;

	nt_start._nodes.clear();
	nt_goal._nodes.clear();

	initConfig.clear();
	DOF_weights.clear();

	int ActiveDoFs = dofSize;
	step_size = 0.10/M_PI*180.0;

	std::string start_string[dofSize];

	for (int i = 0; i < ActiveDoFs; i++) {
		initConfig.push_back(qinit(i) * 180.0 / M_PI);
		DOF_weights.push_back(1.0);
	}

	auto start_node = std::make_shared<RRTNode>(initConfig);

	nt_start.addNode(start_node);
	nt_pointer = &nt_start;
	t_turn = 1;

	// mulitple seed
	for (int i = 0; i < qgoals.size(); i++)
	{
		std::vector<double> goalConfig;
		for (int j = 0; j < ActiveDoFs; j++)
		{
			goalConfig.push_back(qgoals[i][j] * 180.0 / M_PI);
		}
		auto goal_node = std::make_shared<RRTNode>(goalConfig);

		nt_goal.addNode(goal_node);
	}
	/// single seed
	// RRTNode g(goal, 0); // ���� RRT
	// gTree.addNode(g);

	dx_from_init_to_goal.resize(6);
	Vector3d pos_init = CalcBodyToBaseCoordinates(_model.model_, qinit, _model.body_id_vec.back(), _model.body_com_pos.back());
	Matrix3d Rot_init = CalcBodyWorldOrientation(_model.model_, qinit, _model.body_id_vec.back(), true).transpose();
	dx_from_init_to_goal.head(3) = refer_pos - pos_init;
	dx_from_init_to_goal.tail(3) = getPhi(Rot_init, refer_rot);

	dx_from_extend_to_goal.resize(6);

	bool finished = false;
	int iter = 1;
    int a = 1;
    std::vector<double> q_rand;
    RRTNodePtr q_a_near;
    std::vector<double> q_a_reached ;
    RRTNodePtr q_b_near;
    std::vector<double> q_b_reached ;
	int root_tree_index, goal_tree_index;
	while (iter < 50000)
	{
		if (t_turn == 1)
		{
			nt_pointer = &nt_start;								  //		// start tree
			q_rand = generateRandomConfig(); //ConstrainedRandomConfig(_model);					  // last node of goal tree(g2) or random config

			//cout << "start tree -> goal tree" << "\t" <<t_turn << endl;
			q_a_near = nt_pointer->getNearest(q_rand, DOF_weights); // random config�� ���� ����� c_tree�� ��� ã��
			q_a_reached = this->extendNodeWithConstraints(q_rand, q_a_near);   // near��� ���� random config�� Extend

			nt_pointer = &nt_goal;
			q_b_near = nt_pointer->getNearest(q_a_reached, DOF_weights); // random config�� ���� ����� c_tree�� ��� ã��

			//cout << q_b_near->getConfiguration()[0] << "\t" <<  q_b_near->getConfiguration()[1] << endl;
			//cout << "goal tree -> start tree" << endl;
			q_b_reached = this->extendNodeWithConstraints(q_a_reached, q_b_near); // near��� ���� random config�� Extend

			if (getEuclideanNorm(q_a_reached, q_b_reached, dofSize) < step_size)
			{
				nt_pointer = &nt_start;								  //		// start tree
				root_tree_index = nt_pointer->getNearestIndex(q_a_reached, DOF_weights);
				nt_pointer = &nt_goal;
				goal_tree_index = nt_pointer->getNearestIndex(q_b_reached, DOF_weights);

				// q_a_reached가 tree내에 있는 노드가 아니여서 불연속이 생김..
				cout << "t_turn \t" << t_turn << endl;
				cout << q_a_reached[0] << "\t" << q_a_reached[1] << "\t" << q_a_reached[2] << "\t" << q_a_reached[3] << "\t" << q_a_reached[4] << "\t" << q_a_reached[5] << "\t" << q_a_reached[6] << endl;
				cout << q_b_reached[0] << "\t" << q_b_reached[1] << "\t" << q_b_reached[2] << "\t" << q_b_reached[3] << "\t" << q_b_reached[4] << "\t" << q_b_reached[5] << "\t" << q_b_reached[6] << endl;

				cout << getEuclideanNorm(q_a_reached, q_b_reached, dofSize) << endl;
				cout << "reached" << endl;
				a = REACHED;
				finished = true;

				break;
			}
			//t_turn = 2;
		}
		// if (t_turn == 2)
		// {
		// 	c_tree = &gTree;	// goal tree
		// 	q_rand = RandomConfig();// ConstrainedRandomConfig(_model);					  // last node of goal tree(g2) or random config

		// 	//cout << "goal tree -> start _tree " << "\t" <<t_turn << endl;
		// 	q_a_near = c_tree->getNearest(q_rand, DOF_weights); // random config�� ���� ����� c_tree�� ��� ã��

		// 	q_a_reached = this->ConstrainedExtend(q_rand, q_a_near); // near��� ���� random config�� Extend

		// 	c_tree = &tree;
		// 	q_b_near = c_tree->getNearest(q_a_reached, DOF_weights); // random config�� ���� ����� c_tree�� ��� ã��
		// 	//cout << "start tree -> goal tree" << endl;

		// 	q_b_reached = this->ConstrainedExtend(q_a_reached, q_b_near); // near��� ���� random config�� Extend
		// 	if (getEuclideanNorm(q_a_reached, q_b_reached, dofSize) < step_size)
		// 	{
		// 		c_tree = &gTree;								  //		// start tree
		// 		root_tree_index = c_tree->getNearestIndex(q_a_reached, DOF_weights);
		// 		c_tree = &tree;
		// 		goal_tree_index = c_tree->getNearestIndex(q_b_reached, DOF_weights);

		// 		// q_a_reached가 tree내에 있는 노드가 아니여서 불연속이 생김..
		// 		cout << "t_turn \t" << t_turn << endl;
		// 		cout << q_a_reached[0] << "\t" << q_a_reached[1] << "\t" << q_a_reached[2] << "\t" << q_a_reached[3] << "\t" << q_a_reached[4] << "\t" << q_a_reached[5] << "\t" << q_a_reached[6] << endl;
		// 		cout << q_b_reached[0] << "\t" << q_b_reached[1] << "\t" << q_b_reached[2] << "\t" << q_b_reached[3] << "\t" << q_b_reached[4] << "\t" << q_b_reached[5] << "\t" << q_b_reached[6] << endl;

		// 		cout << getEuclideanNorm(q_a_reached, q_b_reached, dofSize) << endl;
		// 		cout << "reached" << endl;
		// 		a = REACHED;
		// 		finished = true;
		// 	}
		// 	t_turn = 1;

		// }
		iter++;
	}

	if (finished) {

		path = nt_start.getPathWithIndex(root_tree_index);
		vector<vector<double>> p2;
		p2 = nt_goal.getPathWithIndex(goal_tree_index);

		reverse(p2.begin(), p2.end());
		path.insert(path.end(), p2.begin(), p2.end());

		std::vector<double> end;
		end = path[path.size()-1];

		for (int i = 0; i < path.size(); i++) {
			std::vector<double> node = path[i];
			for (int j = 0; j < node.size() - 1; j++){
				sout << node[j] << ",";
			}
			sout << node[node.size() - 1] << "\n";

		}//error

	}
	else {

		return false;
	}

	return true;
}
std::vector<double> RRT::generateRandomConfig() {

 	std::vector<double> R;
	do {
		for (int i = 0; i < initConfig.size(); i++) {
			double jointrange = upper_limit(i) - lower_limit(i); // angle
			double r = ((double)rand() / (double)RAND_MAX)*jointrange;
			R.push_back(lower_limit(i) + r);
		}
	} while (R.size() != initConfig.size());

	return R;
}
// std::vector<double> RRT::ConstrainedRandomConfig(Robotmodel model){
// 	std::vector<double> R;
// 	std::vector<double> R_project;

// 	while (R_project.empty())
// 	{
// 		do
// 		{
// 			for (int i = 0; i < start.size(); i++)
// 			{
// 				double jointrange = upper_limit(i) - lower_limit(i); // angle
// 				double r = ((double)rand() / (double)RAND_MAX) * jointrange;
// 				R.push_back(lower_limit(i) + r);
// 			}
// 		} while (R.size() != start.size());

// 		R_project = ProjectConfigForRandomConfig(model, R);
// 	}

// 	return R_project;
//}

std::vector<double> RRT::extendNode(std::vector<double> &node, RRTNodePtr &near){
	std::vector<double> qs, qs_old;
	qs.clear();
	qs_old.clear();
	for (int i = 0; i < dofSize; i++)
	{
		qs.push_back(near->getConfiguration()[i]);
		qs_old.push_back(near->getConfiguration()[i]);
	}

	while(true){
		if (getEuclideanNorm(node, qs, dofSize) < step_size)
		{
			return qs;
		}
		else if (getEuclideanNorm(node, qs, dofSize) - getEuclideanNorm(qs_old, node, dofSize) > 0.0)
		{
			return qs_old; // qs_old is close to node
		}

		for (int i = 0; i < initConfig.size(); i++)
		{
			qs[i] = (qs[i] + min(step_size, getEuclideanNorm(node, qs, dofSize)) * (node[i] - qs[i]) / getEuclideanNorm(node, qs, dofSize));
		}

		if (!checkSelfCollision(_model, qs))
		{
			RRTNodePtr old = near;
			near = std::make_shared<RRTNode>(qs, old);
			nt_pointer->addNode(near);
			qs_old = qs;
		}
		else
		{
			cout << "collision config \t" << qs[0] << "\t" <<qs[1] << "\t" <<qs[2] << "\t" <<qs[3] << "\t" <<qs[4] << "\t" <<qs[5] << "\t" <<qs[6] << "\t" << endl;
			return qs_old;
		}

	}
}

std::vector<double> RRT::extendNodeWithConstraints(std::vector<double> &node, RRTNodePtr &near)
{
	// node = q_rand, near : q_near
	std::vector<double> qs, qs_old;
	qs.clear();
	qs_old.clear();
	bool project_flag = false;
	int iteration_ = 0;

	for (int i = 0; i < dofSize; i++)
	{
		qs.push_back(near->getConfiguration()[i]);
		qs_old.push_back(near->getConfiguration()[i]);
	}
	//cout << "input" << node[0] << "\t" << node[1] << "\t" << node[2] << "\t" << node[3] << "\t" << node[4] << "\t" << node[5] << "\t" << node[6] << endl;
	//cout << "q_near" << qs[0] << "\t"<< qs[1] << "\t"<< qs[2] << "\t"<< qs[3] << "\t"<< qs[4] << "\t"<< qs[5] << "\t"<< qs[6] <<endl;
	while (true)
	{
		if (getEuclideanNorm(node, qs, dofSize) < step_size)
		{
			return qs;
		}
		else if (getEuclideanNorm(node, qs, dofSize) - getEuclideanNorm(qs_old, node, dofSize) > 0.0)
		{
			return qs_old; // qs_old is close to node
		}

		for (int i = 0; i < initConfig.size(); i++) // Move from 'qs' to 'node'
			qs[i] = (qs[i] + min(step_size, getEuclideanNorm(node, qs, dofSize)) * (node[i] - qs[i]) / getEuclideanNorm(node, qs, dofSize));

		//cout << "desired_dir" <<qs[0] << "\t" << qs[1] << "\t" << qs[2] << "\t" << qs[3] << "\t" << qs[4] << "\t" << qs[5] << "\t" << qs[6] << endl;

		// qs_old : near / qs : extend
		qs = projectConfig(_model, qs_old, qs); // project qs onto constraint manifold

		// if (!qs.empty() && (std_norm(qs, qs_old, dofSize) > 3.0*step_size)){
		// 	cout << "output1" <<qs_old[0] << "\t" << qs_old[1] << "\t" << qs_old[2] << "\t" << qs_old[3] << "\t" << qs_old[4] << "\t" << qs_old[5] << "\t" << qs_old[6] << endl;
		// 	cout << "output2" <<qs[0] << "\t" << qs[1] << "\t" << qs[2] << "\t" << qs[3] << "\t" << qs[4] << "\t" << qs[5] << "\t" << qs[6] << endl;

		// 	return qs_old;
		// }

		if (!qs.empty() && !checkExternalCollision(_model, qs))
		{
			//			VectorXd q_project_(7);
			//cout << "extend!" << endl;
			RRTNodePtr old = near;		   //qs_old
			near = std::make_shared<RRTNode>(qs, old) ; //(new RRTNode(qs, old)); // qs : configuration, old : parent node
			nt_pointer->addNode(near);

			if((getEuclideanNorm(qs, qs_old, dofSize) < step_size / 2.0)){
				return qs;
			}
			qs_old = qs;

		}
		else
		{
			//cout << "qs_old" << qs_old[0] << "\t"<< qs_old[1] << "\t"<< qs_old[2] << "\t"<< qs_old[3] << "\t"<< qs_old[4] << "\t"<< qs_old[5] << "\t"<< qs_old[6] <<endl;
			//cout << "return NULL" << endl;
			//cout << "collision config \t" << qs[0] << "\t" <<qs[1] << "\t" <<qs[2] << "\t" <<qs[3] << "\t" <<qs[4] << "\t" <<qs[5] << "\t" <<qs[6] << "\t" << endl;

			return qs_old;
		}
	}
}
std::vector<double> RRT::projectConfig(Robotmodel model, std::vector<double> qold, std::vector<double> &qs)
{
	model.q.resize(dofSize);

	// Tc
	Matrix4d T0_c, T0_obj, Tc_obj;
	T0_c.setIdentity();
	//T0_c.topLeftCorner(3,3) = refer_rot;
	T0_c.topRightCorner(3, 1) = local_rot.transpose()*(refer_pos - local_pos); // global to local pos

	MatrixXd J_temp(6, dofSize), J(6, dofSize);
	VectorXd d_c(6), dx(6), q_error(dofSize), q_old(dofSize);
	dx.setZero();
	Vector3d phi;
	Vector3d s[3], v[3], w[3];

	for (int i = 0; i < dofSize; i++)
	{ // deg -> rad
		qs[i] = qs[i] * M_PI / 180.0;
		qold[i] = qold[i] * M_PI / 180.0;

		q_old(i) = qold[i] * M_PI / 180.0;
	}

	while (true)
	{
		for (int i = 0; i < dofSize; i++)
			model.q(i) = qs[i];

		Matrix3d Rot_temp = CalcBodyWorldOrientation(model.model_, model.q, model.body_id_vec.back(), true).transpose();
		Vector3d pos_temp = CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec.back(), model.body_com_pos.back()); // get position and rotation in EE frame;

		T0_obj.setIdentity();
		T0_obj.topLeftCorner(3, 3) = Rot_temp;
		T0_obj.topRightCorner(3, 1) = pos_temp;
		Tc_obj = T0_c.inverse() * T0_obj;

		d_c.head(3) = Tc_obj.topRightCorner(3, 1);

		Matrix3d Tc_obj_rot = refer_rot.transpose()*Rot_temp;
		d_c(3) = atan2(Tc_obj_rot(2, 1), Tc_obj_rot(2, 2));
		d_c(4) = -asin(Tc_obj_rot(2, 0));
		d_c(5) = atan2(Tc_obj_rot(1, 0), Tc_obj_rot(0, 0));

		// 위치가 베이스 프레임 기준이 아니게 됨.
		for (int i = 0; i < 6; i++)
		{
			if (d_c(i) > C(i, 0)) // max
				dx(i) = d_c(i) - C(i, 0);
			else if (d_c(i) < C(i, 1)) // min
				dx(i) = d_c(i) - C(i, 1) ;
			else
				dx(i) = 0.0;
		}

		// Algorithm 4 - line 3
		if (dx.norm() < 0.01)
		{
			Vector3d q_old_pos = CalcBodyToBaseCoordinates(model.model_, q_old, model.body_id_vec.back(), model.body_com_pos.back());
			double dot_product;

			dx_from_extend_to_goal.head(3) = refer_pos - pos_temp;
			dx_from_extend_to_goal.tail(3) = getPhi(Rot_temp, refer_rot);

			dot_product = (dx_from_init_to_goal.head(3)/dx_from_init_to_goal.head(3).norm()).dot((pos_temp - q_old_pos) / (pos_temp - q_old_pos) .norm());

			// TODO : focal point - init position, final positon.  projected point is inside the ellipsoid
			if ((dx_from_init_to_goal.head(3).norm() < dx_from_extend_to_goal.head(3).norm()) || (dot_product < 0.0))
			{
				qs.clear();
				return qs;
			}

			for (int i = 0; i < dofSize; i++)
			{
				qs[i] = qs[i] / M_PI * 180.0;
			}
				return qs; // convert rad to angle
		}

		// Algorithm 4 - line 4
		CalcPointJacobian6D(model.model_, model.q, model.body_id_vec.back(), model.body_com_pos.back(), J_temp, true);
		J.topLeftCorner(3, dofSize) = J_temp.bottomLeftCorner(3, dofSize);

		rpy(0) = atan2(Rot_temp(2, 1), Rot_temp(2, 2));
		rpy(1) = -asin(Rot_temp(2, 0));
		rpy(2) = atan2(Rot_temp(1, 0), Rot_temp(0, 0));

		E_rpy(0,0) = cos(rpy(2))/cos(rpy(1));
		E_rpy(0,1) = sin(rpy(2))/cos(rpy(1));
		E_rpy(0,2) = 0.0;

		E_rpy(1,0) = sin(rpy(2));
		E_rpy(1,1) = cos(rpy(2));
		E_rpy(1,2) = 0.0;

		E_rpy(2,0) = -cos(rpy(2))*sin(rpy(1))/cos(rpy(1));
		E_rpy(2,1) = sin(rpy(2))*sin(rpy(1))/cos(rpy(1));
		E_rpy(2,2) = 1.0;
		J.bottomLeftCorner(3, dofSize) = E_rpy*J_temp.topLeftCorner(3, dofSize);

		// Algorithm 4 - line 5
		q_error = J.transpose() * (J * J.transpose()).inverse() * dx * 0.3;

		// Algorithm 4 -  line 6
		for (int i = 0; i < dofSize; i++)
			qs[i] -= q_error(i);

		// Algorithm 4 - line 7 : stuck here
		if (!checkJointLimit(qs) || (getEuclideanNorm(qs, qold, dofSize) > 3.0*step_size/180.0*M_PI))
		{
			//cout << "joint limit" << "\t" << qs[0]/M_PI*180.0<< "\t" << qs[1]/M_PI*180.0<< "\t" << qs[2]/M_PI*180.0<< "\t" << qs[3]/M_PI*180.0<< "\t" << qs[4]/M_PI*180.0<< "\t" << qs[5]/M_PI*180.0<< "\t" << qs[6]/M_PI*180.0 << endl;
			qs.clear();
			return qs;
			// flag = false;
			// break;
		}
	}
}

// std::vector<double> RRT::ProjectConfigForRandomConfig(Robotmodel model, std::vector<double> &qs)
// {
// 	model.q.resize(dofSize);

// 	// Tc
// 	Matrix4d T0_c, T0_obj, Tc_obj;
// 	T0_c.setIdentity();
// 	//T0_c.topLeftCorner(3,3) = refer_rot;
// 	T0_c.topRightCorner(3, 1) = local_rot*(refer_pos - local_pos); // global to local pos

// 	MatrixXd J_temp(6, dofSize), J(6, dofSize);
// 	VectorXd d_c(6), dx(6), q_error(dofSize);
// 	dx.setZero();
// 	Vector3d phi;
// 	Vector3d s[3], v[3], w[3];

// 	for (int i = 0; i < dofSize; i++)
// 	{ // deg -> rad
// 		qs[i] = qs[i] * M_PI / 180.0;
// 	}

// 	while (true)
// 	{
// 		for (int i = 0; i < dofSize; i++)
// 			model.q(i) = qs[i];

// 		Matrix3d Rot_temp = CalcBodyWorldOrientation(model.model_, model.q, model.body_id_vec.back(), true).transpose();
// 		Vector3d pos_temp = CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec.back(), model.body_com_pos.back()); // get position and rotation in EE frame;

// 		T0_obj.setIdentity();
// 		T0_obj.topLeftCorner(3, 3) = Rot_temp;
// 		T0_obj.topRightCorner(3, 1) = pos_temp;
// 		Tc_obj = T0_c.inverse() * T0_obj;

// 		d_c.head(3) = Tc_obj.topRightCorner(3, 1);

// 		Matrix3d Tc_obj_rot = refer_rot.transpose()*Rot_temp;
// 		d_c(3) = atan2(Tc_obj_rot(2, 1), Tc_obj_rot(2, 2));
// 		d_c(4) = -asin(Tc_obj_rot(2, 0));
// 		d_c(5) = atan2(Tc_obj_rot(1, 0), Tc_obj_rot(0, 0));

// 		for (int i = 0; i < 6; i++)
// 		{
// 			if (d_c(i) > C(i, 0)) // max
// 				dx(i) = d_c(i) - C(i, 0);
// 			else if (d_c(i) < C(i, 1)) // min
// 				dx(i) = d_c(i) - C(i, 1) ;
// 			else
// 				dx(i) = 0.0;
// 		}

// 		// Algorithm 4 - line 3
// 		if (dx.norm() < 0.01)
// 		{
// 			dx_from_extend_to_goal.head(3) = refer_pos - pos_temp;
// 			dx_from_extend_to_goal.tail(3) = getPhi(Rot_temp, refer_rot);

// 			if (dx_from_init_to_goal.norm() < dx_from_extend_to_goal.norm()){
// 				qs.clear();
// 				return qs;
// 			}

// 			for (int i = 0; i < dofSize; i++)
// 			{
// 				qs[i] = qs[i] / M_PI * 180.0;
// 			}
// 				return qs; // convert rad to angle
// 		}

// 		// Algorithm 4 - line 4
// 		CalcPointJacobian6D(model.model_, model.q, model.body_id_vec.back(), model.body_com_pos.back(), J_temp, true);
// 		J.topLeftCorner(3, dofSize) = J_temp.bottomLeftCorner(3, dofSize);

// 		rpy(0) = atan2(Rot_temp(2, 1), Rot_temp(2, 2));
// 		rpy(1) = -asin(Rot_temp(2, 0));
// 		rpy(2) = atan2(Rot_temp(1, 0), Rot_temp(0, 0));

// 		E_rpy(0,0) = cos(rpy(2))/cos(rpy(1));
// 		E_rpy(0,1) = sin(rpy(2))/cos(rpy(1));
// 		E_rpy(0,2) = 0.0;

// 		E_rpy(1,0) = sin(rpy(2));
// 		E_rpy(1,1) = cos(rpy(2));
// 		E_rpy(1,2) = 0.0;

// 		E_rpy(2,0) = -cos(rpy(2))*sin(rpy(1))/cos(rpy(1));
// 		E_rpy(2,1) = sin(rpy(2))*sin(rpy(1))/cos(rpy(1));
// 		E_rpy(2,2) = 1.0;
// 		J.bottomLeftCorner(3, dofSize) = E_rpy*J_temp.topLeftCorner(3, dofSize);

// 		// Algorithm 4 - line 5
// 		q_error = J.transpose() * (J * J.transpose()).inverse() * dx * 0.3;

// 		// Algorithm 4 -  line 6
// 		for (int i = 0; i < dofSize; i++)
// 			qs[i] -= q_error(i);

// 		// Algorithm 4 - line 7 : stuck here
// 		if (!OutsideJointLimit(qs))
// 		{
// 			//cout << "joint limit" << "\t" << qs[0]/M_PI*180.0<< "\t" << qs[1]/M_PI*180.0<< "\t" << qs[2]/M_PI*180.0<< "\t" << qs[3]/M_PI*180.0<< "\t" << qs[4]/M_PI*180.0<< "\t" << qs[5]/M_PI*180.0<< "\t" << qs[6]/M_PI*180.0 << endl;
// 			qs.clear();
// 			return qs;
// 			// flag = false;
// 			// break;
// 		}
// 	}
// }

double RRT::getDistance(std::vector<double> &a, std::vector<double> &b) {

	//Gets squared Euclidean Distance on C-Space between node and a given
	double accumulated = 0.0f;
	std::vector<double> diff = b;

	for (int i = 0; i< diff.size(); i++) {
		diff[i] -= a[i];
		//double dif = _configuration[i] - (config)[i];
		accumulated += (diff[i] * diff[i] * DOF_weights[i] * DOF_weights[i]);
	}
	return sqrt(accumulated);
}
bool RRT::smoothPath(std::ostream& sout, std::istream& sinput)
{
	std::string fullinput;
	//Parse input
	while (sinput) {
		std::string input;
		sinput >> input;
		fullinput += input;
	}
	int nSmooth = 5; // atoi(fullinput.c_str());

	while (nSmooth) { // while(x) -> while( x != 0 )
		shortcutSmoothing();
		nSmooth--;
	}

	for (int i = 0; i< path.size(); i++) {
		std::vector<double> node = path[i];
		for (int j = 0; j<node.size() - 1; j++) {
			sout << node[j] << ",";
		}
		sout << node[node.size() - 1] << "\n";
	}

	return true;
}
bool RRT::shortcutSmoothing() {
	if (path.size() <= 2) return false;
	int i = ((double)rand() / RAND_MAX)*path.size();
	int j = ((double)rand() / RAND_MAX)*path.size();
	while (abs(j - i)<2) {
		i = ((double)rand() / RAND_MAX)*path.size();
		j = ((double)rand() / RAND_MAX)*path.size();
	}
	if (j < i) { //make sure i is the smaller number
		int a = j;
		j = i;
		i = a;
	}

	if (checkTraj(path[i], path[j])) {
		path.erase(path.begin() + i + 1, path.begin() + j);
		return true;
	}
	return false;
}

bool RRT::checkTraj(std::vector<double> &a, std::vector<double> &b) {
	std::vector<double> u = getUnitVector(a, b); // a->b
	std::vector<double> p = a;
	while (p != b) {
		if (checkExternalCollision(_model, p)) { return false;
		break;
		}
		if (getDistance(p, b) < step_size / 2) {
			p = b;
		}
		else {
			for (int i = 0; i < u.size(); i++) {
				p[i] += u[i] * step_size / 2;
			}
		}
	}
	return true;
}
std::vector<double> RRT::getUnitVector(std::vector<double> &a, std::vector<double> &b) {

	//Gets squared Euclidean Distance on C-Space between node and a given
	double accumulated = 0.0f;
	std::vector<double> diff = b;

	for (int i = 0; i< diff.size(); i++) {
		diff[i] -= a[i];//double dif = _configuration[i] - (config)[i];
		accumulated += (diff[i] * diff[i] * DOF_weights[i] * DOF_weights[i]);
	}
	accumulated = sqrt(accumulated);
	for (int i = 0; i < diff.size(); i++) {
		diff[i] /= accumulated;
	}
	return diff;
}
bool RRT::checkExternalCollision(Robotmodel model, std::vector<double> &config) {
	model.q.resize(dofSize);
	for (int i = 0; i < dofSize; i++)
		model.q(i) = config[i] * M_PI / 180.0;

	// Update Box modeling for link
	for (int i = 0; i < box_num_link; i++)
	{
		Matrix3d Rot_temp3 = Box_link[i].vRot * CalcBodyWorldOrientation(model.model_, model.q, model.body_id_vec[i], true).transpose();
		Box_link[i].vAxis[0] = Rot_temp3.col(0);
		Box_link[i].vAxis[1] = Rot_temp3.col(1);
		Box_link[i].vAxis[2] = Rot_temp3.col(2);
		Box_link[i].vPos = CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec[i], (Box_link[i].vCenter + model.body_com_pos[i]), true);
	}

	bool chk;
	if (box_num_obs == 0)
	{
		chk = false;
		return chk;
	}
	for (int i = 0; i < box_num_obs; i++)
	{
		for (int j = 0; j < box_num_link; j++)
		{
			chk = checkOBBCollision(&Box_obs[i], &Box_link[j]);
			if (chk)
			{
				// cout << "----------------Collision \n"
				// 	 << i << "\t" << j << "\t" << chk << endl;
				// cout << Box_obs[i].vPos.transpose() << endl;
				// cout << Box_link[j].vPos.transpose() << endl;
				// //cout << CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec[j], Vector3d::Zero(), true).transpose() << endl;
				// cout << model.q(0) * 180.0 / M_PI << "\t" << model.q(1) * 180.0 / M_PI << "\t" << model.q(2) * 180.0 / M_PI << "\t" << model.q(3) * 180.0 / M_PI << "\t" << model.q(4) * 180.0 / M_PI << "\t" << model.q(5) * 180.0 / M_PI << "\t" << model.q(6) * 180.0 / M_PI << "\t" << endl;

				//cout << i << "\t" << j << "\t" << chk << endl;

				return chk;
			}
			//cout << "collision" << i << "\t" << j << endl;
		}
	}
	return chk;
}

bool RRT::checkSelfCollision(Robotmodel model, std::vector<double> &config) {
	model.q.resize(dofSize);
	for (int i = 0; i < dofSize; i++)
		model.q(i) = config[i] * M_PI / 180.0;

	// Update Box modeling for link
	for (int i = 0; i < box_num_link; i++)
	{
		Matrix3d Rot_temp3 = Box_link[i].vRot * CalcBodyWorldOrientation(model.model_, model.q, model.body_id_vec[i], true).transpose();
		Box_link[i].vAxis[0] = Rot_temp3.col(0);
		Box_link[i].vAxis[1] = Rot_temp3.col(1);
		Box_link[i].vAxis[2] = Rot_temp3.col(2);
		Box_link[i].vPos = CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec[i], (Box_link[i].vCenter + model.body_com_pos[i]), true);
	}

	bool chk;
	if (box_num_link == 0)
	{
		chk = false;
		return chk;
	}
	for (int i = 0; i < box_num_link; i++)
	{
		for (int j = i + 2; j < box_num_link; j++)
		{
			chk = checkOBBCollision(&Box_link[i], &Box_link[j]);
			if (chk)
			{
				// cout <<  "----------------Collision \n" << i << "\t" << j << "\t" << chk << endl;
				// cout << Box_link[i].vPos.transpose() << endl;
				// cout << Box_link[j].vPos.transpose() << endl;
				// //cout << CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec[j], Vector3d::Zero(), true).transpose() << endl;
				// cout << model.q(0) * 180.0 / M_PI << "\t" << model.q(1) * 180.0 / M_PI << "\t" << model.q(2) * 180.0 / M_PI << "\t" << model.q(3) * 180.0 / M_PI << "\t" << model.q(4) * 180.0 / M_PI << "\t" << model.q(5) * 180.0 / M_PI << "\t" << model.q(6) * 180.0 / M_PI << "\t" << endl;
				return chk;
			}
			// cout << "No collision \n" << i << "\t" << j << "\t" << chk << endl;
			// cout << Box_link[i].vPos.transpose() << endl;
			// cout << Box_link[j].vPos.transpose() << endl;
			// 	cout << CalcBodyToBaseCoordinates(model.model_, model.q, model.body_id_vec[j], Vector3d::Zero(), true).transpose() << endl;
			//cout << model.q(0) * 180.0 / M_PI << "\t" << model.q(1) * 180.0 / M_PI << "\t" << model.q(2) * 180.0 / M_PI << "\t" << model.q(3) * 180.0 / M_PI << "\t" << model.q(4) * 180.0 / M_PI << "\t" << model.q(5) * 180.0 / M_PI << "\t" << model.q(6) * 180.0 / M_PI << "\t" << endl;
		}
	}
	return chk;
}

bool RRT::checkJointLimit(std::vector<double> q) { // radian
	for (int i = 0;i < q.size();i++) {
		if (q[i] > upper_limit(i)/180.0*M_PI) {
			//cout << q[i] << "\t" << upper_limit(i)/180.0*M_PI << "\t" << i << endl;
			return false;
		}
		else if (q[i] < lower_limit(i)/180.0*M_PI) {
			//cout << q[i] << "\t" << lower_limit(i)/180.0*M_PI  << "\t" << i << endl;

			return false;
		}
	}
	return true;

}

bool RRT::checkOBBCollision(ST_OBB* box0, ST_OBB* box1) // Collision-free : False
{
	// compute difference of box centers,D=C1-C0
	Vector3d D = Vector3d(box1->vPos(0) - box0->vPos(0), box1->vPos(1) - box0->vPos(1), box1->vPos(2) - box0->vPos(2));

	float C[3][3];    //matrix C=A^T B,c_{ij}=Dot(A_i,B_j)
	float absC[3][3]; //|c_{ij}|
	float AD[3];      //Dot(A_i,D)
	float R0, R1, R;    //interval radii and distance between centers
	float R01;        //=R0+R1

					  //A0
	C[0][0] = FDotProduct(box0->vAxis[0], box1->vAxis[0]);// vAxis : direction // 3D Dot product
	C[0][1] = FDotProduct(box0->vAxis[0], box1->vAxis[1]);
	C[0][2] = FDotProduct(box0->vAxis[0], box1->vAxis[2]);
	AD[0] = FDotProduct(box0->vAxis[0], D);
	absC[0][0] = (float)fabsf(C[0][0]);
	absC[0][1] = (float)fabsf(C[0][1]);
	absC[0][2] = (float)fabsf(C[0][2]);
	R = (float)fabsf(AD[0]);
	R1 = box1->fAxis(0) * absC[0][0] + box1->fAxis(1) * absC[0][1] + box1->fAxis(2) * absC[0][2];
	R01 = box0->fAxis(0) + R1;
	if (R > R01)
		return 0;

	//A1
	C[1][0] = FDotProduct(box0->vAxis[1], box1->vAxis[0]);
	C[1][1] = FDotProduct(box0->vAxis[1], box1->vAxis[1]);
	C[1][2] = FDotProduct(box0->vAxis[1], box1->vAxis[2]);
	AD[1] = FDotProduct(box0->vAxis[1], D);
	absC[1][0] = (float)fabsf(C[1][0]);
	absC[1][1] = (float)fabsf(C[1][1]);
	absC[1][2] = (float)fabsf(C[1][2]);
	R = (float)fabsf(AD[1]);
	R1 = box1->fAxis(0) * absC[1][0] + box1->fAxis(1) * absC[1][1] + box1->fAxis(2) * absC[1][2];
	R01 = box0->fAxis(1) + R1;
	if (R > R01)
		return 0;

	//A2
	C[2][0] = FDotProduct(box0->vAxis[2], box1->vAxis[0]);
	C[2][1] = FDotProduct(box0->vAxis[2], box1->vAxis[1]);
	C[2][2] = FDotProduct(box0->vAxis[2], box1->vAxis[2]);
	AD[2] = FDotProduct(box0->vAxis[2], D);
	absC[2][0] = (float)fabsf(C[2][0]);
	absC[2][1] = (float)fabsf(C[2][1]);
	absC[2][2] = (float)fabsf(C[2][2]);
	R = (float)fabsf(AD[2]);
	R1 = box1->fAxis(0) * absC[2][0] + box1->fAxis(1) * absC[2][1] + box1->fAxis(2) * absC[2][2];
	R01 = box0->fAxis(2) + R1;
	if (R > R01)
		return 0;

	//B0
	R = (float)fabsf(FDotProduct(box1->vAxis[0], D));
	R0 = box0->fAxis(0) * absC[0][0] + box0->fAxis(1) * absC[1][0] + box0->fAxis(2) * absC[2][0];
	R01 = R0 + box1->fAxis(0);
	if (R > R01)
		return 0;

	//B1
	R = (float)fabsf(FDotProduct(box1->vAxis[1], D));
	R0 = box0->fAxis(0) * absC[0][1] + box0->fAxis(1) * absC[1][1] + box0->fAxis(2) * absC[2][1];
	R01 = R0 + box1->fAxis(1);
	if (R > R01)
		return 0;

	//B2
	R = (float)fabsf(FDotProduct(box1->vAxis[2], D));
	R0 = box0->fAxis(0) * absC[0][2] + box0->fAxis(1) * absC[1][2] + box0->fAxis(2) * absC[2][2];
	R01 = R0 + box1->fAxis(2);
	if (R > R01)
		return 0;

	//A0xB0
	R = (float)fabsf(AD[2] * C[1][0] - AD[1] * C[2][0]);
	R0 = box0->fAxis(1) * absC[2][0] + box0->fAxis(2) * absC[1][0];
	R1 = box1->fAxis(1) * absC[0][2] + box1->fAxis(2) * absC[0][1];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A0xB1
	R = (float)fabsf(AD[2] * C[1][1] - AD[1] * C[2][1]);
	R0 = box0->fAxis(1) * absC[2][1] + box0->fAxis(2) * absC[1][1];
	R1 = box1->fAxis(0) * absC[0][2] + box1->fAxis(2) * absC[0][0];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A0xB2
	R = (float)fabsf(AD[2] * C[1][2] - AD[1] * C[2][2]);
	R0 = box0->fAxis(1) * absC[2][2] + box0->fAxis(2) * absC[1][2];
	R1 = box1->fAxis(0) * absC[0][1] + box1->fAxis(1) * absC[0][0];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A1xB0
	R = (float)fabsf(AD[0] * C[2][0] - AD[2] * C[0][0]);
	R0 = box0->fAxis(0) * absC[2][0] + box0->fAxis(2) * absC[0][0];
	R1 = box1->fAxis(1) * absC[1][2] + box1->fAxis(2) * absC[1][1];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A1xB1
	R = (float)fabsf(AD[0] * C[2][1] - AD[2] * C[0][1]);
	R0 = box0->fAxis(0) * absC[2][1] + box0->fAxis(2) * absC[0][1];
	R1 = box1->fAxis(0) * absC[1][2] + box1->fAxis(2) * absC[1][0];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A1xB2
	R = (float)fabsf(AD[0] * C[2][2] - AD[2] * C[0][2]);
	R0 = box0->fAxis(0) * absC[2][2] + box0->fAxis(2) * absC[0][2];
	R1 = box1->fAxis(0) * absC[1][1] + box1->fAxis(1) * absC[1][0];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A2xB0
	R = (float)fabsf(AD[1] * C[0][0] - AD[0] * C[1][0]);
	R0 = box0->fAxis(0) * absC[1][0] + box0->fAxis(1) * absC[0][0];
	R1 = box1->fAxis(1) * absC[2][2] + box1->fAxis(2) * absC[2][1];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A2xB1
	R = (float)fabsf(AD[1] * C[0][1] - AD[0] * C[1][1]);
	R0 = box0->fAxis(0) * absC[1][1] + box0->fAxis(1) * absC[0][1];
	R1 = box1->fAxis(0) * absC[2][2] + box1->fAxis(2) * absC[2][0];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	//A2xB2
	R = (float)fabsf(AD[1] * C[0][2] - AD[0] * C[1][2]);
	R0 = box0->fAxis(0) * absC[1][2] + box0->fAxis(1) * absC[0][2];
	R1 = box1->fAxis(0) * absC[2][1] + box1->fAxis(1) * absC[2][0];
	R01 = R0 + R1;
	if (R > R01)
		return 0;

	return 1;
}
MatrixXd RRT::MergeRRTResults(MatrixXd joint_target1, MatrixXd joint_target2, int case_){
	// case 0 : RRT / RRT
	// case 1 : RRT / CBiRRT
	// case 2 : CBiRRT / RRT
	// case 3 : CBiRRT / CBiRRT
	MatrixXd merged_joint_target;
	int row[2];
	int reminder[2];
	switch(case_ )
	{
		case 0 :
			row[0] = joint_target1.rows();
			row[1] = joint_target2.rows();
			merged_joint_target.resize(row[0] + row[1] -1, dofSize);
			merged_joint_target.topRows(row[0]) = joint_target1.topRows(row[0]);
			merged_joint_target.bottomRows(row[1]-1) = joint_target2.bottomRows(row[1]-1);
		break;
		case 1 :
			row[0] = joint_target1.rows();
			row[1] = joint_target2.rows() / 20;
			reminder[1] = joint_target2.rows() % 20;
			if (reminder[1] == 0)
			{
				merged_joint_target.resize(row[0] + row[1], dofSize);
				merged_joint_target.topRows(row[0]) = joint_target1.topRows(row[0]);
				for (int i = 0; i < row[1]; i++)
					merged_joint_target.row(row[0] + i) = joint_target2.row(20 * (i + 1) - 1);
			}
			else
			{
				merged_joint_target.resize(row[0] + row[1] + 1,dofSize);
				merged_joint_target.topRows(row[0]) = joint_target1.topRows(row[0]);
				for (int i = 0; i < row[1]; i++)
					merged_joint_target.row(row[0] + i) = joint_target2.row(20 * (i + 1)- 1);

				merged_joint_target.bottomRows(1) = joint_target2.bottomRows(1);
			}
			break;
		case 2 :
			row[0] = joint_target1.rows() /20;
			row[1] = joint_target2.rows() ;
			reminder[0] = joint_target1.rows() % 20;
			if (reminder[0] == 0)
			{
				merged_joint_target.resize(row[0] + row[1],dofSize);
				for (int i = 0; i < row[0]; i++)
					merged_joint_target.row(i) = joint_target1.row(20 * (i));

				merged_joint_target.bottomRows(row[1]) = joint_target2.bottomRows(row[1]);
			}
			else{
				merged_joint_target.resize(row[0] + row[1] + 1,dofSize);
				for (int i = 0; i < row[0] + 1; i++)
					merged_joint_target.row(i) = joint_target1.row(20 * (i));

				merged_joint_target.bottomRows(row[1]) = joint_target2.bottomRows(row[1]);
			}
		break;
		case 3 :
			row[0] = joint_target1.rows() /20;
			row[1] = joint_target2.rows() /20;
			reminder[0] = joint_target1.rows() % 20;
			reminder[1] = joint_target2.rows() % 20;
			if (reminder[0] == 0){
				if (reminder[1] == 0){
					merged_joint_target.resize(row[0] + row[1] + 1,dofSize);
					for (int i = 0; i < row[0] + 1; i++)
						merged_joint_target.row(i) = joint_target1.row(20 * (i));

					for (int i = 0; i < row[1]; i++)
						merged_joint_target.row(row[0] + 1 + i) = joint_target2.row(20 * (i+1));
				}
				else{
					merged_joint_target.resize(row[0] + row[1] + 2,dofSize);
					for (int i = 0; i < row[0] + 1; i++)
						merged_joint_target.row(i) = joint_target1.row(20 * (i));

					for (int i = 0; i < row[1]; i++)
						merged_joint_target.row(row[0] + 1 + i) = joint_target2.row(20 * (i+1));

					merged_joint_target.bottomRows(1) = joint_target2.bottomRows(1);
				}
			}
			else{
				if (reminder[1] == 0){
					merged_joint_target.resize(row[0] + row[1] + 2,dofSize);
					for (int i = 0; i < row[0] + 1; i++)
						merged_joint_target.row(i) = joint_target1.row(20 * (i));

					merged_joint_target.row(row[0]+1) = joint_target1.bottomRows(1);

					for (int i = 0; i < row[1]; i++)
						merged_joint_target.row(row[0] + 2 + i) = joint_target2.row(20 * (i+1));
				}
				else{
					merged_joint_target.resize(row[0] + row[1] + 3,dofSize);
					for (int i = 0; i < row[0] + 1; i++)
						merged_joint_target.row(i) = joint_target1.row(20 * (i));

					merged_joint_target.row(row[0]+1) = joint_target1.bottomRows(1);

					for (int i = 0; i < row[1]; i++)
						merged_joint_target.row(row[0] + 2 + i) = joint_target2.row(20 * (i+1));

					merged_joint_target.bottomRows(1) = joint_target2.bottomRows(1);
				}
			}
		break;
	}

	return merged_joint_target;
}


/*****************************************************************************/
/*  Copyright (c) 2015, Karl Pauwels                                         */
/*  All rights reserved.                                                     */
/*                                                                           */
/*  Redistribution and use in source and binary forms, with or without       */
/*  modification, are permitted provided that the following conditions       */
/*  are met:                                                                 */
/*                                                                           */
/*  1. Redistributions of source code must retain the above copyright        */
/*  notice, this list of conditions and the following disclaimer.            */
/*                                                                           */
/*  2. Redistributions in binary form must reproduce the above copyright     */
/*  notice, this list of conditions and the following disclaimer in the      */
/*  documentation and/or other materials provided with the distribution.     */
/*                                                                           */
/*  3. Neither the name of the copyright holder nor the names of its         */
/*  contributors may be used to endorse or promote products derived from     */
/*  this software without specific prior written permission.                 */
/*                                                                           */
/*  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS      */
/*  "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT        */
/*  LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR    */
/*  A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT     */
/*  HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,   */
/*  SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT         */
/*  LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,    */
/*  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY    */
/*  THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT      */
/*  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE    */
/*  OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.     */
/*****************************************************************************/

#include <numeric>
#include <algorithm>
#include <cstdlib>
#include <math.h>
#include <cstring>
#include <cuda_profiler_api.h>
#include <stdexcept>
#include <Eigen/Dense>
#include <d_multiple_rigid_poses.h>
#include <utility_kernels_pose.h>
#include <multiple_rigid_pose_kernels.h>
#include <utilities.h>

namespace pose {

D_MultipleRigidPoses::D_MultipleRigidPoses(int n_cols, int n_rows,
                                           float nodal_point_x,
                                           float nodal_point_y,
                                           float focal_length_x,
                                           float focal_length_y, float baseline,
                                           Parameters parameters)
    : parameters_{ parameters }, _n_cols{ n_cols }, _n_rows{ n_rows },
      _nodal_point_x{ nodal_point_x }, _nodal_point_y{ nodal_point_y },
      _focal_length_x{ focal_length_x }, _focal_length_y{ focal_length_y },
      _baseline{ baseline }, _currentSparseObject{ 0 },
      _multipleRigidModelsOgre{
        std::unique_ptr<MultipleRigidModelsOgre>{ new MultipleRigidModelsOgre(
            _n_cols, _n_rows, _focal_length_x, _focal_length_y, _nodal_point_x,
            _nodal_point_y, parameters_.near_plane_, parameters_.far_plane_) }
      },
      cub_radix_sorter_{
        std::unique_ptr<util::CubRadixSorter<unsigned int, int> >{
          new util::CubRadixSorter<unsigned int, int>(
              _n_rows * _n_cols * _N_FLOWS, 0, parameters_.getKeyBits())
        }
      },
      _n_objects{ 0 }, _running{ true }, render_state_changed_{ true } {

  camera_pose_.setValid(true);
  previous_camera_pose_.setValid(false);

  /*******************************************************/
  /* Pre-allocate (and initiate )all device space needed */
  /* with maximal possible size                          */
  /* to avoid re-allocation and re-computation           */
  /*******************************************************/

  // linear index to store pixel locations
  // goes from 0 to n_rows*n_cols*n_flows
  // used to identify flow source
  std::vector<int> h_linear_ind(_n_rows * _n_cols * _N_FLOWS);
  std::iota(h_linear_ind.begin(), h_linear_ind.end(), 0);
  d_linear_ind_ =
      util::Device1D<int>::make_unique(_n_rows * _n_cols * _N_FLOWS);
  d_linear_ind_->copyFrom(h_linear_ind);

  // 0,1,2,...,n_objects valid locations
  d_valid_flow_Zbuffer_ =
      util::Device1D<unsigned int>::make_unique(_n_rows * _n_cols * _N_FLOWS);
  d_valid_disparity_Zbuffer_ =
      util::Device1D<unsigned int>::make_unique(_n_rows * _n_cols);
  d_valid_flow_Zbuffer_sub_ =
      util::Device1D<unsigned int>::make_unique(_n_rows * _n_cols * _N_FLOWS);
  d_valid_disparity_Zbuffer_sub_ =
      util::Device1D<unsigned int>::make_unique(_n_rows * _n_cols);
  d_extra_disparity_buffer_ =
      util::Device1D<unsigned int>::make_unique(_n_rows * _n_cols);

  // index of valid locations
  d_ind_flow_Zbuffer_ =
      util::Device1D<int>::make_unique(_n_rows * _n_cols * _N_FLOWS);
  d_ind_disparity_Zbuffer_ =
      util::Device1D<int>::make_unique(_n_rows * _n_cols);
  d_ind_flow_Zbuffer_sub_ =
      util::Device1D<int>::make_unique(_n_rows * _n_cols * _N_FLOWS);
  d_ind_disparity_Zbuffer_sub_ =
      util::Device1D<int>::make_unique(_n_rows * _n_cols);

  // starting index of sorted segments
  d_seg_start_inds_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);

  // Gathering
  d_flow_compact_ =
      util::Device1D<float2>::make_unique(_n_cols * _n_rows * _N_FLOWS);
  d_Zbuffer_flow_compact_ =
      util::Device1D<float>::make_unique(_n_cols * _n_rows * _N_FLOWS);
  d_disparity_compact_ = util::Device1D<float>::make_unique(_n_cols * _n_rows);
  d_Zbuffer_normals_compact_ =
      util::Device1D<float4>::make_unique(_n_cols * _n_rows);
  d_n_values_flow_ = util::Device1D<int>::make_unique(parameters_.max_objects_);
  d_start_ind_flow_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);
  d_n_values_disparity_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);
  d_start_ind_disparity_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);

  // Normal equations
  d_CO_ = util::Device1D<float>::make_unique(
      4 * _MAX_N_VAL_ACCUM * _N_CON_FLOW * parameters_.max_objects_);
  d_CO_reduced_ = util::Device1D<float>::make_unique(_N_CON_FLOW *
                                                     parameters_.max_objects_);
  d_CD_ = util::Device1D<float>::make_unique(
      4 * _MAX_N_VAL_ACCUM * _N_CON_DISP * parameters_.max_objects_);
  d_CD_reduced_ = util::Device1D<float>::make_unique(_N_CON_DISP *
                                                     parameters_.max_objects_);
  d_abs_res_ =
      util::Device1D<float>::make_unique((_N_FLOWS + 1) * _n_rows * _n_cols);

  h_CO_reduced_.resize(_N_CON_FLOW * parameters_.max_objects_);
  h_CD_reduced_.resize(_N_CON_DISP * parameters_.max_objects_);
  segment_normal_eqs_.resize(parameters_.max_objects_);

  // Segment starting indices in single residual structure
  d_offset_ind_res_flow_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);
  d_offset_ind_res_disparity_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);
  d_dTR_ = util::Device1D<float>::make_unique(6 * parameters_.max_objects_);
  d_delta_T_accum_ =
      util::Device1D<float>::make_unique(3 * parameters_.max_objects_);
  d_delta_Rmat_accum_ =
      util::Device1D<float>::make_unique(9 * parameters_.max_objects_);
  d_segment_translation_table_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);

  // Approximate median
  // pre-generate maximum number of random numbers required (uniform between
  // [0,1[ , used for sampling (with replacement) of absolute residuals
  std::vector<float> h_random_numbers(_n_rows * _n_cols * (_N_FLOWS + 1));
  srand(0);
  for (auto &it : h_random_numbers)
    it = (float)((double)rand() / double(RAND_MAX));
  d_random_numbers_ =
      util::Device1D<float>::make_unique(_n_rows * _n_cols * (_N_FLOWS + 1));
  d_random_numbers_->copyFrom(h_random_numbers);
  int pp = int(ceil(double(_n_rows * _n_cols * (_N_FLOWS + 1)) /
                    243.0)); // second stage reduction
  d_median_tmp_ = util::Device1D<float>::make_unique(pp);
  d_abs_res_scales_ =
      util::Device1D<float>::make_unique(parameters_.max_objects_);

  // temps used inside median calculation
  d_median_n_in_ = util::Device1D<int>::make_unique(parameters_.max_objects_);
  d_median_start_inds_ =
      util::Device1D<int>::make_unique(parameters_.max_objects_);

  // Initial Zbuffer
  d_init_Z_ = util::Device1D<float>::make_unique(_n_rows * _n_cols);

  // Residual flow
  d_res_flowx_ = util::Device1D<float>::make_unique(_n_rows * _n_cols);
  d_res_flowy_ = util::Device1D<float>::make_unique(_n_rows * _n_cols);

  // Residual ar flow
  d_res_ar_flowx_ = util::Device1D<float>::make_unique(_n_rows * _n_cols);
  d_res_ar_flowy_ = util::Device1D<float>::make_unique(_n_rows * _n_cols);
}

void D_MultipleRigidPoses::addModel(const char *obj_filename, float obj_scale,
                                    TranslationRotation3D initial_pose) {
  if (_n_objects > (parameters_.max_objects_ - 2))
    throw std::runtime_error("D_MultipleRigidPoses::addModel: Max objects "
                             "exceeded, increase key_bits\n");

  _n_objects++;
  _currentPoses.push_back(initial_pose);
  setSparsePose(initial_pose, _n_objects - 1);

  _multipleRigidModelsOgre->addModel(obj_filename);
}

void D_MultipleRigidPoses::removeAllModels() {
  _multipleRigidModelsOgre->removeAllModels();
  _currentPoses.clear();
  _currentSparseObject = 0;
  _n_objects = 0;
}

void D_MultipleRigidPoses::update(const util::Device1D<float> &d_flowx,
                                  const util::Device1D<float> &d_flowy,
                                  const util::Device1D<float> &d_ar_flowx,
                                  const util::Device1D<float> &d_ar_flowy,
                                  const util::Device2D<float> &d_disparity,
                                  std::bitset<32> segments_to_update) {
// small inaccuracy here is that the optical flow (frame 1 -> frame 2) is
// segmented according to the region estimated at frame 2
// this is necessary in order to update the residual flow according to the
// updated model depth
// have shown (in simulation) that the resulting pose estimates are a lot more
// accurate than those based on initial depth

#ifdef TIME_STEPS
  // Setup timers
  cudaEvent_t start, end;
  float elapsed_time;
  cudaEventCreate(&start);
  cudaEventCreate(&end);
#endif

  if (_running) {

    // reset timers
    int n_timers = 14;
    _compTimes.clear();
    _compTimes.assign(n_timers, 0.0);

    /* 0 preprocess-total
     * 1 ols-total
     * 2 robust-total
     * 3 mark valids
     * 4 radix sort
     * 5 gather
     * 6 compose normal equations
     * 7 reduce normal equations
     * 8 solve normal equations (cpu)
     * 9 compute absolute residuals
     * 10 median absolute residuals
     * 11 render
     * 12 pose accumulation (cpu)
     * 13 residual flow
     */

    std::vector<TranslationRotation3D> delta_poses_accum(
        _n_objects, TranslationRotation3D());

    std::vector<TranslationRotation3D> explained_delta_poses(
        _n_objects, TranslationRotation3D());

#ifdef TIME_STEPS
    cudaEventRecord(start, 0);
#endif

    // first rendering
    render(_currentPoses);

#ifdef TIME_STEPS
    cudaThreadSynchronize();
    cudaEventRecord(end, 0);
    cudaEventSynchronize(end);
    cudaEventElapsedTime(&elapsed_time, start, end);
    _compTimes.at(11) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
    cudaEventRecord(start, 0);
#endif

    // store initial Zbuffer as initial model for residual flow
    convertZbufferToZ(d_init_Z_->data(), _multipleRigidModelsOgre->getZBuffer(),
                      _n_cols, _n_rows, _nodal_point_x, _nodal_point_y,
                      parameters_.near_plane_, parameters_.far_plane_);

#ifdef TIME_STEPS
    cudaThreadSynchronize();
    cudaEventRecord(end, 0);
    cudaEventSynchronize(end);
    cudaEventElapsedTime(&elapsed_time, start, end);
    _compTimes.at(13) += (double)elapsed_time;
#endif

    // start iterations
    std::vector<float> dT_accum(3 * _n_objects);
    std::vector<float> dR_accum(9 * _n_objects);
    std::vector<TranslationRotation3D> poses_accum = _currentPoses;

    for (int it = 1; it <= parameters_.n_icp_outer_it_; it++) {

// update residual flow
#ifdef TIME_STEPS
      cudaEventRecord(start, 0);
#endif

      for (int o = 0; o < _n_objects; o++) {
        auto delta_pose =
            poses_accum.at(o) * _currentPoses.at(o).inverseTransform();
        if (previous_camera_pose_.isValid()) {
          explained_delta_poses.at(o) = camera_pose_.inverseTransform() *
                                        delta_pose * previous_camera_pose_;
        } else {
          explained_delta_poses.at(o) = delta_pose;
        }
      }

      for (int o = 0; o < _n_objects; o++) {
        explained_delta_poses.at(o).getT(&dT_accum[3 * o]);
        explained_delta_poses.at(o).getR_mat(&dR_accum[9 * o]);
      }
      d_delta_T_accum_->copyFrom(dT_accum, 3 * _n_objects);
      d_delta_Rmat_accum_->copyFrom(dR_accum, 9 * _n_objects);

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end, 0);
      cudaEventSynchronize(end);
      cudaEventElapsedTime(&elapsed_time, start, end);
      _compTimes.at(12) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
      cudaEventRecord(start, 0);
#endif
      computeResidualFlow(
          d_res_flowx_->data(), d_res_flowy_->data(), d_res_ar_flowx_->data(),
          d_res_ar_flowy_->data(), d_flowx.data(), d_flowy.data(),
          d_ar_flowx.data(), d_ar_flowy.data(), d_delta_T_accum_->data(),
          d_delta_Rmat_accum_->data(), d_init_Z_->data(),
          _multipleRigidModelsOgre->getSegmentIND(), _n_cols, _n_rows,
          _nodal_point_x, _nodal_point_y, _focal_length_x, _focal_length_y);

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end, 0);
      cudaEventSynchronize(end);
      cudaEventElapsedTime(&elapsed_time, start, end);
      _compTimes.at(13) += (double)elapsed_time;
#endif

      robustPoseUpdates(d_res_flowx_->data(), d_res_flowy_->data(),
                        d_res_ar_flowx_->data(), d_res_ar_flowy_->data(),
                        d_disparity.data(), d_disparity.pitch(),
                        (int)segments_to_update.to_ulong());

      // accumulate delta poses
      for (int s = 0; s < _segment_info.size(); s++) {
        SegmentINFO &tmp = _segment_info.at(s);
        delta_poses_accum.at(tmp.segment_ind) =
            _robustDeltaPoses.at(s) * delta_poses_accum.at(tmp.segment_ind);
      }

#ifdef TIME_STEPS
      cudaEventRecord(start, 0);
#endif

      // update poses, invalidating if magnitude T too large
      poses_accum = _currentPoses;

      // transform to camera frame
      for (auto &it : poses_accum)
        it = camera_pose_.inverseTransform() * it;

      for (int o = 0; o < _n_objects; o++)
        if (delta_poses_accum.at(o).normT2() <
            parameters_.max_t_update_norm_squared_)
          poses_accum.at(o) = delta_poses_accum.at(o) * poses_accum.at(o);
        else // invalidate
          poses_accum.at(o).setValid(false);

      // transform back to world frame
      for (auto &it : poses_accum)
        it = camera_pose_ * it;

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end, 0);
      cudaEventSynchronize(end);
      cudaEventElapsedTime(&elapsed_time, start, end);
      _compTimes.at(12) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
      cudaEventRecord(start, 0);
#endif

      // render at updated pose (unless last iteration)
      if (it < parameters_.n_icp_outer_it_)
        render(poses_accum);

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end, 0);
      cudaEventSynchronize(end);
      cudaEventElapsedTime(&elapsed_time, start, end);
      _compTimes.at(11) += (double)elapsed_time;
#endif
    }

    // update current pose
    _currentPoses = poses_accum;
  }
}

void D_MultipleRigidPoses::updateNormalEquations(
    const util::Device1D<float> &d_flowx, const util::Device1D<float> &d_flowy,
    const util::Device1D<float> &d_ar_flowx,
    const util::Device1D<float> &d_ar_flowy,
    const util::Device2D<float> &d_disparity,
    const std::vector<TranslationRotation3D> &explained_delta_poses,
    std::bitset<32> segments_to_update) {
  // ensure rendering is up-to-date
  render(_currentPoses);

  // store initial Zbuffer as initial model for residual flow
  convertZbufferToZ(d_init_Z_->data(), _multipleRigidModelsOgre->getZBuffer(),
                    _n_cols, _n_rows, _nodal_point_x, _nodal_point_y,
                    parameters_.near_plane_, parameters_.far_plane_);

  // initialize residual flow (will improve later)
  d_res_flowx_->copyFrom(d_flowx);
  d_res_flowy_->copyFrom(d_flowy);
  d_res_ar_flowx_->copyFrom(d_ar_flowx);
  d_res_ar_flowy_->copyFrom(d_ar_flowy);

  // subtract camera-motion-induced flow
  std::vector<float> dT_accum(3 * _n_objects);
  std::vector<float> dR_accum(9 * _n_objects);

  for (int o = 0; o < _n_objects; o++) {
    explained_delta_poses.at(o).getT(&dT_accum[3 * o]);
    explained_delta_poses.at(o).getR_mat(&dR_accum[9 * o]);
  }
  d_delta_T_accum_->copyFrom(dT_accum, 3 * _n_objects);
  d_delta_Rmat_accum_->copyFrom(dR_accum, 9 * _n_objects);

  computeResidualFlow(
      d_res_flowx_->data(), d_res_flowy_->data(), d_res_ar_flowx_->data(),
      d_res_ar_flowy_->data(), d_flowx.data(), d_flowy.data(),
      d_ar_flowx.data(), d_ar_flowy.data(), d_delta_T_accum_->data(),
      d_delta_Rmat_accum_->data(), d_init_Z_->data(),
      _multipleRigidModelsOgre->getSegmentIND(), _n_cols, _n_rows,
      _nodal_point_x, _nodal_point_y, _focal_length_x, _focal_length_y);

  // reset normal equations
  for (auto &it : segment_normal_eqs_)
    it.reset();

  // perform robust iteration
  robustPoseUpdates(d_res_flowx_->data(), d_res_flowy_->data(),
                    d_res_ar_flowx_->data(), d_res_ar_flowy_->data(),
                    d_disparity.data(), d_disparity.pitch(),
                    (int)segments_to_update.to_ulong());

  // prepare normal equations for (external) constrained update
  int min_samples = 1000;
  for (auto &it : _segment_info) {
    // reset if not enough samples or update norm too large
    if (((it.n_values_flow + it.n_values_disparity) < min_samples) &&
        (segment_normal_eqs_.at(it.segment_ind).squaredNormDeltaT() >=
         parameters_.max_t_update_norm_squared_))
      segment_normal_eqs_.at(it.segment_ind).reset();
  }
}

void D_MultipleRigidPoses::setWeights(float w_flow, float w_ar_flow,
                                      float w_disp) {
  parameters_.w_flow_ = w_flow;
  parameters_.w_ar_flow_ = w_ar_flow;
  parameters_.w_disp_ = w_disp;
}

void D_MultipleRigidPoses::getWeights(float &w_flow, float &w_ar_flow,
                                      float &w_disp) const {
  w_flow = parameters_.w_flow_;
  w_ar_flow = parameters_.w_ar_flow_;
  w_disp = parameters_.w_disp_;
}

cudaArray *D_MultipleRigidPoses::getTexture() {
  render(_currentPoses);
  return (_multipleRigidModelsOgre->getTexture());
}

cudaArray *D_MultipleRigidPoses::getSegmentIND() {
  render(_currentPoses);
  return (_multipleRigidModelsOgre->getSegmentIND());
}

cudaArray *D_MultipleRigidPoses::getZbuffer() {
  render(_currentPoses);
  return _multipleRigidModelsOgre->getZBuffer();
}

cudaArray *D_MultipleRigidPoses::getNormalX() {
  render(_currentPoses);
  return _multipleRigidModelsOgre->getNormalX();
}

cudaArray *D_MultipleRigidPoses::getNormalY() {
  render(_currentPoses);
  return _multipleRigidModelsOgre->getNormalY();
}

cudaArray *D_MultipleRigidPoses::getNormalZ() {
  render(_currentPoses);
  return _multipleRigidModelsOgre->getNormalZ();
}

void D_MultipleRigidPoses::setCameraParameters(float focal_length_x,
                                               float focal_length_y,
                                               float nodal_point_x,
                                               float nodal_point_y) {
  render_state_changed_ = true;
  _nodal_point_x = nodal_point_x;
  _nodal_point_y = nodal_point_y;
  _focal_length_x = focal_length_x;
  _focal_length_y = focal_length_y;
  _multipleRigidModelsOgre->updateProjectionMatrix(
      _focal_length_x, _focal_length_y, _nodal_point_x, _nodal_point_y,
      parameters_.near_plane_, parameters_.far_plane_);
}

void
D_MultipleRigidPoses::setCameraPose(const TranslationRotation3D &camera_pose) {
  render_state_changed_ = true;

  // if this is the first camera pose received, make previous and current equal
  previous_camera_pose_ =
      previous_camera_pose_.isValid() ? camera_pose_ : camera_pose;

  camera_pose_ = camera_pose;
  _multipleRigidModelsOgre->updateCameraPose(camera_pose_.rotateX180());
}

bool D_MultipleRigidPoses::isDenseWinner() {

  bool denseWinner = true;
  double thres = parameters_.reliability_threshold_;

  // make sure that sparse does not negatively affect any object's proportion ar
  // valid
  // add some margin to allow minor decrease in reliability (due to noise)
  double margin = parameters_.sparse_intro_allowed_reliability_decrease_;
  bool sparseRelDecrease = false;
  for (int o = 0; o < _n_objects; o++)
    if (_S_ar_flow_prop_valid[o] < (_D_ar_flow_prop_valid[o] - margin))
      sparseRelDecrease = true;

  if (!sparseRelDecrease) {

    // compare reliability at the current sparse object
    int o = _currentSparseObject;

    // we require a pretty high reliability for sparse introduction
    double intro_thres = parameters_.sparse_intro_reliability_threshold_;
    if ((_S_ar_flow_prop_valid[o] > (_D_ar_flow_prop_valid[o] + intro_thres)) ||
        ((_S_ar_flow_prop_valid[o] > intro_thres) &&
         (_D_ar_flow_prop_valid[o] < thres)))
      denseWinner = false;
  }

  return (denseWinner);
}

void D_MultipleRigidPoses::render(
    const std::vector<TranslationRotation3D> &renderPoses) {
  if ((renderPoses != _lastPosesRendered) || (render_state_changed_)) {

    //    util::TimerGPU render_timer;
    _multipleRigidModelsOgre->render(renderPoses);
    //    std::cout << "render time: " << render_timer.read() << " ms\n";
    _lastPosesRendered = renderPoses;
    render_state_changed_ = false;
  }
}

void
D_MultipleRigidPoses::getSegmentLengths(std::vector<int> &lengths,
                                        const std::vector<int> &starting_inds,
                                        int n_segments, int total_length) {

  for (int i = 0; i < (n_segments - 1); i++) {
    int s = starting_inds.at(i);
    if (s >= 0) {

      // find end-point
      int e = starting_inds.at(i + 1);
      for (int j = (i + 2); ((e < 0) && (j < n_segments)); j++)
        e = starting_inds.at(j);

      if (e < 0)
        e = total_length;

      lengths.at(i) = e - s;

    } else

      lengths.at(i) = 0;
  }

  // final segment
  int s = starting_inds.at(n_segments - 1);
  if (s >= 0)
    lengths.at(n_segments - 1) = total_length - s;
  else
    lengths.at(n_segments - 1) = 0;
}

int D_MultipleRigidPoses::getNextSelectedSparseObject(bool dense) {

  std::vector<double> probabilities =
      dense ? _D_ar_flow_prop_valid : _S_ar_flow_prop_valid;

  double sum = 0.0;
  for (int o = 0; o < _n_objects; o++) {
    probabilities.at(o) = 1.0 - probabilities.at(o);
    sum += probabilities.at(o);
  }

  for (int o = 0; o < _n_objects; o++)
    probabilities.at(o) /= sum;

  double r = ((double)rand() / double(RAND_MAX));

  int object = 0;
  double prob_ulim = probabilities.at(object);

  while (r > prob_ulim) {
    object++;
    prob_ulim += probabilities.at(object);
  }

  return (object);
}

void D_MultipleRigidPoses::evaluateARFlowPoseError(
    bool dense, const util::Device1D<float> &d_ar_flowx,
    std::vector<TranslationRotation3D> &poses) {
  if (dense)
    computeARFlowPoseError(d_ar_flowx, poses, _D_ar_flow_prop_valid);
  else
    computeARFlowPoseError(d_ar_flowx, poses, _S_ar_flow_prop_valid);
}

void D_MultipleRigidPoses::computeARFlowPoseError(
    const util::Device1D<float> &d_ar_flowx,
    std::vector<TranslationRotation3D> &poses,
    std::vector<double> &ar_flow_prop_valid) {
  std::vector<double> ar_flow_abs_valid; // discarded
  computeARFlowPoseError(d_ar_flowx, poses, ar_flow_prop_valid,
                         ar_flow_abs_valid);
}

void D_MultipleRigidPoses::computeARFlowPoseError(
    const util::Device1D<float> &d_ar_flowx,
    std::vector<TranslationRotation3D> &poses,
    std::vector<double> &ar_flow_prop_valid,
    std::vector<double> &ar_flow_abs_valid) {
  ar_flow_prop_valid.clear();
  ar_flow_abs_valid.clear();

  render(poses);

  // re-purposing some already allocated data
  auto d_valid_ar_flow_Zbuffer = d_valid_disparity_Zbuffer_.get();
  auto d_valid_Zbuffer = d_valid_disparity_Zbuffer_sub_.get();
  markValidFlowZbufferAndZbufferZeroBased(
      d_valid_ar_flow_Zbuffer->data(), d_valid_Zbuffer->data(),
      d_ar_flowx.data(), _multipleRigidModelsOgre->getSegmentIND(), _n_cols,
      _n_rows, _n_objects);

  // additional buffers required for sorting
  auto d_value = d_ind_disparity_Zbuffer_.get();
  auto d_value_buf = d_ind_disparity_Zbuffer_sub_.get();
  auto d_key_buf = d_extra_disparity_buffer_.get();

  // Radix sort all the indices using the valid marks
  // indices could be ignored here to speed up sorting (but diff is minimal)
  cub_radix_sorter_->sort(*d_valid_ar_flow_Zbuffer, *d_value, *d_key_buf,
                          *d_value_buf);
  pose::extractLabelStartingIndices(d_seg_start_inds_->data(),
                                    d_valid_ar_flow_Zbuffer->data(),
                                    _n_cols * _n_rows, _n_objects);
  std::vector<int> h_seg_start_inds_flow_Zbuffer(_n_objects + 1);
  d_seg_start_inds_->copyTo(h_seg_start_inds_flow_Zbuffer, _n_objects + 1);

  cub_radix_sorter_->sort(*d_valid_Zbuffer, *d_value, *d_key_buf, *d_value_buf);
  pose::extractLabelStartingIndices(d_seg_start_inds_->data(),
                                    d_valid_Zbuffer->data(), _n_cols * _n_rows,
                                    _n_objects);
  std::vector<int> h_seg_start_inds_Zbuffer(_n_objects + 1);
  d_seg_start_inds_->copyTo(h_seg_start_inds_Zbuffer, _n_objects + 1);

  //  for(int o=0;o<(_n_objects+1);o++)
  //    printf("obj %d - flow+z %09d - z
  // %09d\n",o,h_seg_start_inds_flow_Zbuffer[o],h_seg_start_inds_Zbuffer[o]);

  int n_valid_flow_Zbuffer = (h_seg_start_inds_flow_Zbuffer.at(_n_objects) >= 0)
                                 ? h_seg_start_inds_flow_Zbuffer.at(_n_objects)
                                 : (_n_cols * _n_rows);

  _seg_lengths_flow_Zbuffer.assign(_n_objects, 0);

  getSegmentLengths(_seg_lengths_flow_Zbuffer, h_seg_start_inds_flow_Zbuffer,
                    _n_objects, n_valid_flow_Zbuffer);

  int n_valid_Zbuffer = (h_seg_start_inds_Zbuffer.at(_n_objects) >= 0)
                            ? h_seg_start_inds_Zbuffer.at(_n_objects)
                            : (_n_cols * _n_rows);

  _seg_lengths_Zbuffer.assign(_n_objects, 0);
  getSegmentLengths(_seg_lengths_Zbuffer, h_seg_start_inds_Zbuffer, _n_objects,
                    n_valid_Zbuffer);

  //  for(int o=0;o<(_n_objects+1);o++)
  //    printf("obj %d - start+z %09d - length+z %09d - start %09d - length
  // %09d\n",o,h_seg_start_inds_flow_Zbuffer[o],seg_lengths_flow_Zbuffer[o],h_seg_start_inds_Zbuffer[o],seg_lengths_Zbuffer[o]);

  // save proportion valid for each object
  auto bounding_boxes =
      _multipleRigidModelsOgre->getBoundingBoxesInCameraImage(poses);

  for (int o = 0; o < _n_objects; o++) {

    double absolute_count = (double)_seg_lengths_flow_Zbuffer[o];
    double prop = (_seg_lengths_Zbuffer[o] > 0)
                      ? (double)_seg_lengths_flow_Zbuffer[o] /
                            (double)_seg_lengths_Zbuffer[o]
                      : 0.0;

    if (parameters_.check_reliability_) {
      // does the proportion valid AR flow exceed the threshold?
      bool valid = prop > parameters_.reliability_threshold_;

      // is the projected shape 'sufficiently two-dimensional'?
      // the AR flow-based reliability measure fails on near-edge-like object
      // projections
      Eigen::Map<Eigen::Matrix<double, 2, 8> > bb_pixel(
          bounding_boxes.at(o).data());

      // subtract mean
      Eigen::Vector2d mn = bb_pixel.rowwise().mean();
      bb_pixel.colwise() -= mn;

      // covariance
      double ratio = parameters_.max_proportion_projected_bounding_box_;
      Eigen::Matrix2d cov = bb_pixel * bb_pixel.transpose();
      // check validity
      if (std::isfinite(cov(0, 0)) && std::isfinite(cov(0, 1)) &&
          std::isfinite(cov(1, 0)) && std::isfinite(cov(1, 1))) {
        Eigen::EigenSolver<Eigen::Matrix2d> eigen_solver(cov);
        // rotate bounding box
        Eigen::Matrix2d real_eigenvectors = eigen_solver.eigenvectors().real();
        //        Eigen::VectorXd w = eigen_solver.eigenvalues().real();
        bb_pixel = real_eigenvectors.transpose() * bb_pixel;
        // find extent
        auto extent =
            bb_pixel.rowwise().maxCoeff() - bb_pixel.rowwise().minCoeff();
        ratio = extent.maxCoeff() / extent.minCoeff();
      }
      valid =
          valid && (ratio < parameters_.max_proportion_projected_bounding_box_);

      poses.at(o).setValid(valid);
    }

    ar_flow_prop_valid.push_back(prop);
    ar_flow_abs_valid.push_back(absolute_count);
  }
}

void D_MultipleRigidPoses::robustPoseUpdates(
    const float *d_flowx, const float *d_flowy, const float *d_ar_flowx,
    const float *d_ar_flowy, const float *d_disparity, size_t d_disparity_pitch,
    int segments_to_update) {

  cudaArray *d_ZbufferArray = _multipleRigidModelsOgre->getZBuffer();
  cudaArray *d_normalXArray = _multipleRigidModelsOgre->getNormalX();
  cudaArray *d_normalYArray = _multipleRigidModelsOgre->getNormalY();
  cudaArray *d_normalZArray = _multipleRigidModelsOgre->getNormalZ();
  cudaArray *d_segmentINDArray = _multipleRigidModelsOgre->getSegmentIND();

#ifdef TIME_STEPS
  // Setup timers
  cudaEvent_t start, end, start_sub, end_sub;
  float elapsed_time;
  cudaEventCreate(&start);
  cudaEventCreate(&end);
  cudaEventCreate(&start_sub);
  cudaEventCreate(&end_sub);
#endif
#ifdef TIME_STEPS
  cudaEventRecord(start, 0); // preprocessing
#endif

  // initialize normal equations to zero
  std::fill(h_CO_reduced_.begin(), h_CO_reduced_.end(), 0);
  std::fill(h_CD_reduced_.begin(), h_CD_reduced_.end(), 0);
  cudaMemset(d_CO_reduced_->data(), 0,
             _N_CON_FLOW * parameters_.max_objects_ * sizeof(float));
  cudaMemset(d_CD_reduced_->data(), 0,
             _N_CON_DISP * parameters_.max_objects_ * sizeof(float));

  // Determine Zbuffer conversion constants
  // depth = Z_conv1/(Zbuffer+Z_conv2)
  float Z_conv1, Z_conv2;
  get_GL_conv_constants(Z_conv1, Z_conv2, parameters_.far_plane_,
                        parameters_.near_plane_);

// Mark valid locations with segment index
// If a cue's weight equals 0 mark everything invalid
// No initialization required

#ifdef TIME_STEPS
  cudaEventRecord(start_sub, 0);
#endif

  mark_with_zero_based_segmentIND(
      d_valid_flow_Zbuffer_->data(), d_valid_disparity_Zbuffer_->data(),
      d_flowx, d_ar_flowx, (const char *)d_disparity, d_segmentINDArray,
      _n_cols, _n_rows, _n_objects, d_disparity_pitch, parameters_.w_flow_ > 0,
      parameters_.w_ar_flow_ > 0, parameters_.w_disp_ > 0, segments_to_update);

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end_sub, 0);
  cudaEventSynchronize(end_sub);
  cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
  _compTimes.at(3) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
  cudaEventRecord(start_sub, 0);
#endif

  // Radix sort all the indices using the valid marks
  // the _sub_ buffers are only used for double-buffering (ping-pong)
  d_ind_flow_Zbuffer_->copyFrom(*d_linear_ind_);
  d_ind_disparity_Zbuffer_->copyFrom(*d_linear_ind_, _n_cols * _n_rows);

  //  util::TimerGPU sort_timer;
  cub_radix_sorter_->sort(*d_valid_flow_Zbuffer_, *d_ind_flow_Zbuffer_,
                          *d_valid_flow_Zbuffer_sub_, *d_ind_flow_Zbuffer_sub_);
  //  std::cout << "flow sort time: " << sort_timer.read();
  //  sort_timer.reset();
  cub_radix_sorter_->sort(
      *d_valid_disparity_Zbuffer_, *d_ind_disparity_Zbuffer_,
      *d_valid_disparity_Zbuffer_sub_, *d_ind_disparity_Zbuffer_sub_);
//  std::cout << " disp sort time: " << sort_timer.read() << std::endl;

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end_sub, 0);
  cudaEventSynchronize(end_sub);
  cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
  _compTimes.at(4) += (double)elapsed_time;
#endif

  // Get starting indices
  pose::extractLabelStartingIndices(d_seg_start_inds_->data(),
                                    d_valid_flow_Zbuffer_->data(),
                                    _n_cols * _n_rows * _N_FLOWS, _n_objects);
  std::vector<int> h_seg_start_inds_flow(_n_objects + 1);
  d_seg_start_inds_->copyTo(h_seg_start_inds_flow, _n_objects + 1);
  pose::extractLabelStartingIndices(d_seg_start_inds_->data(),
                                    d_valid_disparity_Zbuffer_->data(),
                                    _n_cols * _n_rows, _n_objects);
  std::vector<int> h_seg_start_inds_disparity(_n_objects + 1);
  d_seg_start_inds_->copyTo(h_seg_start_inds_disparity, _n_objects + 1);

  //  printf("before sub\n");
  //  for(int i=0;i<(_n_objects+1);i++)
  //    printf("se %d f %+07d d
  // %+07d\n",i,h_seg_start_inds_flow[i],h_seg_start_inds_disparity[i]);

  // Regularly subsample flow and disparity indices if required when a maximum
  // number of samples has been enforced
  // Relative proportion flow/disp will be maintained approximately

  int n_valid_flow_Zbuffer = (h_seg_start_inds_flow.at(_n_objects) >= 0)
                                 ? h_seg_start_inds_flow.at(_n_objects)
                                 : (_n_cols * _n_rows * _N_FLOWS);
  int n_valid_disparity_Zbuffer =
      (h_seg_start_inds_disparity.at(_n_objects) >= 0)
          ? h_seg_start_inds_disparity.at(_n_objects)
          : (_n_cols * _n_rows);

  //  printf("before sub %s\n",cudaGetErrorString(cudaGetLastError()));
  //  printf("val flow %d val disp
  // %d\n",n_valid_flow_Zbuffer,n_valid_disparity_Zbuffer);

  int n_valid_total = n_valid_flow_Zbuffer + n_valid_disparity_Zbuffer;

  if (n_valid_total > parameters_.max_samples_) {

    // proportionally subsample flow and disparity
    double sub_factor =
        (double)parameters_.max_samples_ / (double)n_valid_total;
    double inv_sub_factor = 1.0 / sub_factor;

    // flow
    int n_valid_flow_Zbuffer_sub =
        (int)floor((double)n_valid_flow_Zbuffer * sub_factor);

    if (n_valid_flow_Zbuffer_sub > 0) {
      subsample_ind_and_labels(
          d_ind_flow_Zbuffer_sub_->data(), d_ind_flow_Zbuffer_->data(),
          d_valid_flow_Zbuffer_sub_->data(), d_valid_flow_Zbuffer_->data(),
          n_valid_flow_Zbuffer_sub, (float)inv_sub_factor);

      // update all regular variables to account for subsampling
      // swap full and subsampled storage!
      n_valid_flow_Zbuffer = n_valid_flow_Zbuffer_sub;
      d_ind_flow_Zbuffer_->swap(*d_ind_flow_Zbuffer_sub_);
      d_valid_flow_Zbuffer_->swap(*d_valid_flow_Zbuffer_sub_);

      // update segment starting indices
      pose::extractLabelStartingIndices(d_seg_start_inds_->data(),
                                        d_valid_flow_Zbuffer_->data(),
                                        n_valid_flow_Zbuffer, _n_objects);
      d_seg_start_inds_->copyTo(h_seg_start_inds_flow, _n_objects + 1);
      //      printf("valsub flow %d\n",n_valid_flow_Zbuffer);
    }

    // disparity
    int n_valid_disparity_Zbuffer_sub =
        (int)floor((double)n_valid_disparity_Zbuffer * sub_factor);

    if (n_valid_disparity_Zbuffer_sub > 0) {
      subsample_ind_and_labels(d_ind_disparity_Zbuffer_sub_->data(),
                               d_ind_disparity_Zbuffer_->data(),
                               d_valid_disparity_Zbuffer_sub_->data(),
                               d_valid_disparity_Zbuffer_->data(),
                               n_valid_disparity_Zbuffer_sub,
                               (float)inv_sub_factor);

      // update all regular variables to account for subsampling
      // swap full and subsampled storage!
      n_valid_disparity_Zbuffer = n_valid_disparity_Zbuffer_sub;
      d_ind_disparity_Zbuffer_->swap(*d_ind_disparity_Zbuffer_sub_);
      d_valid_disparity_Zbuffer_->swap(*d_valid_disparity_Zbuffer_sub_);

      // update segment starting indices
      pose::extractLabelStartingIndices(d_seg_start_inds_->data(),
                                        d_valid_disparity_Zbuffer_->data(),
                                        n_valid_disparity_Zbuffer, _n_objects);
      d_seg_start_inds_->copyTo(h_seg_start_inds_disparity, _n_objects + 1);
      //      printf("valsub disp %d\n",n_valid_disparity_Zbuffer);
    }

    //    printf("after sub\n");
    //    for(int i=0;i<(_n_objects+1);i++)
    //      printf("se %d f %+07d d
    // %+07d\n",i,h_seg_start_inds_flow[i],h_seg_start_inds_disparity[i]);

  } // subsample

  // compute segment lengths
  std::vector<int> seg_lengths_flow(_n_objects);
  getSegmentLengths(seg_lengths_flow, h_seg_start_inds_flow, _n_objects,
                    n_valid_flow_Zbuffer);
  std::vector<int> seg_lengths_disparity(_n_objects);
  getSegmentLengths(seg_lengths_disparity, h_seg_start_inds_disparity,
                    _n_objects, n_valid_disparity_Zbuffer);

//  printf("after subsampling\n");
//  printf("n_f %07d n_d %07d\n", n_valid_flow_Zbuffer,
//         n_valid_disparity_Zbuffer);
//  for (int i = 0; i < _n_objects; i++)
//    printf("se %d f %+07d (%07d) d %+07d (%07d)\n", i,
// h_seg_start_inds_flow[i],
//           seg_lengths_flow[i], h_seg_start_inds_disparity[i],
//           seg_lengths_disparity[i]);

#ifdef TIME_STEPS
  cudaEventRecord(start_sub, 0);
#endif

  // Gather flow and disparity
  if (n_valid_flow_Zbuffer > 0)
    gather_valid_flow_Zbuffer(
        d_flow_compact_->data(), d_Zbuffer_flow_compact_->data(), d_flowx,
        d_flowy, d_ar_flowx, d_ar_flowy, d_ind_flow_Zbuffer_->data(),
        d_ZbufferArray, n_valid_flow_Zbuffer, _n_cols, _n_rows, Z_conv1,
        Z_conv2);

  if (n_valid_disparity_Zbuffer > 0)
    gather_valid_disparity_Zbuffer(
        d_disparity_compact_->data(), d_Zbuffer_normals_compact_->data(),
        (const char *)d_disparity, d_ind_disparity_Zbuffer_->data(),
        d_ZbufferArray, d_normalXArray, d_normalYArray, d_normalZArray,
        n_valid_disparity_Zbuffer, _n_cols, _n_rows, Z_conv1, Z_conv2,
        d_disparity_pitch);

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end_sub, 0);
  cudaEventSynchronize(end_sub);
  cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
  _compTimes.at(5) += (double)elapsed_time;
#endif

  // Gather index information
  _segment_info.clear();

  for (int o = 0; o < _n_objects; o++) {
    int nf = seg_lengths_flow[o];
    int nd = seg_lengths_disparity[o];
    if ((nf > 0) || (nd > 0)) {
      SegmentINFO tmp;
      tmp.segment_ind = o;
      tmp.n_values_flow = nf;
      tmp.n_values_disparity = nd;
      tmp.start_ind_flow = h_seg_start_inds_flow[o];
      tmp.start_ind_disparity = h_seg_start_inds_disparity[o];
      _segment_info.push_back(tmp);
    }
  }

//    printf("------------------------------------\n");
//    for(int o=0;o<_segment_info.size();o++) {
//      SegmentINFO &tmp = _segment_info.at(o);
//      printf("seg %d f_start %06d f_size %06d d_start %06d d_size
// %06d\n",tmp.segment_ind,tmp.start_ind_flow,tmp.n_values_flow,tmp.start_ind_disparity,tmp.n_values_disparity);
//    }
//    printf("------------------------------------\n");

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end, 0);
  cudaEventSynchronize(end);
  cudaEventElapsedTime(&elapsed_time, start, end);
  //  printf("preproc time : %2.4f ms\n",elapsed_time);
  _compTimes.at(0) += (double)elapsed_time;
#endif

  // preparatory book-keeping

  int n_segments = _segment_info.size();

  if (n_segments > 0) {

    std::vector<int> n_values_flow(n_segments);
    std::vector<int> start_ind_flow(n_segments);
    std::vector<int> n_values_disparity(n_segments);
    std::vector<int> start_ind_disparity(n_segments);
    std::vector<int> start_ind_res_flow(n_segments);
    std::vector<int> start_ind_res_disparity(n_segments);

    for (int s = 0; s < n_segments; s++) {
      SegmentINFO &tmp = _segment_info.at(s);
      n_values_flow.at(s) = tmp.n_values_flow;
      start_ind_flow.at(s) = tmp.start_ind_flow;
      n_values_disparity.at(s) = tmp.n_values_disparity;
      start_ind_disparity.at(s) = tmp.start_ind_disparity;

      start_ind_res_flow.at(s) = (s > 0) ? (start_ind_res_disparity.at(s - 1) +
                                            n_values_disparity.at(s - 1))
                                         : 0;
      start_ind_res_disparity.at(s) =
          start_ind_res_flow.at(s) + n_values_flow.at(s);
    }

    // Compute offset between residual input and output indices. Residuals are
    // stored as [seg1_flow seg1_disp seg2_flow seg2_disp ... ] to facilitate
    // median extraction

    std::vector<int> offset_ind_res_flow(n_segments);
    std::vector<int> offset_ind_res_disparity(n_segments);

    for (int s = 0; s < n_segments; s++) {
      offset_ind_res_flow.at(s) =
          start_ind_res_flow.at(s) - start_ind_flow.at(s);
      offset_ind_res_disparity.at(s) =
          start_ind_res_disparity.at(s) - start_ind_disparity.at(s);
    }

    // Create segment translation table (map zero-based original segment indices
    // to compressed segment indices)

    std::vector<int> segment_translation_table(_n_objects);
    for (int s = 0; s < n_segments; s++)
      segment_translation_table.at(_segment_info.at(s).segment_ind) = s;

    d_segment_translation_table_->copyFrom(segment_translation_table,
                                           _n_objects);
    d_n_values_flow_->copyFrom(n_values_flow, n_segments);
    d_start_ind_flow_->copyFrom(start_ind_flow, n_segments);
    d_n_values_disparity_->copyFrom(n_values_disparity, n_segments);
    d_start_ind_disparity_->copyFrom(start_ind_disparity, n_segments);
    d_offset_ind_res_flow_->copyFrom(offset_ind_res_flow, n_segments);
    d_offset_ind_res_disparity_->copyFrom(offset_ind_res_disparity, n_segments);
  }

  // take care here since _d_CO and _d_CD's maximum sizes are not exactly
  // determined in this way (especially the times 4)

  dim3 threadBlock_normal(64, 1, 1);
  int gridDim_x_normal_equations =
      (n_segments > 0)
          ? divUp(_MAX_N_VAL_ACCUM, n_segments * threadBlock_normal.x) * 4
          : 0;
  dim3 blockGrid_normal(gridDim_x_normal_equations, n_segments);
  //  dim3 threadBlock_reduce_64(64,1);
  //  dim3 blockGrid_reduce_flow(_N_CON_FLOW, n_segments);
  //  dim3 blockGrid_reduce_disparity(_N_CON_DISP, n_segments);

  dim3 threadBlock_reduce_64_mult(64, 4);
  dim3 blockGrid_reduce_flow_mult(divUp(_N_CON_FLOW, 4), n_segments);
  dim3 blockGrid_reduce_disparity_mult(divUp(_N_CON_DISP, 4), n_segments);

/****************/
/* OLS Estimate */
/****************/

#ifdef TIME_STEPS
  cudaEventRecord(start, 0); // ols-time
#endif

  if (n_segments > 0) {

#ifdef TIME_STEPS
    cudaEventRecord(start_sub, 0);
#endif

    normal_eqs_flow(blockGrid_normal, threadBlock_normal, d_CO_->data(),
                    d_flow_compact_->data(), d_Zbuffer_flow_compact_->data(),
                    d_ind_flow_Zbuffer_->data(), _focal_length_x,
                    _focal_length_y, _nodal_point_x, _nodal_point_y, _n_rows,
                    _n_cols, d_n_values_flow_->data(),
                    d_start_ind_flow_->data());
    normal_eqs_disparity(blockGrid_normal, threadBlock_normal, d_CD_->data(),
                         d_disparity_compact_->data(),
                         d_Zbuffer_normals_compact_->data(),
                         d_ind_disparity_Zbuffer_->data(), _focal_length_x,
                         _focal_length_y, _nodal_point_x, _nodal_point_y,
                         _baseline, _n_cols, d_n_values_disparity_->data(),
                         d_start_ind_disparity_->data(), parameters_.w_disp_);

#ifdef TIME_STEPS
    cudaThreadSynchronize();
    cudaEventRecord(end_sub, 0);
    cudaEventSynchronize(end_sub);
    cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
    _compTimes.at(6) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
    cudaEventRecord(start_sub, 0);
#endif

    reduce_normal_eqs_64_mult_constr(blockGrid_reduce_flow_mult,
                                     threadBlock_reduce_64_mult,
                                     d_CO_reduced_->data(), d_CO_->data(),
                                     gridDim_x_normal_equations, _N_CON_FLOW);
    reduce_normal_eqs_64_mult_constr(blockGrid_reduce_disparity_mult,
                                     threadBlock_reduce_64_mult,
                                     d_CD_reduced_->data(), d_CD_->data(),
                                     gridDim_x_normal_equations, _N_CON_DISP);

//    reduce_normal_eqs_64_GPU<<<blockGrid_reduce_flow,threadBlock_reduce_64>>>(_d_CO_reduced,
// _d_CO, gridDim_x_normal_equations);
//    reduce_normal_eqs_64_GPU<<<blockGrid_reduce_disparity,threadBlock_reduce_64>>>(_d_CD_reduced,
// _d_CD, gridDim_x_normal_equations);

#ifdef TIME_STEPS
    cudaThreadSynchronize();
    cudaEventRecord(end_sub, 0);
    cudaEventSynchronize(end_sub);
    cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
    _compTimes.at(7) += (double)elapsed_time;
#endif
  }

  d_CO_reduced_->copyTo(h_CO_reduced_);
  d_CD_reduced_->copyTo(h_CD_reduced_);

#ifdef TIME_STEPS
  cudaEventRecord(start_sub, 0);
#endif

  // Solve systems
  std::vector<float> dTdR(6 * n_segments);
  for (int s = 0; s < n_segments; s++) {
    int curr_part = _segment_info.at(s).segment_ind;
    segment_normal_eqs_.at(curr_part).compose(
        &h_CO_reduced_.at(_N_CON_FLOW * s), &h_CD_reduced_.at(_N_CON_DISP * s));
    segment_normal_eqs_.at(curr_part).solve(&dTdR.at(6 * s));
  }

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end_sub, 0);
  cudaEventSynchronize(end_sub);
  cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
  _compTimes.at(8) += (double)elapsed_time;
#endif

  // Copy to device memory
  d_dTR_->copyFrom(dTdR, 6 * n_segments);

  //  for(int s=0;s<n_segments;s++) {
  //    printf("seg %d - ",_segment_info.at(s).segment_ind);
  //    for(int i=0;i<6;i++)
  //      printf("%+03.4f ",dTdR[s*6+i]);
  //    printf("\n");
  //  }

  // Store OLS estimates
  _OLSDeltaPoses.clear();
  for (int i = 0; i < n_segments; i++)
    _OLSDeltaPoses.push_back(TranslationRotation3D(&dTdR[i * 6]));

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end, 0);
  cudaEventSynchronize(end);
  cudaEventElapsedTime(&elapsed_time, start, end);
  //  printf("OLS time : %2.4f ms\n",elapsed_time);
  _compTimes.at(1) += (double)elapsed_time;
#endif

/****************/
/* M-estimation */
/****************/

#ifdef TIME_STEPS
  cudaEventRecord(start, 0); // robust-time
#endif

  for (int it = 0; it < parameters_.n_icp_inner_it_; it++) {

    // Compute flow and disparity residuals

    if (n_segments > 0) {

#ifdef TIME_STEPS
      cudaEventRecord(start_sub, 0);
#endif

      dim3 threadBlock_res(256, 1);

      if (n_valid_flow_Zbuffer > 0) {
        dim3 blockGrid_res(divUp(n_valid_flow_Zbuffer, threadBlock_res.x), 1);
        flow_absolute_residual_scalable(
            blockGrid_res, threadBlock_res, d_abs_res_->data(),
            d_flow_compact_->data(), d_Zbuffer_flow_compact_->data(),
            d_ind_flow_Zbuffer_->data(), d_valid_flow_Zbuffer_->data(),
            _focal_length_x, _focal_length_y, _nodal_point_x, _nodal_point_y,
            _n_rows, _n_cols, n_valid_flow_Zbuffer,
            d_offset_ind_res_flow_->data(),
            d_segment_translation_table_->data(), parameters_.w_flow_,
            parameters_.w_ar_flow_, d_dTR_->data());
      }

      if (n_valid_disparity_Zbuffer > 0) {
        dim3 blockGrid_res(divUp(n_valid_disparity_Zbuffer, threadBlock_res.x),
                           1);
        disp_absolute_residual_scalable(
            blockGrid_res, threadBlock_res, d_abs_res_->data(),
            d_disparity_compact_->data(), d_Zbuffer_normals_compact_->data(),
            d_ind_disparity_Zbuffer_->data(),
            d_valid_disparity_Zbuffer_->data(), _focal_length_x,
            _focal_length_y, _nodal_point_x, _nodal_point_y, _baseline, _n_cols,
            n_valid_disparity_Zbuffer, d_offset_ind_res_disparity_->data(),
            d_segment_translation_table_->data(), parameters_.w_disp_,
            d_dTR_->data());
      }

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end_sub, 0);
      cudaEventSynchronize(end_sub);
      cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
      _compTimes.at(9) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
      cudaEventRecord(start_sub, 0);
#endif

      // Median absolute residuals
      std::vector<float> abs_res_scales(n_segments);
      int pp[n_segments];
      for (int s = 0; s < n_segments; s++) {
        SegmentINFO &tmp = _segment_info.at(s);
        pp[s] = tmp.n_values_flow + tmp.n_values_disparity;
      }
      approx_multiple_medians_shuffle_cuda(
          abs_res_scales.data(), d_abs_res_->data(), d_median_tmp_->data(),
          d_random_numbers_->data(), pp, n_segments, d_median_n_in_->data(),
          d_median_start_inds_->data());

      for (int s = 0; s < n_segments; s++)
        abs_res_scales.at(s) *= 6.9460f;

      d_abs_res_scales_->copyFrom(abs_res_scales, n_segments);

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end_sub, 0);
      cudaEventSynchronize(end_sub);
      cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
      _compTimes.at(10) += (double)elapsed_time;
#endif

//  for(int s=0;s<n_segments;s++)
//    printf("%2.5f ",abs_res_scales[s]);
//  printf("\n");

// Weighted constraints

#ifdef TIME_STEPS
      cudaEventRecord(start_sub, 0);
#endif

      normal_eqs_flow_weighted(
          blockGrid_normal, threadBlock_normal, d_CO_->data(),
          d_flow_compact_->data(), d_Zbuffer_flow_compact_->data(),
          d_ind_flow_Zbuffer_->data(), _focal_length_x, _focal_length_y,
          _nodal_point_x, _nodal_point_y, _n_rows, _n_cols,
          d_n_values_flow_->data(), d_start_ind_flow_->data(),
          d_abs_res_scales_->data(), parameters_.w_flow_,
          parameters_.w_ar_flow_, d_dTR_->data());

      normal_eqs_disparity_weighted(
          blockGrid_normal, threadBlock_normal, d_CD_->data(),
          d_disparity_compact_->data(), d_Zbuffer_normals_compact_->data(),
          d_ind_disparity_Zbuffer_->data(), _focal_length_x, _focal_length_y,
          _nodal_point_x, _nodal_point_y, _baseline, _n_cols,
          d_n_values_disparity_->data(), d_start_ind_disparity_->data(),
          d_abs_res_scales_->data(), parameters_.w_disp_, d_dTR_->data());

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end_sub, 0);
      cudaEventSynchronize(end_sub);
      cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
      _compTimes.at(6) += (double)elapsed_time;
#endif

#ifdef TIME_STEPS
      cudaEventRecord(start_sub, 0);
#endif

      reduce_normal_eqs_64_mult_constr(blockGrid_reduce_flow_mult,
                                       threadBlock_reduce_64_mult,
                                       d_CO_reduced_->data(), d_CO_->data(),
                                       gridDim_x_normal_equations, _N_CON_FLOW);
      reduce_normal_eqs_64_mult_constr(blockGrid_reduce_disparity_mult,
                                       threadBlock_reduce_64_mult,
                                       d_CD_reduced_->data(), d_CD_->data(),
                                       gridDim_x_normal_equations, _N_CON_DISP);

//      reduce_normal_eqs_64_GPU<<<blockGrid_reduce_flow,threadBlock_reduce_64>>>(_d_CO_reduced,
// _d_CO, gridDim_x_normal_equations);
//      reduce_normal_eqs_64_GPU<<<blockGrid_reduce_disparity,threadBlock_reduce_64>>>(_d_CD_reduced,
// _d_CD, gridDim_x_normal_equations);

#ifdef TIME_STEPS
      cudaThreadSynchronize();
      cudaEventRecord(end_sub, 0);
      cudaEventSynchronize(end_sub);
      cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
      _compTimes.at(7) += (double)elapsed_time;
#endif
    }

    d_CO_reduced_->copyTo(h_CO_reduced_);
    d_CD_reduced_->copyTo(h_CD_reduced_);

#ifdef TIME_STEPS
    cudaEventRecord(start_sub, 0);
#endif

    // Solve systems
    for (int s = 0; s < n_segments; s++) {
      int curr_part = _segment_info.at(s).segment_ind;
      segment_normal_eqs_.at(curr_part)
          .compose(&h_CO_reduced_.at(_N_CON_FLOW * s),
                   &h_CD_reduced_.at(_N_CON_DISP * s));
      segment_normal_eqs_.at(curr_part).solve(&dTdR.at(6 * s));
    }

#ifdef TIME_STEPS
    cudaThreadSynchronize();
    cudaEventRecord(end_sub, 0);
    cudaEventSynchronize(end_sub);
    cudaEventElapsedTime(&elapsed_time, start_sub, end_sub);
    _compTimes.at(8) += (double)elapsed_time;
#endif

    // Copy to device memory
    d_dTR_->copyFrom(dTdR, 6 * n_segments);

    //    printf("ROBUST ITERATION %d\n",it);
    //    printf("-------------------\n");
    //    for(int s=0;s<n_segments;s++) {
    //      printf("seg %d - ",_segment_info.at(s).segment_ind);
    //      for(int i=0;i<6;i++)
    //        printf("%+03.4f ",dTdR[s*6+i]);
    //      printf("\n");
    //    }
  }

  // Store robust estimates
  _robustDeltaPoses.clear();
  for (int i = 0; i < n_segments; i++)
    _robustDeltaPoses.push_back(TranslationRotation3D(&dTdR[i * 6]));

#ifdef TIME_STEPS
  cudaThreadSynchronize();
  cudaEventRecord(end, 0);
  cudaEventSynchronize(end);
  cudaEventElapsedTime(&elapsed_time, start, end);
  //  printf("robust time : %2.4f ms\n",elapsed_time);
  _compTimes.at(2) += (double)elapsed_time;
#endif
}
}