#include "avi.h"

/* 
 * avi.c 
 *
 * This MEX-file is an interface to the Microsoft AVIFILE routines.  
 * There are three modes of operation for this MEX-file; 'open',
 * 'addframe' and 'close'.  The syntax for calling this MEX-file is as follows:
 *
 * ID = AVI('open',FILENAME) initializes the interface to the AVIFILE routines
 *		and returns ID, a unique integer ID corresponding to the open file
 *		FILENAME.  The MATLAB M-file AIVFILE should be used to call this routine.
 *
 *  AVI('addframe',FRAME,BITMAPINFO,FRAMENUM,FPS,QUALITY,ID, STREAMNAME) adds FRAME 
 *		number FRAMENUM to the stream in the AVI file represented by ID.  The 
 *		BITMAPINFO is the bitmapheader structure of the AVIFILE object.  
 *		The FPS (frames per second) and QUALITY parameters are required by the
 *		AVIFILE routines. STREAMNAME is a string describing the video stream.
 *      The MATLAB M-file ADDFRAME should be used to call this routine.
 *
 *  AVI('close',ID) finishes writing the AVI file represented by ID.  This will 
 *		close the stream and file handles. This routine should be called by
 *		the MATLAB M-file AVICLOSE.
 *
 *	aviutils.c contains all the routines called by AVI.  
 */

/* 
 * Copyright 1984-2000 The MathWorks, Inc. 
 * $Revision: 1.2 $  $Date: 2000/02/23 19:38:47 $
*/

static char rcsid[] = "$Id: avi.c,v 1.2 2000/02/23 19:38:47 clawton Exp $";

static int mexAtExitIsSet = 0;

void mexFunction(int nlhs,
                 mxArray *plhs[],
                 int nrhs,
                 const mxArray *prhs[])
{


    char *mode;
    void (*opMode)(int nlhs,
                   mxArray *plhs[],
                   int nrhs,
                   const mxArray *prhs[]);

    if( (nrhs < 2) && (nrhs > 8) )
        mexErrMsgTxt("Invalid number of input arguments.");

    if(!mexAtExitIsSet)
        {
            InitNode();
            mexAtExit(exitMex);
            mexAtExitIsSet = 1;
        }

    if ( mxIsChar(prhs[0]) ) 
        mode = mxArrayToString(prhs[0]);
    else
        mexErrMsgTxt("First input to AVI must be a string.");

    if (mode == NULL)
        mexErrMsgTxt("Memory allocation failed.");
	
    if(!strcmp(mode,"open"))
    	opMode = aviopen;
    else if(!strcmp(mode,"addframe"))
    	opMode = addframe;
	else if(!strcmp(mode,"close"))
    	opMode = aviclose;
	else
		mexErrMsgTxt("Unrecognized mode for AVI.");

	mxFree(mode);
		
	(*opMode)(nlhs,
                 plhs,
                 nrhs,
                 prhs);

}

/*
 * AVIOPEN opens the AVI file and returns a unique file ID in plhs
*/

void aviopen(int nlhs,
                 mxArray *plhs[],
                 int nrhs,
                 const mxArray *prhs[])
{
	PAVIFILE pfile;
	HRESULT hr;
    char *filename;
	int id;
	
	/* Validate inputs */
	if (nrhs != 2)
		mexErrMsgTxt("Invalid number of inputs to AVIOPEN.");

	if ( !mxIsChar(prhs[1]) )
		mexErrMsgTxt("The first input to AVIOPEN must be the filename.");
/*
    filename = malloc((mxGetN(prhs[1]) + 1)*sizeof(char));
	*/
	filename = mxArrayToString(prhs[1]);
	if( filename == NULL )
		mexErrMsgTxt("Memory allocation failed.");
/*
	mxGetString(prhs[1], filename, mxGetN(prhs[1])+1);
*/

    AVIFileInit();
    hr = AVIFileOpen(&pfile,		    /* returned file pointer */
                     filename,		            /* file name */
                     OF_WRITE | OF_CREATE,	    /* mode to open file with */
                     0L);
	mxFree(filename);
    if (hr != AVIERR_OK)
        mexErrMsgTxt("Failed to open file.");
	
	/* Keep track of stream handles */
	id = addNodetoList(pfile,NULL);

	/* Return a unique number */
	plhs[0] = mxCreateDoubleMatrix(1,1,mxREAL);
	*((double *) mxGetPr(plhs[0])) = id;
}

/* 
 * ADDFRAME  Adds a frame to the video stream in the avifile opened by aviopen. 
 *	 If this is the first stream, the stream is created, otherwise it is 
 *	 appended to the end.  
 */
 
void addframe(int nlhs,
                 mxArray *plhs[],
                 int nrhs,
                 const mxArray *prhs[])

