/* lsq_util.c */

// Least squares utilities.

#include "lsq.h"

extern char errstr[200];  // error string

int calculate_dose_diff(GRID_ARRAY *beamlets, float weight_diff, int weight_bixel_num, GRID *dose)
{
// only change the summed dose by the dose change of a single beamlet
 static int i;

 // dose calculation with sparse beamlet matrices
 for (i=0; i<beamlets->array[weight_bixel_num].size; i++)
	 dose->matrix[beamlets->array[weight_bixel_num].indices[i]] 
		 += weight_diff*beamlets->array[weight_bixel_num].data[i];
 return(SUCCESS);
}

// put a full dose calculation file here
int calculate_dose(GRID_ARRAY *beamlets, GRID *dose)
// do a full dose calculation
{
	static int i,j;

	// clear the dose grid
	for (i=0;i<beamlets->Nvox;i++)
		dose->matrix[i] = 0.0;

	// dose calculation with sparse beamlet matrices
	for (j=0;j<beamlets->num_beamlets;j++)  // loop through beams
		for (i=0;i<beamlets->array[j].size;i++)  // loop through non-zero dose voxels
			dose->matrix[beamlets->array[j].indices[i]]
			+= beamlets->beam_weight[j]*beamlets->array[j].data[i];

	return(SUCCESS);
}

int calc_dvh(GRID *dose, int Nind, int *ind, DVH *dvh)
// Calculates a dose volume histogram using the dose grid based on the parameters in DVH.
// Memory for the dvh array must be pre-allocated before calling this function.
// The resulting DVH is given in units of percent volume [0..100].
{
	int i,n,m;

	// fill the dvh with zeros
	for (m=0;m<dvh->Nbins;m++)
		dvh->dvh[m] = 0.0;
		
	// quit if del_d is less than or equal to zero
	if (dvh->del_d <= 0.0)
	{
		sprintf(errstr,"del_d was less than or equal to zero");
		printf("del_d error\n");
		return(FAILURE);
	}

	// calculate the DVH
	for (n=0;n<Nind;n++)
	{
		// dose index in the dose grid
		i = ind[n];
		// find the dose index in the dvh
		m = (int)(dose->matrix[i]/dvh->del_d);
		if (m >= 0 && m < dvh->Nbins)
			dvh->dvh[m] += (float)1.0;
	}
		
	// integrate the dvh from d to infinity for all d, then normalize to 100% at d=0
	for (m=dvh->Nbins-2;m>=0;m--)
		dvh->dvh[m] += dvh->dvh[m+1];

	// normalization step
	if (Nind != 0)  // avoid divide-by-zero errors
		for (m=0;m<dvh->Nbins;m++)
			dvh->dvh[m] *= (float)100.0/(float)Nind;
	else  // no voxels, but 100% still gets 0 Gy or more
		dvh->dvh[0] = 100;

	return(SUCCESS);
}

int output_dose_weights(GRID_ARRAY *beamlets, GRID *dose, 
		                char *dosefilename, char *weightfilename)
// Writes the current dose distribution and beamlet weights to dosefilename
// and weightfilename, respectively.
{
	FILE *dosefile, *weightfile;

	dosefile = fopen(dosefilename,"wb");
	weightfile = fopen(weightfilename,"wb");

	// print dose and weight files
	if ((int)fwrite(dose->matrix,sizeof(float),beamlets->Nvox,dosefile) != beamlets->Nvox)
	{
		sprintf(errstr,"Unable to write dose data to file \n%s.",dosefilename);
		return(FAILURE);
	}

	if ((int)fwrite(beamlets->beam_weight,sizeof(float),beamlets->num_beamlets,weightfile) != beamlets->num_beamlets)
	{
		sprintf(errstr,"Unable to write beamlet weight data to file \n%s.",weightfilename);
		return(FAILURE);
	}
	fclose(dosefile);
	fclose(weightfile);
	return(SUCCESS);
}

