#include "RxProcessor.h"

RxProcessor::RxProcessor() {
	iHandle = 0;
	uID = 0;
	dFocus = 0.5;
	pxImages = NULL;
	displayRate = 20;
	haveCameraCal = false;
	imageUpdated = false;
	cudaReady = false;
	downsample = false;
	displayType = PARTICLE;
	int uBufferSize = 512;
	bool bOverwrite = false;
	imageOverLCM = false;
	m_xImageBuffer.Initialize(uBufferSize, bOverwrite);

	std::stringstream ss;
	Poco::Timestamp t;
	// set a default output file
	ss << "E:\\" << "RxProc_" << t.utcTime() << ".csv";
	outputFile = ss.str();

	// Authenticate MGPU runtime
	printf("Authenticate LFR...\n");
	Rx::LFR::CLightFieldRuntime::Authenticate();
	printf("Initializing CLUVizTool...\n");
	CLUViz::Tool::Init();
}

RxProcessor::~RxProcessor() {
	haveCameraCal = false;
}

void RxProcessor::setVideoOutputFile(string name) {
	this->outputFile = name + ".mp4";
	outputVideo.open(this->outputFile, cv::VideoWriter::fourcc('D', 'I', 'V', 'X'), 30, cv::Size(1024, 1024));
	if (!outputVideo.isOpened())
	{
		cout << "Could not open the output video for write: " << name << endl;
		return;
	}
}

RxProcessor::RxProcessor(unsigned int uID, Rx::LFR::CCalibration& xDefaultCalibration) {

	// Authenticate MGPU runtime
	printf("Authenticate LFR...\n");
	Rx::LFR::CLightFieldRuntime::Authenticate();

	this->uID = uID;
	setCameraCalibration(uID, xDefaultCalibration);

}

void RxProcessor::setCameraCalibration(unsigned int uID, Rx::LFR::CCalibration& xDefaultCalibration) {
	
	// For this example just work with one camera and one GPU
	xCudaCompute.SetCudaDevice(Rx::LFR::CCuda::GetDevice(uID));
	xCudaCompute.ApplyCalibration(xDefaultCalibration, true);
	xCudaCompute.GetParams().ImportParameterFromFile("C:\\Users\\proberts\\Desktop\\CurrentSettings-2019-08-09.rxset");


	// Get the image access interface.
	pxImages = static_cast<Rx::LFR::ICudaDataImages*>(xCudaCompute.GetInterface(Rx::LFR::Interfaces::ECudaCompute::Images));
	haveCameraCal = true;

}

void RxProcessor::setImageFormat(Rx::CRxImageFormat fmt) {
	this->imageFormat = fmt;
}

void RxProcessor::initCudaWithImage() {
	std::cout << " Loading first, slow image on to GPU..." << endl;
	Rx::CRxImage tmp;
	tmp.Create(imageFormat);
	xCudaCompute.UploadRawImage(tmp);
	std::cout << " Done loading first, slow image on to GPU..." << endl;
	cudaReady = true;
}

void RxProcessor::addImage(Rx::CRxImage xImage)
{

	try
	{

		if (!m_xImageBuffer.MoveIn(std::move(xImage))) {
			std::cout << "proc: dropped Frame!" << std::endl;
		}


	}
	catch (Rx::CRxException& ex)
	{
		printf("Exception occured:\n%s\n\n", ex.ToString(true).ToCString());
		printf("Press any key to end program...\n");
		_getch();
	}
}


int RxProcessor::start() {

	try
	{

		if (!haveCameraCal) {
			std::cout << "Can't start processor thread without calibration." << std::endl;
			return 1;
		}

		// load the window
		if (!imageOverLCM) {
			CLUViz::Tool::CreateViewImage(iHandle, 0, 0, 1000, 800, "RxProcessorView");
		}

		// Start the image processing thread
		startThread();

	}

	catch (Rx::CRxException& ex)
	{
		printf("Exception occured:\n%s\n\n", ex.ToString(true).ToCString());
		printf("Press any key to end program...\n");
		_getch();
	}

	return 0;

}

int RxProcessor::stop() {

	// Finalize Cuda and the runtime
	Rx::LFR::CLightFieldRuntime::End();

	// Stop thread
	stopThread();

	outputVideo.release();

	return 0;

}

