#----------------------------------------------------------------------
#  GFDM method test - 3d heat sink
#  Same test case of PrePoMax tutorial #21
#  Author: Andrea Pavan
#  Date: 08/01/2023
#  License: GPLv3-or-later
#----------------------------------------------------------------------
using ElasticArrays;
using LinearAlgebra;
using SparseArrays;
using PyPlot;
include("utils.jl");


#problem definition
l1 = 54.4e-3;       #domain x size
l2 = 43e-3;     #domain y size
l3 = 60e-3;     #domain z size
tb = 3e-3;		#bottom thickness
tf = 4e-3;		#fins thickness
sf = 12.8e-3;		#fins spacing
kconst = 237;       #thermal conductivity
qb = 18382.4;       #bottom heat flux
Tamb = 20;      #ambient temperature
hconst = 23;        #convective film coefficient

meshSize = 1.65e-3;
surfaceMeshSize = 1.65e-3;
minNeighbors = 60;
minSearchRadius = meshSize/2;


#read pointcloud from a SU2 file
time1 = time();
pointcloud = ElasticArray{Float64}(undef,3,0);      #2xN 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("14_3d_heat_sink_mesh_32082.su2");
#pointcloud = parseSU2mesh("14_3d_heat_sink_mesh_83267.su2");
#pointcloud = parseSU2mesh("14_3d_heat_sink_mesh_169186.su2");
N = size(pointcloud,2);

for i=1:N
    if pointcloud[2,i]<=1e-6
        #bottom surface
        push!(boundaryNodes, i);
        append!(normals, [0,-1,0]);
    elseif pointcloud[2,i]>=l2-1e-6
        #fins top
        push!(boundaryNodes, i);
        append!(normals, [0,1,0]);
    elseif pointcloud[3,i]<=1e-6
        #lateral surface
        push!(boundaryNodes, i);
        append!(normals, [0,0,-1]);
    elseif pointcloud[3,i]>=l3-1e-6
        #lateral surface
        push!(boundaryNodes, i);
        append!(normals, [0,0,1]);
    elseif tb-1e-6<=pointcloud[2,i]<=tb+1e-6 && (l1-tf-sf-1e-6<=pointcloud[1,i]<=l1-tf+1e-6)
        #bottom surface between the fins
        push!(boundaryNodes, i);
        append!(normals, [0,1,0]);
    elseif tb-1e-6<=pointcloud[2,i]<=tb+1e-6 && (l1-2*tf-2*sf-1e-6<=pointcloud[1,i]<=l1-2*tf-sf+1e-6)
        #bottom surface between the fins
        push!(boundaryNodes, i);
        append!(normals, [0,1,0]);
    elseif tb-1e-6<=pointcloud[2,i]<=tb+1e-6 && (l1-3*tf-3*sf-1e-6<=pointcloud[1,i]<=l1-3*tf-2*sf+1e-6)
        #bottom surface between the fins
        push!(boundaryNodes, i);
        append!(normals, [0,1,0]);
    elseif pointcloud[1,i]<=1e-6
        #vertical lateral surface
        push!(boundaryNodes, i);
        append!(normals, [-1,0,0]);
    elseif pointcloud[1,i]>=l1-1e-6
        #vertical lateral surface
        push!(boundaryNodes, i);
        append!(normals, [1,0,0]);
    elseif l1-tf-1e-6<=pointcloud[1,i]<=l1-tf+1e-6 && pointcloud[2,i]>=tb-1e-6
        #vertical surface between fins
        push!(boundaryNodes, i);
        append!(normals, [-1,0,0]);
    elseif l1-tf-sf-1e-6<=pointcloud[1,i]<=l1-tf-sf+1e-6 && pointcloud[2,i]>=tb-1e-6
        #vertical surface between fins
        push!(boundaryNodes, i);
        append!(normals, [1,0,0]);
    elseif l1-2*tf-sf-1e-6<=pointcloud[1,i]<=l1-2*tf-sf+1e-6 && pointcloud[2,i]>=tb-1e-6
        #vertical surface between fins
        push!(boundaryNodes, i);
        append!(normals, [-1,0,0]);
    elseif l1-2*tf-2*sf-1e-6<=pointcloud[1,i]<=l1-2*tf-2*sf+1e-6 && pointcloud[2,i]>=tb-1e-6
        #vertical surface between fins
        push!(boundaryNodes, i);
        append!(normals, [1,0,0]);
    elseif l1-3*tf-2*sf-1e-6<=pointcloud[1,i]<=l1-3*tf-2*sf+1e-6 && pointcloud[2,i]>=tb-1e-6
        #vertical surface between fins
        push!(boundaryNodes, i);
        append!(normals, [-1,0,0]);
    elseif l1-3*tf-3*sf-1e-6<=pointcloud[1,i]<=l1-3*tf-3*sf+1e-6 && pointcloud[2,i]>=tb-1e-6
        #vertical surface between fins
        push!(boundaryNodes, i);
        append!(normals, [1,0,0]);
    else
        push!(internalNodes, i);
    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
