/* lsq is a linear least squares optimization algorithm for finding
optimal energy fluence weighting factors for radiation therapy.  Tomotherapy
has a patent on the algorithm.  The algorithm is explained in detail by
Dave Shepard in the publication: PMB 45 (2000) 69-90.  */

#include "lsq.h" 

char errstr[200];  // globally available string for passing errors
int firstIterFlag;  // flag that is activated only on the first iteration

#define OPT_INPUT_FILE "optInput.txt"

int main(int argc, char *argv[])
//  Expects one or zero command line input arguments. If a lsq header file is not 
//  included as an input argument, then OPT_INPUT_FILE is used.
{
 float weight_diff, start_time, end_time, update_factor;
 float clockStart, clockStop;

 int *usedVoxels;

 int i, j, iter_num;  

 // filename strings
 char opt_input_file[200];

 float *E;   // objective function values

 char doseFileStr[200];
 char weightFileStr[200];

 // structures for the optimization
 PRESC presc;
 PRESCGRID prescGrid;
 OPT_PARAM op;
 GRID dose;
 GRID doseTot;
 GRID_ARRAY beamlets;

 // dummy file
 FILE *fid;

 // check to see if an input argument was supplied
 if (argc > 2)  // calling the executable counts as one argument
 {
	printf("Expecting one or zero input command line arguments, received %d.\n",argc);
	return(FAILURE);
 }

 if (argc == 1)  // only the executable was called, no other arguments
	 strcpy(opt_input_file,OPT_INPUT_FILE);
 else  // one arguments was supplied
	 strcpy(opt_input_file,argv[1]);

 // read in optimization parameters
 if (read_in_opt_parameters(&op, opt_input_file) == FAILURE)
 {
	 printf("Failed to read in optimization properties:\n%s\n",errstr);
	 return(FAILURE);
 }
 else
	 printf("Successfully read-in optimization parameters.\n");

 // print optimization parameters for verification of readin
 printf("Niterations = %d\n",op.Niterations);
 printf("prescription_filename = %s\n",op.prescription_filename);
 printf("initial_beam_weights_filename = %s\n",op.initial_beam_weights_filename);
 printf("beamlet_header_file = %s\n",op.beamlet_header_file);
 printf("dose_batch_base_name = %s\n",op.dose_batch_base_name);
 printf("dose_batch_extension = %s\n",op.dose_batch_extension);
 printf("weight_batch_base_name = %s\n",op.weight_batch_base_name);
 printf("weight_batch_extension = %s\n",op.weight_batch_extension);
 printf("modFactor = %d\n",op.modFactor);

 // read in the prescription
 if (read_in_prescription(&presc, op.prescription_filename) == FAILURE)
 {
	 printf("Failed to read in prescription:\n%s\n",errstr);
	 return(FAILURE);
 }
 else
	 printf("Successfully read-in prescription.\n");

 fillPrescriptionGrid(&presc, &prescGrid);

 if ( (usedVoxels = (int *)malloc(sizeof(int)*presc.Nvox)) == NULL)
 {
   printf("Memory allocation for usedVoxels array failed.\n");
   return(FAILURE);
 }

 fid = fopen("prescGridTest.bin","wb");
 fwrite(prescGrid.numNodes,sizeof(int),prescGrid.Nvox,fid);
 fclose(fid);

 if (SMARTBEAMLETS == 1)
	 find_used_voxels(&presc,usedVoxels);
 else
	 for (i=0;i<presc.Nvox;i++) usedVoxels[i] = 1;
		 
 // count number of used voxels
 i = 0;
 for (j=0;j<presc.Nvox;j++)
	 if (usedVoxels[j] == 1)
		 i++;

 printf("total number of used voxels = %d\n",i);

  // allocate memory for dose matrix
 if ((dose.matrix = (float *)malloc(sizeof(float)*presc.Nvox)) == NULL) 
 {
   printf("Memory allocation for dose matrix failed.\n");
   return(FAILURE);
 }

 dose.x_count = presc.x_count;
 dose.y_count = presc.y_count;
 dose.z_count = presc.z_count;
 dose.Nvox = presc.Nvox;

 // allocate memory for the total dose matrix
 if ((doseTot.matrix = (float *)malloc(sizeof(float)*presc.Nvox)) == NULL) 
 {
   printf("Memory allocation for dose matrix failed.\n");
   return(FAILURE);
 }
 
 doseTot.x_count = presc.x_count;
 doseTot.y_count = presc.y_count;
 doseTot.z_count = presc.z_count;
 doseTot.Nvox = presc.Nvox;

 // load the beamlets into memory, calculate full initial guess
 if (read_in_beamlets(&beamlets,op.beamlet_header_file,op.initial_beam_weights_filename,usedVoxels,&op,&doseTot) == FAILURE)
 {
	 printf("Read in beamlets failed: %s\n",errstr);
	 return(FAILURE);
 }

 // ensure that dimensions of beamlets and prescription masks are consistent
 if (beamlets.x_count != presc.x_count 
  || beamlets.y_count != presc.y_count 
  || beamlets.z_count != presc.z_count)
 {
	 printf("Prescription dimensions (%d,%d,%d) do not match beamlet dimensions (%d,%d,%d)\n",
		 presc.x_count,presc.y_count,presc.z_count,
		 beamlets.x_count,beamlets.y_count,beamlets.z_count);
	 return(FAILURE);
 }

 // calculate initial dose for the purposes of the optimization
 clockStart = clock()/(float)CLOCKS_PER_SEC;
 calculate_dose(&beamlets, &dose);
 clockStop = clock()/(float)CLOCKS_PER_SEC;

 printf("initial dose successfully calculated\n");
 printf("Initial dosecalc time = %f\n",clockStop - clockStart);
 calcAllDvhs(&dose, &presc);       // recalculate dvhs
 printf("initial dvhs successfully calculated\n");

 // start the clock
 start_time = clock()/(float) CLOCKS_PER_SEC;
 printf ("Starting iterations\n" );
 
 E = (float *)malloc(sizeof(float)*op.Niterations+1);

 // write the initial guess to a file
 sprintf(doseFileStr,"%s%d.%s",op.dose_batch_base_name,0,op.dose_batch_extension);
 fid = fopen(doseFileStr,"wb");
 if (SMARTBEAMLETS == 1)
	fwrite(doseTot.matrix,sizeof(float),beamlets.Nvox,fid);
 else
	fwrite(dose.matrix,sizeof(float),beamlets.Nvox,fid);

 fclose(fid);

// write initial weights guess to a file
 sprintf(weightFileStr,"%s%d.%s",op.weight_batch_base_name,0,op.weight_batch_extension);
 fid = fopen(weightFileStr,"wb");
 fwrite(beamlets.beam_weight,sizeof(float),beamlets.num_beamlets,fid);
 fclose(fid);

 // calculate the initial objective function value
 E[0] = calc_obj_func(&beamlets,&dose,&presc);

 printf("Initial objective function value = %f\n",E[0]);

 // calculate one optimization iteration
 for (iter_num = 0; iter_num < op.Niterations; iter_num++)
 {
	 if (iter_num == 0)
		 firstIterFlag = 0;
	 else
		 firstIterFlag = 0;
	
   for (j=0;j<beamlets.num_beamlets;j++)
	if (beamlets.beam_weight[j] != 0.0)  // proceed only if the beam is not turned off
	{
		if (calc_update_factor(&beamlets, &dose, &presc, j, &update_factor, &prescGrid) == FAILURE)
		{
			printf("Failed in update factor calculation for iteration %d, beamlet number %d:\n%s\n",iter_num,j,errstr);
			return(FAILURE);
		}

		// difference between new and old weights
		weight_diff = (update_factor - (float)1.0)*beamlets.beam_weight[j];
		// new beamlet weight
		beamlets.beam_weight[j] *= update_factor;
	}
	
	// apply modulation factor 
    truncateBeamWeights(beamlets.beam_weight,beamlets.num_beamlets,op.modFactor);

	calculate_dose(&beamlets, &dose); // recalculate the dose distribution
	calcAllDvhs(&dose, &presc);       // recalculate dvhs

	// calculate the objective function
	E[iter_num+1] = calc_obj_func(&beamlets,&dose,&presc);
	printf("Objective function for iteration %d of %d was %f\n",iter_num+1,op.Niterations,E[iter_num+1]);
	
	if ((int)fmod((double)iter_num+1.0,op.Nperbatch) == 0 || iter_num + 1 == op.Niterations)
	{
		printf("Printing dose and weights.\n");
		
		if (SMARTBEAMLETS == 1)
		{
			printf("Loading full beamlets for dose calculation ...\n");
			// load the beamlets and calculate dose
			if (loadBeamletsCalcDose(&op, &doseTot, &beamlets) != SUCCESS)
			{
				printf("Failed to calculate total dose for output of results: %s\n",errstr);
				return(FAILURE);
			}
		}

		// create dose filename
		sprintf(doseFileStr,"%s%d.%s",op.dose_batch_base_name,iter_num+1,op.dose_batch_extension);
		fid = fopen(doseFileStr,"wb");
		if (SMARTBEAMLETS == 1)
			fwrite(doseTot.matrix,sizeof(float),beamlets.Nvox,fid);
		else
			fwrite(dose.matrix,sizeof(float),beamlets.Nvox,fid);
		fclose(fid);

		// create weights filename
		sprintf(weightFileStr,"%s%d.%s",op.weight_batch_base_name,iter_num+1,op.weight_batch_extension);
		fid = fopen(weightFileStr,"wb");
		fwrite(beamlets.beam_weight,sizeof(float),beamlets.num_beamlets,fid);
		fclose(fid);
	}
 } 
 
 end_time = clock()/((float) CLOCKS_PER_SEC);

 // write the objective function vector to a file
 fid = fopen(op.obj_func_name,"wb");
 fwrite(E,sizeof(float),op.Niterations+1,fid);
 fclose(fid);

 free(usedVoxels);

 // free the memory allocated for the dose grid
 free(dose.matrix);
 free(doseTot.matrix);

 for (i=0; i<beamlets.num_beamlets;i++)
 {
	free(beamlets.array[i].indices);
	free(beamlets.array[i].data);
 } 

 free(beamlets.beam_weight);
 free(beamlets.initial_beam_weight);

 // free prescription arrays
 for (i=0;i<presc.Ntissue;i++)
 {
	free(presc.array[i].ind);
	free(presc.array[i].dMinus);
	free(presc.array[i].dPlus);
 }

 free(presc.array);

 // free prescription grid
 for (i=0;i<prescGrid.Nvox;i++)
	 if (prescGrid.numNodes[i] > 0)
		free(prescGrid.array[i]);

 free(prescGrid.array);
 free(prescGrid.numNodes);

 printf("Elapsed optimization time %f\n",end_time-start_time);

 // create a file that signals that the optimization is done
 fid = fopen("optDone.txt","w");
 fprintf(fid,"Optimization complete\n");
 fclose(fid);

 return(SUCCESS);
}