#----------------------------------------------------------------------
#  GFDM method test - 2d beam thermal analysis + mechanical post-proc
#  Reference: https://reference.wolfram.com/language/PDEModels/tutorial/Multiphysics/ModelCollection/ThermalLoad.html
#  Author: Andrea Pavan
#  Date: 24/02/2023
#  License: GPLv3-or-later
#----------------------------------------------------------------------
using ElasticArrays;
using LinearAlgebra;
using SparseArrays;
using PyPlot;
include("utils.jl");


#problem definition
l1 = 5.0;       #domain x size
l2 = 0.5;       #domain y size
E = 210e9;      #Young modulus
ν = 0.3;        #Poisson ratio
α = 12.0e-6;        #thermal expansion coefficient
k = 44;     #thermal conductivity
ρ = 7500;       #material density
cp = 470;       #specific heat capacity
Tref = 20;      #thermal expansion reference temperature
Tcold = 0;      #left and lower surfaces temperature
qtop = 1e4;     #top surface heat flux
t0 = 0;     #starting time
Δt = 60;       #timestep
tf = 10800;     #ending time

meshSize = 0.05;
#meshSize = 0.025;
surfaceMeshSize = meshSize;
minNeighbors = 10;
minSearchRadius = meshSize;


#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:surfaceMeshSize:l2-surfaceMeshSize
    append!(pointcloud, [0,i]);
    append!(normals, [-1,0]);
    append!(pointcloud, [l1,i]);
    append!(normals, [1,0]);
end
#for i=0+surfaceMeshSize:surfaceMeshSize:l1-surfaceMeshSize
for i=0:surfaceMeshSize:l1
    append!(pointcloud, [i,0]);
    append!(normals, [0,-1]);
    append!(pointcloud, [i,l2]);
    append!(normals, [0,1]);
end
boundaryNodes = collect(range(1,size(pointcloud,2)));
for x=0+meshSize:meshSize:l1-meshSize
    for y=0+meshSize:meshSize:l2-meshSize
        newP = [x,y]+(rand(Float64,2).-0.5).*meshSize/5;
        insertP = true;
        for j in boundaryNodes
            if (newP[1]-pointcloud[1,j])^2+(newP[2]-pointcloud[2,j])^2<(0.75*meshSize)^2
                insertP = false;
            end
        end
        if insertP
            append!(pointcloud, newP);
            append!(normals, [0,0]);
        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));

