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

#include <opencv/cv.h>

#include <sstream>

using namespace cv;
using namespace std;

class Traite_image {
	public:
		Mat prev;
		bool first = true;
		int resize_f = 8;

		ros::NodeHandle n;

		image_transport::ImageTransport it;
		image_transport::Publisher      pub;
		image_transport::Subscriber     sub;


		Traite_image() : n("~"),it(n) {
			pub = it.advertise("/image_out", 1);
			sub = it.subscribe("/usb_cam/image_raw", 1, [this](const sensor_msgs::ImageConstPtr& img) -> void { this->on_image(img);},ros::VoidPtr(),image_transport::TransportHints("compressed"));
		}


		// This processes an image and publishes the result.
		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 (Exception& e) {
				std::ostringstream errstr;
				errstr << "cv_bridge exception caught: " << e.what();
				return;
			}

			//Mat& input  = const_cast<Mat&>(bridge_input->image);
			const Mat& input = bridge_input->image;
			Mat next;
			resize(input, next, Size(input.size().width/resize_f, input.size().height/resize_f));
			cvtColor(next, next, CV_BGR2GRAY);
			Mat output; //   (input.rows, input.cols, CV_32FC2);
			ROS_INFO("got input");
			if (first) {
				prev = next.clone();
				first = false;
				ROS_INFO("first done");
			}

			//unsigned int size           = input.rows * input.cols * 3;
			//unsigned char* begin_input  = (unsigned char*)(input.data);
			//unsigned char* end_input    = (unsigned char*)(input.data) + size;
			//unsigned char* out          = (unsigned char*)(output.data);
			//unsigned char* in           = begin_input;

			// This is an efficient way to process each channel in each pixel,
			// with an iterator taste.
			//while(in != end_input) {
			//  *(out++) = *(ptr_prev) - *(in);
			//  *(ptr_prev++) = *(in++);
			//}

			Mat_<Point2f> flow;
			Ptr<DenseOpticalFlow> tvl1 = createOptFlow_DualTVL1();

			tvl1->calc(prev, next, flow);

			drawOpticalFlow(flow, output);

			pub.publish(cv_bridge::CvImage(msg->header, "rgb8", output).toImageMsg());
			// bridge_input is handled by a smart-pointer. No explicit delete needed.

			ROS_INFO("pub");

			prev = next.clone();
		}

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

		Vec3b computeColor(float fx, float fy)
		{
			static bool first = true;

			// relative lengths of color transitions:
			// these are chosen based on perceptual similarity
			// (e.g. one can distinguish more shades between red and yellow
			//  than between yellow and green)
			const int RY = 15;
			const int YG = 6;
			const int GC = 4;
			const int CB = 11;
			const int BM = 13;
			const int MR = 6;
			const int NCOLS = RY + YG + GC + CB + BM + MR;
			static Vec3i colorWheel[NCOLS];

			if (first)
			{
				int k = 0;

				for (int i = 0; i < RY; ++i, ++k)
					colorWheel[k] = Vec3i(255, 255 * i / RY, 0);

				for (int i = 0; i < YG; ++i, ++k)
					colorWheel[k] = Vec3i(255 - 255 * i / YG, 255, 0);

				for (int i = 0; i < GC; ++i, ++k)
					colorWheel[k] = Vec3i(0, 255, 255 * i / GC);

				for (int i = 0; i < CB; ++i, ++k)
					colorWheel[k] = Vec3i(0, 255 - 255 * i / CB, 255);

				for (int i = 0; i < BM; ++i, ++k)
					colorWheel[k] = Vec3i(255 * i / BM, 0, 255);

				for (int i = 0; i < MR; ++i, ++k)
					colorWheel[k] = Vec3i(255, 0, 255 - 255 * i / MR);

				first = false;
			}

			const float rad = sqrt(fx * fx + fy * fy);
			const float a = atan2(-fy, -fx) / (float)CV_PI;

			const float fk = (a + 1.0f) / 2.0f * (NCOLS - 1);
			const int k0 = static_cast<int>(fk);
			const int k1 = (k0 + 1) % NCOLS;
			const float f = fk - k0;

			Vec3b pix;

			for (int b = 0; b < 3; b++)
			{
				const float col0 = colorWheel[k0][b] / 255.f;
				const float col1 = colorWheel[k1][b] / 255.f;

				float col = (1 - f) * col0 + f * col1;

				if (rad <= 1)
					col = 1 - rad * (1 - col); // increase saturation with radius
				else
					col *= .75; // out of range

				pix[2 - b] = static_cast<uchar>(255.f * col);
			}

			return pix;
		}

		void drawOpticalFlow(const Mat_<Point2f>& flow, Mat& dst, float maxmotion = -1)
		{
			dst.create(flow.size(), CV_8UC3);
			dst.setTo(Scalar::all(0));

			// determine motion range:
			float maxrad = maxmotion;

			if (maxmotion <= 0)
			{
				maxrad = 1;
				for (int y = 0; y < flow.rows; ++y)
				{
					for (int x = 0; x < flow.cols; ++x)
					{
						Point2f u = flow(y, x);

						if (!isFlowCorrect(u))
							continue;

						maxrad = max(maxrad, sqrt(u.x * u.x + u.y * u.y));
					}
				}
			}

			for (int y = 0; y < flow.rows; ++y)
			{
				for (int x = 0; x < flow.cols; ++x)
				{
					Point2f u = flow(y, x);

					if (isFlowCorrect(u))
						dst.at<Vec3b>(y, x) = computeColor(u.x / maxrad, u.y / maxrad);
				}
			}
		}

};


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

	return 0;
}