#include "stereo.hpp"
#include "stereosgm.hpp"
#include "../filters/salient_gradient.hpp"
#include "../filters/focal_cluster.hpp"

#include <opencv2/highgui.hpp>
#include <opencv2/imgproc.hpp>

struct StereoCSF::Impl {
	Array2D<uchar> l;
	Array2D<uchar> r;
	Array2D<uchar> gl;
	Array2D<uchar> gr;
	Array2D<uchar> temp;
	Bucket1D<short2, 64> buckets_l;
	Bucket2D<ushort, 64> buckets_r;
	Array1D<int> focal;

	Impl(int width, int height) :
		l(width, height), r(width, height),
		gl(width, height), gr(width, height), temp(width, height),
		buckets_l(height), buckets_r(16, height), focal(1024) {}
};

StereoCSF::StereoCSF() : impl_(nullptr) {
	impl_ = new Impl(0, 0);
}

void StereoCSF::compute(cv::InputArray l, cv::InputArray r, cv::OutputArray disparity) {

	//cudaSetDevice(0);

	if (l.rows() != impl_->l.height || r.cols() != impl_->l.width) {
		delete impl_; impl_ = nullptr;
		impl_ = new Impl(l.cols(), l.rows());
	}

	mat2gray(l, impl_->l);
	mat2gray(r, impl_->r);

	Array2D<float> disp_array(l.cols(), l.rows());
	disp_array.toGpuMat().setTo(cv::Scalar(0.0f));

	SalientGradientGrouped sgr = {impl_->r.data(), impl_->gr.data(), impl_->temp.data(), impl_->buckets_r.data(), impl_->r.width, impl_->r.height};
	parallel1DWarpSM(sgr, r.rows(), r.cols());

	for (int fx = 200; fx < l.cols()-200; fx += 50) {
	for (int fy = 200; fy < l.rows()-200; fy += 50) {
		short2 focal_pt = {short(fx), short(fy)};
		SalientGradient sgl = {focal_pt, 1000, impl_->l.data(), impl_->gl.data(), impl_->temp.data(), impl_->buckets_l.data(), impl_->l.width, impl_->l.height};
		parallel1DWarpSM(sgl, l.rows(), l.cols());
		impl_->focal.toGpuMat().setTo(cv::Scalar(0));	

		FocalCluster fc = {focal_pt, impl_->buckets_l.data(), impl_->buckets_r.data(), impl_->focal.data(), 1024};
		parallel1DWarp(fc, l.rows(), 1);

		FocalSelector fs = {focal_pt, impl_->buckets_l.data(), impl_->buckets_r.data(), impl_->focal.data(), disp_array.data(), 1024};
		parallel1DWarp(fs, l.rows(), 1);
	}
	}

	disp_array.toGpuMat().download(disparity);

	cv::Mat gradtmp;
	impl_->gl.toGpuMat().download(gradtmp);
	//cv::imshow("Gradient Left", gradtmp);

	cv::Mat tmp;
	impl_->focal.toGpuMat().download(tmp);

	double minval;
	double maxval;
	int minloc[2];
	int maxloc[2];
	cv::minMaxIdx(tmp, &minval, &maxval, minloc, maxloc);

	std::cout << "Focal Disparity = " << maxloc[1] << std::endl;

	tmp.convertTo(tmp, CV_8UC1, 255.0/maxval);
	cv::resize(tmp,tmp, cv::Size(tmp.cols, 100));
	cv::applyColorMap(tmp, tmp, cv::COLORMAP_TURBO);
	cv::imshow("Disparity Hist", tmp);

	cv::Mat imgleft, imgright;
	impl_->l.toGpuMat().download(imgleft);
	impl_->r.toGpuMat().download(imgright);

	cv::cvtColor(imgleft,imgleft, cv::COLOR_GRAY2BGR);
	cv::cvtColor(imgright,imgright, cv::COLOR_GRAY2BGR);
	cv::drawMarker(imgleft, cv::Point(1000,800), cv::Scalar(0,0,255));
	cv::drawMarker(imgright, cv::Point(1000-maxloc[1],800), cv::Scalar(0,0,255));

	cv::resize(imgleft,imgleft, cv::Size(imgleft.cols/2, imgleft.rows/2));
	cv::resize(imgright,imgright, cv::Size(imgright.cols/2, imgright.rows/2));
	cv::imshow("Left Focal", imgleft);
	cv::imshow("Right Focal", imgright);

	//impl_->gr.toGpuMat().download(tmp);
	//cv::resize(tmp,tmp, cv::Size(tmp.cols/2, tmp.rows/2));
	//cv::imshow("Gradients Right", tmp);

	cudaSafeCall(cudaDeviceSynchronize());
	disparity.create(l.rows(), l.cols(), CV_32F);
	cv::waitKey(10000);
}

StereoCSF::~StereoCSF() {
	if (impl_) {
		delete impl_;
		impl_ = nullptr;
	}
}