void RxProcessor::threadedFunction() {

	frameCounter = 0;

	while (isThreadRunning())
	{

		// Wait for the image buffer to be not empty
		if (!m_xImageBuffer.WaitForNotEmpty(1000))
		{
			yield();
			continue;
		}

		// Check if a new image should be get from the image buffer or the previous image sould be updated
		if (!m_xImageBuffer.MoveOut(xCapturedImage))
		{
			{
				// Wait again! This functions says that it waits until a frame is available
				yield();
				continue;
			}
		}

		// Check display rate
		auto finish = std::chrono::high_resolution_clock::now();
		std::chrono::duration<double> elapsed = finish - displayClockStart;

		try {

			if (cudaReady) {

				switch (displayType) {
				case RAW:
					rawImageProc();
					break;
				case TOTAL_FOCUS:
					totalFocusProc();
					break;
				case DEPTH_MAP:
					depthMapProc();
					break;
				case REFOCUS:
					refocusProc();
					break;
				case PARTICLE:
					particleStats();
					break;
				default:
					rawImageProc();
					break;
				}
				frameCounter++;
			}

		}
		catch (Rx::CRxException& ex)
		{
			printf("Exception occured:\n%s\n\n", ex.ToString(true).ToCString());
			printf("Press any key to end program...\n");
			_getch();
		}

	}


}

int RxProcessor::queueSize() {
	return m_xImageBuffer.GetSize();
}

void RxProcessor::setDisplayType(IMAGE_DISPLAY_TYPE t) {
	if (t != displayType) {
		displayType = t;
		imageUpdated = false;
	}
}

void RxProcessor::drawImage(Rx::CRxImage& img) {

	if (imageOverLCM) {

		img.SetTimestamp(xCapturedImage.GetTimestamp());

		if (img.GetFormat().m_iWidth > 1280 || img.GetFormat().m_iHeight > 1280) {
			Rx::CRxImage dstImage;
			xCudaCompute.ResizeImage(&dstImage, &img, img.GetFormat().m_iWidth/1280);
			img = dstImage;
		}

		// check for color image
		int pixeltype = 1;
		if (img.GetFormat().m_ePixelType != Rx::Interop::Runtime28::EPixelType::Lum)
			pixeltype = 3;

		if (img.GetFormat().GetBitDepth() > 8) {
			Rx::CRxImage dstImage;
			img.ConvertType(dstImage, Rx::Interop::Runtime28::EPixelType::Lum, Rx::Interop::Runtime28::EDataType::Byte);
			img = dstImage;
		}

		lcmImage.timestamp = xCapturedImage.GetTimestamp();
		lcmImage.width = img.GetFormat().m_iWidth;
		lcmImage.height = img.GetFormat().m_iHeight;
		lcmImage.dataType = 1;
		lcmImage.dataSize = lcmImage.width * lcmImage.height*pixeltype;
		lcmImage.pixelType = pixeltype;
		lcmImage.data.resize(lcmImage.dataSize);

		memcpy(lcmImage.data.data(), img.GetDataPtr(), lcmImage.dataSize);

		//totalFocusImage.data = testImage;

		lcm.publish("EYERISIMAGE", &lcmImage);

	}

	else {

		// Flip image vertically
		cv::Mat flipImage;
		if (img.GetFormat().m_ePixelType == Rx::Interop::Runtime28::EPixelType::Lum) {
			flipImage = cv::Mat(img.GetFormat().m_iHeight, img.GetFormat().m_iWidth, CV_16UC1, img.GetDataPtr());
		}
		else {
			flipImage = cv::Mat(img.GetFormat().m_iHeight, img.GetFormat().m_iWidth, CV_8UC3, img.GetDataPtr());
		}

		if (flipImage.empty()) {
			cout << "image empty" << endl;
		}

		cv::rotate(flipImage, flipImage, cv::ROTATE_180);

		cout << frameCounter << endl;

		CLUViz::Tool::ViewSetImage(iHandle, &img, imageUpdated);

		try {
			cv::Mat outImage(img.GetFormat().m_iHeight, img.GetFormat().m_iWidth, CV_8UC1);
			cv::Mat color(img.GetFormat().m_iHeight, img.GetFormat().m_iWidth, CV_8UC3);
			cv::Mat resized(1024, 1024, CV_8UC3);
			flipImage.convertTo(outImage, CV_8UC1, 0.015625);
			cv::cvtColor(outImage, color, cv::COLOR_GRAY2BGR);
			cv::resize(color, resized, cv::Size(1024, 1024));
			outputVideo.write(resized);
		}
		catch (exception& e) {
			cout << e.what() << endl;
		}

		imageUpdated = true;
	}

}

