/* kernel_readin.c */
#include "mex.h"
#include "defs.h"

#define MATLAB_KERNELS  (prhs[0])
#define GEOMETRY        (prhs[1])
#define BEAMLETS        (plhs[0])   // output structure
/* #define TERMA           (plhs[1])
#define KERMAC			(plhs[2])
#define DEFF            (plhs[3])
#define TERMA_MASK      (plhs[4])
#define DOSE_MASK       (plhs[5]) */

//function prototypes
void check_MATLAB_KERNELS(const mxArray *);
void check_GEOMETRY(const mxArray *);
void load_kernels(const mxArray *, MONO_KERNELS *);
void load_Geometry(const mxArray *, FLOAT_GRID *);
void load_beam(const mxArray *, int, BEAM *);
int calc_deff(FLOAT_GRID *,FLOAT_GRID *, FLOAT_GRID *, BEAM *);
int terma_kerma(FLOAT_GRID *,FLOAT_GRID *,FLOAT_GRID *,MONO_KERNELS *,FLOAT_GRID *);
int make_poly_kernel(MONO_KERNELS *, KERNEL *);
int calc_dose(FLOAT_GRID *,FLOAT_GRID *,FLOAT_GRID *,KERNEL *,BEAM *, FLOAT_GRID *);
int terma_dose_masks(FLOAT_GRID *, FLOAT_GRID *, BEAM *);

char errstr[200];  // error string that all routines have access to