float calc_obj_func(GRID_ARRAY *beamlets, GRID *dose, PRESC *presc)
// Calculate the objective function value.
{
	int i, j, k;
	float E, Eplus, Eminus, Edvh_plus, Edvh_minus;
	float diff, del_d_plus, del_d_minus;

	E = 0.0;

	for (k=0;k<presc->Ntissue;k++)  // loop through all tissue types
	 if (presc->array[k].Nind != 0 && presc->array[k].alpha > 0.0)  
	 // continue only if the ROI contains voxels and has an importance above zero
	 {
		// initialize objective function terms for tissue k to zero
		Eplus = 0.0;
		Eminus = 0.0;
		Edvh_plus = 0.0;
		Edvh_minus = 0.0;

		// calculate the over- and underdose penalty terms
		for (j=0;j<presc->array[k].Nind;j++)
		{
			i = presc->array[k].ind[j];   // current voxel index

			// overdose penalty term
			diff = dose->matrix[i] - presc->array[k].dPlus[j];
			if (diff > 0.0)
				Eplus += diff*diff;

			// underdose penalty term
			diff = dose->matrix[i] - presc->array[k].dMinus[j];
			if (diff < 0.0)
				Eminus += diff*diff;
		}

		// calculate the dvh penalties if they apply
		if ((presc->array[k].betaVPlus > 0.0) || (presc->array[k].betaVMinus > 0.0))
		{
			if (presc->array[k].betaVPlus > 0.0) 
			// penalize overdosed dose volume considerations
			{
				del_d_plus = presc->array[k].dvh.del_d_plus;

				// dvh_plus constraints
				for (j=0;j<presc->array[k].Nind;j++)
				{
					i = presc->array[k].ind[j];  // current voxel of interest
					diff = dose->matrix[i] - presc->array[k].dVPlus; 
					if (diff >= 0 && diff <= del_d_plus)  // see if this voxel should be penalized
						Edvh_plus += diff*diff;
				}
			}

			if (presc->array[k].betaVMinus > 0.0)
			// penalize underdosed dose volume considerations
			{
				del_d_minus = presc->array[k].dvh.del_d_minus;
			
				// dvh_minus constraints
				for (j=0;j<presc->array[k].Nind;j++)
				{
					i = presc->array[k].ind[j];  // current voxel of interest
					diff = dose->matrix[i] - presc->array[k].dVMinus; 
					if (diff >= del_d_minus && diff <= 0)  // see if this voxel should be penalized
						Edvh_minus += diff*diff; 
				}
			}
		}

		// update the objective function
		E += presc->array[k].alpha/(float)presc->array[k].Nind
			    *(presc->array[k].betaPlus*Eplus
				+ presc->array[k].betaMinus*Eminus
				+ presc->array[k].betaVPlus*Edvh_plus
				+ presc->array[k].betaVMinus*Edvh_minus);
	 }
	 return(E);
} 

int compare(void *arg1, void *arg2)
// Compare two floats to each other.
{
   if (*(float *)arg1 < *(float *)arg2)
	   return -1;
   else if (*(float *)arg1 == *(float *)arg2)
	   return 0;
   else
	   return 1;
}

int calcAllDvhs(GRID *dose, PRESC *presc)
// Calculates the dvh for each tissue type in the prescription
{
	int i, j, k;
	float d_max;
	char tmpstr[200];

	for (k=0;k<presc->Ntissue;k++)
	{
		// find the current maximum dose in the ROI
		d_max = 0.0;
		for (j=0;j<presc->array[k].Nind;j++)
		{
			i = presc->array[k].ind[j];  // current voxel index
			if (dose->matrix[i] > d_max)
				d_max = dose->matrix[i];
		}

		// set up dose bins for the current dvh
		presc->array[k].dvh.Nbins = Ndvh_bins;
		presc->array[k].dvh.dmax = d_max*(float)1.1;  // go over the max dose by a small amount
		if (d_max != 0.0)  // zero dose delivered to this tissue
			presc->array[k].dvh.del_d = presc->array[k].dvh.dmax/(float)presc->array[k].dvh.Nbins;
		else
			presc->array[k].dvh.del_d = (float)100.0/(float)presc->array[k].dvh.Nbins;

		// calculate the dvh for this tissue
		if (calc_dvh(dose, presc->array[k].Nind, presc->array[k].ind, &presc->array[k].dvh) == FAILURE)
		{
			strcpy(tmpstr,errstr);  // copy the error string passed from the dvh calculator
			sprintf(errstr,"Failed in DVH calculation for tissue %d:\n%s\n",k,tmpstr);
			return(FAILURE);
		} 

		// if applicable in the optimization, find del_d_plus and del_d_minus
		if (presc->array[k].betaVPlus > 0.0) 
		{
			// find where vPlus occurs in the DVH
			for (j=0;j<presc->array[k].dvh.Nbins;j++)
				if (presc->array[k].dvh.dvh[j] < presc->array[k].vPlus)
					break;

			// j is the dvh index that marks the crossing from above to below vPlus

			// use j to see if any dvh penalties will be applied here
			if ((float)j*presc->array[k].dvh.del_d >= presc->array[k].dVPlus)
				presc->array[k].dvh.del_d_plus 
					   = (float)j*presc->array[k].dvh.del_d - presc->array[k].dVPlus;
			else
				presc->array[k].dvh.del_d_plus = 0.0;
		}

		if (presc->array[k].betaVMinus > 0.0)
		{
			// find where vMinus occurs in the DVH
			for (j=0;j<presc->array[k].dvh.Nbins;j++)
				if (presc->array[k].dvh.dvh[j] < presc->array[k].vMinus)
					break;

			// j is the dvh index that marks the crossing from above to below vMinus

			// use j to see if any dvh penalties will be applied here
			if ((float)j*presc->array[k].dvh.del_d <= presc->array[k].dVMinus)
				presc->array[k].dvh.del_d_minus
				   = (float)j*presc->array[k].dvh.del_d - presc->array[k].dVMinus;
			else
				presc->array[k].dvh.del_d_minus = 0.0;
		}
	}
	return(SUCCESS);
}

