#----------------------------------------------------------------------
#  GFDM method test - 2d heat equation, Dirichlet BC, Steady solver
#  Author: Andrea Pavan
#  Date: 15/12/2022
#  License: GPLv3-or-later
#----------------------------------------------------------------------
using ElasticArrays;
using LinearAlgebra;
using SparseArrays;
using Printf;
using PyPlot;
include("utils.jl");


#problem definition
l1 = 1.0;       #domain x size
l2 = 0.1;       #domain y size
uL = 400;       #left border temperature
uR = 300;       #right border temperature
kcost = 0.1;        #thermal conductivity
rho = 8.0;      #density
ccost = 3.8;        #specific heat capacity
t0 = 0.0;       #starting time
tend = 10.0;     #ending time
dt = 0.2;       #time step
u0(x,y) = 0.0;      #initial conditions
uD(x,y) = uL+x*(uR-uL)/l1;      #Dirichlet boundary conditions

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


#pointcloud generation
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
normals = ElasticArray{Float64}(undef,2,0);     #2xN matrix containing the components [nx;ny] of the normal of each boundary node
for i=0:surfaceMeshSize:l1
    append!(pointcloud, [i,0.0]);
    append!(normals, [0,-1]);
end
for i=surfaceMeshSize:surfaceMeshSize:l2
    append!(pointcloud, [l1,i]);
    append!(normals, [1,0]);
end
for i=l1-surfaceMeshSize:-surfaceMeshSize:0
    append!(pointcloud, [i,l2]);
    append!(normals, [0,1]);
end
for i=l2-surfaceMeshSize:-surfaceMeshSize:surfaceMeshSize
    append!(pointcloud, [0,i]);
    append!(normals, [-1,0]);
end
boundaryNodes = collect(range(1,size(pointcloud,2)));
for i=0:meshSize:l1
    for j=0:meshSize:l2
        if i>0 && i<l1 && j>0 && j<l2
            append!(pointcloud, [i,j]);
        end
    end
end
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));


#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 *= 1.1;
    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
r2 = 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]);
    r2[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];
        r2[i][j] = P[i][:,j]'P[i][:,j];
    end
    r2max = maximum(r2[i]);
    for j=1:Nneighbors[i]
        w2[i][j] = exp(-1*r2[i][j]/r2max)^2;
    end
end


#least square matrix inversion
A = Vector{Matrix}(undef,N);        #least-squares matrices
condA = Vector{Float64}(undef,N);       #condition number
invA = Vector{Matrix}(undef,N);     #inverse matrices
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,:];
    A[i] = [sum(xj.^2 .*w2[i])  sum(xj.*yj.*w2[i])  0.5*sum(xj.^3 .*w2[i])  0.5*sum(xj.*yj.^2 .*w2[i])  sum(xj.^2 .*yj.*w2[i]);
        sum(xj.*yj.*w2[i])  sum(yj.^2 .*w2[i])  0.5*sum(xj.^2 .*yj.*w2[i])  0.5*sum(yj.^3 .*w2[i])  sum(xj.*yj.^2 .*w2[i]);
        0.5*sum(xj.^3 .*w2[i])  0.5*sum(xj.^2 .*yj.*w2[i])  0.25*sum(xj.^4 .*w2[i])  0.25*sum(xj.^2 .*yj.^2 .*w2[i])  0.5*sum(xj.^3 .*yj.*w2[i]);
        0.5*sum(xj.*yj.^2 .*w2[i])  0.5*sum(yj.^3 .*w2[i])  0.25*sum(xj.^2 .*yj.^2 .*w2[i])  0.25*sum(yj.^4 .*w2[i])  0.5*sum(xj.*yj.^3 .*w2[i]);
        sum(xj.^2 .*yj.*w2[i])  sum(xj.*yj.^2 .*w2[i])  0.5*sum(xj.^3 .*yj.*w2[i])  0.5*sum(xj.*yj.^3 .*w2[i])  sum(xj.^2 .*yj.^2 .*w2[i])];
    condA[i] = cond(A[i]);
    invA[i] = inv(A[i]);
    B[i] = zeros(Float64,5,1+Nneighbors[i]);
    B[i][:,1] = [-sum(xj.*w2[i]); -sum(yj.*w2[i]); -0.5*sum(xj.^2 .*w2[i]); -0.5*sum(yj.^2 .*w2[i]); -sum(xj.*yj.*w2[i])];
    B[i][1,2:end] = xj.*w2[i];
    B[i][2,2:end] = yj.*w2[i];
    B[i][3,2:end] = 0.5*xj.^2 .*w2[i];
    B[i][4,2:end] = 0.5*yj.^2 .*w2[i];
    B[i][5,2:end] = xj.*yj.*w2[i];
    C[i] = invA[i]*B[i];
end
println("Inverted least-squares matrices in ", round(time()-time3,digits=2), " s");
println("Matrices properties:");
println("  Max condition number: ",round(maximum(condA[internalNodes]),digits=2));
println("  Avg condition number: ",round(sum(condA[internalNodes])/length(internalNodes),digits=2));
println("  Min condition number: ",round(minimum(condA[internalNodes]),digits=2));


#implicit euler matrix assembly
time4 = time();
rows = Int[];
cols = Int[];
vals = Float64[];
for i in boundaryNodes
    push!(rows, i);
    push!(cols, i);
    push!(vals, 1);
end
for i in internalNodes
    push!(rows, i);
    push!(cols, i);
    push!(vals, (C[i][3,1]+C[i][4,1])*kcost/(rho*ccost));
    for j=1:lastindex(neighbors[i])
        push!(rows, i);
        push!(cols, neighbors[i][j]);
        push!(vals, (C[i][3,1+j]+C[i][4,1+j])*kcost/(rho*ccost));
    end
end
M = sparse(rows,cols,vals,N,N);
println("Completed matrix assembly in ", round(time()-time4,digits=2), " s");


#solve linear system
time5 = time();
u = u0.(pointcloud[1,:],pointcloud[2,:]);       #numerical solution
ue = uD.(pointcloud[1,:],pointcloud[2,:]);      #exact analytical solution
erru = ue-u;        #numerical solution error
b = zeros(N);       #rhs vector
for i in boundaryNodes
    b[i] = uD(pointcloud[1,i],pointcloud[2,i]);
end
for i in internalNodes
    b[i] = 0;
end
@printf("%6s | %10s\n","Step","max(err(u))");
@printf("%6i | %10.4e\n",0,maximum(erru));
u = M\b;

#error calculation
erru = ue-u;        #numerical solution error
maxerru = maximum(abs.(erru));
@printf("%6i | %10.4e\n",1,maxerru);
println("Linear system solved in ", round(time()-time5,digits=2), " s");

#solution plot
figure(4);
scatter(pointcloud[1,:],pointcloud[2,:],c=u,cmap="inferno");
colorbar();
title("Numerical solution");
axis("equal");
display(gcf());
