#include <ipa_room_exploration/energy_functional_explorator.h>

// Constructor
EnergyFunctionalExplorator::EnergyFunctionalExplorator()
{

}

// Function that computes the energy functional for each pair of nodes.
double EnergyFunctionalExplorator::E(const EnergyExploratorNode& location, const EnergyExploratorNode& neighbor, const double cell_size_in_pixel, const double previous_travel_angle)
{
	float energy_functional = 0.0;

	// 1. translational distance
	cv::Point diff = neighbor.center_ - location.center_;
	energy_functional += cv::norm(diff)/cell_size_in_pixel;

	// 2. rotational distance
	const double travel_angle_to_node = std::atan2(diff.y, diff.x);
	double diff_angle = travel_angle_to_node - previous_travel_angle;
	while (diff_angle < -PI)
		diff_angle += 2*PI;
	while (diff_angle > PI)
		diff_angle -= 2*PI;
	energy_functional += std::abs(diff_angle)*PI_2_INV;	// 1.01 for punishing turns a little bit more on a tie

	// 3. neighboring function, determining how many neighbors of the neighbor have been visited
	int visited_neighbors = 0;
	for(std::vector<EnergyExploratorNode*>::const_iterator n=neighbor.neighbors_.begin(); n!=neighbor.neighbors_.end(); ++n)
		if ((*n)->obstacle_ == false && (*n)->visited_ == true)
			++visited_neighbors;
	energy_functional += 4. - 0.5*visited_neighbors;

	int wall_points = 0;
	for(std::vector<EnergyExploratorNode*>::const_iterator n=neighbor.neighbors_.begin(); n!=neighbor.neighbors_.end(); ++n)
		if ((*n)->obstacle_ == true)
			++wall_points;
	energy_functional += 0.72 - 0.09*wall_points;

	//std::cout << "E: " << cv::norm(diff)/cell_size << " + " << std::abs(diff_angle)*PI_2_INV << " + " << 4. - 0.5*visited_neighbors << " + " << 0.72 - 0.09*wall_points << "                        angles: " << travel_angle_to_node << ", " << previous_travel_angle << "   diff ang: " << diff_angle << std::endl;

	return energy_functional;
}

