%--------------------------
 % @Author: Jingqiao Hu
 % @Date: 2022-01-14 21:23:30
 % @LastEditTime: 2022-05-09 20:33:32
%--------------------------
dbstop if error
feature accel on
addpath util util_simulation util_bezier util_micro util_mtimesx util_output
addpath util_cvt mexCVT util_tetmesh util_MMC util_intersection

%% adjustable paras
opt_load_rho = 1;
opt_load_cvt = 1;
opt_load_mmc = 0;
opt_load_def = 0;

volfrac = 0.3;

t = 8; t0 = 0.1; t1 = 15; % for MMC
pb = 0.1; % for blending
% epsilon = 5.8; % for Heaviside
% step = 1.6; % for Heaviside
dgt0 = 5;

e1 = 1; nu0 = 0.3; e0 = 1e-5; penal = 3;
vF = 1;

nbridge = 0;

scalarE0 = 0.2;

optDesign = 'cantilever';
% optDesign = 'cantilever_hole';

nelx = 4;
nely = 2; 
nelz = 1;
microx = 64;
dx = 1;

%% init
lx = dx * microx;
globalx = nelx * microx;
globaly = nely * microx;
globalz = nelz * microx;
maxx = globalx / 2; minx = -maxx;
maxy = globaly / 2; miny = -maxy;
maxz = globalz / 2; minz = -maxz;

datapath = "data/mat/cvt_mesh_"+optDesign+".mat";
if opt_load_cvt == 0
    [seeds, rho, cpnts] = generateCVT(opt_load_rho, optDesign, volfrac, vF, ...
        nelx, nely, nelz, lx, microx);
            
    if strcmp(optDesign, 'cantilever_hole')
        idx = find(seeds(:,1) < -32 & seeds(:,1) > -96 & ...
                    seeds(:,2) < 32 & seeds(:,2) > -32 & ...
                    seeds(:,3) < 32 & seeds(:,3) > -32);
        seeds(idx,:) = [];
        figure; scatter3(seeds(:,1), seeds(:,2), seeds(:,3), 'filled'); axis equal;
    end
    [c_polyMesh, c_nodesFine, c_tets, c_seg, c_polygon] = ...
        generateMesh(seeds, nelx, nely, nelz, microx, optDesign);
%     write_cin(seeds, optDesign);    
    
    save(datapath, "c_nodesFine", "c_tets", "c_seg", "seeds", "c_polygon", ...
        "c_polyMesh", "cpnts", "rho");
else
    load(datapath);
end
% figure; scatter3(seeds(:,1), seeds(:,2), seeds(:,3), 'filled'); axis equal;
% load seeds % cvt init version

nvar = numel(seeds) + size(seeds, 1);
npnts = size(seeds, 1);

epsilon = compute_epsilon(c_tets, c_nodesFine);
step = epsilon * 2;

[c_bnid, c_dofid, c_bnidSaved, c_bnidSavedID, nbnidFine] = border_check(c_nodesFine, c_polygon);

%% MMC
T = repmat(t, npnts, 1);

[xval, xold1, xold2, xmin, xmax, low, upp, m_mma, c_mma, d_mma, a0_mma, ...
    a_mma] = MMA_paras(t, t0, t1, seeds, maxx, minx, maxy, miny, maxz, minz);

%% coarse mesh
[c_edofMat, c_nodeCoarse, cbns, c_nid, c_dupB] = assemble_cvt(c_polygon);
% figure; scatter3(cbns(:,1), cbns(:,2), cbns(:,3)); hold on;  

c_phi = parametric_B_cvt(c_polygon, c_nodesFine, c_bnid, c_bnidSaved, ...
    c_bnidSavedID, c_edofMat, c_nid, c_dupB, nbnidFine, cbns, c_nodeCoarse);

[freedofs, fixeddofs, fext, alldofs] = design_domain(optDesign, maxx, minx, ...
    maxy, miny, maxz, minz, vF, cbns, 0);
U = zeros(alldofs, 1);

%% fine mesh
[c_sK, c_vol, vol0] = prepare_sK(c_nodesFine, c_tets, 1, nu0);

rho1 = rho.^3;

scalarE = scalarE0;
% scalarE = 1;

