#include "ros/ros.h"
#include <image_transport/image_transport.h>
#include <cv_bridge/cv_bridge.h>
#include <sensor_msgs/image_encodings.h>
#include <papillon/BoundingBox.h>

#include <opencv/cv.h>

#include <sstream>

using namespace std;

class Traite_image {
	public:
		const static int SENSITIVITY_VALUE = 20;
		const static int BLUR_SIZE = 9;
		const static int CLOSE_SIZE = 30;
		const static int ERODE_SIZE = 2;
		const static int NB_FRAME_DROP = 1;


		vector<cv::Mat> prevs;
		cv::Mat last_T;
		bool first = true;
		int resize_f = 1;

		ros::NodeHandle n;

		image_transport::ImageTransport it;
		image_transport::Subscriber     sub;
		image_transport::Publisher      pub_img;
		image_transport::Publisher      pub_thres;
		ros::Publisher                  pub_cmd;


		Traite_image() : n("~"),it(n) {
			pub_img = it.advertise("/papillon/image_out", 1);
			pub_thres = it.advertise("/papillon/thres_out", 1);
			pub_cmd = n.advertise<papillon::BoundingBox>("/papillon/bbox", 1);
			sub = it.subscribe("/bebop/image_raw", 1, [this](const sensor_msgs::ImageConstPtr& img) -> void { this->on_image(img);},ros::VoidPtr(),image_transport::TransportHints("compressed"));
		}


		// Callback function executedwhen recieving an image from the camera
		// Publishes to /papillon/image_out
		void on_image(const sensor_msgs::ImageConstPtr& msg) {

			cv_bridge::CvImageConstPtr bridge_input;
			try {
				bridge_input = cv_bridge::toCvShare(msg,sensor_msgs::image_encodings::RGB8);
			}
			catch (cv::Exception& e) {
				std::ostringstream errstr;
				errstr << "cv_bridge exception caught: " << e.what();
				return;
			}

			const cv::Mat& input = bridge_input->image;
			cv::Mat next;
			resize(input, next, cv::Size(input.size().width/resize_f, input.size().height/resize_f));
			cv::Mat output;
			next.copyTo(output);
			
			cv::cvtColor(next, next, cv::COLOR_BGR2GRAY);
			
			if (first) {
				for (int i = 0; i < NB_FRAME_DROP; ++i) {
					prevs.push_back(next.clone());
				}
				first = false;
				ROS_INFO("Ready");
			}

			cv::Mat next_stab;
			stabiliseImg(prevs.back(), next, next_stab);
			int crop_ratio = 6;
			float crop_x = next_stab.size().width/crop_ratio;
			float crop_y = next_stab.size().height/crop_ratio;
			float crop_w = next_stab.size().width*(1-2.0/crop_ratio);
			float crop_h = next_stab.size().height*(1-2.0/crop_ratio);
			cv::Rect myROI(crop_x, crop_y, crop_w, crop_h);
			cv::Mat next_stab_cropped = next_stab(myROI);
			cv::Mat prev_cropped = prevs.back()(myROI);
			cv::Mat closed_thres;
			searchForMovement(prev_cropped, next_stab_cropped, output, closed_thres);


			pub_img.publish(cv_bridge::CvImage(msg->header, "rgb8", output).toImageMsg());
			pub_thres.publish(cv_bridge::CvImage(msg->header, "mono8", closed_thres).toImageMsg());

			prevs.pop_back();
			prevs.insert(prevs.begin(), next.clone());
		}

		void stabiliseImg(cv::Mat prev_grey, cv::Mat cur_grey, cv::Mat &output){
			// vector from prev to cur
			vector <cv::Point2f> prev_corner, cur_corner;
			vector <cv::Point2f> prev_corner2, cur_corner2;
			vector <uchar> status;
			vector <float> err;

			cv::goodFeaturesToTrack(prev_grey, prev_corner, 200, 0.01, 30);
			cv::calcOpticalFlowPyrLK(prev_grey, cur_grey, prev_corner, cur_corner, status, err);

			// weed out bad matches
			for(size_t i=0; i < status.size(); i++) {
				if(status[i]) {
					prev_corner2.push_back(prev_corner[i]);
					cur_corner2.push_back(cur_corner[i]);
				}
			}

			cv::Mat T = estimateRigidTransform(prev_corner2, cur_corner2, true); // false = rigid transform, no scaling/shearing

			if(T.data == NULL) {
				last_T.copyTo(T);
			}
			T.copyTo(last_T);

			cv::warpAffine(cur_grey, output, T, cur_grey.size(),cv::INTER_CUBIC|cv::WARP_INVERSE_MAP);
		}

