#include <mex.h>
#include <string.h>
#include "VidHeader.h"

typedef unsigned char uchar; 

static void print_usage(const char * error_str){
	mexPrintf("%s\n", error_str);
	mexPrintf("Usage: vidCropMex, \n"
	"Inputs: \n"
	"   < input video file name : string >\n"
	"   < output video file name : string >\n"
	"	< [empty] or [2 x 1 : double] >	crop start and end frames \n"
	"	< [0 : double] > placeholder argument \n"
	"	< scalar : double >	gap in pixels between merged videos \n"
	"	[ [empty] or [4 x 1 : double] ] bbox crop [l t w h] \n"
	"	[ [empty] or [3 x 1 : unit8] ] RGB background fill color\n"
	" OR (alternate version) \n"
	"   < {K x 1 : string} > cell array of input video file names \n"
	"   < output video file name : string >\n"
	"	< [empty] or [2 x K : double] > crop start and end frames for each video \n"
	"	< [ 0 or 1 or 2 : double] > merge cropped input videos horizontally or vertically or sequentially \n"
	"	< scalar : double >	gap in pixels between merged videos \n"
	"	[ [empty] or [4 x K : double] ] bbox crop for all video [l t w h] \n"
	"	[ [empty] or [3 x 1 : unit8] RGB background fill color \n"
	"Outputs: \n"
	"   None.\n"
	);
	mexPrintf("\n"); mexErrMsgTxt(error_str);
}

#define my_max(a, b) (((a) > (b))? (a) : (b))
#define my_min(a, b) (((a) < (b))? (a) : (b))

void set_bounds(double  * ptr, int vidi, VidHeader * vid_header, int & l, int & t, int & w, int & h){
	if( ptr != NULL ) {
		ptr += 4*vidi;
		l = (int)ptr[0];
		t = (int)ptr[1];
		w = (int)ptr[2];
		h = (int)ptr[3];
		l = my_max(l, 0);
		t = my_max(t, 0);
		w = my_min(vid_header -> width - l, w);
		h = my_min(vid_header -> height - t, h);
	} else {
		l = 0, t = 0, w = vid_header -> width, h = vid_header -> height;
	}
}

