#----------------------------------------------------------------------
#  Direct GFDM method test - 2d Laplace equation, Dirichlet BC
#  Validation: https://en.wikipedia.org/wiki/Laplace%27s_equation
#  Author: Andrea Pavan
#  Date: 21/12/2022
#  License: GPLv3-or-later
#----------------------------------------------------------------------
using ElasticArrays;
using LinearAlgebra;
using SparseArrays;
using PyPlot;
include("utils.jl");


#problem definition
r1 = 2.0;       #inner radius
r2 = 4.0;       #outer radius
function uD(x,y)        #Dirichlet boundary conditions
    r = sqrt(x^2+y^2);
    θ = atan(y,x);
    if abs(r-r2)<1e-9
        return 4*sin(5*θ);
    end
    return 0;
end

meshSize = 0.25;     #distance target between internal nodes
surfaceMeshSize = 0.25;      #distance target between boundary nodes
minNeighbors = 8;       #minimum number of neighbors allowed
minSearchRadius = meshSize/2;       #starting search radius


#pointcloud generation - boundary nodes
time1 = time();
pointcloud = ElasticArray{Float64}(undef,2,0);      #2xN matrix containing the coordinates [X;Y] of each node
boundaryNodes = Vector{Int};        #indices of the boundary nodes
dirichletNodes = Vector{Int}(undef,0);       #indices of the Dirichlet nodes
normals = ElasticArray{Float64}(undef,2,0);     #2xN matrix containing the components [nx;ny] of the normal of each boundary node
dθ = surfaceMeshSize/r1;
for i=0:dθ:2*pi-dθ
    append!(pointcloud, r1.*[cos(i),sin(i)]);
    append!(normals, -[cos(i),sin(i)]);
    push!(dirichletNodes, size(pointcloud,2));
end
dθ = surfaceMeshSize/r2;
for i=0:dθ:2*pi-dθ
    append!(pointcloud, r2.*[cos(i),sin(i)]);
    append!(normals, [cos(i),sin(i)]);
    push!(dirichletNodes, size(pointcloud,2));
end
boundaryNodes = collect(range(1,size(pointcloud,2)));

#pointcloud generation - internal nodes (cartesian)
#=for x=-r2+meshSize:meshSize:r2-meshSize
    for y=-r2+meshSize:meshSize:r2-meshSize
        if x^2+y^2>r1^2 && x^2+y^2<r2^2
            #P = [x,y];
            P = [x,y]+(rand(Float64,2).-0.5).*meshSize/5;

            #check if P is too close to the boundary
            insertP = true;
            for j in boundaryNodes
                if (x-pointcloud[1,j])^2+(y-pointcloud[2,j])^2<(0.75*meshSize)^2
                    insertP = false;
                end
            end
            if insertP
                append!(pointcloud, P);
            end
        end
    end
end=#

#pointcloud generation - internal nodes (flowmesher)
NinternalPoints = 0;
surfaceMargin = surfaceMeshSize/2;
while NinternalPoints<650
    x = rand(-r2:1e-6:r2);
    y = rand(-r2:1e-6:r2);
    if x^2+y^2>(r1+surfaceMargin)^2 && x^2+y^2<(r2-surfaceMargin)^2
        append!(pointcloud, [x,y]);
        global NinternalPoints += 1;
    end
