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 = 20.0
R2 = R^2

# Charge distribution
rho = 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
            rho[i, j] = 1.0
        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), :])

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

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

contourf(phi)

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