// astratest.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"

astra_status_t check( const astra_status_t  status )
{
	switch ( status )
	{
		case ASTRA_STATUS_SUCCESS:
			break;

		default:
			printf( "Some error (%d) has occured!\n", status );
	}

	return status;
}

cv::Mat DepthFrame2Mat( const astra::DepthFrame & depth_frame )
{
	if ( !depth_frame.is_valid() )
	{
		return cv::Mat();
	}
	
	cv::Mat depth_map = cv::Mat( cv::Size( depth_frame.width(), depth_frame.height() ), CV_16U );

#ifndef _DEBUG
#pragma omp parallel for
#endif
	for ( int y = 0; y < depth_map.rows; ++y )
	{
		for ( int x = 0; x < depth_map.cols; ++x )
		{
			const auto value = ( unsigned __int16 )( depth_frame.data()[x + y*depth_frame.width()] );
			depth_map.at<unsigned __int16>( y, x ) = value;
		}
	}

	return depth_map;
}

cv::Mat ColorFrame2Mat( const astra::ColorFrame & color_frame )
{
	if ( !color_frame.is_valid() )
	{
		return cv::Mat();
	}

	cv::Mat image = cv::Mat( cv::Size( color_frame.width(), color_frame.height() ), CV_8UC3 );

#ifndef _DEBUG
#pragma omp parallel for
#endif
	for ( int y = 0; y < image.rows; ++y )
	{
		for ( int x = 0; x < image.cols; ++x )
		{
			const auto value = color_frame.data()[x + y*color_frame.width()];
			image.at<cv::Vec3b>( y, x ) = cv::Vec3b( value.b, value.g, value.r );
		}
	}

	return image;
}

cv::Mat DepthMap2XYZ( const cv::Mat depth_map, const astra::CoordinateMapper & coordinate_mapper )
{
	cv::Mat xyz = cv::Mat( depth_map.size(), CV_32FC3 );

	for ( int y = 0; y < depth_map.rows; ++y )
	{
		for ( int x = 0; x < depth_map.cols; ++x )
		{
			const auto depth = depth_map.at<unsigned __int16>( y, x );
			cv::Vec3f point_ws;			
			coordinate_mapper.convert_depth_to_world(
				( float )x, ( float )y, ( float )depth,
				&point_ws[0], &point_ws[1], &point_ws[2] );
			xyz.at<cv::Vec3f>( y, x ) = point_ws;
		}
	}

	return xyz;
}

cv::Mat IRFrame2Mat( const astra::InfraredFrame16 & ir_frame )
{
	if ( !ir_frame.is_valid() )
	{
		return cv::Mat();
	}

	cv::Mat ir = cv::Mat( cv::Size( ir_frame.width(), ir_frame.height() ), CV_16U );

#ifndef _DEBUG
#pragma omp parallel for
#endif
	for ( int y = 0; y < ir.rows; ++y )
	{
		for ( int x = 0; x < ir.cols; ++x )
		{
			const auto value = ( unsigned __int16 )( ir_frame.data()[x + y*ir_frame.width()] );
			ir.at<unsigned __int16>( y, x ) = value;
		}
	}

	return ir;
}

cv::Mat Point2Mat( const astra::PointFrame & point_frame )
{
	if ( !point_frame.is_valid() )
	{
		return cv::Mat();
	}

	cv::Mat point = cv::Mat( cv::Size( point_frame.width(), point_frame.height() ), CV_32FC3 );

#ifndef _DEBUG
#pragma omp parallel for
#endif
	for ( int y = 0; y < point.rows; ++y )
	{
		for ( int x = 0; x < point.cols; ++x )
		{
			const auto value = ( astra::Vector3f )( point_frame.data()[x + y*point_frame.width()] );
			point.at<cv::Vec3f>( y, x ) = cv::Vec3f( value.x, value.y, value.z );
		}
	}

	return point;
}

