package visUtils;

/** Java routines to do various simple tasks
 * surrounding data visualization. To be called
 * in a single short line only.
 */

public class VisUtils {
	
	/** 
	  * given a sample array of free parameters, this returns the log of the
	  * sample density at each sample point (relative to the parameter space)
	  * 
	  * @param samples: 1st dim -> sample index; 2nd dim -> free parameter index 
	  * @param regularizer: width of heat kernel used to turn samples into
	  * 					density
	  * @return array corresponding to the sample index with values of the density
	  * 		at each sample point (in the parameter space). The final scale is
	  * 		linearly normalized so that the max value is one and the min value
	  * 		is 0.
	  */
	 public final static double[] sampleLogDensity(double samples[][]) {
		 int numberOfSamples = samples.length, numberOfParameters = samples[0].length;
		 double dummy = 0, densityMaximum = 0, densityMinimum = 0;
		 double[] sampleDensity=new double[numberOfSamples];
			
		 // Zero all the posterior values
		for (int i=0;i<numberOfSamples;i++){
			sampleDensity[i]=0;			
		}
			
		// Here we use a Gaussian kernel density function to turn the sample data into
		// a sample density. No underlying grid needs to be used to perform
		// this calculation.
		for (int l=0;l<numberOfSamples;l++){
			for (int i=l+1;i<numberOfSamples;i++){
				dummy = 0;
				for (int k=0;k<numberOfParameters;k++)
					dummy+=Math.pow((samples[l][k]-samples[i][k]),2);
				sampleDensity[i]+=dummy;
				sampleDensity[l]+=dummy;
			}
		}
			
		// Normalise the posterior values to have a max of one
		densityMinimum = sampleDensity[0];
		for (int i=0;i<numberOfSamples;i++){
			if(sampleDensity[i]>densityMaximum)
				densityMaximum=sampleDensity[i];
			if(sampleDensity[i]<densityMinimum)
				densityMinimum=sampleDensity[i];
		}
		
		for (int i=0;i<numberOfSamples;i++)
			sampleDensity[i]=(sampleDensity[i]-densityMaximum)/(densityMinimum-densityMaximum);
		
		return sampleDensity;
	 }
	 
	 /** 
	  * given a sample array of free parameters, this returns the
	  * sample density at each sample point (relative to the parameter space)
	  * 
	  * @param samples: 1st dim -> sample index; 2nd dim -> free parameter index 
	  * @return array corresponding to the sample index with values of the density
	  * 		at each sample point (in the parameter space). The final scale is
	  * 		linearly normalized so that the max value is one and the min value
	  * 		is 0.
	  */
	 public final static double[] sampleDensity(double samples[][]) {
		 int numberOfSamples = samples.length, numberOfParameters = samples[0].length;
		 double dummy = 0, densityMaximum = 0, densityMinimum = 0;
		 double[] sampleDensity=new double[numberOfSamples];
			
		 // Zero all the posterior values
		for (int i=0;i<numberOfSamples;i++){
			sampleDensity[i]=0;			
		}
			
		// Here we use a Gaussian kernel density function to turn the sample data into
		// a sample density. No underlying grid needs to be used to perform
		// this calculation.
		for (int l=0;l<numberOfSamples;l++){
			for (int i=l+1;i<numberOfSamples;i++){
				dummy = 0;
				for (int k=0;k<numberOfParameters;k++)
					dummy+=Math.pow((samples[l][k]-samples[i][k]),2);
				sampleDensity[i]+=dummy;
				sampleDensity[l]+=dummy;
			}
		}
		dummy=0;
		for (int l=0;l<numberOfSamples;l++){
			dummy+=sampleDensity[l];
		}
		dummy=dummy/numberOfSamples;
		
		for (int l=0;l<numberOfSamples;l++){
			sampleDensity[l]=Math.exp(-sampleDensity[l]/(2*dummy));
		}
			
		// Normalise the posterior values to have a max of one
		densityMinimum = sampleDensity[0];
		densityMaximum = sampleDensity[0];
		for (int i=0;i<numberOfSamples;i++){
			if(sampleDensity[i]>densityMaximum)
				densityMaximum=sampleDensity[i];
			if(sampleDensity[i]<densityMinimum)
				densityMinimum=sampleDensity[i];
		}
		
		if (densityMinimum == densityMaximum)
		      throw new RuntimeException("Sample density function is constant and not able to be normalized!");
		
		for (int i=0;i<numberOfSamples;i++)
			sampleDensity[i]=sampleDensity[i]-densityMinimum/(densityMaximum-densityMinimum);
		
		return sampleDensity;
	 }
	 
	 /**
	  * This sorts a given array of doubles and returns the index
	  * remapping. That is the original index of the value placed
	  * in the ith position in the sort is returned, NOT the original
	  * value
	  * 
	  * @param sortArray array to be sorted
	  * @return returns the index remapping that occurred to
	  * 		resort the array.
	  */
	 
