package ie.dcu.image.dt;

import ie.dcu.matrix.*;


/**
 * Compute distance transforms using Meijster's algorithm:
 * 
 * For details on the algorithm see:
 * 
 * <pre>
 *    A General Algorithm for Computing Distance Transforms in Linear Time.
 *    Meijster et al. Computational Imaging and Vision (2000)
 * </pre>
 * 
 * Meijster's algorithm appears to be the most efficient exact Euclidean 
 * distance transform algorithm in most situations. For more details see:
 * 
 * <pre>
 *    2D Euclidean distance transform algorithms: A comparative survey. 
 *    Fabbri et al. ACM Computing Surveys (2008) vol. 40 (1) pp. 1-44
 * </pre>
 * 
 * <b>Note:</b> If the matrix given for transform contains no background
 * pixels, then each pixel will have a value of Integer.MAX_VALUE - rows**2 
 * -cols**2
 * 
 * 
 * @author Kevin McGuinness
 */
public class DistanceTransform {
	private IntMatrix matrix;
	private int inf;
	private boolean done;
	
	/**
	 * Create distance transform object
	 */
	public DistanceTransform() {
		done = false;
	}
	
	/**
	 * Returns true if the computation has been done.
	 */
	public boolean isComputed() {
		return done;
	}
	
	/**
	 * Returns true if the transform has been initialized.
	 */
	public boolean isInitialized() {
		return matrix != null;
	}
	
	
	/**
	 * Private initializer
	 */
	private void init(Matrix m) {
		this.matrix = new IntMatrix(m.rows, m.cols);
		this.inf = Integer.MAX_VALUE - m.rows * m.rows - m.cols * m.cols;
		this.done = false;
	}
	
	/**
	 * Private routine to check if the transform is possible.
	 */
	private boolean isTransformPossible() {
		// Cannot be called when done
		assert (!done);
		
		for (int i = 0; i < matrix.size; i++) {
			if (matrix.values[i] == 0) {
				return true;
			}
		}
		return false;
	}
	
	/**
	 * Initialize transform with an integer matrix.
	 * 
	 * @param m
	 *        An integer matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(IntMatrix m, int foregroundValue) {
		init(m);
		
		for (int i = 0; i < m.size; i++) {
			matrix.values[i] = m.values[i] == foregroundValue ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Initialize transform with an short matrix.
	 * 
	 * @param m
	 *        An short matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(ShortMatrix m, short foregroundValue) {
		init(m);
		
		for (int i = 0; i < m.size; i++) {
			matrix.values[i] = m.values[i] == foregroundValue ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Initialize transform with an byte matrix.
	 * 
	 * @param m
	 *        An byte matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(ByteMatrix m, byte foregroundValue) {
		init(m);
		
		for (int i = 0; i < m.size; i++) {
			matrix.values[i] = m.values[i] == foregroundValue ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Initialize transform with an long matrix.
	 * 
	 * @param m
	 *        An long matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(LongMatrix m, long foregroundValue) {
		init(m);
		
		for (int i = 0; i < m.size; i++) {
			matrix.values[i] = m.values[i] == foregroundValue ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Initialize transform with an double matrix.
	 * 
	 * @param m
	 *        An double matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(DoubleMatrix m, double foregroundValue) {
		init(m);
		
		for (int i = 0; i < m.size; i++) {
			matrix.values[i] = m.values[i] == foregroundValue ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Initialize transform with an float matrix.
	 * 
	 * @param m
	 *        An float matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(FloatMatrix m, float foregroundValue) {
		init(m);
		
		for (int i = 0; i < m.size; i++) {
			matrix.values[i] = m.values[i] == foregroundValue ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Initialize transform with an arbitrary matrix.
	 * 
	 * @param m
	 *        A matrix
	 * @param foregroundValue
	 *        The value of a foreground pixel
	 * @return 
	 *        true if the transform is possible
	 */
	public boolean init(Matrix m, Number foregroundValue) {
		init(m);
		
		double fgv = foregroundValue.doubleValue();
		for (int i = 0; i < m.size; i++) {
			
			matrix.values[i] = m.valueAt(i).doubleValue() == fgv ? inf : 0;
		}
		
		return isTransformPossible();
	}
	
	/**
	 * Performs the column transform
	 */
	private final void transformCols() {
				
		for (int j = 0; j < matrix.cols; j++) {

			// Forward
			for (int i = 1, b = 1; i < matrix.rows; i++) {
				int idx = i * matrix.cols + j;
				
				if (matrix.values[idx] > matrix.values[idx-matrix.cols]) {
					matrix.values[idx] = matrix.values[idx-matrix.cols] + b;
					b += 2;
				} else {
					b = 1;
				}
			}
			
			// Backward
			for (int i = matrix.rows - 2, b = 1; i >= 0; i--) {
				int idx = i * matrix.cols + j;
				
				if (matrix.values[idx] > matrix.values[idx+matrix.cols] + b) {
					matrix.values[idx] = matrix.values[idx+matrix.cols] + b;
					b += 2;
				} else {
					b = 1;
				}
			}
		}
	}
	