end
P = pointcloud;
#v = rand(-meshSize/5:1e-6:meshSize/5,2,size(P,2));
v = zeros(Float64,2,size(P,2));
F = zeros(Float64,2,size(P,2));
#maxV = maximum(abs.(v));
maxV = 1.0;
dt = 0.1;
h = 3*meshSize;
iter = 1;
while maxV>5e-3
    for i=1+length(boundaryNodes):size(P,2)
        global F[:,i] = [0,0];
        #if v[:,i]'v[:,i]<1e-2
            for j=1:size(P,2)
                rd2 = transpose(P[:,j]-P[:,i])*(P[:,j]-P[:,i]);
                if i!=j && rd2<h^2
                    global F[:,i] += exp(-6.5*rd2/(h^2))*(P[:,i]-P[:,j])/sqrt(rd2);
                end
            end
        #end
        ang = atan(P[2,i],P[1,i]);
        #F[:,i] .+= 1e-1/(transpose(r1.*[cos(ang),sin(ang)]-P[:,i])*(r1.*[cos(ang),sin(ang)]-P[:,i]));
        #F[:,i] .-= 1e-1./(transpose(r2.*[cos(ang),sin(ang)]-P[:,i])*(r2.*[cos(ang),sin(ang)]-P[:,i]));
    end
    global v += dt*F;
    for i=1+length(boundaryNodes):size(P,2)
        modV = sqrt(v[1,i]^2+v[2,i]^2);
        if modV>0.5*meshSize/dt
            v[:,i] /= modV;
            v[:,i] *= 0.5*meshSize/dt;
        end
    end
    global v .*= 0.5;
    global maxV = sum(abs.(v[:]))/length(v[:]);
    println("Flowmesher: iter = ", iter, ", meanV = ",round(maxV,digits=4));
    global P += dt.*v;
    for i=1+length(boundaryNodes):size(P,2)
        if P[1,i]^2+P[2,i]^2<(r1+surfaceMargin)^2
            ang = atan(P[2,i],P[1,i]);
            rad = sqrt(P[:,i]'P[:,i]);
            P[:,i] += 2*((r1+surfaceMargin).*[cos(ang),sin(ang)]-P[:,i]);
            v[:,i] = [0,0];
        end
        if P[1,i]^2+P[2,i]^2>(r2-surfaceMargin)^2
            ang = atan(P[2,i],P[1,i]);
            rad = sqrt(P[:,i]'P[:,i]);
            P[:,i] += 2*((r2-surfaceMargin).*[cos(ang),sin(ang)]-P[:,i]);
            v[:,i] = [0,0];
        end
    end
    global iter += 1;
    #=if mod(iter,10)==0
        figure();
        #plot(pointcloud[1,boundaryNodes],pointcloud[2,boundaryNodes],"r.");
        plot(P[1,:],P[2,:],"k.");
        title("Pointcloud plot");
        axis("equal");
        display(gcf());
        sleep(0.1);
    end=#
end
pointcloud = P;

internalNodes = collect(range(1+length(boundaryNodes),size(pointcloud,2)));
println("Generated pointcloud in ", round(time()-time1,digits=2), " s");
println("Pointcloud properties:");
println("  Boundary nodes: ",length(boundaryNodes));
println("  Internal nodes: ",length(internalNodes));
println("  Memory: ",memoryUsage(pointcloud,boundaryNodes));

#pointcloud plot
figure();
plot(pointcloud[1,dirichletNodes],pointcloud[2,dirichletNodes],"r.");
plot(pointcloud[1,internalNodes],pointcloud[2,internalNodes],"k.");
title("Pointcloud plot");
axis("equal");
display(gcf());


#neighbor search
time2 = time();
N = size(pointcloud,2);     #number of nodes
neighbors = Vector{Vector{Int}}(undef,N);       #vector containing N vectors of the indices of each node neighbors
Nneighbors = zeros(Int,N);      #number of neighbors of each node
for i=1:N
    searchradius = minSearchRadius;
    while Nneighbors[i]<minNeighbors
        neighbors[i] = Int[];
        #check every other node
        for j=1:N
            if j!=i && all(abs.(pointcloud[:,j]-pointcloud[:,i]).<searchradius)
                push!(neighbors[i],j);
            end
        end
        unique!(neighbors[i]);
        Nneighbors[i] = length(neighbors[i]);
        searchradius += minSearchRadius/2;
    end
end
println("Found neighbors in ", round(time()-time2,digits=2), " s");
println("Connectivity properties:");
println("  Max neighbors: ",maximum(Nneighbors)," (at index ",findfirst(isequal(maximum(Nneighbors)),Nneighbors),")");
println("  Avg neighbors: ",round(sum(Nneighbors)/length(Nneighbors),digits=2));
println("  Min neighbors: ",minimum(Nneighbors)," (at index ",findfirst(isequal(minimum(Nneighbors)),Nneighbors),")");


#neighbors distances and weights
time3 = time();
P = Vector{Array{Float64}}(undef,N);        #relative positions of the neighbors
d2 = Vector{Vector{Float64}}(undef,N);      #relative distances of the neighbors
w2 = Vector{Vector{Float64}}(undef,N);      #neighbors weights
for i=1:N
    P[i] = Array{Float64}(undef,2,Nneighbors[i]);
    d2[i] = Vector{Float64}(undef,Nneighbors[i]);
    w2[i] = Vector{Float64}(undef,Nneighbors[i]);
    for j=1:Nneighbors[i]
        P[i][:,j] = pointcloud[:,neighbors[i][j]]-pointcloud[:,i];
        d2[i][j] = P[i][:,j]'P[i][:,j];
    end
    d2max = maximum(d2[i]);
    for j=1:Nneighbors[i]
        w2[i][j] = exp(-1*d2[i][j]/d2max)^2;
    end
end
w2pde = 2.0;        #least squares weight for the pde
w2bc = 2.0;     #least squares weight for the boundary conditions


#least square matrix inversion
A = Vector{Matrix}(undef,N);        #least-squares matrices
condA = Vector{Float64}(undef,N);       #condition number
B = Vector{Matrix}(undef,N);        #least-squares decomposition matrices
C = Vector{Matrix}(undef,N);        #derivatives coefficients matrices
for i in internalNodes
    xj = P[i][1,:];
    yj = P[i][2,:];
    V = zeros(Float64,1+Nneighbors[i],6);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], xj[j]^2, yj[j]^2, xj[j]*yj[j]];
    end
    V[end,:] = [0, 0, 0, 2, 2, 0];
    W = Diagonal(vcat(w2[i],w2pde));
    A[i] = transpose(V)*W*V;
    condA[i] = cond(A[i]);
    (Q,R) = qr(A[i]);
    C[i] = inv(R)*transpose(Q)*transpose(V)*W;