%% iteration
datapath = "data/mat/cvt_"+optDesign+"_v"+volfrac+".mat";
if opt_load_mmc == 0

    loop = 0; objVr5 = 1; maxloop = 100;
    vars_loop = zeros(npnts, 4, maxloop);
    obj_loop = zeros(maxloop, 2);

    while loop < maxloop && objVr5 > 1e-3
        loop = loop + 1; tic   

        vars_loop(:,:,loop) = [seeds, T];
        
        %% cvt
        conn_list = dt_connection(seeds);

        indices = 1 : int32(npnts);
        [c_edges, cenP, vedges, c_nnbEdgeIDX, c_nnbSeedIDX] = intersection_info(seeds, ...
            indices, nelx, nely, nelz, microx,optDesign);   

        c_interPolySeed = do_intersection(vedges, c_polyMesh, c_nnbEdgeIDX, seeds);

        [E, dE] = cvt_energy(seeds, cenP);

        %% MMC
        MMCs = update_MMCs(vedges, T, c_nnbSeedIDX);
        if mod(loop,5) ==0 || loop==1 
            plot_mmc(MMCs, minx, maxx, miny, maxy, minz, maxz);
            % output_mmc(MMCs, maxx, minx, maxy, miny, maxz, minz, loop);
        end

        [c_rho, c_dist] = projection_independent(c_nodesFine, c_tets, MMCs, ...
            epsilon, c_nnbEdgeIDX, c_interPolySeed);
        
        [c_dH, c_interEles] = sensitivity_mmc_align(c_nodesFine, c_tets, seeds, ...
            step, T, epsilon, conn_list, c_rho, c_dist, c_interPolySeed, ...
            nelx, nely, nelz, microx,optDesign);

        %% simulation
        c_psi = preparation_cvt(e1, e0, penal, c_phi, c_dofid, c_rho, c_tets, c_sK);

        K = stiffness_cvt(c_rho, c_sK, c_psi, c_edofMat, alldofs, e0, e1, penal, c_vol);
        
        U(freedofs) = K(freedofs, freedofs) \ fext(freedofs);
        c = fext' * U;

%         extract_def_nodes(U, c_polygon, cbns)

        [v_local, v] = volume_cvt(c_rho, c_vol, vol0);

        %% optimization   
        [dc, dv] = sensitivity_cvt(c_dH, c_interEles, c_vol, c_rho, c_psi, c_sK, ...
            c_edofMat, U, e0, e1, penal, vol0);

        % significant digits for sensitivity truncation 
        dgt = dgt0 - floor(log10([max(abs(dc(:))), max(abs(dE(:))), max(abs(dv(:)))])); 
        % truncated objective and cons sensitivity
        dc = round(dc*10^dgt(1))/10^dgt(1);
        dE = round(dE*10^dgt(2))/10^dgt(2);
        dv = round(dv*10^dgt(3))/10^dgt(3);
        
        f0val= (1-scalarE) * c + E* scalarE; 
        df0dx= - dc * (1-scalarE) + dE * scalarE; 
        fval = v / volfrac - 1;
        dfdx = dv' / volfrac; 
        
        [xmma, ~, ~, ~, ~, ~, ~, ~, ~, low, upp] = mmasub_mmc(m_mma, nvar, ...
            loop, xval, xmin, xmax, xold1, xold2, f0val, df0dx, fval, dfdx,...
            low, upp, a0_mma, a_mma, c_mma, d_mma);

        if sum(isnan(xmma)) > 0
            xmma = xval;
            scalarE = 1;
            disp('NAN optimization!');
        else
            scalarE = scalarE0;
        end
        xold2 = xold1; xold1 = xval;         
    
        %% update seeds
        xval = xmma;
        % xval = round(xmma*1e4)/1e4; 
        vars = reshape(xval, 4, [])'; % [n,3]
        seeds = vars(:,1:3);
        T = vars(:,4);

        OBJ(loop) = f0val;
        if loop >= 5 && (v-volfrac)/volfrac < 1e-3
            objVr5 = abs(max(abs(OBJ(loop-4:loop)-mean(OBJ(loop-4:loop))))/mean(OBJ(loop-4:loop)));
        end
        
        obj_loop(loop, :) = [c,E];

        toc;
        fprintf(' It.:%3i C.:%6.3f E.:%6.3f, V.:%6.4f ch.:%6.2f \n',loop,c,E,v,objVr5);
    end 
    save(datapath, 'seeds', 'T', 'MMCs','c','E','vars_loop','obj_loop', 'U');
    output_mmc(MMCs, maxx, minx, maxy, miny, maxz, minz, loop);
else
    load(datapath);
end