	/**
	 * Meijsters f function
	 */
	private static int f(int x, int i, int pixel) {
		return (x - i)*(x - i) + pixel;
	}
	
	/**
	 * Performs the row transform
	 */
	private final void transformRows() {
		// Meijster row transform
		
		int[] s = new int[matrix.cols];
		int[] t = new int[matrix.cols];
		int[] r = new int[matrix.cols];
		
		for (int i = 0; i < matrix.rows; i++) {
			int q = s[0] = t[0] = 0;
			int offset = i * matrix.cols;
			
			for (int u = 1; u < matrix.cols; ++u) {
				int im_r_u = matrix.values[offset + u];
				
				while (q != -1 && 
					f(t[q], s[q], matrix.values[offset + s[q]]) > 
					f(t[q], u, im_r_u)) {
					--q;
				}
				
				if (q == -1) {
					q = 0;
					s[0] = u;
				} else {
					int w = 1 + ((u * u - (s[q] * s[q]) + 
							matrix.values[offset + u] - matrix.values[offset + s[q]]) 
							/ (2 * (u - s[q])));
					if (w < matrix.cols) {
						++q;
						s[q] = u;
						t[q] = w;
					}
				}
			}
			
			System.arraycopy(matrix.values, offset, r, 0, matrix.cols);

		    for (int u = matrix.cols - 1; u != -1; --u) {
		    	matrix.values[offset + u] = f(u, s[q], r[s[q]]);

				if (u == t[q]) {
					--q;
				}
			}
		}
	}
	
	/**
	 * Carries out the transform if it has not already been completed
	 */
	private void squareTransform() {
		if (matrix == null) {
			throw new IllegalStateException("transform object not initialized");
		}
		
		if (!done) {
			if (matrix.rows > 0 && matrix.cols > 0) {
				transformCols();
				transformRows();
			}
			done = true;
		}
	}
	
	/**
	 * Compute the square euclidean distance transform and return the result as
	 * a flattened integer array of square distance values, that has a size of 
	 * rows * columns.
	 * 
	 * @return flat array of square distance values.
	 */
	public IntMatrix computeSquareTransform() {
		squareTransform();
		return matrix.clone();
	}
	
	/**
	 * Compute the euclidean distance transform and return the result as a 
	 * flattened array of double precision floating point values. The returned
	 * array has a dimension of rows * columns.
	 * 
	 * @return flattened array of euclidean distance values.
	 */
	public DoubleMatrix computeTransform() {
		squareTransform();
		
		DoubleMatrix result = new DoubleMatrix(matrix.rows, matrix.cols);
		for (int i = 0; i < matrix.size; i++) {
			result.values[i] = Math.sqrt(matrix.values[i]);
		}
		
		return result;
	}
	
	/**
	 * Compute the euclidean distance transform and return the result as a 
	 * flattened array of single precision floating point values. The returned
	 * array has a dimension of rows * columns.
	 * 
	 * @return flattened array of euclidean distance values.
	 */
	public FloatMatrix computeTransformFloat() {
		squareTransform();
		
		FloatMatrix result = new FloatMatrix(matrix.rows, matrix.cols);
		for (int i = 0; i < matrix.size; i++) {
			result.values[i] = (float) Math.sqrt(matrix.values[i]);
		}
		
		return result;
	}
	
	// Some test code
	public static void main(String[] args) {
		
		// Create rectangular mask
		int[][] pixels = new int[11][11];
		for (int i = 0; i < pixels.length; i++) {
			for (int j = 0; j < pixels[i].length; j++) {
				pixels[i][j] = 0;
			}
		}
		pixels[5][5] = 1;
		
		// Print mask
		System.out.println("Pixels: ");
		for (int i = 0; i < pixels.length; i++) {
			for (int j = 0; j < pixels[i].length; j++) {
				System.out.printf("%3d", pixels[i][j]);
				System.out.print(' ');
			}
			System.out.println();
		}
		
		// Transform
		DistanceTransform dt = new DistanceTransform();
		dt.init(new IntMatrix(pixels), 0);
		DoubleMatrix result = dt.computeTransform();
		
		// Print transform
		System.out.println("Result: ");
		for (int i = 0; i < result.rows; i++) {
			for (int j = 0; j < result.cols; j++) {
				System.out.printf("%4.2f", result.doubleAt(i, j));
				System.out.print(' ');
			}
			System.out.println();
		}
	}
}