void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[])
{
	int dims[3], dims2[2];  // dimensions of input arrays
	const int *celldims;
	int b, B;     // beam indices
	const int one = 1;
	// matlab arrays for easy output if need be (just comment these lines out and uncomment
	// the #define lines above and the mxDestroy lines below).
	mxArray *TERMA, *KERMAC, *DEFF, *TERMA_MASK, *DOSE_MASK;
	mxArray *DOSE, *cell, *struc, *field, *ind_field, *data_field;
	
	// sparse storage parameters
	int Nind, Ntotal, i, j;
	int *intPtr;
	float *floatPtr;

	// information for sparse output data:
	int nfields = 5;
	const char *field_names[5] = {"x_count","y_count","z_count",
							"non_zero_indices","non_zero_values"};

	char tmpstr[200];

	FLOAT_GRID density;
	FLOAT_GRID terma_mask;
	FLOAT_GRID dose_mask;
	FLOAT_GRID deff;
	FLOAT_GRID terma;
	FLOAT_GRID kermac;
	FLOAT_GRID dose;

	MONO_KERNELS mono_kernels;
	KERNEL poly_kernel;

	BEAM beam;

	// check and load-in input data
	check_MATLAB_KERNELS(MATLAB_KERNELS); // see function comments for proper input format
	check_GEOMETRY(GEOMETRY);             // see function comments for proper input format

	load_kernels(MATLAB_KERNELS, &mono_kernels);
	load_Geometry(GEOMETRY,&density);

	// copy density grid geometry information to all other grids of interest

	copy_grid_geometry(&density,&terma_mask);
	copy_grid_geometry(&density,&dose_mask);
	copy_grid_geometry(&density,&deff);
	copy_grid_geometry(&density,&terma);
	copy_grid_geometry(&density,&kermac);
	copy_grid_geometry(&density,&dose);

	// Allocate memory for all of the grids
	
	// dimensions of all the grids
	dims[0] = density.x_count;
	dims[1] = density.y_count;
	dims[2] = density.z_count;
	Ntotal = dims[0]*dims[1]*dims[2];

	// Allocate memory for all of the data grids using Matlab mxArray functions. 
	// The idea behind this is for the outputs to send the outputs directly to Matlab.
 	if ((TERMA_MASK = mxCreateNumericArray(3,dims,mxSINGLE_CLASS,mxREAL))==NULL)
	  mexErrMsgTxt("Could not allocate memory for terma_mask grid");

 	terma_mask.matrix = (float *)mxGetData(TERMA_MASK);

	if ((DOSE_MASK = mxCreateNumericArray(3,dims,mxSINGLE_CLASS,mxREAL))==NULL)
	  mexErrMsgTxt("Could not allocate memory for dose_mask grid"); 

	dose_mask.matrix = (float *)mxGetData(DOSE_MASK);

	if ((DEFF = mxCreateNumericArray(3,dims,mxSINGLE_CLASS,mxREAL))==NULL)
	  mexErrMsgTxt("Could not allocate memory for deff grid"); 

	deff.matrix = (float *)mxGetData(DEFF); 

	if ((TERMA = mxCreateNumericArray(3,dims,mxSINGLE_CLASS,mxREAL))==NULL)
	  mexErrMsgTxt("Could not allocate memory for terma grid"); 

	terma.matrix = (float *)mxGetData(TERMA);

 	if ((KERMAC = mxCreateNumericArray(3,dims,mxSINGLE_CLASS,mxREAL))==NULL)
	  mexErrMsgTxt("Could not allocate memory for kermac grid"); 

	kermac.matrix = (float *)mxGetData(KERMAC);

	if ((DOSE = mxCreateNumericArray(3,dims,mxSINGLE_CLASS,mxREAL))==NULL)
	  mexErrMsgTxt("Could not allocate memory for dose matrix"); 
	
	// Point dose matrix pointer at DOSE Matlab matrix data
	dose.matrix = (float *)mxGetData(DOSE);

	//create polyenergetic kernel from mono kernels and fluence, mu data
 	if (SUCCESS != make_poly_kernel(&mono_kernels,&poly_kernel) )
	{
		sprintf(tmpstr,"Failed making polyenergetic kernel!\n");
		strcat(tmpstr,errstr);
		strcpy(errstr,tmpstr);
		mexErrMsgTxt(errstr);
	} 

	// create the output structure
	celldims = mxGetDimensions(mxGetField(GEOMETRY,0,"beam"));
	B = celldims[0]*celldims[1];  // number of beams in the calculation
	BEAMLETS = mxCreateCellArray(2,celldims);

	// do calculations for all beams
 	for (b=0;b<B;b++)
	{
		load_beam(GEOMETRY,b,&beam);
		// create terma and dose masks
		if (SUCCESS != terma_dose_masks(&terma_mask,&dose_mask,&beam))
		{
			sprintf(tmpstr,"Failed in terma_dose_masks!\n");
			strcat(tmpstr,errstr);
			strcpy(errstr,tmpstr);
			mexErrMsgTxt(errstr);
		}

		//create effective depth array from density array
		if (SUCCESS != calc_deff(&density,&deff,&terma_mask,&beam))
		{
			sprintf(tmpstr,"Failed in calc_deff!\n");
			strcat(tmpstr,errstr);
			strcpy(errstr,tmpstr);
			mexErrMsgTxt(errstr);
		}

		//create kerma and terma arrays
		//note kerma is collision kerma and is used for a kernel hardening correction
 		if (SUCCESS != terma_kerma(&deff,&terma,&kermac,&mono_kernels,&terma_mask))
		{
			sprintf(tmpstr,"Failed in terma_kerma calculation!\n");
			strcat(tmpstr,errstr);
			strcpy(errstr,tmpstr);
			mexErrMsgTxt(errstr);
		}

		//use all this stuff to calculate dose
		if ( (SUCCESS != calc_dose(&density,&terma,&dose,&poly_kernel,&beam,&dose_mask)) ) 
		{
			sprintf(tmpstr,"Failed calculating dose!\n");
			strcat(tmpstr,errstr);
			strcpy(errstr,tmpstr);
			mexErrMsgTxt(errstr);
		}

		FILE *fid;
		fid = fopen("dose.bin","wb");
		fwrite(dose.matrix,sizeof(float),Ntotal,fid);
		fclose(fid);

		// count the number of non-zero dose values
		Nind = 0;
		for (i=0;i<Ntotal;i++)
			if (dose.matrix[i] > 0.0)
				Nind++;

		dims2[0] = 1;
		dims2[1] = Nind;

		// fill in the sparse matrix structure for this beam

		// create the structure that will hold the sparse beamlet
		struc = mxCreateStructArray(1,&one,nfields,field_names);

        // indices field
        ind_field = mxCreateNumericArray(2,dims2,mxINT32_CLASS,mxREAL);

		// data field
		data_field = mxCreateNumericArray(2,dims2,mxSINGLE_CLASS,mxREAL);

		// x, y, and z counts
		dims2[0] = 1;
		dims2[1] = 1;

		field = mxCreateNumericArray(2,dims2,mxINT32_CLASS,mxREAL);
		intPtr = (int *)mxGetData(field);
		intPtr[0] = dims[0];
		mxSetField(struc,0,"x_count",mxDuplicateArray(field));

		field = mxCreateNumericArray(2,dims2,mxINT32_CLASS,mxREAL);
		intPtr = (int *)mxGetData(field);
		intPtr[0] = dims[1];
		mxSetField(struc,0,"y_count",mxDuplicateArray(field));

		field = mxCreateNumericArray(2,dims2,mxINT32_CLASS,mxREAL);
		intPtr = (int *)mxGetData(field);
		intPtr[0] = dims[2];
		mxSetField(struc,0,"z_count",mxDuplicateArray(field));

		intPtr = (int *)mxGetData(ind_field);
		floatPtr = (float *)mxGetData(data_field);

		// extract the sparse data
		j = 0;  // index just for the sparse data
		for (i=0;i<Ntotal;i++)
			if (dose.matrix[i] > 0.0)
			{
				intPtr[j] = i + 1;
				floatPtr[j] = dose.matrix[i];
				j++;
			}
		
		mxSetField(struc,0,"non_zero_indices",mxDuplicateArray(ind_field));
		mxSetField(struc,0,"non_zero_values",mxDuplicateArray(data_field));

		mxSetCell(BEAMLETS,b,mxDuplicateArray(struc));
		printf("Finished calculating dose with beam %d\n",b+1);

		// reset the matrices to zeros
		for (i=0;i<Ntotal;i++)
		{
			terma_mask.matrix[i] = 0.0;
			dose_mask.matrix[i] = 0.0;
			deff.matrix[i] = 0.0;
			terma.matrix[i] = 0.0;
			kermac.matrix[i] = 0.0;
			dose.matrix[i] = 0.0;
		}
	} 

	// destroy the arrays created with mxCreate that are not output matrices
	mxDestroyArray(TERMA);
	mxDestroyArray(KERMAC);
	mxDestroyArray(DEFF);
	mxDestroyArray(TERMA_MASK);
	mxDestroyArray(DOSE_MASK);
	mxDestroyArray(DOSE); 

	// free the kernels
	free(poly_kernel.angular_boundary);
	free(poly_kernel.radial_boundary);
	// free(poly_kernel.total_matrix);
	for (j=0;j<N_KERNEL_CATEGORIES;j++)
		free(poly_kernel.matrix[j]);

	// cannot free mono_kernels because they are part of a Matlab input structure
}