	 public final static int[] sortIndex(double sortArray[]) {
		 double[][] indexArray = new double[sortArray.length][2];
		 int[] sortedIndex = new int[sortArray.length];
		 for (int i = 0; i<sortArray.length; i++) {
			 indexArray[i][0]=sortArray[i];
			 indexArray[i][1]=i;
		 }
		 QuickSort(indexArray,0,indexArray.length-1);
		 
		 for (int i = 0; i<sortArray.length; i++) {
			 sortedIndex[i]=(int) indexArray[i][1];
		 }
		 
		 return sortedIndex;
	 }
	 
	 // Standard quicksort
	 private static void QuickSort(double a[][], int l, int r) { 
		 int M = 4; 
		 int i; 
		 int j; 
		 double v; 
		 if ((r-l)>M) { 
			 i = (r+l)/2; 
			 if (a[l][0]>a[i][0]) swap(a,l,i); // Tri-Median Methode! 
			 if (a[l][0]>a[r][0]) swap(a,l,r); 
			 if (a[i][0]>a[r][0]) swap(a,i,r); 
			 j = r-1; 
			 swap(a,i,j); 
			 i = l; 
			 v = a[j][0]; 
			 for(;;) { 
				 while(a[++i][0]<v); 
				 while(a[--j][0]>v); 
				 if (j<i) break; 
				 swap (a,i,j); 
			 } 
			 swap(a,i,r-1); 
			 QuickSort(a,l,j); QuickSort(a,i+1,r); 
		 }
	 }
	 
	 // value swap for in-place quicksort
	 private static void swap(double a[][], int i, int j) { 
		 double[] T = new double[2]; 
		 T[0] = a[i][0];
		 T[1] = a[i][1];
		 a[i][0] = a[j][0];
		 a[i][1] = a[j][1];
		 a[j][0] = T[0];
		 a[j][1] = T[1];
	 } 

	/**
	  * For a given array returns an rgb color array that can be 
	  * imported into matlab. This particular mapping corresponds
	  * to the default matlab color map
	  * 
	  * @param arr for which the matlab RGB color array is constructed
	  * @return an array of rgb values for each value of arr, the first
	  * 		index corresponds to the index of arr
	  */
	 
	 public final static double[][] matlabRGBMatrix(double arr[]) {
		 
		 // Find min and max of array
		 double min,max;
		 min = arr[0];
		 max = arr[0];
		 
		 for (int i=1;i<arr.length;i++) {
			 if (arr[i]>max)
				 max=arr[i];
			 if (arr[i]<min)
				 min=arr[i];
		 }
		 
		 // Normalize via linear mapping so that min -> 0, max -> 1
		 for (int i=1;i<arr.length;i++)
			 arr[i]=(arr[i]-min)/(max-min);
		 
		 // Now we construct the colour function for the multi-line plot visualisation
		double[][] rgbMatrix = new double[arr.length][3];
		
		// This specific colour function is taken from the default height
		// colour function in MATLAB for surface plots.
		for (int i=0;i<arr.length;i++) {
			arr[i]=3.9392*arr[i]; // This scaling makes the following code
												// easier to read
			if(arr[i]<=.5) {
				rgbMatrix[i][0]=.5+arr[i];
				rgbMatrix[i][1]=0;
				rgbMatrix[i][2]=0;
			}
			else if(arr[i]<=1.5) {
				rgbMatrix[i][0]=1;
				rgbMatrix[i][1]=(arr[i]-.5);
				rgbMatrix[i][2]=0;
			}
			else if(arr[i]<=2.5) {
				rgbMatrix[i][0]=((2.5-arr[i]));
				rgbMatrix[i][1]=1;
				rgbMatrix[i][2]=((arr[i]-1.5));
			}
			else if(arr[i]<=3.5) {
				rgbMatrix[i][0]=0;
				rgbMatrix[i][1]=((3.5-arr[i]));
				rgbMatrix[i][2]=1;
			}
			else{
				rgbMatrix[i][0]=0;
				rgbMatrix[i][1]=0;
				rgbMatrix[i][2]=((4.5-arr[i]));
			}
		}
		
		return rgbMatrix;
	 }
	 
