using DelimitedFiles
using Plots

# Solves the Eikonal equation using the Fast Sweeping Method
# Input is in file with name stored in source
# Input domain has ni x nj values, with a spacing of h
# Default tolerance is 1E-6
# Returns ni x nj solution matrix without extra layer of points
function solveEikonal(ni, nj, h, source, tol = 1E-6)
	F = zeros(ni+2, nj+2)
	U = ones(ni+2, nj+2) * Inf
	initialize!(F, U, ni, nj, source)
	maxerr = 1
	iters = 0
	while maxerr > tol
		maxerr = 0
		maxerr = sweep!(U, ni+1, 2, 2, nj+1, F, h, maxerr) # North-East
		maxerr = sweep!(U, ni+1, 2, nj+1, 2, F, h, maxerr) # North-West
		maxerr = sweep!(U, 2, ni+1, nj+1, 2, F, h, maxerr) # South-West
		maxerr = sweep!(U, 2, ni+1, 2, nj+1, F, h, maxerr) # South-East
		iters += 1
	end
	println("$iters iterations, maxerr = $(Float64(maxerr))")
	U
end

# Perform one set of sweeps on matrix U using directions specified in
# i1,ib,ja,jb, and speed function F, and current value of maximum
# absolute error (err = (U[i, j]-Unew)/U[i, j]), where Unew
# is new value of U[i,j] calculated using solveQuadratic()
# h is the spacing between points
# Returns updated value of maxerr
function sweep!(U, ia, ib, ja, jb, F, h, maxerr)
	#stepi = sign(ib - ia)
	#stepj = sign(jb - ja)
	stepi = 0
	stepj = 0
	if ib < ia
		stepi = -1
	else
		stepi = 1
	end
	if jb < ja
		stepj = -1
	else
		stepj = 1
	end
	for i in ia:stepi:ib
		for j in ja:stepj:jb
			Unew = solveQuadratic(U, i, j, F, h)
			if Unew < U[i, j]
				err = (U[i, j] - Unew) / U[i, j]
				if err > maxerr
					maxerr = err
				end
			U[i, j] = Unew
			end
		end
	end
	maxerr
end

# Solve the discretized Eikonal equation at (i,j), given
# speed function F, and spacing between points h
# Returns the new value of U[i,j]
function solveQuadratic(U, i, j, F, h)
	if U[i, j] == 0
		return 0
	end
	a = min(U[i-1, j], U[i+1, j])
	b = min(U[i, j-1], U[i, j+1])
	if abs(a - b) >= h / F[i, j]
		return min(a, b) + h / F[i, j]
	else
		return (a + b + sqrt(2* h^2 / F[i, j]^2 - (a - b)^2)) / 2
	end
end

# Reads from source a speed function F defined a ni x nj grid
# as well as coordinates of boundary of front
# (input is terminated with a negative number on last line)
# Also initializes solution matrix U
# U and F are (ni+2) x (nj+2) 2D matrices
# Requires the DelimitedFiles package
function initialize!(F, U, ni, nj, source)
    temp = readdlm(source)
    for j in 1:nj, i in 1:ni
        F[i+1, j+1] = Float64(temp[i, j])
    end
    for i in ni+1:size(temp, 1)-1 # skip last line that has -1 terminator
        # +2: skip border and convert input from 0-based indexing
        U[temp[i,1]+2, temp[i,2]+2] = 0
    end
	# println(U)
	# println(F)
    nothing
end

# Call on file "ex1.txt" with the following format:
display(view(solveEikonal(7, 7, 1, "ex1.txt"), 2:8, 2:8))

# Call on file "test100.txt" with the following format:
heatmap(solveEikonal(100, 100, 1, "test100.txt"), color=:gist_yarg)