#pointcloud plot
#=figure();
plot(pointcloud[1,boundaryNodes],pointcloud[2,boundaryNodes],"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
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,2,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


#thermal - boundary conditions
g1 = zeros(Float64,N);
g2 = zeros(Float64,N);
g3 = zeros(Float64,N);
for i in boundaryNodes
    if pointcloud[1,i]<=0+1e-6
        #left surface
        g1[i] = 1.0;
        g2[i] = 0.0;
        g3[i] = Tcold;
    elseif pointcloud[2,i]<=0+1e-6
        #bottom surface
        g1[i] = 1.0;
        g2[i] = 0.0;
        g3[i] = Tcold;
    elseif pointcloud[2,i]>=l2-1e-6
        #upper surface
        g1[i] = 0.0;
        g2[i] = k;
        g3[i] = qtop;
    else
        #everywhere else
        g1[i] = 0.0;
        g2[i] = 1.0;
        g3[i] = 0.0;
    end
end
Fx = zeros(Float64,N);      #volumetric loads
Fy = zeros(Float64,N);


#thermal - least square matrix inversion
C = Vector{Matrix}(undef,N);        #stencil 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[1+Nneighbors[i],:] = [0, 0, 0, 2*k/(ρ*cp), 2*k/(ρ*cp), 0];
    W = Diagonal(vcat(w[i],wpde));
    (Q,R) = qr(W*V);
    C[i] = inv(R)*transpose(Matrix(Q))*W;
end
for i in boundaryNodes
    #println("Boundary node: ",i);
    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[1+Nneighbors[i],:] = [0, 0, 0, 2*k/(ρ*cp), 2*k/(ρ*cp), 0];
    V[2+Nneighbors[i],:] = [g1[i], g2[i]*normals[1,i], g2[i]*normals[2,i], 0, 0, 0];
    W = Diagonal(vcat(w[i],wpde,wbc));
    (Q,R) = qr(W*V);
    C[i] = inv(R)*transpose(Matrix(Q))*W;
end
println("Thermal - Inverted least-squares matrices in ", round(time()-time3,digits=2), " s");


#mechanical - boundary conditions
g1u = zeros(Float64,N);
g2u = zeros(Float64,N);
g3u = zeros(Float64,N);
g1v = zeros(Float64,N);
g2v = zeros(Float64,N);
g3v = zeros(Float64,N);
for i in boundaryNodes
    if pointcloud[1,i]<=0+1e-6
        #left side - u=0, v=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;
    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;
    end
end


#mechanical - least square matrix inversion
Cm = Vector{Matrix}(undef,N);       #stencil coefficients matrices
for i in internalNodes
    xj = P[i][1,:];
    yj = P[i][2,:];
    V = zeros(Float64,2+2*Nneighbors[i],12);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], xj[j]^2, yj[j]^2, xj[j]*yj[j], 0, 0, 0, 0, 0, 0];
        V[j+Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 1, xj[j], yj[j], xj[j]^2, yj[j]^2, xj[j]*yj[j]];
    end
    V[1+2*Nneighbors[i],:] = [0, 0, 0, 2*1/(1-ν^2), 2*0.5*1/(1+ν), 0,  0, 0, 0, 0, 0, 1*ν/(1-ν^2)+0.5*1/(1+ν)];
    V[2+2*Nneighbors[i],:] = [0, 0, 0, 0, 0, 1*ν/(1-ν^2)+0.5*1/(1+ν),  0, 0, 0, 2*0.5*1/(1+ν), 2*1/(1-ν^2), 0];
    W = Diagonal(vcat(w[i],w[i],wpde,wpde));
    VF = svd(W*V);
    Cm[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,:];
    nx = normals[1,i];
    ny = normals[2,i];
    V = zeros(Float64,4+2*Nneighbors[i],12);
    for j=1:Nneighbors[i]
        V[j,:] = [1, xj[j], yj[j], xj[j]^2, yj[j]^2, xj[j]*yj[j], 0, 0, 0, 0, 0, 0];
        V[j+Nneighbors[i],:] = [0, 0, 0, 0, 0, 0, 1, xj[j], yj[j], xj[j]^2, yj[j]^2, xj[j]*yj[j]];
    end
    V[1+2*Nneighbors[i],:] = [0, 0, 0, 2*1/(1-ν^2), 2*0.5*1/(1+ν), 0,  0, 0, 0, 0, 0, 1*ν/(1-ν^2)+0.5*1/(1+ν)];
    V[2+2*Nneighbors[i],:] = [0, 0, 0, 0, 0, 1*ν/(1-ν^2)+0.5*1/(1+ν),  0, 0, 0, 2*0.5*1/(1+ν), 2*1/(1-ν^2), 0];
    V[3+2*Nneighbors[i],:] = [g1u[i], g2u[i]*nx*1/(1-ν^2), g2u[i]*ny*0.5*1/(1+ν), 0, 0, 0,  0, g2u[i]*ny*0.5*1/(1+ν), g2u[i]*nx*1*ν/(1-ν^2), 0, 0, 0];
    V[4+2*Nneighbors[i],:] = [0, g2v[i]*ny*1*ν/(1-ν^2), g2v[i]*nx*0.5*1/(1+ν), 0, 0, 0,  g1v[i], g2v[i]*nx*0.5*1/(1+ν), g2v[i]*ny*1/(1-ν^2), 0, 0, 0];
    W = Diagonal(vcat(w[i],w[i],wpde,wpde,wbc,wbc));
    VF = svd(W*V);
    Cm[i] = transpose(VF.Vt)*inv(Diagonal(VF.S))*transpose(VF.U)*W;
end
println("Mechanical - Inverted mechanical least-squares matrices");


#thermal - matrix assembly
time4 = time();
rows = Int[];
cols = Int[];
vals = Float64[];
for i=1:N
    push!(rows, i);
    push!(cols, i);
    push!(vals, 1-C[i][1,1+Nneighbors[i]]/Δt);
    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("Thermal - Completed matrix assembly in ", round(time()-time4,digits=2), " s");


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

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

    #zu equation
    push!(rows, 2*N+i);
    push!(cols, 2*N+i);
    push!(vals, 1);
    push!(rows, 2*N+i);
    push!(cols, i);
    push!(vals, -1/Δt);

    #zv equation
    push!(rows, 3*N+i);
    push!(cols, 3*N+i);
    push!(vals, 1);
    push!(rows, 3*N+i);
    push!(cols, N+i);
    push!(vals, -1/Δt);
end
Mm = sparse(rows,cols,vals,4*N,4*N);
println("Mechanical - Completed matrix assembly");


#time propagation
time5 = time();
T = Tcold*ones(N);      #solution at the previous step
dTdt = zeros(N);        #temperature time derivative
u = zeros(N);       #initial x displacement
v = zeros(N);       #initial y displacement
zu = zeros(N);      #initial x speed
zv = zeros(N);      #initial y speed
dTdx = zeros(N);        #temperature gradient along x
dTdy = zeros(N);        #temperature gradient along y
qx = zeros(N);      #heat flux along x
qy = zeros(N);      #heat flux along y

#initial plot
#=global iter = 0;
titlestr = "Numerical solution - t = 0";
filenamestr = "vid/000.jpg";=#

