#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 = REFOCUS;
	int uBufferSize = 512;
	bool bOverwrite = false;
	imageOverLCM = false;
	m_xImageBuffer.Initialize(uBufferSize, bOverwrite);

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

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

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();

	return 0;

}

void RxProcessor::threadedFunction() {

	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) {
				frameCounter++;

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

		}
		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());
		}

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

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

}

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);

}