void RxProcessor::particleStats() {

	
	try {
		// Upload the image as the new raw image of all further CUDA computations
		//printf("Uploading image to Cuda compute instance...\n");
		xCudaCompute.UploadRawImage(xCapturedImage);

		// Pre-process raw light field image before calling any processing functions
		//printf("Pre-processing light field image...\n");
		xCudaCompute.Compute_PreProcess();
		xCudaCompute.Compute_DepthRay();
		xCudaCompute.Compute_DepthMap();
		xCudaCompute.Compute_TotalFocus();
		pxImages->Download(Rx::LFR::EImage::TotalFocus_View_Object_Pinhole, &xTFImage);

		xCudaCompute.Compute_Depth3D();
		xCudaCompute.Compute_DepthColorCode(Rx::LF::ESpace::View_Object_Pinhole);
		pxImages->Download(Rx::LFR::EImage::DepthMapColored_View_Object_Pinhole, &xDepthImage);

		// Get total focus image and count particles
		cv::Mat particleImage;
		cv::Mat edgeImage;
		int dilation_size = 3;
		cv::Mat sel = cv::getStructuringElement(cv::MORPH_RECT,
			cv::Size(2 * dilation_size + 1, 2 * dilation_size + 1),
			cv::Point(dilation_size, dilation_size));
		particleImage = cv::Mat(xTFImage.GetFormat().m_iHeight, xTFImage.GetFormat().m_iWidth, CV_16UC1, xTFImage.GetDataPtr());

		//cout << "Rx to Mat." << endl;

		// Convert the 8-bit scale
		particleImage.convertTo(particleImage, CV_8UC1, 0.015625);// 12-bit to 8-bit scaling

		edgeImage = particleImage.clone();

		cv::Canny(particleImage, edgeImage, 5, 75, 3);
		cv::dilate(edgeImage, edgeImage, sel);
		cv::erode(edgeImage, edgeImage, sel);

		//cout << "Canny." << endl;

		vector<vector<cv::Point> > contours;
		vector<cv::Vec4i> hierarchy;

		// Find contours in the binary image
		findContours(edgeImage, contours, hierarchy, cv::RETR_EXTERNAL, CV_CHAIN_APPROX_SIMPLE);

		cv::Mat displayImage = particleImage.clone();
		cv::cvtColor(displayImage, displayImage, cv::COLOR_GRAY2BGR);

		displayImage = displayImage * 4;

		//cout << "Contours." << endl;

		// Limit the number of ROIs to process
		int roiCount = 0;
		int allRoiCount = 0;
		double averageRoiArea = 0.0;
		double averageRoiMal = 0.0;
		double maxRoiArea = -1.0;
		double maxRoiMal = -1.0;

		int downsampleFactor = 1;

		int minArea = 3;
		int maxArea = 400;

		// Examine and save valid contours as individual ROIs
		for (unsigned int i = 0; i < contours.size(); i++) {

			cv::Rect bb = boundingRect(contours[i]);

			// Important! Compute area in raw pixels, not downsampled pixels
			double area = contourArea(contours[i]) * downsampleFactor * downsampleFactor;
			cv::RotatedRect minEllipse;

			averageRoiArea += area;

			// Fit an ellipse to get major axis length
			double mal = 0.0;
			if (contours[i].size() >= 6) {
				minEllipse = fitEllipse(contours[i]);
				if (minEllipse.size.width > minEllipse.size.height)
					mal = minEllipse.size.width * downsampleFactor;
				else
					mal = minEllipse.size.height * downsampleFactor;
				
			}

			if (mal < 20 && area > minArea && area < maxArea) {

				if (contours[i].size() >= 6) {
					cv::ellipse(displayImage, minEllipse, cv::Scalar(0, 255, 0));
				}
				else {
					int rad = sqrt(bb.height*bb.width)/2;
					cv::circle(displayImage, cv::Point((int)(bb.x + bb.width / 2), (int)(bb.y + bb.height / 2)), rad, cv::Scalar(0, 0, 255));
				}
				roiCount++;

			}

			


		}

		//cout << "Stats." << endl;

		if (roiCount > 0)
			averageRoiArea /= roiCount;
		char buffer[256];
		sprintf(buffer, "%d,%14.6f,%d,%d,%f", frameCounter,xCapturedImage.GetTimestamp(), xCapturedImage.GetID(), roiCount, averageRoiArea);
		std::ofstream myfile;

		myfile.open(outputFile.c_str(), ios::app);
		myfile << buffer << endl;
		myfile.close();
		cout << buffer << endl;

		//displayImage.convertTo(displayImage, CV_16UC1, 4 / 0.015625);

		
		Rx::CRxImage disp;
		disp.Create(xTFImage.GetFormat().m_iWidth,xTFImage.GetFormat().m_iHeight, Rx::Interop::Runtime28::EPixelType::BGR, Rx::Interop::Runtime28::EDataType::UByte,displayImage.data);

		drawImage(disp);
	}
	catch (Rx::CRxException& ex)
	{
		printf("Exception occured:\n%s\n\n", ex.ToString(true).ToCString());
		printf("Press any key to end program...\n");
		_getch();
	}

}

