using Pkg
Pkg.add("Plots")
using Plots

function del2_5(a::Matrix{Float64}, dx = 1.0)
	#=
	Returns a finite-difference approximation of the laplacian of the array a, 
	with lattice spacing dx, using the five-point stencil: 

					0 1 0 
					1 -4 1 
					0 1 0 
	=#

	del2 = zeros(size(a))
	del2[2:end-1, 2:end-1] .= (a[2:end-1, 3:end] + a[2:end-1, 1:end-2] +
							   a[3:end, 2:end-1] + a[1:end-2, 2:end-1] -
							   4.0 * a[2:end-1, 2:end-1]) ./ (dx^2)
	return del2
end

function del2_9(a::Matrix{Float64}, dx = 1.0)
	#=
	Returns a finite-difference approximation of the laplacian of the array a, 
	with lattice spacing dx, using the nine-point stencil: 

					1/6 2/3 1/6 
					2/3 -10/3 2/3 
					1/6 2/3 1/6 
	=#

	del2 = zeros(size(a))
	del2[2:end-1, 2:end-1] .= (4.0 * (a[2:end-1, 3:end] + a[2:end-1, 1:end-2] +
									  a[3:end, 2:end-1] + a[1:end-2, 2:end-1]) +
							   (a[1:end-2, 1:end-2] + a[1:end-2, 3:end] +
								a[3:end, 1:end-2] + a[3:end, 3:end]) -
							   20.0 * a[2:end-1, 2:end-1]) ./ (6.0 * dx^2)
	return del2
end

function invDel2_5(b::Matrix{Float64}, dx = 1.0, N = 10000)
	#=
	Relaxes over N steps to a discrete approximation of the inverse laplacian 
	of the source term b, each step is a weighted average over the four neighboring 
	points and the source. This is the Jacobi algorithm.
	=#

	invDel2 = zeros(size(b))
	newInvDel2 = zeros(size(b))

	for m in 1:N
		newInvDel2[2:end-1, 2:end-1] .= 0.25 * (invDel2[2:end-1, 3:end] +
												invDel2[2:end-1, 1:end-2] +
												invDel2[3:end, 2:end-1] +
												invDel2[1:end-2, 2:end-1] -
												(dx^2) * b[2:end-1, 2:end-1])
		invDel2 .= newInvDel2
	end

	diff = del2_5(invDel2, dx) - b
	diffSq = diff .* diff
	error = sqrt(sum(diffSq))

	println("\nerror = ", error)

	return invDel2
end


# # Define variables
# L = 10
# dx = 1.0

# # Parabola of revolution has constant Laplacian
# phi = zeros((L, L))
# for i ∈ 1:L
# 	x = (i + 0.5 - 0.5 * L) * dx
# 	for j ∈ 1:L
# 		y = (j + 0.5 - 0.5 * L) * dx
# 		phi[i, j] = x^2 + y^2
# 	end
# end

# println(size(phi))

# println(del2_5(phi, dx))

# println("\n\n")

# println(del2_9(phi, dx))

# Electrostatics example: uniformly charged cylinder (rho = 1) of radius R
L = 100
dx = 1.0
R = 1.0
R2 = R^2

# Charge distribution
rho = zeros((L, L))
for i ∈ 1:L
	y = (i - 1 + 0.5 - 0.5 * L) * dx
	println("y = ", y)
	for j ∈ 1:L
		x = (j - 1 + 0.5 - 0.5 * L) * dx
		# this is the rod here
		if y < 0 && abs(x) < R
			rho[i, j] = 1.0 # high voltage
		elseif y > 0 && abs(x) < R
			rho[i, j] = -1.0 # conducting plane
		else
			rho[i, j] = 0.0
		end
	end
end

rhoPlot = plot(rho[Int(L / 2), :])

# Exact (analytical) electric potential
# phi = zeros((L, L))
# for i ∈ 1:L
# 	x = (i + 0.5 - 0.5 * L) * dx
# 	for j ∈ 1:L
# 		y = (j + 0.5 - 0.5 * L) * dx
# 		r2 = x^2 + y^2
# 		if r2 < R2
# 			phi[i, j] = -π * r2
# 		else
# 			phi[i, j] = -π * (R2 + R2 * log(r2 / R2))
# 		end
# 	end
# end

# phiPlot = plot(phi[Int(L / 2), :])

# Charge density obtained from exact potential
# rho = -1.0 / (4.0 * π) * del2_5(phi, dx)
# rhoPlotLattice = plot(rho[Int(L / 2), :])

function calc_field(phi::Matrix{Float64}, dx = 1.0)
	E_x = zeros(size(phi))
	E_y = zeros(size(phi))

	E_x[2:end-1, 2:end-1] .= -(phi[3:end, 2:end-1] - phi[1:end-2, 2:end-1]) / (2 * dx)
	E_y[2:end-1, 2:end-1] .= -(phi[2:end-1, 3:end] - phi[2:end-1, 1:end-2]) / (2 * dx)

	# for i ∈ 2:size(phi, 1)-1
	# 	for j ∈ 2:size(phi, 2)-1
	# 		E_x[i, j] = -(phi[i+1, j] - phi[i-1, j]) / (2 * dx)
	# 		E_y[i, j] = -(phi[i, j+1] - phi[i, j-1]) / (2 * dx)
	# 	end
	# end

	return E_x, E_y
end

phi = invDel2_5(-4.0 * π * rho, dx, 1000)
# phi = invDelSOR(-4.0 * π * rho, dx, 500)
phi .-= phi[Int(L / 2), Int(L / 2)]

phiPlotInvDel = plot(phi[Int(L / 2), :])

# phi_contour = contourf(phi, title = "Electric Potential", color = :viridis, aspect_ratio = :equal, colorbar_title = "V")
# savefig(phi_contour, "hw7/5.6-phi.png")

# plot the arrows of the magnatic field on the contour plot
xxs = [x for x in 1:L for y in 1:L]
xxy = [y for x in 1:L for y in 1:L]
d_x, d_y = calc_field(phi, dx)

function df(i, j)
	norm = d_x[Int64(i)] * d_x[Int64(i)] + d_y[Int64(j)] * d_y[Int64(j)]
	if norm == 0
		return (0, 0)
	end
	return (d_x[Int64(i)] / norm, d_y[Int64(j)] / norm)
end

quiver(xxs, xxy, quiver = df, aspect_ratio = :equal, legend = false)
savefig("hw7/5.6-test.png")


# d_x, d_y = calc_field(phi, dx)
# field_y = contourf(d_y, title = "Electric Field (y)", color = :viridis, aspect_ratio = :equal, colorbar_title = "V/m")

# field_x = contourf(d_x, title = "Electric Field (x)", color = :viridis, aspect_ratio = :equal, colorbar_title = "V/m")


# savefig(field_y, "hw7/5.6-fy.png")
# savefig(field_x, "hw7/5.6-fx.png")

#plot(rhoPlot, rhoPlotLattice)
#plot(phiPlot, phiPlotInvDel)