// Function that plans a coverage path trough the given map, using the method proposed in
//
//	Bormann Richard, Joshua Hampp, and Martin Hägele. "New brooms sweep clean-an autonomous robotic cleaning assistant for
//	professional office cleaning." Robotics and Automation (ICRA), 2015 IEEE International Conference on. IEEE, 2015.
//
// This method discretizes the free space, that should be covered, into several nodes. Each of the node has to be covered, in order
// to cover the whole area. The path starts at the node that is closest to the given starting position and chooses the next node as
// the one that minimizes the energy functional, provided in the paper above. To do this here the following steps are done.
//	I.	The free area gets discretized into several nodes, using the given cell_size parameter, starting at the upper left white pixel of
//		the room. Whenever the overlaid grid then hits a white pixel, this point is added as a node. Then after all nodes have been found
//		the direct 8 neighbors for each node are found, which will be used later in the energy functional.
//	II.	After all nodes have been found, the coverage path is computed.
//			i.	The start node gets chosen as the one that is closest to the given starting position and is an edge of the given room, i.e
//				a node that has less than 4 neighbors.
//			ii.	The next node is then iteratively chosen from the directly neighboring ones, by finding the node that minimizes the given
//				energy functional and wasn't visited before.
//			iii.If in the neighborhood no accessible point could be found, search for the next node in the whole grid to continue the path.
//			iv.	This procedure is repeated, until all created nodes have been covered.
// III.	If wanted use the given vector from the robot middlepoint to the fov middlepoint to map the fov poses to the
//		robot footprint poses. To do so simply a vector transformation is applied. If the computed robot pose is not in the
//		free space, another accessible point is generated by finding it on the radius around the fov middlepoint s.t.
//		the distance to the last robot position is minimized. If this is not wanted one has to set the corresponding
//		Boolean to false (shows that the path planning should be done for the robot footprint).
//
void EnergyFunctionalExplorator::getExplorationPath(const cv::Mat& room_map, std::vector<geometry_msgs::Pose2D>& path, const float map_resolution,
			const cv::Point starting_position, const cv::Point2d map_origin, const double grid_spacing_in_pixel,
			const bool plan_for_footprint, const Eigen::Matrix<float, 2, 1> robot_to_fov_vector)
{
	const int grid_spacing_as_int = std::floor(grid_spacing_in_pixel);
	const int half_grid_spacing_as_int = std::floor(grid_spacing_in_pixel*0.5);

	// *********************** I. Find the main directions of the map and rotate it in this manner. ***********************
	cv::Mat R;
	cv::Rect bbox;
	cv::Mat rotated_room_map;
	RoomRotator room_rotation;
	room_rotation.computeRoomRotationMatrix(room_map, R, bbox, map_resolution);
	room_rotation.rotateRoom(room_map, rotated_room_map, R, bbox);

	// compute min/max room coordinates
	cv::Point min_room(1000000, 1000000), max_room(0, 0);
	for (int v=0; v<rotated_room_map.rows; ++v)
	{
		for (int u=0; u<rotated_room_map.cols; ++u)
		{
			if (rotated_room_map.at<uchar>(v,u)==255)
			{
				min_room.x = std::min(min_room.x, u);
				min_room.y = std::min(min_room.y, v);
				max_room.x = std::max(max_room.x, u);
				max_room.y = std::max(max_room.y, v);
			}
		}
	}
	cv::Mat inflated_rotated_room_map;
	cv::erode(rotated_room_map, inflated_rotated_room_map, cv::Mat(), cv::Point(-1, -1), half_grid_spacing_as_int);

	// *********************** II. Find the nodes and their neighbors ***********************
	// get the nodes in the free space
	std::vector<std::vector<EnergyExploratorNode> > nodes; // 2-dimensional vector to easily find the neighbors
	int number_of_nodes = 0;

//	// compute min/max room coordinates
//	int min_y = 1000000, min_x = 1000000;
//	for (int v=0; v<rotated_room_map.rows; ++v)
//	{
//		for (int u=0; u<rotated_room_map.cols; ++u)
//		{
//			if (rotated_room_map.at<uchar>(v,u)==255)
//			{
//				min_x = std::min(min_x, u);
//				min_y = std::min(min_y, v);
//			}
//		}
//	}
//	if (min_x-grid_spacing_as_int > 0)
//		min_x -= grid_spacing_as_int;
//	if (min_y-grid_spacing_as_int > 0)
//		min_y -= grid_spacing_as_int;

	// todo: create grid in external class - it is the same in all approaches
	// todo: if first/last row or column in grid has accessible areas but center is inaccessible, create a node in the accessible area
	for(int y=min_room.y+half_grid_spacing_as_int; y<max_room.y; y+=grid_spacing_as_int)
	{
		// for the current row create a new set of neurons to span the network over time
		std::vector<EnergyExploratorNode> current_row;
		for(int x=min_room.x+half_grid_spacing_as_int; x<max_room.x; x+=grid_spacing_as_int)
		{
			// create node if the current point is in the free space
			EnergyExploratorNode current_node;
			current_node.center_ = cv::Point(x,y);
			//if(rotated_room_map.at<uchar>(y,x) == 255)				// could make sense to test all pixels of the cell, not only the center
			if (GridGenerator::completeCellTest(inflated_rotated_room_map, current_node.center_, grid_spacing_as_int) == true)
			{
				current_node.obstacle_ = false;
				current_node.visited_ = false;
				++number_of_nodes;
			}
			// add the obstacle nodes as already visited
			else
			{
				current_node.obstacle_ = true;
				current_node.visited_ = true;
				++number_of_nodes;
			}
			current_row.push_back(current_node);
		}

		// insert the current row into grid
		nodes.push_back(current_row);
	}
	std::cout << "found " << number_of_nodes <<  " nodes" << std::endl;

	// find the neighbors for each node
	EnergyExploratorNode* first_accessible_node = 0;
	std::vector<EnergyExploratorNode*> corner_nodes; // vector that stores the corner nodes, i.e. nodes with 3 or less neighbors
	for(size_t row=0; row<nodes.size(); ++row)
	{
		for(size_t column=0; column<nodes[row].size(); ++column)
		{
			std::vector<EnergyExploratorNode*>& current_neighbors = nodes[row][column].neighbors_;
			for(int dy=-1; dy<=1; ++dy)
			{
				// don't exceed the current row
				if(row+dy < 0 || row+dy >= nodes.size())
					continue;

				// get the neighbors left from the current neuron
				if(column > 0)
					current_neighbors.push_back(&nodes[row+dy][column-1]);

				// get the nodes on the same column as the current neuron
				if(dy != 0)
					current_neighbors.push_back(&nodes[row+dy][column]);

				// get the nodes right from the current neuron
				if(column < nodes[row].size()-1)
					current_neighbors.push_back(&nodes[row+dy][column+1]);
			}

			// check if the current node is a corner, i.e. nodes that have 3 or less neighbors that are not obstacles
			int non_obstacle_neighbors = nodes[row][column].countNonObstacleNeighbors();
			if(non_obstacle_neighbors<=3 && nodes[row][column].obstacle_==false)
				corner_nodes.push_back(&nodes[row][column]);

			if (first_accessible_node==0 && nodes[row][column].obstacle_==false)
				first_accessible_node = &nodes[row][column];
		}
	}
	std::cout << "found neighbors, corners: " << corner_nodes.size() << std::endl;
	if (first_accessible_node == 0)
	{
		std::cout << "Warning: there are no accessible points in this room." << std::endl;
		return;
	}

//	// testing
//	cv::Mat test_map = rotated_room_map.clone();
//	for (size_t i=0; i<nodes.size(); ++i)
//		for (size_t j=0; j<nodes[i].size(); ++j)
//			if (nodes[i][j].obstacle_==false)
//				cv::circle(test_map, nodes[i][j].center_, 2, cv::Scalar(127), CV_FILLED);
//	cv::imshow("grid", test_map);
//	cv::waitKey();
//	for(size_t i=0; i<nodes.size(); ++i)
//	{
//		for(size_t j=0; j<nodes[i].size(); ++j)
//		{
//			cv::Mat test_map = rotated_room_map.clone();
//
//			std::vector<EnergyExploratorNode*> neighbors = nodes[i][j].neighbors_;
//			for(std::vector<EnergyExploratorNode*>::iterator n=neighbors.begin(); n!=neighbors.end(); ++n)
//				cv::circle(test_map, (*n)->center_, 2, cv::Scalar(127), CV_FILLED);
//
//			cv::imshow("neighbors", test_map);
//			cv::waitKey();
//		}
//	}

	// *********************** III. Plan the coverage path ***********************
	// i. find the start node of the path as a corner that is closest to the starting position
	std::vector<cv::Point> starting_point_vector(1, starting_position); // opencv syntax
	cv::transform(starting_point_vector, starting_point_vector, R);
	const cv::Point rotated_starting_point = starting_point_vector[0]; // Point that keeps track of the last point after the boustrophedon path in each cell
	EnergyExploratorNode* start_node = first_accessible_node;
	double min_distance = 1e10;
	for(std::vector<EnergyExploratorNode*>::iterator corner=corner_nodes.begin(); corner!=corner_nodes.end(); ++corner)
	{
		cv::Point diff = (*corner)->center_ - rotated_starting_point;
		double current_distance = diff.x*diff.x+diff.y*diff.y;
		if(current_distance<=min_distance)
		{
			start_node = *corner;
			min_distance = current_distance;
		}
	}
	std::cout << "start node: " << start_node->center_ << std::endl;

	// insert start node into coverage path
	std::vector<cv::Point2f> fov_coverage_path;
	fov_coverage_path.push_back(cv::Point2f(start_node->center_.x, start_node->center_.y));
	start_node->visited_ = true;	// mark visited nodes as obstacles

	// ii. starting at the start node, find the coverage path, by choosing the node that min. the energy functional
	EnergyExploratorNode* last_node = start_node;
	double previous_travel_angle = 0;  //always use x-direction in the rotated map  //std::atan2(rotated_starting_point.y-start_node->center_.y, rotated_starting_point.x-start_node->center_.x);
	for(std::vector<EnergyExploratorNode*>::iterator neighbor=last_node->neighbors_.begin(); neighbor!=last_node->neighbors_.end(); ++neighbor)
	{
		if ((*neighbor)->obstacle_==false && (*neighbor)->center_.y==start_node->center_.y && (*neighbor)->center_.x>start_node->center_.x)
		{
			previous_travel_angle = 0;
			break;
		}
		if ((*neighbor)->obstacle_==false && (*neighbor)->center_.y==start_node->center_.y && (*neighbor)->center_.x<start_node->center_.x)
		{
			previous_travel_angle = PI;
			break;
		}
		if ((*neighbor)->obstacle_==false && (*neighbor)->center_.y<start_node->center_.y && (*neighbor)->center_.x==start_node->center_.x)
		{
			previous_travel_angle = -0.5*PI;
		}
		if ((*neighbor)->obstacle_==false && (*neighbor)->center_.y>start_node->center_.y && (*neighbor)->center_.x==start_node->center_.x)
		{
			previous_travel_angle = 0.5*PI;
		}
	}
//	cv::Mat path_map = rotated_room_map.clone();
//	cv::circle(path_map, fov_coverage_path[0], 2, cv::Scalar(100), CV_FILLED);
	do
	{
		//std::cout << "Point: " << last_node->center_ << std::endl;
		// check the direct neighbors, if at least one is not already visited
		std::vector<EnergyExploratorNode*> not_visited_neighbors;
		for(std::vector<EnergyExploratorNode*>::iterator neighbor=last_node->neighbors_.begin(); neighbor!=last_node->neighbors_.end(); ++neighbor)
			if ((*neighbor)->obstacle_ == false && (*neighbor)->visited_ == false)
				not_visited_neighbors.push_back(*neighbor);

		// if there are not visited direct neighbors, find the one of them that minimizes the energy functional
		double min_energy = 1e10;
		EnergyExploratorNode* next_node = 0;
		if (not_visited_neighbors.size() > 0)
		{
			// find best neighbor
			for (std::vector<EnergyExploratorNode*>::iterator candidate=not_visited_neighbors.begin(); candidate!=not_visited_neighbors.end(); ++candidate)
			{
				const double current_energy = E(*last_node, **candidate, grid_spacing_in_pixel, previous_travel_angle);
				//std::cout << "Neighbor: " << (*candidate)->center_ << "    energy: " << current_energy << std::endl;
				if(current_energy < min_energy)
				{
					min_energy = current_energy;
					next_node = *candidate;
				}
			}
		}
		// if no direct neighbor is unvisited, search for the next node in all unvisited nodes
		else
		{
			// find best next node
			for (size_t row=0; row<nodes.size(); ++row)
			{
				for (size_t col=0; col<nodes[row].size(); ++col)
				{
					// only check free nodes and not visited ones
					if (nodes[row][col].obstacle_==false && nodes[row][col].visited_==false)
					{
						// check if current node has a better energy
						const double current_energy = E(*last_node, nodes[row][col], grid_spacing_in_pixel, previous_travel_angle);
						if (current_energy < min_energy)
						{
							min_energy = current_energy;
							next_node = &nodes[row][col];
						}
					}
				}
			}
			if (next_node == 0)
				break;				// stop if all nodes have been visited
		}
		// add next node to path and set robot location
		previous_travel_angle = std::atan2(next_node->center_.y-last_node->center_.y, next_node->center_.x-last_node->center_.x);
		fov_coverage_path.push_back(cv::Point2f(next_node->center_.x, next_node->center_.y));
		next_node->visited_ = true;	// mark visited nodes as obstacles

//		cv::circle(path_map, next_node->center_, 2, cv::Scalar(100), CV_FILLED);
//		cv::line(path_map, next_node->center_, last_node->center_, cv::Scalar(127));
//		cv::imshow("path", path_map);
//		cv::waitKey();

		last_node = next_node;
	} while (true);

	// transform the calculated path back to the originally rotated map
	std::vector<geometry_msgs::Pose2D> fov_poses;
	room_rotation.transformPathBackToOriginalRotation(fov_coverage_path, fov_poses, R);

//	// go trough the found fov-path and compute the angles of the poses s.t. it points to the next pose that should be visited
//	for(unsigned int point_index=0; point_index<fov_coverage_path.size(); ++point_index)
//	{
//		// get the vector from the current point to the next point
//		geometry_msgs::Pose2D current_point = fov_coverage_path[point_index];
//		geometry_msgs::Pose2D next_point = fov_coverage_path[(point_index+1)%(fov_coverage_path.size())];
//
//		float angle = std::atan2(next_point.y-current_point.y, next_point.x-current_point.x);
//
//		// save the found angle
//		fov_coverage_path[point_index].theta = angle;
//	}

	// if the path should be planned for the footprint, transform the image points to the map coordinates
	if(plan_for_footprint==true)
	{
		for(std::vector<geometry_msgs::Pose2D>::iterator pose=fov_poses.begin(); pose!=fov_poses.end(); ++pose)
		{
			geometry_msgs::Pose2D current_pose;
			current_pose.x = (pose->x * map_resolution) + map_origin.x;
			current_pose.y = (pose->y * map_resolution) + map_origin.y;
			current_pose.theta = pose->theta;
			path.push_back(current_pose);
		}
		return;
	}

//	// testing
//	cv::Mat path_map = rotated_room_map.clone();
//	cv::circle(path_map, fov_coverage_path[0], 2, cv::Scalar(100), CV_FILLED);
//	for(std::vector<cv::Point>::iterator path_node=fov_coverage_path.begin()+1; path_node!=fov_coverage_path.end(); ++path_node)
//	{
//		cv::circle(path_map, *path_node, 2, cv::Scalar(100), CV_FILLED);
//		cv::line(path_map, *path_node, *(path_node-1), cv::Scalar(127));
//		cv::imshow("path", path_map);
//		cv::waitKey();
//	}
//	cv::imshow("path", path_map);
//	cv::waitKey();

	// ****************** IV. Map the found fov path to the robot path ******************
	//mapPath(room_map, path, fov_poses, robot_to_fov_vector, map_resolution, map_origin, starting_position);
	ROS_INFO("Starting to map from field of view pose to robot pose");
	cv::Point robot_starting_position = (fov_poses.size()>0 ? cv::Point(fov_poses[0].x, fov_poses[0].y) : starting_position);
	cv::Mat inflated_room_map;
	cv::erode(room_map, inflated_room_map, cv::Mat(), cv::Point(-1, -1), half_grid_spacing_as_int);
	mapPath(inflated_room_map, path, fov_poses, robot_to_fov_vector, map_resolution, map_origin, robot_starting_position);
}
