#  poisson.jl -- Solve the Poisson Equation
#
#  del^2 y = f for x such that domain(x)>c
#  y = psi for x such that domain(x)=c

using SparseArrays, LinearAlgebra, Plots

xmin=[-2,-2]
xmax=[2,2]
mres=100

# specify the boundary condition
function psi(x)
    if x[1]^2+x[2]^2>1.7^2
        return 3.0
    end
    return 5.0
end
# specify the distribution
function f(x)
    return 0.0
end

# create spatial grid to fit the desired resolution
if xmax[1]-xmin[1]<xmax[2]-xmin[2]
    dx=(xmax[2]-xmin[2])/(mres-1)
else
    dx=(xmax[1]-xmin[1])/(mres-1)
end
nu=Int.(round.((xmax-xmin)/dx).+1)
image=zeros(Int,nu...)
xmax=xmin+dx*(nu.-1)
x1s=[xmin[1]+(i1-1)*dx for i1=1:size(image)[1]]
x2s=[xmin[2]+(i2-1)*dx for i2=1:size(image)[2]]
function xi(i)
    return [x1s[i[1]],x2s[i[2]]]
end

# specify the domain using level sets
function domain(x)
    if x[1]*x[1]+x[2]*x[2]<1.8^2
        if (x[1]-.5)*(x[1]-.5)+x[2]^2<0.5^2
            return 0.0
        end
        return 2.0
    end
    return 0.0
end

# count the grid points inside the domain
function enumerate()
    count=0
    for i1=2:size(image)[1]-1
        for i2=2:size(image)[2]-1
            if domain(xi([i1,i2]))>1.0
                count+=1
                image[i1,i2]=count
            end
        end
    end
    return count
end

# create A and F using the 4-point stencil
function stencil(count)
    xs=[]; ys=[]; as=[]
    fv=Vector{Float64}(undef,count)
    for i1=1:size(image)[1]
        for i2=1:size(image)[2]
            function look(j1,j2)
                if image[j1,j2]>0 # still inside the domain
                    push!(xs,image[i1,i2])
                    push!(ys,image[j1,j2])
                    push!(as,1.0)
                else # subtract the boundary to the other side in F
                    fv[image[i1,i2]]-=psi(xi([j1,j2]))
                end
            end
            if image[i1,i2]>0
                fv[image[i1,i2]]=dx^2*f(xi([i1,i2]))
                push!(xs,image[i1,i2])
                push!(ys,image[i1,i2])
                push!(as,-4.0)
                look(i1,i2-1) # down
                look(i1-1,i2) # left
                look(i1+1,i2) # right
                look(i1,i2+1) # up
            end
        end
    end
    return (sparse(Int.(xs),Int.(ys),Float64.(as)),fv)
end

# by construction the exact solution is psi
function exact(count)
    xv=zeros(count)
    for i1=1:size(image)[1]
        for i2=1:size(image)[2]
            if image[i1,i2]>0
                xv[image[i1,i2]]=psi(xi([i1,i2]))
            end
        end
    end
    return xv
end

# convert a solution back to the grid
function physical(x)
    black=minimum(x)
    p=zeros(nu...)
    for i1=1:size(image)[1]
        for i2=1:size(image)[2]
            if image[i1,i2]>0
                p[i1,i2]=x[image[i1,i2]]
            else
                p[i1,i2]=black
            end
        end
    end
    return p
end

# various plotting functions
function pltdomain()
    heatmap(x1s,x2s,(x->min(1,x)).(image)',c=:hot)
end
function plterror()
    heatmap(x1s,x2s,physical(AP-EX)'*1000,c=:hot)
end
function pltexact()
    heatmap(x1s,x2s,physical(EX)',c=:hot)
end
function pltapprox()
    heatmap(x1s,x2s,physical(AP)',c=:hot)
end

# solve the problem and print the error
count=enumerate()
EX=exact(count)
A,F=stencil(count)
AP=A\F  # solving Au=F and call u=AP using sparse matrices...
println(mres," ",norm(AP-EX)*dx)