void mexFunction( int nlhs, mxArray *plhs[],
                  int nrhs, const mxArray *prhs[] ){
	
	if( nrhs < 5 || mxGetClassID(prhs[1]) != mxCHAR_CLASS 	|| mxGetClassID(prhs[2]) != mxDOUBLE_CLASS 	|| 
					mxGetClassID(prhs[3]) != mxDOUBLE_CLASS || mxGetClassID(prhs[4]) != mxDOUBLE_CLASS	)
		print_usage("Incorrect number inputs and/or input datatype for some arguments.");
	
	int num_input_videos = 0;
	int buflen = 2000;
	char * in_fname = new char[buflen];
	VidHeader ** in_vid_header = NULL;
	double * input_video_crop_frames = NULL;
	if( mxGetNumberOfElements(prhs[2]) != 0 ) {
		input_video_crop_frames = new double[mxGetNumberOfElements(prhs[2])];
		memcpy( input_video_crop_frames, mxGetPr(prhs[2]), mxGetNumberOfElements(prhs[2]) * sizeof(double) );
	}

	if ( mxGetClassID(prhs[0]) == mxCELL_CLASS ){
		
		num_input_videos = mxGetNumberOfElements(prhs[0]);
		if( mxGetNumberOfElements(prhs[2]) != 0 && ( num_input_videos != mxGetN(prhs[2]) || mxGetM(prhs[2]) != 2 ) )
			print_usage("Crop frames input array should be size [Number input videos x 2].");
		
		in_vid_header = new VidHeader*[num_input_videos];
		for( int vidi = 0; vidi < num_input_videos; vidi++ ){
			if( mxGetClassID(mxGetCell(prhs[0],vidi)) != mxCHAR_CLASS )
				print_usage("One of the cell entries for filenames array is not a string.");
			
			mxGetString(mxGetCell(prhs[0],vidi), in_fname, buflen);
			in_vid_header[vidi] = new VidHeader(in_fname, stderr);
		}
		
	} else {
		if( mxGetClassID(prhs[0]) != mxCHAR_CLASS || mxGetNumberOfElements(prhs[2]) != 2 )
			print_usage("Input filename is not a string or crop frames array is not length 2.");
		
		num_input_videos = 1;
		in_vid_header = new VidHeader*[num_input_videos];
		mxGetString(prhs[0], in_fname, buflen);
		in_vid_header[0] = new VidHeader(in_fname, stderr);
	}

	if( input_video_crop_frames == NULL ){
		input_video_crop_frames = new double[2*num_input_videos];
		for( int vidi = 0; vidi < num_input_videos; vidi++ ){
			input_video_crop_frames[2*vidi] 	= 0;
			input_video_crop_frames[2*vidi + 1] = in_vid_header[vidi] -> num_frames;
		}
	}
	int max_num_input_frames = 0, num_all_input_frames = 0;
	for( int vidi = 0; vidi < num_input_videos; vidi++ ){

		int start 	= (int)input_video_crop_frames[2*vidi];
		int end 	= (int)input_video_crop_frames[2*vidi+1];
		start 	= ( start < 0 )? 0 : start;
		end 	= ( end < in_vid_header[vidi] -> num_frames )? end : in_vid_header[vidi] -> num_frames;
		input_video_crop_frames[2*vidi] 	= start;
		input_video_crop_frames[2*vidi + 1] = end;
		max_num_input_frames 				= my_max( max_num_input_frames, end - start );
		num_all_input_frames 				+= end - start;
	}
	
	if(	mxGetNumberOfElements(prhs[3]) != 1 )
		print_usage("Fourth input argument should be scalar.");
	int merge_videos_horiz_or_vert	= (int)( mxGetScalar(prhs[3]) + 0.5 );

	if(	mxGetNumberOfElements(prhs[4]) != 1 )
		print_usage("Gap between videos input argument should be scalar.");
	int gap_between_videos = (int)( mxGetScalar(prhs[4]) + 0.5 );

	double * crop_bboxes 	= NULL;
	int crop_bbox_argi 		= 5;
	if( nrhs > crop_bbox_argi && mxGetNumberOfElements(prhs[crop_bbox_argi]) > 0 ){
		if( mxGetClassID(prhs[crop_bbox_argi]) != mxDOUBLE_CLASS || mxGetM(prhs[crop_bbox_argi]) != 4 || mxGetN(prhs[crop_bbox_argi]) != num_input_videos )
			print_usage("Incorrect crop bbox array size or datatype");

		crop_bboxes = mxGetPr( prhs[crop_bbox_argi] );
	}

	int background_fill_rgb_argi 	= 6;
	unsigned char background_fill_rgb[ 	3 ];
	if( nrhs > background_fill_rgb_argi && mxGetNumberOfElements(prhs[background_fill_rgb_argi]) > 0 ){
		if( mxGetClassID(prhs[background_fill_rgb_argi]) != mxUINT8_CLASS || mxGetNumberOfElements(prhs[background_fill_rgb_argi]) != 3 )
			print_usage("Incorrect background_fill_rgb array size or datatype");
		
		unsigned char * input_background_fill_rgb	= ( unsigned char * ) mxGetData( prhs[background_fill_rgb_argi] );
		for( int i	= 0;	i	<	3;	i++ )	background_fill_rgb[ i ]	= 	input_background_fill_rgb[ i ];
	}	else
		for( int i	= 0;	i	<	3;	i++ )	background_fill_rgb[ i ]	= 	(uchar)255;

	int out_ww = 0, out_hh = 0;
	for( int vidi = 0; vidi < num_input_videos; vidi++ ) {
		int l, t, w, h; 
		set_bounds( crop_bboxes, vidi, in_vid_header[vidi], l, t, w, h );

		switch( merge_videos_horiz_or_vert ){
			case 0:
				out_ww 	+= w + gap_between_videos;
				out_hh 	= my_max( out_hh, h );
				break;
			case 1:
				out_ww	= my_max( out_ww, w );
				out_hh 	+= h + gap_between_videos;
				break;
			case 2:
				out_hh 	= my_max( out_hh, h );
				out_ww	= my_max( out_ww, w );
				break;
		}
	}
	switch( merge_videos_horiz_or_vert ){
		case 0: out_ww 	-= gap_between_videos; break;
		case 1: out_hh 	-= gap_between_videos; break;
	}
	out_ww = (out_ww / 8 + (out_ww % 8 > 0 )) * 8;
	out_hh = (out_hh / 8 + (out_hh % 8 > 0 )) * 8;
	
	char * out_fname = new char[buflen];
	mxGetString(prhs[1], out_fname, buflen);
	VidHeader * out_vid_header = NULL;
	int num_out_frames 	= ( merge_videos_horiz_or_vert < 2 ) ? max_num_input_frames : num_all_input_frames;
	out_vid_header = new VidHeader( out_fname, stderr, out_ww, out_hh, num_out_frames );
	
	if( merge_videos_horiz_or_vert < 2 ){
		for( int framei = 0; framei < max_num_input_frames; framei++ ){

			int	out_vid_header_pixmapi	=	0;
			for(		int xy			=	0;	xy		< 	out_ww	* out_hh;	xy++ 	)
				for(	int	rgbi		=	0;	rgbi	<	3;					rgbi++	)
						out_vid_header	->	pixmap[	out_vid_header_pixmapi++ ]	=	background_fill_rgb[ rgbi ];
			// memset(out_vid_header->pixmap, (uchar)255, out_vid_header->pixmap_size);

			int out_x = 0, out_y = 0;
			for( int vidi = 0; vidi < num_input_videos; vidi++ ){
				int start 	= (int)input_video_crop_frames[2*vidi];
				int end 	= (int)input_video_crop_frames[2*vidi+1];
				int l, t, w, h; 
				set_bounds( crop_bboxes, vidi, in_vid_header[vidi], l, t, w, h );
				if( framei < end - start ){
					in_vid_header[vidi] -> readFrameFromVideoFile( framei + start );
					for( int y = t; y < t + h; y++ ){
						int in = 3 * ( y*in_vid_header[vidi] -> width + l );
						int out; 
						switch( merge_videos_horiz_or_vert ){
							case 0: out = 3 * ( (y-t) * out_vid_header -> width + out_x ); break;
							case 1: out = 3 * ( (y-t) + out_y ) * out_vid_header -> width; break;
						}
							
						for( int x = 3*l; x < 3*(l + w); x++ )
							out_vid_header -> pixmap[out++] = in_vid_header[vidi] -> pixmap[in++];	
					}
				}
				switch( merge_videos_horiz_or_vert ){
					case 0: out_x  += w + gap_between_videos; break;
					case 1: out_y  += h + gap_between_videos; break;
				}
			}
			out_vid_header -> writeNextFrameToVideoFile();
			if( framei % 200 == 199 )
				fprintf( stdout,  "%6d of %6d frames in output written.\n", framei, num_out_frames );
		}

	} else {
		int out_vid_framei = 0;	
		for( int vidi = 0; vidi < num_input_videos; vidi++ ){
			int start 	= (int)input_video_crop_frames[2*vidi];
			int end 	= (int)input_video_crop_frames[2*vidi+1];
			for( int framei = start; framei < end; framei++ ){

				memset(out_vid_header->pixmap, (uchar)255, out_vid_header->pixmap_size);

				in_vid_header[vidi] -> readFrameFromVideoFile( framei );
				int l, t, w, h, out_x, out_y;
				set_bounds( crop_bboxes, vidi, in_vid_header[vidi], l, t, w, h );
				out_x = ( out_vid_header -> width - w )/2;
				out_y = ( out_vid_header -> height - h )/2;

				for( int y = t; y < t + h; y++ ){
					int in 	= 3 * ( y*in_vid_header[vidi] -> width + l );
					int out = 3 * ( ( y - t + out_y ) * out_vid_header -> width + out_x );
					for( int x = 3*l; x < 3*(l + w); x++ )
						out_vid_header -> pixmap[out++] = in_vid_header[vidi] -> pixmap[in++];
				}
				out_vid_header -> writeNextFrameToVideoFile();
				if( out_vid_framei % 200 == 199 )
					fprintf( stdout,  "%6d of %6d frames in output written.\n", out_vid_framei, num_out_frames );
				out_vid_framei++;
			}
		}
	}
	
	for( int vidi = 0; vidi < num_input_videos; vidi++ )
		delete in_vid_header[vidi];
	delete [] in_vid_header;
	delete out_vid_header;
	
	delete [] in_fname;
	delete [] out_fname;
	delete [] input_video_crop_frames;
}