int find_used_voxels(PRESC *presc, int *usedVoxels)
// Finds all voxels that are affected by the prescription.  Voxels that are
// affected are marked as 1, those that aren't are marked as 0.
{
	int i,j,k;

	// voxels are unused until proven otherwise
	for (i=0;i<presc->x_count*presc->y_count*presc->z_count;i++)
		usedVoxels[i] = 0;

	// loop through all tissues, finding voxels that are used
	for (k=0;k<presc->Ntissue;k++)
		if (   (presc->array[k].alpha > 0.0)
			&& (   presc->array[k].betaMinus > 0.0 || presc->array[k].betaPlus
			    || presc->array[k].betaVMinus || presc->array[k].betaVPlus       ))
			for (j=0;j<presc->array[k].Nind;j++)
				usedVoxels[presc->array[k].ind[j]] = 1;

	return(SUCCESS);
}

int fillPrescriptionGrid(PRESC *presc, PRESCGRID *prescGrid)
{
	int i,j,k;
	int *currTiss;

	prescGrid->x_count = presc->x_count;
	prescGrid->y_count = presc->y_count;
	prescGrid->z_count = presc->z_count;
	prescGrid->Nvox = presc->x_count*presc->y_count*presc->z_count;
	
	prescGrid->array = (PRESCNODE **)malloc(sizeof(PRESCNODE *)*prescGrid->Nvox);
	prescGrid->numNodes = (int *)malloc(sizeof(int)*prescGrid->Nvox);

	for (i=0;i<prescGrid->Nvox;i++) prescGrid->numNodes[i] = 0;

	// Count the number of tissues that fall inside each voxel
	for (k=0;k<presc->Ntissue;k++)
		if (   (presc->array[k].alpha > 0.0)
		    && (   presc->array[k].betaMinus > 0.0 || presc->array[k].betaPlus
			    || presc->array[k].betaVMinus || presc->array[k].betaVPlus       ))
			for (j=0;j<presc->array[k].Nind;j++)
				prescGrid->numNodes[presc->array[k].ind[j]]++;

	// allocate memory for the prescription grid nodes
	for (i=0;i<prescGrid->Nvox;i++)
		if (prescGrid->numNodes[i] > 0)
			prescGrid->array[i] = (PRESCNODE *)malloc(sizeof(PRESCNODE)*prescGrid->numNodes[i]);

	// tally of the current number of tissues that have been placed in a voxel
	currTiss = (int *)malloc(sizeof(int)*prescGrid->Nvox);
	for (i=0;i<prescGrid->Nvox;i++) currTiss[i] = 0;

	// fill the prescription grid nodes
	for (k=0;k<presc->Ntissue;k++)
		if (   (presc->array[k].alpha > 0.0)
			&& (   presc->array[k].betaMinus > 0.0 || presc->array[k].betaPlus
			    || presc->array[k].betaVMinus || presc->array[k].betaVPlus       ))
			for (j=0;j<presc->array[k].Nind;j++)
			{
				i = presc->array[k].ind[j];  // current voxel
				prescGrid->array[i][currTiss[i]].tissueInd = k;
				prescGrid->array[i][currTiss[i]].dMinus = presc->array[k].dMinus[j];
				prescGrid->array[i][currTiss[i]].dPlus = presc->array[k].dPlus[j];
				currTiss[i]++;
			}

	return(SUCCESS);
}