void RxProcessor::totalFocusProc() {

	// Upload the image as the new raw image of all further CUDA computations
	//printf("Uploading image to Cuda compute instance...\n");
	xCudaCompute.UploadRawImage(xCapturedImage);

	// Pre-process raw light field image before calling any processing functions
	//printf("Pre-processing light field image...\n");
	xCudaCompute.Compute_PreProcess();
	xCudaCompute.Compute_DepthRay();
	xCudaCompute.Compute_DepthMap();
	xCudaCompute.Compute_TotalFocus();
	pxImages->Download(Rx::LFR::EImage::TotalFocus_View_Object_Pinhole, &xOutputImage);


	drawImage(xOutputImage);

}

void RxProcessor::refocusProc() {

	// Upload the image as the new raw image of all further CUDA computations
	//printf("Uploading image to Cuda compute instance...\n");
	xCudaCompute.UploadRawImage(xCapturedImage);

	// Pre-process raw light field image before calling any processing functions
	//printf("Pre-processing light field image...\n");
	xCudaCompute.Compute_PreProcess();
	xCudaCompute.GetParams().SetValue(Rx::LFR::Params::ECudaCompute::Focus_RelativeFocusPlane, dFocus);
	xCudaCompute.Compute_RefocusBasic();
	// Get refocus image from CUDA device
	//printf("Download image from CUDA device...\n");
	pxImages->Download(Rx::LFR::EImage::RefocusBasic, &xOutputImage);

	drawImage(xOutputImage);

}

void RxProcessor::depthMapProc() {

	// Upload the image as the new raw image of all further CUDA computations
	//printf("Uploading image to Cuda compute instance...\n");
	xCudaCompute.UploadRawImage(xCapturedImage);

	// Pre-process raw light field image before calling any processing functions
	//printf("Pre-processing light field image...\n");
	xCudaCompute.Compute_PreProcess();
	xCudaCompute.Compute_DepthRay();
	xCudaCompute.Compute_DepthMap();
	xCudaCompute.Compute_Depth3D();

	xCudaCompute.Compute_DepthColorCode(Rx::LF::ESpace::View_Object_Pinhole);
	pxImages->Download(Rx::LFR::EImage::DepthMapColored_View_Object_Pinhole, &xOutputImage);



	drawImage(xOutputImage);

}

void RxProcessor::rawImageProc() {

	// Upload the image as the new raw image of all further CUDA computations
	//printf("Uploading image to Cuda compute instance...\n");
	xCudaCompute.UploadRawImage(xCapturedImage);

	// Pre-process raw light field image before calling any processing functions
	//printf("Pre-processing light field image...\n");
	xCudaCompute.Compute_PreProcess();
	pxImages->Download(Rx::LFR::EImage::Raw, &xOutputImage);

	drawImage(xOutputImage);

}