end
for i in dirichletNodes
    xj = P[i][1,:];
    yj = P[i][2,:];
    V = zeros(Float64,2+Nneighbors[i],6);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], xj[j]^2, yj[j]^2, xj[j]*yj[j]];
    end
    V[end-1,:] = [0, 0, 0, 2, 2, 0];
    V[end,:] = [1, 0, 0, 0, 0, 0];
    W = Diagonal(vcat(w2[i],[w2pde,w2bc]));
    A[i] = transpose(V)*W*V;
    condA[i] = cond(A[i]);
    (Q,R) = qr(A[i]);
    C[i] = inv(R)*transpose(Q)*transpose(V)*W;
end
println("Inverted least-squares matrices in ", round(time()-time3,digits=2), " s");
println("Matrices properties:");
println("  Max condition number (internal nodes): ",round(maximum(condA[internalNodes]),digits=2));
println("  Avg condition number (internal nodes): ",round(sum(condA[internalNodes])/length(internalNodes),digits=2));
println("  Min condition number (internal nodes): ",round(minimum(condA[internalNodes]),digits=2));


#matrix assembly
time4 = time();
rows = Int[];
cols = Int[];
vals = Float64[];
for i in dirichletNodes
    push!(rows, i);
    push!(cols, i);
    push!(vals, 1);
    for j=1:lastindex(neighbors[i])
        push!(rows, i);
        push!(cols, neighbors[i][j]);
        push!(vals, -C[i][1,j]);
    end
end
for i in internalNodes
    push!(rows, i);
    push!(cols, i);
    push!(vals, 1);
    for j=1:Nneighbors[i]
        push!(rows, i);
        push!(cols, neighbors[i][j]);
        push!(vals, -C[i][1,j]);
    end
end
M = sparse(rows,cols,vals,N,N);
println("Completed matrix assembly in ", round(time()-time4,digits=2), " s");


#linear system solution
time5 = time();
b = zeros(N);       #rhs vector
for i in dirichletNodes
    b[i] = C[i][1,end]*uD(pointcloud[1,i],pointcloud[2,i]);
end
for i in internalNodes
    b[i] = 0;
end
u = M\b;
println("Linear system solved in ", round(time()-time5,digits=2), " s");


#solution plot
figure();
scatter(pointcloud[1,:],pointcloud[2,:],c=u,cmap="hsv");
colorbar();
title("Laplace equation - annular domain");
axis("equal");
display(gcf());