int loadBeamletsCalcDose(OPT_PARAM *op, GRID *dose, GRID_ARRAY *beamlets)
// Loads all beamlets and calculates total dose using the current set of beamlet
// weights.  This function is necessary when the beamlets used for the optimization
// have been scaled down to only include voxels that are used in the optimization.
{
	int i,j,k,m,beamletNum,Nbeamlets,dims[3],Nind,num;
	char bixbatchfile[1000];
	int *ind;
	float *data;
	FILE *fid;

	beamletNum = 0;
	
	// reset the dose distribution to all zeros
	for (i=0;i<beamlets->Nvox;i++) dose->matrix[i] = 0.0;

	// load one beamlet at at a time, calculating total dose along the way
	for (k=0;k<beamlets->num_beamlet_batches;k++)
	{
		// open the beamlet batch file
		sprintf(bixbatchfile,"%s/%s%d.%s",op->bix_dir,op->bix_batch_base_filename,k,op->bix_batch_extension);
		if ((fid = fopen(bixbatchfile,"rb")) == NULL)
		{
			sprintf(errstr,"Unable to open file for beamlet batch %d.\n,filename: %s",k,bixbatchfile);
			return(FAILURE);
		}

		// read-in the number of beamlets that are in this file
		if((int)fread(&Nbeamlets,sizeof(int),1,fid) != 1)
		{
			sprintf(errstr,"Unable to read number of beamlets in file:\n%s",bixbatchfile);
			return(FAILURE);
		}

		// read-in all of the beamlets in the current batch
		for (j=0;j<Nbeamlets;j++)
		{
			if((int)fread(&num,sizeof(int),1,fid) != 1)
			{
				sprintf(errstr,"Unable to read beamlet number for beamlet %d in file:\n%s",j,bixbatchfile);
				return(FAILURE);
			}

			// read the beamlet dimensions
			if((int)fread(dims,sizeof(int),3,fid) != 3)
			{
				sprintf(errstr,"Unable to read beamlet dimensions for beamlet %d in file:\n%s",j,bixbatchfile);
				return(FAILURE);
			}

			// ensure that the beamlet has the same grid size as the prescription
			if((dims[0] != beamlets->x_count) || (dims[1] != beamlets->y_count) || (dims[2] != beamlets->z_count))
			{
				sprintf(errstr,"Mismatch in grid dimensions for beamlet %d in file:\n%s",j,bixbatchfile);
				return(FAILURE);
			}
	
			// read the number of non-zero values in the current beamlet
			if((int)fread(&Nind,sizeof(int),1,fid) != 1)
			{
				sprintf(errstr,"Unable to read number of non-zero values for beamlet %d in file:\n%s",j,bixbatchfile);
				return(FAILURE);
			}

			// allocate memory and read-in indices for the current beamlet
			if((ind = (int *)malloc(sizeof(int)*Nind)) == NULL)
			{
				sprintf(errstr,"Unable to allocate memory for indices of beamlet %d\n",num);
				return(FAILURE);
			}

			if((int)fread(ind,sizeof(int),Nind,fid) != Nind)
			{
				sprintf(errstr,"Unable to read-in indices for beamlet %d\n",num);
				return(FAILURE);
			}

			// allocate memory and read-in indices for the current beamlet
			if((data = (float *)malloc(sizeof(float)*Nind)) == NULL)
			{
				sprintf(errstr,"Unable to allocate memory for data of beamlet %d\n",num);
				return(FAILURE);
			}	

			if((int)fread(data,sizeof(float),Nind,fid) != Nind)
			{
				sprintf(errstr,"Unable to read-in dose data for beamlet %d\n",num);
				return(FAILURE);
			}

			// add the current beamlet to the dose distribution
			for (m=0;m<Nind;m++)
				dose->matrix[ind[m]] += data[m]*beamlets->beam_weight[beamletNum];

			// free indices and data for the full beamlet
			free(ind);
			free(data);

			beamletNum++;  // increment the beamlet number
		}
		
		printf("Successfully loaded beamlet batch %d of %d.\n",k+1,beamlets->num_beamlet_batches);

		fclose(fid);
	}
	return(SUCCESS);
}

int truncateBeamWeights(float *beamWeights, int Nbeamlets, float modFactor)
// Applies modulation factor to all of the beamlet weights, including
// the ones that are not flagged.
{
	int j, Nnonzero;
	float meanInt, maxInt;
	
	Nnonzero = 0;
	maxInt = 0.0;
	meanInt = 0.0;

	for (j=0;j<Nbeamlets;j++)
		if (beamWeights[j] > 0.0)
		{
			Nnonzero++;
			meanInt += beamWeights[j];
		}
	
	meanInt /= (float)Nnonzero;  // calculate the mean of the intensity weights
	
	maxInt = modFactor*meanInt;  // maximum allowable intensity

	// set any weights that are greater than the maximum to the maximum intensity
	for (j=0;j<Nbeamlets;j++)
		if (beamWeights[j] > maxInt)
			beamWeights[j] = maxInt;

	return(SUCCESS);
}