	 /**
	  * This function aides in the plotting of a density profile on a 2D grid defined
	  * by xGrid and yGrid. Basically, this function samples at even intervals along
	  * a curve defined by 'line'. At each sample point a heat kernel is applied with
	  * a width determined by the length between samples. The heat kernel's value is
	  * then summed at each point on the xGrid. Then the histogram values are normalized
	  * to 1 along angular poloidal slices eminating from centre.
	  * @param line defines the curve that will be sampled on
	  * @param xGrid defines 2D grid
	  * @param yGrid defines 2D grid
	  * @param centre centre of poloidal slice normalization
	  * @param numberLinePoints number of samples on the curve
	  * @param numberPoloidalSlices number of angular slices to normalize on
	  * @return 2D array of values reflecting the density plot on the 2D grid corresponding
	  * 		to the line plot.
	  */
	 public final static double[][] lineDensityPlot(double line[][], double xGrid[], double yGrid[], double centre[], int numberLinePoints, int numberPoloidalSlices) {		
		 double[] lineIncrementalLength = new double[line[0].length];

		// First, we calculate the total length of the line sample in configuration
		// space.
		
		lineIncrementalLength[0]=0;
		for (int i = 1; i<line[0].length;i++) {
			lineIncrementalLength[i]=lineIncrementalLength[i-1];
			lineIncrementalLength[i]+=Math.sqrt(Math.pow(line[0][i]-line[0][i-1],2)+Math.pow(line[0][i]-line[0][i-1],2));
		}
		
		double[][] lineResamples = new double[2][numberLinePoints];
		double currentLength = 0, currentInterp = 0;
		int currentInnerIndex = 0;
		
		double lcfsLineIncrement = lineIncrementalLength[lineIncrementalLength.length-1]/(numberLinePoints-1);
		for (int i=0;i<lineResamples[0].length;i++) {
			currentLength = i*lcfsLineIncrement;
			for (int j=currentInnerIndex;j<lineIncrementalLength.length-1;j++) {
				if (currentLength>=lineIncrementalLength[j] && currentLength<lineIncrementalLength[j+1]) {
					currentInterp = (currentLength-lineIncrementalLength[j])/(lineIncrementalLength[j+1]-lineIncrementalLength[j]);
					lineResamples[0][i] = line[0][j+1]*currentInterp+line[0][j]*(1-currentInterp);
					lineResamples[1][i] = line[1][j+1]*currentInterp+line[1][j]*(1-currentInterp);
					currentInnerIndex=j;
					break;
				}
			}
		}
		
		int plotXGridSize = xGrid.length, plotYGridSize = yGrid.length;
		
		// set the array to hold posterior values (again histograms on x-slices)		
		double[][] densityValues = new double[plotXGridSize][plotYGridSize];
				
		double[][] polodialAngle = new double[plotXGridSize][plotYGridSize];
		                                
		// construct the 2D grid of the configuration space
		for (int i=0; i<plotXGridSize; i++) {
			for (int j=0; j<plotYGridSize; j++) {
				densityValues[i][j]=0;
				polodialAngle[i][j]=Math.atan(Math.abs(centre[0]-xGrid[i])/Math.abs(centre[1]-yGrid[j]));
				// Pick the branch for atan now
				if (xGrid[i]<centre[0] && yGrid[j]>=centre[1])
					polodialAngle[i][j]+=Math.PI/2;
				else if (xGrid[i]<centre[0] && yGrid[j]<=centre[1])
					polodialAngle[i][j]+=Math.PI;
				else if (xGrid[i]>=centre[0] && yGrid[j]<centre[1])
					polodialAngle[i][j]+=3*Math.PI/2;
			}
		}
		
		// Here we use a Gaussian kernel density function to turn the sample data into
		// a posterior density on a x-slice. 
		double currentYSampleValue;
		double currentXSampleValue;

		for(int i=0;i<numberLinePoints;i++) {
			currentYSampleValue=lineResamples[1][i];
			currentXSampleValue=lineResamples[0][i];
			for (int j=0; j<plotXGridSize; j++) {
				for (int k=0; k<plotYGridSize; k++)
					densityValues[j][k] += Math.exp(-((currentYSampleValue-yGrid[k])*(currentYSampleValue-yGrid[k])+(currentXSampleValue-xGrid[j])*(currentXSampleValue-xGrid[j]))/(.5*lcfsLineIncrement*lcfsLineIncrement));
			}
		}
		
		// determine the maximum density on each polodial-slice
		for (int k=0;k<numberPoloidalSlices;k++) {
			double maxSliceValue=-1;
			for (int i=0; i<plotXGridSize; i++) {
				for (int j=0; j<plotYGridSize; j++) {
					if (densityValues[i][j]>maxSliceValue && (polodialAngle[i][j]>=2*k*Math.PI/(numberPoloidalSlices) && polodialAngle[i][j]<2*(k+1)*Math.PI/(numberPoloidalSlices)))
						maxSliceValue=densityValues[i][j];
				}
			}
			for (int i=0; i<plotXGridSize; i++) {
				for (int j=0; j<plotYGridSize; j++) {
					if (polodialAngle[i][j]>=2*k*Math.PI/(numberPoloidalSlices) && polodialAngle[i][j]<2*(k+1)*Math.PI/(numberPoloidalSlices))
						densityValues[i][j]=densityValues[i][j]/maxSliceValue;
				}
			}
		}
		
		return densityValues;
	 }
}