g1 = zeros(Float64,N);
g2 = zeros(Float64,N);
g3 = zeros(Float64,N);
for i in boundaryNodes
    if pointcloud[2,i]<=1e-6
        #bottom surface
        g1[i] = 0.0;
        g2[i] = kconst;
        g3[i] = qb;
    else
        #everywhere else
        g1[i] = hconst;
        g2[i] = kconst;
        g3[i] = hconst*Tamb;
    end
end

#=
#neighbor search (brute force)
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=#

#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),")");

#connectivity plot
#=figure();
plot1Idx = rand(1:N,5);
#plot1Idx = rand(1+length(boundaryNodes):N,5);
#plot1Idx = [40, 9, 30064];
plot3D(pointcloud[1,:],pointcloud[2,:],pointcloud[3,:],marker=".",linestyle="None",color="lightgray");
for i in plot1Idx
    connColor = rand(3);
    plot3D(pointcloud[1,neighbors[i]],pointcloud[2,neighbors[i]],pointcloud[3,neighbors[i]],marker=".",linestyle="None",color=connColor);
    for j in neighbors[i]
        plot3D([pointcloud[1,i],pointcloud[1,j]],[pointcloud[2,i],pointcloud[2,j]],[pointcloud[3,i],pointcloud[3,j]],"-",color=connColor);
    end
end
plot3D(pointcloud[1,plot1Idx],pointcloud[2,plot1Idx],pointcloud[3,plot1Idx],"k.");
title("Connectivity plot");
axis("equal");
display(gcf());=#


#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,3,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(-6*r2[i][j]/r2max)^2;
    end
end
w2pde = 2.0;        #least squares weight for the pde
w2bc = 2.0;     #least squares weight for the boundary condition


#least square matrix inversion
A = Vector{Matrix}(undef,N);        #least-squares matrices
C = Vector{Matrix}(undef,N);        #derivatives coefficients matrices
for i in internalNodes
    xj = P[i][1,:];
    yj = P[i][2,:];
    zj = P[i][3,:];
    V = zeros(Float64,1+Nneighbors[i],10);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, xj[j]*yj[j], xj[j]*zj[j], yj[j]*zj[j]];
    end
    V[1+Nneighbors[i],:] = [0, 0, 0, 0, 2, 2, 2, 0, 0, 0];
    W = Diagonal(vcat(w2[i],w2pde));
    A[i] = transpose(V)*W*V;
    #(Q,R) = qr(A[i]);
    #C[i] = inv(R)*transpose(Q)*transpose(V)*W;
    (Ui,Si,Vi) = svd(A[i]);
    C[i] = Vi*inv(Diagonal(Si))*transpose(Ui)*transpose(V)*W;
end
for i in boundaryNodes
    #println("Boundary node: ",i);
    xj = P[i][1,:];
    yj = P[i][2,:];
    zj = P[i][3,:];
    V = zeros(Float64,2+Nneighbors[i],10);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], zj[j], xj[j]^2, yj[j]^2, zj[j]^2, xj[j]*yj[j], xj[j]*zj[j], yj[j]*zj[j]];
    end
    V[1+Nneighbors[i],:] = [0, 0, 0, 0, 2, 2, 2, 0, 0, 0];
    V[2+Nneighbors[i],:] = [g1[i], g2[i]*normals[1,i], g2[i]*normals[2,i], g2[i]*normals[3,i], 0, 0, 0, 0, 0, 0];
    W = Diagonal(vcat(w2[i],w2pde,w2bc));
    A[i] = transpose(V)*W*V;
    #(Q,R) = qr(A[i]);
    #C[i] = inv(R)*transpose(Q)*transpose(V)*W;
    (Ui,Si,Vi) = svd(A[i]);
    C[i] = Vi*inv(Diagonal(Si))*transpose(Ui)*transpose(V)*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
    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 internalNodes
    b[i] = 0;
end
for i in boundaryNodes
    b[i] = C[i][1,end]*g3[i];
end
u = M\b;
println("Linear system solved in ", round(time()-time5,digits=2), " s");

#solution plot
figure();
#plt = scatter3D(pointcloud[1,internalNodes],pointcloud[2,internalNodes],pointcloud[3,internalNodes],c=u[internalNodes],cmap="inferno");
plt = scatter3D(pointcloud[1,:],pointcloud[2,:],pointcloud[3,:],c=u,cmap="jet");
title("Numerical solution");
axis("equal");
colorbar(plt);
display(gcf());