{
	HRESULT hr;
	PBITMAPINFO bi;
	BITMAPINFOHEADER bih;
	AVICOMPRESSOPTIONS opts;
	AVISTREAMINFO strhdr;
	PAVIFILE pfile;
	PAVISTREAM ps,
		psCompressed;
	FOURCC CompCode;
	NodeType *handle;

	uint8_T KeyFrameRate,
		*frame,
		*colormap;
	mxArray *mxWidth,
		*mxHeight;

	char *compression,
		*StreamName;
	double *FrameNumber;
	int identifier,
		ImageSize;
	unsigned int i,
		j;

	/* Validate Inputs */
	if (!mxIsUint8(prhs[1]))
		mexErrMsgTxt("Frame must be a uint8 matrix.");
	if(!mxIsStruct(prhs[2]))
		mexErrMsgTxt("Third input to addframe must be a structure representing BITMAPINFO.");
	if(!mxIsNumeric(prhs[3]))
		mexErrMsgTxt("Fourth input to addframe must be the frame number.");
	if(!mxIsNumeric(prhs[4]))
		mexErrMsgTxt("Fifth input to addframe must be the frames per second.");
	if(!mxIsNumeric(prhs[5]))
		mexErrMsgTxt("Sixth input to addframe must be the Quality.");
	if(!mxIsNumeric(prhs[6]))
		mexErrMsgTxt("Seventh input to addframe must be the Id number.");
	if(!mxIsChar(prhs[7]))
		mexErrMsgTxt("Eighth input to addframe must be a stream name.");
	if(!mxIsNumeric(prhs[8]))
		mexErrMsgTxt("Last input to addframe must be the key frame rate.");

	/* Find the stream handle for this particular file. */
	identifier = (int) *mxGetPr(prhs[6]);
	handle = FindNodeInList(identifier);
	if( handle == NULL)
		mexErrMsgTxt("Failed to find the open file handle.");
	psCompressed = handle->psCompressed;
	pfile = handle->pfile;

	compression = mxArrayToString(mxGetField(prhs[2],0,"biCompression"));
	if (compression == NULL)
		mexErrMsgTxt("Failed to allocate memory.");

	frame = (uint8_T *)mxGetData(prhs[1]);
	FrameNumber = mxGetPr(prhs[3]);

	KeyFrameRate = (int) *mxGetPr(prhs[8]);

	CompCode = mmioFOURCC(compression[0],compression[1],compression[2],
                              compression[3]);
	mxFree(compression);

	mxWidth = mxGetField(prhs[2],0,"biWidth");
	mxHeight = mxGetField(prhs[2],0,"biHeight");

	/* Populate bih structure */
	bih.biSize = sizeof(BITMAPINFOHEADER);
	bih.biWidth = (long) *mxGetPr(mxWidth);
	bih.biHeight = (long) *mxGetPr(mxHeight);
	bih.biPlanes = 1;
	bih.biBitCount = (int) *mxGetPr(mxGetField(prhs[2],0,"biBitCount")); 
	bih.biCompression = 0; /* 0 For uncompressed */
	bih.biSizeImage = (int) *mxGetPr(mxGetField(prhs[2],0,"biSizeImage"));

	bih.biXPelsPerMeter = 0;
	bih.biYPelsPerMeter = 0;
	bih.biClrUsed = (int) *mxGetPr(mxGetField(prhs[2],0,"biClrUsed")); 
	bih.biClrImportant = 0;

	colormap = (uint8_T *)mxGetData(mxGetField(prhs[2],0,"Colormap"));
	ImageSize = (int) *mxGetPr(mxGetField(prhs[2],0,"biSizeImage"));

	bi = (PBITMAPINFO) mxMalloc(sizeof(BITMAPINFOHEADER) + 
                    sizeof(RGBQUAD) * bih.biClrUsed);

	bi->bmiHeader = bih;

	if( colormap != NULL )
	{
		j=0;
		for(i=0;i<bih.biClrUsed; i++)
		{
			bi->bmiColors[i].rgbBlue = colormap[j++];
			bi->bmiColors[i].rgbGreen = colormap[j++];
			bi->bmiColors[i].rgbRed = colormap[j++];
			bi->bmiColors[i].rgbReserved = colormap[j++];
		}
	}

	memset(&strhdr, 0, sizeof(strhdr));
	strhdr.fccType                = streamtypeVIDEO;/* stream type */
	strhdr.fccHandler             = 0;
	strhdr.dwScale                = 1;
	strhdr.dwRate                 = (int) *mxGetPr(prhs[4]);		
	strhdr.dwSuggestedBufferSize  = ImageSize;

	StreamName = mxArrayToString(prhs[7]);
	if (StreamName == NULL)
		mexErrMsgTxt("Memory allocation failed.");

	strcpy(strhdr.szName, StreamName);
	mxFree(StreamName);

	SetRect(&strhdr.rcFrame, 0, 0,		    /* rectangle for stream */
	   (int) *mxGetPr(mxWidth),
	    (int) *mxGetPr(mxHeight));

	if (*FrameNumber == 0)
	{

	/* 	For first frame, create the stream.	*/
		hr = AVIFileCreateStream(pfile,		    /* file pointer */
			     &ps,		    /* returned stream pointer */
		         &strhdr);	    /* stream header */
		if (hr != AVIERR_OK)
			mexErrMsgTxt("Failed to create video stream.");

		memset(&opts, 0, sizeof(opts));
		if (CompCode ==0)				/* Uncompressed AVI */
		{
			opts.fccType = streamtypeVIDEO;
			opts.fccHandler = CompCode;
			opts.dwKeyFrameEvery = 0;
			opts.dwQuality = 0;
			opts.dwFlags = 0;
			opts.lpFormat = bi;
			opts.cbFormat = bi->bmiHeader.biSize + 
			bi->bmiHeader.biClrUsed * sizeof(RGBQUAD);
			
		}
		else							/* Compressed AVI */
		{
			opts.fccType = 0; /* streamtypeVIDEO doesn't seem to work for all compressors. */
			opts.fccHandler = CompCode;
			opts.dwKeyFrameEvery = KeyFrameRate;
			opts.dwQuality = (int) *mxGetPr(prhs[5]);
			opts.dwFlags = AVICOMPRESSF_KEYFRAMES;
			opts.lpFormat = bi;
			opts.cbFormat = bi->bmiHeader.biSize + 
                            bi->bmiHeader.biClrUsed * sizeof(RGBQUAD);
		}

		hr = AVIMakeCompressedStream(&psCompressed, ps, &opts, NULL);
		if (hr == AVIERR_NOCOMPRESSOR)
                    {		
                        if (ps)
			{
                            AVIStreamClose(ps);
                            ps = NULL;
			}
                    mexErrMsgTxt("Can not locate compressor.");
                    }
		else if (hr == AVIERR_MEMORY)
                    {	
                        if (ps)
			{
                            AVIStreamClose(ps);
                            ps = NULL;
			}
                    mexErrMsgTxt("Not enough memory available for compressor.");
                    }
		else if (hr == AVIERR_UNSUPPORTED)
                    {	
                        if (ps)
			{
                            AVIStreamClose(ps);
                            ps = NULL;
			}
                    mexErrMsgTxt("Compressor can not compress this data type.");
                    }
		else if (hr != AVIERR_OK)
                    {	
                        if (ps)
			{
                            AVIStreamClose(ps);
                            ps = NULL;
			}
                    mexErrMsgTxt("Failed to make compressed stream. Unable to determine reason.");
                    }
		/* Assign the stream to the list */
		handle->psCompressed = psCompressed;
		if (ps)
		{
			AVIStreamClose(ps);
			ps = NULL;
		}

		hr = AVIStreamSetFormat(psCompressed, 0,
			bi,	    /* stream format */
			bi->bmiHeader.biSize + 
			bi->bmiHeader.biClrUsed * sizeof(RGBQUAD));
		if (hr != AVIERR_OK)
		mexErrMsgTxt("Failed to set stream format.");

	} /* End of code required for first frame */

	mxFree(bi);
	
	hr = AVIStreamWrite(psCompressed,	/* stream pointer */
		(int) *FrameNumber,				/* time of this frame */
		1,				/* number to write */
		frame,/* pointer to data */
		ImageSize,	/* size of this frame */
		0,			 /* flags.... */
		NULL,
		NULL);
	if (hr != AVIERR_OK)
		mexErrMsgTxt("Failed to write stream data.");
}


/* 
 * AVICLOSE close the stream and the file then remove this information from
 *	the list.
 */

void aviclose(int nlhs,
                 mxArray *plhs[],
                 int nrhs,
                 const mxArray *prhs[])
{

	PAVIFILE pfile;
	PAVISTREAM psCompressed;
	int identifier;
	NodeType *handle;

	/*Need to find the stream handle for this particular file. */
	identifier = (int) *mxGetPr(prhs[1]);
	handle = FindNodeInList(identifier);
	
	if(handle == NULL)
		mexErrMsgTxt("The file is not open.");

	psCompressed = handle->psCompressed;
	pfile = handle->pfile;

	if (psCompressed)
	{
		AVIStreamClose(psCompressed);
		psCompressed = NULL;
	}

	if (pfile)
	{
	    AVIFileClose(pfile);
		pfile = NULL;
	}

	deleteNodeFromList(identifier);	     

}