#----------------------------------------------------------------------
#  GFDM method test - 3d beam under flexural load
#  Author: Andrea Pavan
#  Date: 14/03/2023
#  License: GPLv3-or-later
#----------------------------------------------------------------------
using ElasticArrays;
using LinearAlgebra;
using SparseArrays;
using PyPlot;
include("utils.jl");


#problem definition
l1 = 20.0;      #domain x size
l2 = 1.0;       #domain y size
l3 = 1.0;       #domain z size
E = 71.7e9;     #Young modulus
ν = 0.33;       #Poisson ratio
qy = 1e5;       #distributed load on the top
t0 = 0;         #starting time
Δt = 0.01;       #timestep
tf = 0.05;       #ending time

meshSize = 0.2;
#meshSize = 0.1;
#meshSize = 0.075;
surfaceMeshSize = meshSize;
minNeighbors = 60;
minSearchRadius = meshSize;

μ = 0.5*E/(1+ν);        #Lamè coefficients
λ = E*ν/((1+ν)*(1-2*ν));


#read pointcloud from a SU2 file
time1 = time();
pointcloud = ElasticArray{Float64}(undef,3,0);      #3xN matrix containing the coordinates [X;Y;Z] of each node
boundaryNodes = Vector{Int}(undef,0);       #indices of the boundary nodes
internalNodes = Vector{Int}(undef,0);       #indices of the internal nodes
normals = ElasticArray{Float64}(undef,3,0);     #3xN matrix containing the components [nx;ny;nz] of the normal of each boundary node

pointcloud = parseSU2mesh("18b_3d_flexural_beam_mesh_4459.su2");
#pointcloud = parseSU2mesh("18b_3d_flexural_beam_mesh_24932.su2");
#pointcloud = parseSU2mesh("18b_3d_flexural_beam_mesh_51728.su2");
#cornerPoint = findall((pointcloud[2,:].<=1e-6).*(pointcloud[1,:].<=1e-6));
#pointcloud = pointcloud[:, setdiff(1:end,cornerPoint)];
N = size(pointcloud,2);
for i=1:N
    if pointcloud[1,i]<=0+1e-5
        #left surface
        push!(boundaryNodes, i);
        append!(normals, [-1,0,0]);
    elseif pointcloud[2,i]>=l2-1e-5
        #top surface
        push!(boundaryNodes, i);
        append!(normals, [0,1,0]);
    elseif pointcloud[1,i]>=l1-1e-5
        #right surface
        push!(boundaryNodes, i);
        append!(normals, [1,0,0]);
    elseif pointcloud[2,i]<=0+1e-5
        #bottom surface
        push!(boundaryNodes, i);
        append!(normals, [0,-1,0]);
    elseif pointcloud[3,i]<=0+1e-5
        #front surface
        push!(boundaryNodes, i);
        append!(normals, [0,0,-1]);
    elseif pointcloud[3,i]>=l3-1e-5
        #rear surface
        push!(boundaryNodes, i);
        append!(normals, [0,0,1]);
    else
        push!(internalNodes, i);
        append!(normals, [0,0,0]);
    end
end

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));


#boundary conditions
N = size(pointcloud,2);     #number of nodes
g1u = zeros(Float64,N);
g2u = zeros(Float64,N);
g3u = zeros(Float64,N);
g1v = zeros(Float64,N);
g2v = zeros(Float64,N);
g3v = zeros(Float64,N);
g1w = zeros(Float64,N);
g2w = zeros(Float64,N);
g3w = zeros(Float64,N);
for i in boundaryNodes
    if pointcloud[1,i]<=0+1e-5
        #left surface - u=0, v=0, w=0
        g1u[i] = 1.0;
        g2u[i] = 0.0;
        g3u[i] = 0.0;
        g1v[i] = 1.0;
        g2v[i] = 0.0;
        g3v[i] = 0.0;
        g1w[i] = 1.0;
        g2w[i] = 0.0;
        g3w[i] = 0.0;
    #=elseif pointcloud[2,i]>=l2-1e-5
        #top side - ty=qy
        g1u[i] = 0.0;
        g2u[i] = 1.0;
        g3u[i] = 0.0;
        g1v[i] = 0.0;
        g2v[i] = 1.0;
        g3v[i] = qy;
        g1w[i] = 0.0;
        g2w[i] = 1.0;
        g3w[i] = 0.0;=#
    elseif pointcloud[1,i]>l1-1e-5
        #right surface - u=0.5, v=0, w=0
        g1u[i] = 1.0;
        g2u[i] = 0.0;
        g3u[i] = 0.5μ;
        g1v[i] = 1.0;
        g2v[i] = 0.0;
        g3v[i] = 0.0;
        g1w[i] = 1.0;
        g2w[i] = 0.0;
        g3w[i] = 0.0;
    else
        #everywhere else
        g1u[i] = 0.0;
        g2u[i] = 1.0;
        g3u[i] = 0.0;
        g1v[i] = 0.0;
        g2v[i] = 1.0;
        g3v[i] = 0.0;
        g1w[i] = 0.0;
        g2w[i] = 1.0;
        g3w[i] = 0.0;
    end