		void searchForMovement(cv::Mat prev_grey, cv::Mat cur_grey, cv::Mat &output, cv::Mat &out2){
			cv::GaussianBlur(prev_grey, prev_grey, cv::Size(BLUR_SIZE,BLUR_SIZE), 3.0);
			cv::GaussianBlur(cur_grey, cur_grey, cv::Size(BLUR_SIZE,BLUR_SIZE), 3.0);

			// Subtract the 2 last frames and threshold them
			cv::Mat thres;
			cv::absdiff(prev_grey,cur_grey,thres);

			cv::Mat element = getStructuringElement( cv::MORPH_ELLIPSE,
                                   cv::Size( 2*ERODE_SIZE + 1, 2*ERODE_SIZE+1 ),
                                   cv::Point( ERODE_SIZE, ERODE_SIZE ) );
			// Apply the erode operation
			cv::erode(thres, thres, element );

			cv::threshold(thres, thres, SENSITIVITY_VALUE, 255, cv::THRESH_BINARY);

			// Intermediate output
			thres.copyTo(out2);

			cv::Mat closed_thres;
			cv::Mat structuringElement = getStructuringElement(cv::MORPH_ELLIPSE, cv::Size(CLOSE_SIZE, CLOSE_SIZE));
			cv::morphologyEx( thres, closed_thres, cv::MORPH_CLOSE, structuringElement );

			bool objectDetected = false;
			cv::Mat temp;
			closed_thres.copyTo(temp);
			vector< vector<cv::Point> > contours;
			//find contours of filtered image using openCV findContours function
			cv::findContours(temp,contours,CV_RETR_EXTERNAL,CV_CHAIN_APPROX_SIMPLE );// retrieves external contours

			//if contours vector is not empty, we have found some objects
			if(contours.size()>0){
				vector<cv::Rect> nc_rects; // Non connected rectangles
				for(size_t i=0; i<contours.size();i++)
				{
					nc_rects.push_back(cv::boundingRect(contours[i]));
				}

				vector<cv::Rect> c_rects; // Connected rectangles
				cleanBBoxes(nc_rects, cur_grey.size(), c_rects);
				if (c_rects.size() > 0) {
					for (const auto& rect : c_rects)
						cv::rectangle(output, rect, cv::Scalar(0, 255, 0), 2);

					cv::Rect objBRect = c_rects.front();
					papillon::BoundingBox bbox = papillon::BoundingBox();
					bbox.x = objBRect.x / (float)cur_grey.size().width;
					bbox.y = objBRect.y / (float)cur_grey.size().height;
					bbox.width = objBRect.width / (float)cur_grey.size().width;
					bbox.height = objBRect.height / (float)cur_grey.size().height;
					pub_cmd.publish(bbox);
				}
			}
		}

		void cleanBBoxes(vector<cv::Rect> nc_rects, cv::Size img, vector<cv::Rect> &c_rects) {
			int max = 0;
			for (auto r : nc_rects) {
				cv::Point center = rectCenter(r);
				if (r.width > img.width / 2 || r.height > img.height / 8)
					continue;

				if (center.y > img.height * 2 / 3 || center.y < img.height / 3)
					continue;

				int r_area = r.area();
				if (max < r_area) {
					max = r_area;
					c_rects.insert(c_rects.begin(), r);
				} else {
					c_rects.push_back(r);
				}
			}
		}

		cv::Point rectCenter(cv::Rect r) {
			return cv::Point(r.x + r.width/2, r.y + r.height / 2);
		}

		inline bool isFlowCorrect(cv::Point2f u)
		{
			return !cvIsNaN(u.x) && !cvIsNaN(u.y) && fabs(u.x) < 1e9 && fabs(u.y) < 1e9;
		}
};


int main(int argc, char **argv)
{
	ros::init(argc, argv, "test_opencv");
	Traite_image dataset=Traite_image();
	ros::spin();

	return 0;
}