for t=t0+Δt:Δt:tf
    #thermal problem
    b = zeros(N);       #rhs vector
    for i=1:N
        b[i] = -C[i][1,1+Nneighbors[i]]*T[i]/Δt;
    end
    for i in boundaryNodes
        b[i] += C[i][1,2+Nneighbors[i]]*g3[i];
    end
    Tnew = M\b;
    global dTdt = (Tnew-T)./Δt;
    global T = Tnew;

    #heat fluxes
    global dTdx = zeros(N);
    global dTdy = zeros(N);
    for i=1:N
        for j=1:Nneighbors[i]
            global dTdx[i] += C[i][2,j]*T[neighbors[i][j]];        #C[i][j]*u[j]
            global dTdy[i] += C[i][3,j]*T[neighbors[i][j]];
        end
        global dTdx[i] += C[i][2,1+Nneighbors[i]]*dTdt[i];     #C[i][1+N]*gPDE
        global dTdy[i] += C[i][3,1+Nneighbors[i]]*dTdt[i];
    end
    for i in boundaryNodes
        global dTdx[i] += C[i][2,2+Nneighbors[i]]*g3[i];       #C[i][2+N]*gBC
        global dTdy[i] += C[i][3,2+Nneighbors[i]]*g3[i];
    end
    global qx = -k*dTdx;
    global qy = -k*dTdy;

    #mechanical problem
    bm = zeros(4*N);       #rhs vector
    for i=1:N
        bm[i] = 0;
        bm[N+i] = 0;
        bm[i] += Cm[i][1,1+2*Nneighbors[i]]*(-(ρ/E)*zu[i]/Δt - Fx[i]/E + 1*α*dTdx[i]/(1-ν)) + Cm[i][1,2+2*Nneighbors[i]]*(-(ρ/E)*zv[i]/Δt - Fy[i]/E + 1*α*dTdy[i]/(1-ν));
        bm[N+i] += Cm[i][7,1+2*Nneighbors[i]]*(-(ρ/E)*zu[i]/Δt - Fx[i]/E + 1*α*dTdx[i]/(1-ν)) + Cm[i][7,2+2*Nneighbors[i]]*(-(ρ/E)*zv[i]/Δt - Fy[i]/E + 1*α*dTdy[i]/(1-ν));
        bm[2*N+i] = -u[i]/Δt;
        bm[3*N+i] = -v[i]/Δt;
    end
    for i in boundaryNodes
        nx = normals[1,i];
        ny = normals[2,i];
        bm[i] += Cm[i][1,3+2*Nneighbors[i]]*(g3u[i]/E + g2u[i]*nx*1*α*(T[i]-Tref)/(1-ν)) + Cm[i][1,4+2*Nneighbors[i]]*(g3v[i]/E + g2v[i]*ny*1*α*(T[i]-Tref)/(1-ν));
        bm[N+i] += Cm[i][7,3+2*Nneighbors[i]]*(g3u[i]/E + g2u[i]*nx*1*α*(T[i]-Tref)/(1-ν)) + Cm[i][7,4+2*Nneighbors[i]]*(g3v[i]/E + g2v[i]*ny*1*α*(T[i]-Tref)/(1-ν));
    end
    sol = Mm\bm;
    global u = sol[1:N];
    global v = sol[N+1:2*N];
    global zu = sol[2*N+1:3*N];
    global zv = sol[3*N+1:4*N];
    println("t = ",t,"; Tmax = ",round(maximum(T),digits=2),"; umax = ",round(maximum(abs.(u)),digits=4),"; vmax = ",round(maximum(abs.(v)),digits=4));

    #transient plot
    #=global iter += 1;
    global titlestr = "Numerical solution - t = "*string(t);
    global filenamestr = "vid/"*string(lpad(iter,3,"0"))*".jpg";
    figure();
    scatter(pointcloud[1,:]+20*u,pointcloud[2,:]+20*v,c=T,cmap="inferno");
    title(titlestr);
    colorbar();
    axis("equal");
    savefig(filenamestr);=#
end
println("Simulation completed in ", round(time()-time5,digits=2), " s");

#thermal solution plot
#=figure();
plt = scatter(pointcloud[1,:],pointcloud[2,:],c=T,cmap="inferno");
title("Numerical solution - temperature");
axis("equal");
colorbar(plt);
display(gcf());=#


#heat fluxes plot
#=figure();
plt = scatter(pointcloud[1,:],pointcloud[2,:],c=qx,cmap="jet");
title("Numerical solution - qx heat flux");
axis("equal");
colorbar(plt);
display(gcf());
figure();
plt = scatter(pointcloud[1,:],pointcloud[2,:],c=qy,cmap="jet");
title("Numerical solution - qy heat flux");
axis("equal");
colorbar(plt);
display(gcf());=#


#displacement plot
#=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="Oranges");
scatter(pointcloud[1,:],pointcloud[2,:],c=v,cmap="jet");
colorbar();
title("Numerical solution - y displacement");
axis("equal");
display(gcf());=#

#deformed beam plot
figure();
scatter(pointcloud[1,:]+20*u,pointcloud[2,:]+20*v,c=T,cmap="inferno");
colorbar();
title("Numerical solution");
axis("equal");
display(gcf());

println("\nWolfram reference solution:");
println("> umax = 0.0049009");
println("> vmax = 0.0243852");
#println("CalculiX FEM:");
#println("> umax = 0.0050147");
#println("> vmax = 0.0258213");
#println("> Tmax = 88.29");
println("GFDM:");
println("> umax = ",maximum(abs.(u)));
println("> vmax = ",maximum(abs.(v)));
println("> Tmax = ",maximum(abs.(T)));