end
Fx = zeros(Float64,N);      #volumetric loads
Fy = zeros(Float64,N);
Fz = zeros(Float64,N);

#boundary conditions plot
#=figure();
plt = scatter3D(pointcloud[1,:],pointcloud[2,:],pointcloud[3,:],c=g2v);
title("Numerical error");
axis("equal");
colorbar(plt);
display(gcf());=#


#neighbor search (cartesian cells)
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
(neighbors,Nneighbors,cell) = cartesianNeighborSearch(pointcloud,meshSize,minNeighbors);

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
w = Vector{Vector{Float64}}(undef,N);       #neighbors weights
for i=1:N
    P[i] = Array{Float64}(undef,3,Nneighbors[i]);
    r2[i] = Vector{Float64}(undef,Nneighbors[i]);
    w[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]
        w[i][j] = exp(-6*r2[i][j]/r2max);
        #w[i][j] = 1.0;
    end
end
wpde = 2.0;       #least squares weight for the pde
wbc = 2.0;        #least squares weight for the boundary condition


#least square matrix inversion
C = Vector{Matrix}(undef,N);       #stencil coefficients matrices
for i in internalNodes
    xj = P[i][1,:];
    yj = P[i][2,:];
    zj = P[i][2,:];
    V = zeros(Float64,3+3*Nneighbors[i],30);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, yj[j]*zj[j], xj[j]*zj[j], xj[j]*yj[j],  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        V[j+Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, yj[j]*zj[j], xj[j]*zj[j], xj[j]*yj[j],  0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        V[j+2*Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, yj[j]*zj[j], xj[j]*zj[j], xj[j]*yj[j]];
    end
    V[1+3*Nneighbors[i],:] = [0, 0, 0, 0, 2*(2+λ/μ), 2, 2, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ,  0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0];
    V[2+3*Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ,  0, 0, 0, 0, 2, 2*(2+λ/μ), 2, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0, 0];
    V[3+3*Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0,  0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0, 0,  0, 0, 0, 0, 2, 2, 2*(2+λ/μ), 0, 0, 0];
    W = Diagonal(vcat(w[i],w[i],w[i],wpde,wpde,wpde));
    VF = svd(W*V);
    C[i] = transpose(VF.Vt)*inv(Diagonal(VF.S))*transpose(VF.U)*W;
end
for i in boundaryNodes
    xj = P[i][1,:];
    yj = P[i][2,:];
    zj = P[i][2,:];
    V = zeros(Float64,6+3*Nneighbors[i],30);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, yj[j]*zj[j], xj[j]*zj[j], xj[j]*yj[j],  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        V[j+Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, yj[j]*zj[j], xj[j]*zj[j], xj[j]*yj[j],  0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        V[j+2*Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 0, 0, 0,  1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, yj[j]*zj[j], xj[j]*zj[j], xj[j]*yj[j]];
    end
    V[1+3*Nneighbors[i],:] = [0, 0, 0, 0, 2*(2+λ/μ), 2, 2, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ,  0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0];
    V[2+3*Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ,  0, 0, 0, 0, 2, 2*(2+λ/μ), 2, 0, 0, 0,  0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0, 0];
    V[3+3*Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0,  0, 0, 0, 0, 0, 0, 0, 1+λ/μ, 0, 0,  0, 0, 0, 0, 2, 2, 2*(2+λ/μ), 0, 0, 0];
    V[4+3*Nneighbors[i],:] = [g1u[i], g2u[i]*normals[1,i]*(2+λ/μ), g2u[i]*normals[2,i], g2u[i]*normals[3,i], 0, 0, 0, 0, 0, 0,  0, g2u[i]*normals[2,i], g2u[i]*normals[1,i]*λ/μ, 0, 0, 0, 0, 0, 0, 0,  0, g2u[i]*normals[3,i], 0, g2u[i]*normals[1,i]*λ/μ, 0, 0, 0, 0, 0, 0];
    V[5+3*Nneighbors[i],:] = [0, g2v[i]*normals[2,i]*λ/μ, g2v[i]*normals[1,i], 0, 0, 0, 0, 0, 0, 0,  g1v[i], g2v[i]*normals[1,i], g2v[i]*normals[2,i]*(2+λ/μ), g2v[i]*normals[3,i], 0, 0, 0, 0, 0, 0,  0, 0, g2v[i]*normals[3,i], g2v[i]*normals[2,i]*λ/μ, 0, 0, 0, 0, 0, 0];
    V[6+3*Nneighbors[i],:] = [0, g2w[i]*normals[3,i]*λ/μ, 0, g2w[i]*normals[1,i], 0, 0, 0, 0, 0, 0,  0, 0, g2w[i]*normals[3,i]*λ/μ, g2w[i]*normals[2,i], 0, 0, 0, 0, 0, 0,  g1w[i], g2w[i]*normals[1,i], g2w[i]*normals[2,i], g2w[i]*normals[3,i]*(2+λ/μ), 0, 0, 0, 0, 0, 0];
    W = Diagonal(vcat(w[i],w[i],w[i],wpde,wpde,wpde,wbc,wbc,wbc));
    VF = svd(W*V);
    C[i] = transpose(VF.Vt)*inv(Diagonal(VF.S))*transpose(VF.U)*W;
    #C[i] = transpose(VF.Vt)*inv(Diagonal(VF.S.+1e-12))*transpose(VF.U)*W;
end
println("Inverted least-squares matrices in ", round(time()-time3,digits=2), " s");


#matrix assembly
time4 = time();
rows = Int[];
cols = Int[];
vals = Float64[];
for i=1:N
    #u equation
    push!(rows, i);
    push!(cols, i);
    push!(vals, 1-C[i][1,1+3*Nneighbors[i]]/Δt);
    for j=1:Nneighbors[i]
        push!(rows, i);
        push!(cols, neighbors[i][j]);
        push!(vals, -C[i][1,j]);
        push!(rows, i);
        push!(cols, N+neighbors[i][j]);
        push!(vals, -C[i][1,j+Nneighbors[i]]);
        push!(rows, i);
        push!(cols, 2*N+neighbors[i][j]);
        push!(vals, -C[i][1,j+2*Nneighbors[i]]);
    end
    push!(rows, i);
    push!(cols, N+i);
    push!(vals, -C[i][1,2+3*Nneighbors[i]]/Δt);
    push!(rows, i);
    push!(cols, 2*N+i);
    push!(vals, -C[i][1,3+3*Nneighbors[i]]/Δt);

    #v equation
    push!(rows, N+i);
    push!(cols, N+i);
    push!(vals, 1-C[i][11,2+3*Nneighbors[i]]/Δt);
    for j=1:Nneighbors[i]
        push!(rows, N+i);
        push!(cols, neighbors[i][j]);
        push!(vals, -C[i][11,j]);
        push!(rows, N+i);
        push!(cols, N+neighbors[i][j]);
        push!(vals, -C[i][11,j+Nneighbors[i]]);
        push!(rows, N+i);
        push!(cols, 2*N+neighbors[i][j]);
        push!(vals, -C[i][11,j+2*Nneighbors[i]]);
    end
    push!(rows, N+i);
    push!(cols, i);
    push!(vals, -C[i][11,1+3*Nneighbors[i]]/Δt);
    push!(rows, N+i);
    push!(cols, 2*N+i);
    push!(vals, -C[i][11,3+3*Nneighbors[i]]/Δt);

    #w equation
    push!(rows, 2*N+i);
    push!(cols, 2*N+i);
    push!(vals, 1-C[i][21,3+3*Nneighbors[i]]/Δt);
    for j=1:Nneighbors[i]
        push!(rows, 2*N+i);
        push!(cols, neighbors[i][j]);
        push!(vals, -C[i][21,j]);
        push!(rows, 2*N+i);
        push!(cols, N+neighbors[i][j]);
        push!(vals, -C[i][21,j+Nneighbors[i]]);
        push!(rows, 2*N+i);
        push!(cols, 2*N+neighbors[i][j]);
        push!(vals, -C[i][21,j+2*Nneighbors[i]]);
    end
    push!(rows, 2*N+i);
    push!(cols, i);
    push!(vals, -C[i][21,1+3*Nneighbors[i]]/Δt);
    push!(rows, 2*N+i);
    push!(cols, N+i);
    push!(vals, -C[i][21,2+3*Nneighbors[i]]/Δt);
end
M = sparse(rows,cols,vals,3*N,3*N);
println("Completed matrix assembly in ", round(time()-time4,digits=2), " s");


#fictious time propagation
time5 = time();
uprev = zeros(3*N);     #solution at the previous step
for t=t0:Δt:tf
    b = zeros(3*N);         #rhs vector
    for i=1:N
        b[i] = -C[i][1,1+3*Nneighbors[i]]*uprev[i]/Δt -C[i][1,2+3*Nneighbors[i]]*uprev[N+i]/Δt -C[i][1,3+3*Nneighbors[i]]*uprev[2*N+i]/Δt;
        b[N+i] = -C[i][11,1+3*Nneighbors[i]]*uprev[i]/Δt -C[i][11,2+3*Nneighbors[i]]*uprev[N+i]/Δt -C[i][11,3+3*Nneighbors[i]]*uprev[2*N+i]/Δt;
        b[2*N+i] = -C[i][21,1+3*Nneighbors[i]]*uprev[i]/Δt -C[i][21,2+3*Nneighbors[i]]*uprev[N+i]/Δt -C[i][21,3+3*Nneighbors[i]]*uprev[2*N+i]/Δt;
        b[i] += -C[i][1,1+3*Nneighbors[i]]*Fx[i]/μ -C[i][1,2+3*Nneighbors[i]]*Fy[i]/μ -C[i][1,3+3*Nneighbors[i]]*Fz[i]/μ;     #volumetric loads
        b[N+i] += -C[i][11,1+3*Nneighbors[i]]*Fx[i]/μ -C[i][11,2+3*Nneighbors[i]]*Fy[i]/μ -C[i][11,3+3*Nneighbors[i]]*Fz[i]/μ;
        b[2*N+i] += -C[i][21,1+3*Nneighbors[i]]*Fx[i]/μ -C[i][21,2+3*Nneighbors[i]]*Fy[i]/μ -C[i][21,3+3*Nneighbors[i]]*Fz[i]/μ;
    end
    for i in boundaryNodes
        b[i] += C[i][1,4+3*Nneighbors[i]]*g3u[i]/μ + C[i][1,5+3*Nneighbors[i]]*g3v[i]/μ + C[i][1,6+3*Nneighbors[i]]*g3w[i]/μ;
        b[N+i] += C[i][11,4+3*Nneighbors[i]]*g3u[i]/μ + C[i][11,5+3*Nneighbors[i]]*g3v[i]/μ + C[i][11,6+3*Nneighbors[i]]*g3w[i]/μ;
        b[2*N+i] += C[i][21,4+3*Nneighbors[i]]*g3u[i]/μ + C[i][21,5+3*Nneighbors[i]]*g3v[i]/μ + C[i][21,6+3*Nneighbors[i]]*g3w[i]/μ;
    end
    sol = qr(M)\b;
    println("t = ",t,"; vmax = ",maximum(abs.(sol[N+1:2*N])));
    global uprev = sol;
end
println("Linear system solved in ", round(time()-time5,digits=2), " s");


#displacement plot
u = uprev[1:N];
v = uprev[N+1:2*N];
w = uprev[2*N+1:3*N];
figure();
scatter(pointcloud[1,:],pointcloud[2,:],c=u,cmap="jet");
colorbar();
title("Numerical solution - x displacement");
axis("equal");
display(gcf());
figure();
scatter(pointcloud[1,:],pointcloud[2,:],c=v,cmap="jet");
colorbar();
title("Numerical solution - y displacement");
axis("equal");
display(gcf());
figure();
scatter(pointcloud[1,:],pointcloud[2,:],c=w,cmap="jet");
colorbar();
title("Numerical solution - z displacement");
axis("equal");
display(gcf());


#deformed beam 3d plot
#=figure();
plt = scatter3D(pointcloud[1,:]+2*u,pointcloud[2,:]+2*v,pointcloud[3,:]+2*w,c=sqrt.(u.^2+v.^2+w.^2),cmap="Oranges");
title("Numerical solution - deflection");
axis("equal");
colorbar(plt);
display(gcf());=#

println("Euler-Bernoulli beam theory:");
println("> Max deflection: vmax = ",qy*(l1^4)/(8*E*(l2^4)/12));
println("GFDM:");
println("> Max deflection: vmax = ",maximum(abs.(v)));
