#----------------------------------------------------------------------
#  GFDM method test - 3d printing of a straight oblong line
#  Author: Andrea Pavan
#  Date: 19/03/2023
#  License: GPLv3-or-later
#----------------------------------------------------------------------
using Base.Threads;
using ElasticArrays;
using LinearAlgebra;
using SparseArrays;
using PyPlot;
include("utils.jl");


#problem definition
l1 = 0.55e-3;       #xextrusion width
l2 = 0.2e-3;        #layer height
l3 = 1e-3;          #extrusion segment length
L = 10e-3;          #total domain length
v = 40e-3;          #extrusion feed rate

k = 0.24;           #thermal conductivity
ρ = 1370;           #material density
cp = 1050;          #specific heat capacity
Tamb = 20;          #ambient temperature
Tbed = 50;          #bed temperature
h = 20;             #convective film coefficient
T0 = 230;           #initial temperature

meshSize = 0.05e-3;
surfaceMeshSize = meshSize;
minNeighbors = 40;
minSearchRadius = meshSize;

#mkdir("vid1");
#mkdir("vid2");
println("Number of threads: ",nthreads());


#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
segmentBoundaryNodes = Vector{Int}(undef,0);               #indices of the boundary nodes
segmentInternalNodes = Vector{Int}(undef,0);               #indices of the internal nodes
segmentNormals = ElasticArray{Float64}(undef,3,0);         #3xN matrix containing the components [nx;ny;nz] of the normal of each boundary node

segmentPointcloud = parseSU2mesh("17_3d_heat_filament_1613.su2");
#cornerPoint = findall(pointcloud[3,:].<=0+1e-6);
#pointcloud = pointcloud[:, setdiff(1:end,cornerPoint)];
#cornerPoint = findall(pointcloud[3,:].>=l3-1e-6);
#pointcloud = pointcloud[:, setdiff(1:end,cornerPoint)];
Nsegment = size(segmentPointcloud,2);
for i=1:Nsegment
    if segmentPointcloud[2,i]>=l2/2-1e-6
        #top surface
        push!(segmentBoundaryNodes, i);
        append!(segmentNormals, [0,1,0]);
    elseif segmentPointcloud[2,i]<=-l2/2+1e-6
        #bottom surface
        push!(segmentBoundaryNodes, i);
        append!(segmentNormals, [0,-1,0]);
    elseif segmentPointcloud[1,i]>=(l1-l2)/2-1e-6 && (segmentPointcloud[1,i]-(l1-l2)/2)^2+segmentPointcloud[2,i]^2>=(l2/2)^2
        #right bulb
        push!(segmentBoundaryNodes, i);
        append!(segmentNormals, [segmentPointcloud[1,i]-(l1-l2)/2,segmentPointcloud[2,i],0]./(l2/2));
    elseif segmentPointcloud[1,i]<=-(l1-l2)/2+1e-6 && (segmentPointcloud[1,i]+(l1-l2)/2)^2+segmentPointcloud[2,i]^2>=(l2/2)^2
        #left bulb
        push!(segmentBoundaryNodes, i);
        append!(segmentNormals, [segmentPointcloud[1,i]+(l1-l2)/2,segmentPointcloud[2,i],0]./(l2/2));
    elseif segmentPointcloud[3,i]<=1e-6
        #rear section
        push!(segmentBoundaryNodes, i);
        append!(segmentNormals, [0,0,-1]);
    elseif segmentPointcloud[3,i]>=l3-1e-6
        #front section
        push!(segmentBoundaryNodes, i);
        append!(segmentNormals, [0,0,1]);
    else
        push!(segmentInternalNodes, i);
        append!(segmentNormals, [0,0,0]);
    end
end

println("Generated pointcloud in ", round(time()-time1,digits=2), " s");
println("Pointcloud properties:");
println("  Boundary nodes: ",length(segmentBoundaryNodes));
println("  Internal nodes: ",length(segmentInternalNodes));
println("  Memory: ",memoryUsage(pointcloud,segmentBoundaryNodes));