%% recover deformation
if opt_load_def == 0
    indices = 1 : int32(npnts);
    [c_edges, cenP, vedges, c_nnbEdgeIDX, c_nnbSeedIDX] = intersection_info(seeds, ...
        indices, nelx, nely, nelz, microx, optDesign);   
    c_interPolySeed = do_intersection(vedges, c_polyMesh, c_nnbEdgeIDX, seeds);

    MMCs = update_MMCs(vedges, T, c_nnbSeedIDX);

    [c_rho, c_dist] = projection_independent(c_nodesFine, c_tets, MMCs, ...
            epsilon, c_nnbEdgeIDX, c_interPolySeed);

    [c_psi, c_proj] = preparation_cvt(e1, e0, penal, c_phi, c_dofid, c_rho, c_tets, c_sK);
    
%     U = U / 10;

    ref_nodes = cell(size(c_edofMat, 1), 1);
    def_nodes = cell(size(c_edofMat, 1), 1);
    for ele = 1 : size(c_edofMat, 1)
        edof = c_edofMat{ele};
        ue = U(edof);
        he = c_proj{ele} * ue; % he order: [4-borders, inner]
        
        dofid = c_dofid{ele};
        ufine = he(dofid);

        ref_nodes{ele} = c_nodesFine{ele};
        def_nodes{ele} = ref_nodes{ele} + reshape(ufine, 3, [])';
    end
    
    recover_def_mmc_surf(ref_nodes, def_nodes, MMCs, optDesign);
   
else
    load(datapath);
end


%% utils
function plot_mmc(MMCs, minx, maxx, miny, maxy, minz, maxz)
    figure(1); clf;
%     for i = 1:size(MMCs, 1)
%         t = MMCs(i, 1);
        x = MMCs(:,[2,5]);
        y = MMCs(:,[3,6]);
        z = MMCs(:,[4,7]);
        plot3(x',y',z','LineWidth',2, 'Color','k'); hold on;  
%     end
    view([10 25]);
    xlabel('x axis');
    ylabel('y axis');
    zlabel('z axis');
    axis equal;
    axes1 = gca;
    set(axes1,'Xlim',[minx, maxx],'Ylim',[miny, maxy], 'Zlim',[minz, maxz]);
%     axis off;
end

function defn1 = extract_def_nodes(U, c_polygon, cbns)

    defn = reshape(U, 3, [])' + cbns;
    defn1 = [];
    for i = 1:size(c_polygon)
        segs = c_polygon{i};
        for j = 1:size(segs, 1)
            nodes = segs{j}(:, 1:3);
    
            [~, idx] = ismembertol(nodes, cbns, 1e-3, 'Byrows', true); 
            idx1 = idx(idx > 0);
            defn1 = [defn1; defn(idx1, :)];
        end
    end
    figure; scatter3(defn1(:,1), defn1(:,2), defn1(:,3));
end

function [xval, xold1, xold2, xmin, xmax, low, upp, m_mma, c_mma, d_mma, a0_mma, ...
    a_mma] = MMA_paras(t, t0, t1, seeds, maxx, minx, maxy, miny, maxz, minz)

    np = size(seeds, 1);

    T  = repmat(t, np, 1);

    xval = reshape([seeds, T]', [], 1); 
    xold1= xval; xold2 = xval; 

    xmin = repmat([minx, miny, minz, t0], np, 1)'; 
    xmax = repmat([maxx, maxy, maxz, t1], np, 1)'; 
    xmin = xmin(:);
    xmax = xmax(:);

    low  = xmin; upp = xmax; 

    m_mma = 1; % number of constrain
    c_mma = 1000*ones(m_mma,1); d_mma = zeros(m_mma,1); a0_mma = 1; a_mma = zeros(m_mma,1); 
end

function epsilon = compute_epsilon(c_tets, c_nodesFine)

    len2 = zeros(size(c_tets,1), 1);
    for ele = 1:size(c_tets,1)
        tets = c_tets{ele};
        nodes = c_nodesFine{ele};
        x = nodes(:,1);
        y = nodes(:,2);
        z = nodes(:,3);
        x = x(tets);
        y = y(tets);
        z = z(tets);
    
        len1 = sqrt((x(:,1) - x(:,2)) .^ 2 + (y(:,1) - y(:,2)) .^ 2 + (z(:,1) - z(:,2)) .^ 2);
        len2(ele) = mean(len1);
    end
    epsilon = mean(len2) / 2;
end