int test_depth_ir()
{	
	const uint8_t fps = 30;
	const uint32_t width = 640;
	const uint32_t height = 480;

	check( astra::initialize() ); // prepare Astra to do our bidding
	// causes weird exception: cpptoml::parse_exception
	
	astra::StreamSet device; // encapsulates a group of related data sources		
	astra::StreamReader reader = device.create_reader(); // tap into one of the streams			

	astra::DepthStream depth_stream = reader.stream<astra::DepthStream>(); // create the low-level depth stream 	
	astra::ImageStreamMode depth_mode( width, height, fps, astra_pixel_formats::ASTRA_PIXEL_FORMAT_DEPTH_MM );
	depth_stream.set_mode( depth_mode );
	depth_stream.enable_mirroring( false );
	depth_stream.enable_registration( false );
	depth_stream.start(); // starting depth stream tells Astra that we’re interested in getting depth data	
	
	if ( depth_stream.hFov() > 0.0f )
	{
		printf( "depth_stream: fov_x = %0.3f deg, fov_y = %0.3f deg\n", RAD2DEG( depth_stream.hFov() ), RAD2DEG( depth_stream.vFov() ) );
	}
	else
	{	
		printf( "Camera not found! Try to reconnect the camera or plug it in other usb port.\n" );
		check( astra::terminate() );		

		return -1;
	}

	astra::ColorStream color_stream = reader.stream<astra::ColorStream>(); // color works similarly	
	astra::ImageStreamMode color_mode( width, height, fps, astra_pixel_formats::ASTRA_PIXEL_FORMAT_RGB888 );
	color_stream.set_mode( color_mode );
	color_stream.enable_mirroring( false );
	color_stream.start();

	astra::InfraredStream ir_stream = reader.stream<astra::InfraredStream>();
	astra::ImageStreamMode ir_mode( width, height, fps, astra_pixel_formats::ASTRA_PIXEL_FORMAT_GRAY16 );
	ir_stream.set_mode( ir_mode );
	ir_stream.enable_mirroring( false );
	//ir_stream.start();

	astra::PointStream point_stream = reader.stream<astra::PointStream>(); // high-level stream
	//point_stream.start();

	cv::namedWindow( "Depth" );	
	cv::moveWindow( "Depth", 0, 0 );	
	cv::namedWindow( "RGB" );
	cv::moveWindow( "RGB", 650, 0 );
	//cv::namedWindow( "IR" );
	//cv::moveWindow( "IR", 650 * 2, 0 );

	char key;
	int frame_index = 0;

	do
	{
		const double t0 = omp_get_wtime();

		astra::Frame frame = reader.get_latest_frame(); // retrieve the latest...
		const auto depth_frame = frame.get<astra::DepthFrame>(); //... depth frame
		const auto color_frame = frame.get<astra::ColorFrame>(); //... color frame
		const auto ir_frame = frame.get<astra::InfraredFrame16>(); //... ir frame
		const auto point_frame = frame.get<astra::PointFrame>(); // ... point frame

		frame_index = depth_frame.frame_index();
		
		cv::Mat depth_map = DepthFrame2Mat( depth_frame ); // values are in the range <0, 10000>
		cv::Mat color = ColorFrame2Mat( color_frame ); // values are in the range <0, 1023> ???
		cv::Mat ir = IRFrame2Mat( ir_frame ); // values are in the range <0, 1023> ???		
		cv::Mat xyz = Point2Mat( point_frame );

		//cv::Mat xyz_data = DepthMap2XYZ( depth_map( cv::Rect( 260, 180, window, window ) ), depth_stream.coordinateMapper() );
		if ( point_frame.is_valid() )
		{
			astra::Vector3f center_xyz = ( astra::Vector3f )( point_frame.data()[point_frame.width() / 2 +
				( point_frame.height() / 2 )*point_frame.width()] );
			printf( "%0.3f m, ", center_xyz.z * 1e-3 ); // LHS, x-right, y-up, z-forward
		}

			
		//ir *= ( 65535.0 / 1023.0 );
		//cv::imwrite( "depth.png", depth_map );

		if ( !depth_map.empty() )
		{
			// up to now, depth_map values are in mm's
			depth_map *= ( 65535.0 / 10000.0 );
			cv::imshow( "Depth", depth_map );
			//cv::imwrite( "depth.png", depth_map );
		}
		if ( !color.empty() )
		{
			cv::imshow( "RGB", color );
			//cv::imwrite( "rgb.png", color );
		}
		if ( !ir.empty() ) cv::imshow( "IR", ir );		
		key = cv::waitKey( 10 );

		const double t1 = omp_get_wtime();
		printf( "%0.3f ms (%0.1f fps)\n", ( t1 - t0 )*1e+3, 1.0 / ( t1 - t0 ) );		
	} while ( key != 'q' );

	point_stream.stop();
	ir_stream.stop();
	color_stream.stop();
	depth_stream.stop();

	check( astra::terminate() ); // properly turn off attached devices

	return 0;
}

// occasionally causes memory leak (Debug + VLD) in opencv_videoio310d.dll!cv::Algorithm::clear()
// this works for Orbbec Astra Pro and ZED stereo UVC cameras
int test_color()
{
	cv::VideoCapture cap( 0 ); // open the default camera, how to enumerate?

	if ( cap.isOpened() )
	{
		cap.set( CV_CAP_PROP_FRAME_WIDTH, 1280 ); // also lower res works well
		cap.set( CV_CAP_PROP_FRAME_HEIGHT, 720 );
		cap.set( CV_CAP_PROP_FPS, 10.0 );
		printf( "(%d x %d) px @ %0.1f fps\n",
			( int )cap.get( CV_CAP_PROP_FRAME_WIDTH ),
			( int )cap.get( CV_CAP_PROP_FRAME_HEIGHT ),
			cap.get( CV_CAP_PROP_FPS ) );		

		cv::namedWindow( "Image" );
		char key;
		
		do {			
			const double t0 = omp_get_wtime();			

			cv::Mat image;
			cap >> image; // get a new frame from Astra camera in ~30 ms (30 fps) on 10yo hw, ZED runs at 4 fps
			
			cv::imshow( "Image", image );
			key = cv::waitKey( 10 );
			if ( key == 's' )
			{
				cv::imwrite( "image.png", image );
			}
			
			const double t1 = omp_get_wtime();
			printf( "%0.3f ms (%0.1f fps)\n", ( t1 - t0 )*1e+3, 1.0 / ( t1 - t0 ) );
		} while ( key != 'q' );

		cap.release();
	}
	else
	{
		printf( "Camera not found!\n" );
	}

	return 0;
}

int main()
{
	return test_depth_ir();
    //return test_color();
}
