#include <vector>
#include <queue>
#include <math.h>
#include <stdexcept>
#include "marcher.h"
#include "heap.cpp"

// we don't want the function names to get 'mangled' by the compiler
extern "C"
{

	Marcher::Marcher(const double* const c, const int ndim, const long* const shape, const double* const dx, const int order):
		c(c),
		ndim(ndim),
		shape(shape),
		dx(dx),
		order(order)
	{
		size = 1;
		dx_sq_inv = new double[ndim];
		shift = new long[ndim];

		for (int i = ndim - 1; i >= 0; i--)
		{
			dx_sq_inv[i] = 1.0 / pow(dx[i], 2);

			shift[i] = size;
			size *= shape[i];
		}

		flags = new char[size];

		alpha_sq = new double[ndim];
		beta = new double[ndim];
		skip = new bool[ndim];
	}

	Marcher::~Marcher()
	{
		delete[] shift;
		delete[] flags;

		delete[] dx_sq_inv;
		delete[] alpha_sq;
		delete[] beta;
		delete[] skip;
	}

	void Marcher::initialize(const size_t x_s, double* tau)
	{
		for (size_t i = 0; i < size; i++)
		{
			flags[i] = UNKNOWN;
			tau[i] = INF;
		}

		tau[x_s] = 0;
	}

	void Marcher::solve(const size_t x_s, double* const tau)
	{
		initialize(x_s, tau);

		auto heap_comp = [&tau](const size_t e1, const size_t e2){ return tau[e1] < tau[e2]; };
		Heap<decltype(heap_comp)> front(heap_comp, size);
		front.push(x_s);

		// list of points with minimal tau-values for each while-iteration
		std::vector<size_t> minima;

		// main loop for fast marching
		while (!front.empty())
		{
			minima.clear();

			// add the point with minimal tau-value in front
			minima.push_back(front.top());
			// mark it as known
			flags[front.top()] = KNOWN;
			double value = tau[front.top()];
			// remove it from the heap
			front.pop();

			// append all points that have the same (minimal) tau-value -> front.push will get faster by this
			bool done = false;
			while (!done)
			{
				if (!front.empty() && tau[front.top()] == value)
				{
					minima.push_back(front.top());
					flags[front.top()] = KNOWN;
					front.pop();
				}
				else
					done = true;
			}

			for (size_t x_i : minima)
			{
				// now we check all neighbors of x_i
				size_t rem = x_i;
				for (int d = 0; d < ndim; d++)
				{
					// index in the current axis
					long dim_i = rem / shift[d];
					rem -= dim_i * shift[d];

					size_t x_n = x_i - shift[d];
					// valid neighbor to the 'left'
					if (dim_i > 0 && flags[x_n] != KNOWN)
					{
						tau[x_n] = solve_quadratic(x_n, tau);
						// add neighbor to front-heap if it is not yet on it
						if (flags[x_n] == UNKNOWN)
						{
							front.push(x_n);
							flags[x_n] = FRONT;
						}
						// otherwise update the heap
						else
							front.update(x_n);
					}

					x_n = x_i + shift[d];
					// valid neighbor to the 'right'
					if (dim_i < shape[d] - 1 && flags[x_n] != KNOWN)
					{
						tau[x_n] = solve_quadratic(x_n, tau);
						// add neighbor to front-heap if it is not yet on it
						if (flags[x_n] == UNKNOWN)
						{
							front.push(x_n);
							flags[x_n] = FRONT;
						}
						// otherwise update the heap
						else
							front.update(x_n);
					}
				}
			}
		}
	}

	const double nine_fourths = 9.0 / 4.0;
	const double one_third = 1.0 / 3.0;

	double Marcher::solve_quadratic(const size_t x, double* const tau) const
	{
		size_t rem = x;
		for (int d = 0; d < ndim; d++)
		{
			// index in the current axis
			long dim_i = rem / shift[d];
			rem -= dim_i * shift[d];

			double tau_n = INF;
			char direction = 0;
			// choose the direction considering the non-factored first order approximation (tau_i-1 < tau_i+1 -> backward direction)
			// valid known neighbor to the 'left'
			if (dim_i > 0 && flags[x - shift[d]] == KNOWN)
			{
				tau_n = tau[x - shift[d]];
				direction = -1;
			}
			// valid known neighbor to the 'right' (with smaller value)
			if (dim_i < shape[d] - 1 && flags[x + shift[d]] == KNOWN && tau[x + shift[d]] < tau_n)
			{
				tau_n = tau[x + shift[d]];
				direction = 1;
			}

			// only include summands that are valid, skip the rest
			if (direction != 0)
			{
				skip[d] = false;

				// valid known second neighbor + second order condition -> use second order
				if (order == 2 && ((direction == -1 && dim_i > 1) || (direction == 1 && dim_i < shape[d] - 2)) && flags[x + 2 * direction * shift[d]] == KNOWN
					&& tau_n > tau[x + 2 * direction * shift[d]])
				{
					alpha_sq[d] = nine_fourths * dx_sq_inv[d];
					beta[d] = one_third * (4.0 * tau_n - tau[x + 2 * direction * shift[d]]);
				}
				// otherwise fall back to first order
				else
				{
					alpha_sq[d] = dx_sq_inv[d];
					beta[d] = tau_n;
				}
			}
			else
				skip[d] = true;
		}


		double c_base = -1.0 / pow(this->c[x], 2);
		double a, b, c;
		double disc;

		// the currently biggest "beta" and it's dimension
		double biggest;
		int biggest_d;

		do
		{
			a = 0.0, b = 0.0;
			c = c_base;

			biggest = N_INF;
			biggest_d = -1;

			for (int d = 0; d < ndim; d++)
			{
				if (!skip[d])
				{
					a += alpha_sq[d];
					b -= 2.0 * alpha_sq[d] * beta[d];
					c += alpha_sq[d] * pow(beta[d], 2);

					if (beta[d] > biggest)
					{
						biggest_d = d;
						biggest = beta[d];
					}
				}
			}

			if (biggest_d == -1)
				throw std::runtime_error("Negative discriminant in solve_quadratic.");

			disc = pow(b, 2) - 4.0 * a * c;
			skip[biggest_d] = true;
		} while (disc < 0);

		// a is always positive, so the '+' solution is larger (causality)
		return (-b + sqrt(disc)) / (2.0 * a);
	}
}