
/* FMM.cs - 2d Fast Marching Method (FMM) implementation.
 * Copyright (C) 2004-2005 -- Sebastian Nowozin
 *
 * This program is free software released under the GNU General Public
 * License, which is included in this software package (doc/LICENSE).
 */

/* FMM.cs
 *
 * A two dimensional Fast Marching Method (FMM) implementation using first
 * order difference approximations.  The references used to produce this
 * implementation:
 *
 * J.A. Sethian, "Level Set Methods and Fast Marching Methods"
 * Stanley Osher, Ronald Fedkiw, "Level Set Methods and Dynamic Implicit
 *     Surfaces"
 * ITK, The Insight Segmentation Toolkit, which has a n-dimensional generic
 *     FMM implementation (although there are some known flaws in it, its an
 *     excellent implementation and the quadratic equation solver is from that
 *     implementation).
 *
 * Within autopano-sift it is used for fast distance map creation when
 * smart-pruning control points.
 */

using System;
using System.Text;
using System.Collections;


public class
FMM
{
	public enum State {
		Unknown,
		Known,
		Trial,
	};

	/* The constant propagation speed.  For a distance map its 1.0 everywhere.
	 */
	private double speed = 1.0;

	public FMM (int x, int y)
	{
		mhIndexMap = new int[y, x];
		Map = new State[y, x];
		TMap = new double[y, x];
		CauseMap = new int[y, x];
		xdim = x;
		ydim = y;

		for (int py = 0 ; py < ydim ; ++py) {
			for (int px = 0 ; px < xdim ; ++px) {
				mhIndexMap[py, px] = -1;
				Map[py, px] = State.Unknown;
				TMap[py, px] = 0.0;
				CauseMap[py, px] = -1;
			}
		}

		/* Create minheap and pass a delegate to be called whenever the heap
		 * order changes.  This is used to keep the heap backpointer map up to
		 * date.
		 */
		minheap = new MinHeap (512, this.HeapNotifyHandler);
	}

	public readonly int xdim, ydim;

	/* The heap backpointer map and the heap itself.
	 */
	private int[,] mhIndexMap;
	private MinHeap minheap;

	/* The set map, managing the disjoint KNOWN, TRIAL and UNKNOWN sets.
	 */
	public State[,] Map;

	/* The arrival time map.
	 */
	public double[,] TMap;

	public int[,] CauseMap;

	public override string ToString ()
	{
		StringBuilder sb = new StringBuilder ();

		sb.Append ("FMM (");
		sb.AppendFormat ("minheap = {0}, ", minheap);
		sb.Append (")");

		return (sb.ToString ());
	}

	/* Calculate the total sum of all points.  All points are assumed to be
	 * KNOWN points.
	 */
	public double CalculateTotalSum ()
	{
		double sum = 0.0;

		for (int y = 0 ; y < ydim ; ++y) {
			for (int x = 0 ; x < xdim ; ++x) {
				if (Map[y, x] != State.Known)
					throw (new FMMNotFinishedException
						("There are points with a unknown or trial state."));

				sum += TMap[y, x];
			}
		}

		return (sum);
	}

	public class FMMNotFinishedException : ApplicationException
	{
		public FMMNotFinishedException (string message) :
			base (message)
		{
		}
	}

	public void HeapNotifyHandler (int heapIndex, object val)
	{
		MapPoint mp = (MapPoint) val;
		mhIndexMap[mp.Y, mp.X] = heapIndex;
	}

	/* cause: an integer that will be carried with the front as it propagates.
	 */
	public void AddKnown (int x, int y, double val, int cause)
	{
		if (x < 0 || y < 0 || x >= xdim || y >= ydim)
			throw (new ArgumentException ("Coordinates out of dimension"));

		/* If there is already a KNOWN point with a lower arrival time, skip
		 * this assignment.
		 */
		if (Map[y, x] == State.Known && TMap[y, x] < val)
			return;

		Map[y, x] = State.Known;
		TMap[y, x] = val;
		CauseMap[y, x] = cause;
	}

	public void InitializeBand ()
	{
		for (int y = 0 ; y < ydim ; ++y) {
			for (int x = 0 ; x < xdim ; ++x) {
				if (Map[y, x] == State.Known)
					continue;

				bool hasKnownNeighbor = false;
				if (y > 0 && Map[y-1, x] == State.Known)
					hasKnownNeighbor = true;
				else if (y < (ydim-1) && Map[y+1, x] == State.Known)
					hasKnownNeighbor = true;
				else if (x > 0 && Map[y, x-1] == State.Known)
					hasKnownNeighbor = true;
				else if (x < (xdim-1) && Map[y, x+1] == State.Known)
					hasKnownNeighbor = true;

				if (hasKnownNeighbor)
					Update (x, y);
			}
		}
	}

	private class MapPoint
	{
		public double Time;
		public int X;
		public int Y;

		public MapPoint ()
		{
			X = Y = -1;
			Time = -1.0;
		}

		public MapPoint (int x, int y, double time)
		{
			this.X = x;
			this.Y = y;
			this.Time = time;
		}
	}

	private readonly int[] syOffset = new int[] { 0, 0, -1, 1, -1, -1, 1, 1, };
	private readonly int[] sxOffset = new int[] { -1, 1, 0, 0, -1, 1, -1, 1, };

	/* Do a given number of iterations by converting UNKNOWN and TRIAL values to
	 * KNOWN values.
	 *
	 * @param iterationCount If greater than zero, then as maximum there are as
	 * many iterations done as specified.  If the final state is reached before,
	 * no further iterations are done.
	 * @param T_step If iteration_count is less than or equal to zero and T_step
	 * is larger than zero, then as many steps are done so we progress at least
	 * T_step time units forward.
	 *
	 * If both iterationCount and T_step are zero or less than zero, the FMM
	 * is marching until no further marching is possible.
	 *
	 * @returns True if there are further iterations possible, false if the
	 * final state has been reached.
	 */
	public bool DoIterations (int iterationCount, double T_step)
	{
		double T_start = -1.0;
		double T_now = -1.0;
		int iter = 0;

		while ((iterationCount <= 0 || iter < iterationCount) &&
			(T_step <= 0.0 || ((T_now - T_start) < T_step)) &&
			minheap.IsEmpty == false)
		{
			MapPoint minTrial = (MapPoint) minheap.DeleteMin (out T_now);

			iter += 1;
			//T_now = minTrial.Time;

			// XXX: fmm->causing_T = T_now;
			if (T_start < 0.0)
				T_start = T_now;

			/* Descriptions as in [Malladi, Sethian], "An O(N log N) algorithm for
			 * shape modelling".
			 *
			 * Step 2a, get minimal T element: is in minTrial already
			 * Step 2b, add the point to KNOWN
			 */
			Map[minTrial.Y, minTrial.X] = State.Known;

			/* Step 2c, tag as neighbors any points that are not KNOWN, if the
			 * neighbor is in UNKNOWN, remove it from that set and add it to
			 * the TRIAL set.
			 */
			bool[] isNeighbor = new bool[8];
			for (int n = 0 ; n < 4 ; ++n) {
				int py = minTrial.Y + syOffset[n];
				int px = minTrial.X + sxOffset[n];

				isNeighbor[n] = false;

				// Out of computational domain?
				if (px < 0 || px >= xdim || py < 0 || py >= ydim)
					continue;

				// Only use UNKNOWN and TRIAL
				if (Map[py, px] == State.Known)
					continue;

				// Tag as neighbor
				isNeighbor[n] = true;

				if (Map[py, px] == State.Unknown)
					Map[py, px] = State.Trial;
			}

			/* Step 2d, recompute the values of T at all neighbors by solving
			 * the quadratic equation.
			 */
			for (int n = 0 ; n < 4 ; ++n) {
				if (isNeighbor[n] == false)
					continue;

				Update (minTrial.X + sxOffset[n], minTrial.Y + syOffset[n]);
			}

			if (minheap.IsEmpty)
				return (false);
		}

		return (true);
	}

	private void Update (int x, int y)
	{
		MapPoint[] mv = new MapPoint[2];
		for (int i = 0 ; i < mv.Length ; ++i)
			mv[i] = new MapPoint ();

		/* We only have two dimensions to consider, and here is what we do: we
		 * find the direction for each dimension in which the interface is
		 * nearest.  Nearest means, the point in the direction is already
		 * known or a trial point and has an older time (smaller value) than
		 * the point in the exact opposite direction (diametral position).
		 */
		for (int d = 0 ; d < 2 ; ++d) {
			mv[d].Time = -1.0;

			// Left/Right, in the given dimension
			for (int s = -1 ; s <= 1 ; s += 2) {
				int px, py;

				px = x + (d == 0 ? s : 0);
				py = y + (d == 1 ? s : 0);
				if (px < 0 || px >= xdim || py < 0 || py >= ydim)
					continue;

				// We only want KNOWN values
				if (Map[py, px] != State.Known)
					continue;

				// Now we got the value, lets check if its minimal
				if (mv[d].Time < 0.0 || TMap[py, px] < mv[d].Time) {
					mv[d].Time = TMap[py, px];
					mv[d].X = px;
					mv[d].Y = py;
				}
			}
		}

		/* We have to order the list of quadrants starting with the minimal T
		 * value in increasing order.
		 *
		 * The only two cases to check: first value is invalid or both are
		 * valid and the second is lower.  In either case, swap the elements.
		 */
		if (mv[0].Time < 0.0 || (mv[1].Time >= 0.0 && mv[1].Time < mv[0].Time)) {
			MapPoint mvTemp = mv[1];

			mv[1] = mv[0];
			mv[0] = mvTemp;
		}

		/* If you want a real flexible FMM with a position-dependent speed
		 * function, the line to go here is:
		 * speed = speedFunction (x, y);
		 *
		 * We use a constant speed to obtain a distance map.
		 */
		double T = SolveQuadratic (mv, speed);

		// Write this value back and make the point a TRIAL point.
		if (T >= 0.0) {
			TMap[y, x] = T;
			Map[y, x] = State.Trial;

			/* Propagate the cause label.
			 *
			 * Of course the cause is not always a 1.0 or 0.0 relationship
			 * when two fronts merge, but we treat it as such to basically
			 * build a Voronoi cell label map.
			 */
			CauseMap[y, x] = CauseMap[mv[0].Y, mv[0].X];

			/* If the new point was already a TRIAL point, i.e. its time got
			 * reduced, remove the old trial element from the heap by looking
			 * up the backpoint index map (mhIndexMap).
			 */
			if (mhIndexMap[y, x] >= 0)
				minheap.Delete (mhIndexMap[y, x]);

			/* Now in either case add the new trial value to the heap.
			 */
			minheap.Insert (T, new MapPoint (x, y, T));
		}
	}

	/* Solve the quadratic equation:
	 * \left(\frac{\phi_{i,j} - \phi_1}{\Delta x}\right)^2 +
	 *     \left(\frac{\phi_{i,j} - \phi_2}{\Delta y}\right)^2 =
	 *     \frac{1}{F_{i,j}}
	 *
	 * Where F_{i,j} is the speed at the location (i,j).  We further assume
	 * \Delta x = \Delta y = 1.0.
	 *
	 * This function is derived on the ITK FMM implementation.
	 */
	private double SolveQuadratic (MapPoint[] mv, double speed)
	{
		double a, b, c;
		double solution = -1.0;
		double dis;

		a = b = 0.0;
		if (speed > 0.0) {
			c = -1.0 * (1.0 / (speed * speed));
		} else
			return (-1.0);

		for (int d = 0 ; d < mv.Length ; ++d) {
			// This direction has no T value.
			if (mv[d].Time < 0.0)
				continue;

			/* In case this is true it means that the arrival time calculated is
			 * _lower_ than the time already attached to a point we consider TRIAL
			 * or KNOWN.  Infact, it should only happen for TRIAL points.  As the
			 * values are ordered with increasing times, we break out of the
			 * dimension loop.
			 */
			if (solution >= 0.0 && solution < mv[d].Time)
				break;

			a += 1.0;
			b += mv[d].Time;
			c += mv[d].Time * mv[d].Time;

			dis = b * b - a * c;

			solution = (Math.Sqrt (dis) + b) / a;
		}

		return (solution);
	}

	public class MinHeap
	{
		private struct MHElem {
			public double Value;
			public object ElementObject;
		}
		private MHElem[] array;

		/* The number of array cells used within the array, whereas the actual
		 * array can be larger to minimize the number of necessary
		 * reallocations.
		 */
		int len;

		/* If non-null, this function will be called whenever an element
		 * position is changed, an element is added or removed.  The first
		 * parameter is the element in the heap, the second is the object
		 * stored in the heap.
		 * An first parameter of -1 means the element is not part of the heap
		 * anymore.
		 */
		public delegate void HeapNotify (int idx, object val);
		private HeapNotify notify = null;

		public MinHeap (int initialAllocationLength, HeapNotify notify)
		{
			array = new MHElem[initialAllocationLength];
			len = 0;
			this.notify = notify;
		}

		public override string ToString ()
		{
			return (String.Format ("MinHeap (len = {0}, allocated = {1})",
				len, array.Length));
		}

		private void Exchange (int i1, int i2)
		{
			MHElem mhTemp = array[i1];
			array[i1] = array[i2];
			array[i2] = mhTemp;

			if (notify != null) {
				notify (i1, array[i1].ElementObject);
				notify (i2, array[i2].ElementObject);
			}
		}

		private void SiftUp (int i)
		{
			if (i == 0)
				return;

			do {
				if (array[i].Value < array[i >> 1].Value) {
					Exchange (i, i >> 1);
				} else
					break;

				i >>= 1;
			} while (i != 0);
		}

		private void SiftDown (int i)
		{
			if ((i << 1) > len)
				return;

			do {
				i <<= 1;
				if ((i + 1) <= len && array[i + 1].Value < array[i].Value)
					i += 1;

				if (array[i >> 1].Value > array[i].Value) {
					Exchange (i, i >> 1);
				} else
					break;
			} while ((i << 1) <= len);
		}

		public void Insert (double val, object obj)
		{
			len += 1;
			if (len > array.Length) {
				// Allocate and copy the current elements into a double sized
				// array.
				MHElem[] newArray = new MHElem[array.Length << 1];
				Array.Copy (array, newArray, array.Length);
				array = newArray;
			}

			// Guaranteed space available
			array[len - 1].Value = val;
			array[len - 1].ElementObject = obj;

			if (notify != null)
				notify (len - 1, obj);

			SiftUp (len - 1);
		}

		public bool IsEmpty {
			get {
				return (len == 0);
			}
		}

		public void Delete (int idx)
		{
			if (idx < 0 || idx >= len)
				throw (new ArgumentException ("Index " + idx +
					" out of length " + len + "."));

			MHElem x = array[idx];
			MHElem y = array[len - 1];

			len -= 1;
			if (notify != null)
				notify (-1, x.ElementObject);

			// If we removed the last element, just return early
			if (idx == len)
				return;

			// Move last element to removed place.
			array[idx] = y;
			if (notify != null)
				notify (idx, y.ElementObject);

			if (y.Value <= x.Value)
				SiftUp (idx);
			else
				SiftDown (idx);
		}

		public object DeleteMin (out double Value)
		{
			Value = array[0].Value;
			object retObj = array[0].ElementObject;

			Delete (0);

			return (retObj);
		}
	}
}