#3d printing process simulation
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
N = size(pointcloud,2);
uprev = T0*ones(N);                                 #solution at the previous step
t = 0;                                              #simulation time
tf = 0.5;                                           #ending time
iter = 0;                                           #step number
idxCI = 0;                                          #reference internal node index
idxCB = 0;                                          #reference boundary node index
TCI = Float64[];                                    #reference internal node temperature history
TCB = Float64[];                                    #reference internal node temperature history
for s=l3:l3:L
    #generate pointcloud
    append!(pointcloud, segmentPointcloud);
    pointcloud[3,N+1:end] .+= s-l3;
    append!(normals, segmentNormals);
    append!(boundaryNodes, segmentBoundaryNodes.+N);
    append!(internalNodes, segmentInternalNodes.+N);
    append!(uprev,T0*ones(Nsegment));
    if N==0
        global idxCI = rand(segmentInternalNodes);
    end
    global N = size(pointcloud,2);

    #neighbor search (cartesian cells)
    time2 = time();
    neighbors = Vector{Vector{Int}}(undef,N);       #vector containing N vectors of the indices of each node segmentNeighbors
    Nneighbors = zeros(Int,N);                      #number of segmentNeighbors of each node
    (neighbors,Nneighbors,cell) = cartesianNeighborSearch(pointcloud,meshSize,minNeighbors);
    time3 = time();

    #neighbors distances and weights
    P = Vector{Array{Float64}}(undef,N);            #relative positions of the segmentNeighbors
    r2 = Vector{Vector{Float64}}(undef,N);          #relative distances of the segmentNeighbors
    w = Vector{Vector{Float64}}(undef,N);           #segmentNeighbors weights
    @threads 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

    #check boundary nodes
    for idx=1:lastindex(boundaryNodes)
        i = boundaryNodes[idx];
        for j in neighbors[i]
            if dot(pointcloud[:,j]-pointcloud[:,i],normals[:,i])>1e-9
                #the node i is no longer a boundary node
                boundaryNodes[idx] = 0;
                append!(internalNodes,i);
                normals[:,i] = [0,0,0];
            end
        end
    end
    deleteat!(boundaryNodes,findall(boundaryNodes.==0));

    #boundary conditions
    g1 = zeros(Float64,N);
    g2 = zeros(Float64,N);
    g3 = zeros(Float64,N);
    for i in boundaryNodes
        if pointcloud[2,i]<=-l2/2+1e-6
            #bottom surface
            g1[i] = 1.0;
            g2[i] = 0.0;
            g3[i] = Tbed;
        else
            #everywhere else
            g1[i] = h;
            g2[i] = k;
            g3[i] = h*Tamb;
            if idxCB==0
                global idxCB = i;
            end
        end
    end

    #least square matrix inversion
    C = Vector{Matrix}(undef,N);        #stencil coefficients matrices
    @threads 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*k/(ρ*cp), 2*k/(ρ*cp), 2*k/(ρ*cp), 0, 0, 0];
        W = Diagonal(vcat(w[i],wpde));
        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
    @threads for i in boundaryNodes
        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*k/(ρ*cp), 2*k/(ρ*cp), 2*k/(ρ*cp), 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(w[i],wpde,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
    time4 = time();

    #matrix assembly
    Δt = l3/v;                                      #timestep
    global iter += 1;
    global t += Δt;
    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);

    #linear system solution
    b = zeros(N);       #rhs vector
    for i in internalNodes
        b[i] = -C[i][1,1+Nneighbors[i]]*uprev[i]/Δt;
    end
    for i in boundaryNodes
        b[i] = -C[i][1,1+Nneighbors[i]]*uprev[i]/Δt + C[i][1,2+Nneighbors[i]]*g3[i];
    end
    u = M\b;
    dudt = (u-uprev)./Δt;
    println("t = ",round(t,digits=5),"; N = ",N,"; umax = ",round(maximum(u),digits=2));
    global uprev = u;
    push!(TCI,u[idxCI]);
    push!(TCB,u[idxCB]);
    time5 = time();
    println("Timings:");
    println("  Neighbors search: ",round(time3-time2,digits=2)," s");
    println("  LS matrix inversion: ",round(time4-time3,digits=2)," s");
    println("  Linear system: ",round(time5-time4,digits=2)," s");

    #temperature plot
    #=figure();
    plt = scatter3D(pointcloud[3,:],pointcloud[1,:],pointcloud[2,:],c=uprev,cmap="jet");
    title("Temperature - t = "*string(round(t,digits=5)));
    colorbar(plt);
    xlim(-2e-3,12e-3);
    ylim(-7e-3,7e-3);
    zlim(-7e-3,7e-3);
    savefig("vid1/"*string(lpad(iter,3,"0"))*".jpg");=#

    #heat fluxes
    #=dudx = zeros(N);
    dudy = zeros(N);
    dudz = zeros(N);
    for i=1:N
        for j=1:Nneighbors[i]
            dudx[i] += C[i][2,j]*u[neighbors[i][j]];    #C[i][j]*u[j]
            dudy[i] += C[i][3,j]*u[neighbors[i][j]];
            dudz[i] += C[i][4,j]*u[neighbors[i][j]];
        end
        dudx[i] += C[i][2,1+Nneighbors[i]]*dudt[i];     #C[i][1+N]*gPDE
        dudy[i] += C[i][3,1+Nneighbors[i]]*dudt[i];
        dudz[i] += C[i][4,1+Nneighbors[i]]*dudt[i];
    end
    for i in boundaryNodes
        dudx[i] += C[i][2,2+Nneighbors[i]]*g3[i];       #C[i][2+N]*gBC
        dudy[i] += C[i][3,2+Nneighbors[i]]*g3[i];
        dudz[i] += C[i][4,2+Nneighbors[i]]*g3[i];
    end
    qx = -k*dudx;
    qy = -k*dudy;
    qz = -k*dudz;

    #heatflux plot
    figure();
    plt = scatter3D(pointcloud[3,:],pointcloud[1,:],pointcloud[2,:],c=qz,cmap="jet");
    title("qz heatflux - t = "*string(round(t,digits=5)));
    colorbar(plt);
    axis("equal");
    #display(gcf());
    savefig("vid2/"*string(lpad(iter,3,"0"))*".jpg");=#

    #cooling simulation
    if s==L
        #temperature plot
        figure();
        plt = scatter3D(pointcloud[3,:],pointcloud[1,:],pointcloud[2,:],c=uprev,cmap="jet");
        title("Temperature - t = "*string(round(t,digits=5)));
        colorbar(plt);
        xlim(-2e-3,12e-3);
        ylim(-7e-3,7e-3);
        zlim(-7e-3,7e-3);
        display(gcf());
        
        while t<tf
            global t += Δt;
            global iter += 1;
            #linear system solution
            b = zeros(N);       #rhs vector
            for i in internalNodes
                b[i] = -C[i][1,1+Nneighbors[i]]*uprev[i]/Δt;
            end
            for i in boundaryNodes
                b[i] = -C[i][1,1+Nneighbors[i]]*uprev[i]/Δt + C[i][1,2+Nneighbors[i]]*g3[i];
            end
            u = M\b;
            dudt = (u-uprev)./Δt;
            println("t = ",round(t,digits=5),"; N = ",N,"; umax = ",round(maximum(u),digits=2));
            global uprev = u;
            push!(TCI,u[idxCI]);
            push!(TCB,u[idxCB]);

            #temperature plot
            #=figure();
            plt = scatter3D(pointcloud[3,:],pointcloud[1,:],pointcloud[2,:],c=uprev,cmap="jet");
            title("Temperature - t = "*string(round(t,digits=5)));
            colorbar(plt);
            xlim(-2e-3,12e-3);
            ylim(-7e-3,7e-3);
            zlim(-7e-3,7e-3);
            savefig("vid1/"*string(lpad(iter,3,"0"))*".jpg");=#
        end
    end
end
println("Simulation completed in ",round(time()-time1,digits=2)," s");

#reference node time history
Δt = l3/v;
figure();
plot(LinRange(Δt,t,length(TCI)),TCI,"k.",label="Random internal node");
plot(LinRange(Δt,t,length(TCB)),TCB,"r.",label="Random boundary node");
title("Temperature evolution for different nodes");
legend(loc="upper right");
xlabel("Time (s)");
ylabel("Temperature (°C)");
display(gcf());
