%--------------------------
 % @Author: Jingqiao Hu
 % @Date: 2022-01-14 21:23:30
 % @LastEditTime: 2022-05-10 19:41:13
%--------------------------
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_mmc = 0;

volfrac = 0.3;

t = 0.1; t0 = 1e-5; t1 = 0.3; % for MMC
dgt0 = 5;

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

num_add = 0;

scalarE = 0.2;
scalarE0 = 0.2;

% vF = 0.0001;
vF = 0.001;
optDesign = 'cantilever';

nelx = 8;
nely = 2; 
nelz = 2;
microx = 16;
dx = 1 / microx;

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

[cx,cy,cz] = meshgrid(minx+0.5:0.5:maxx-0.5, minz+0.5:0.5:maxz-0.5, minz+0.5:maxz-0.5);
seeds = [cx(:), cy(:), cz(:)];
% figure; scatter3(seeds(:,1), seeds(:,2), seeds(:,3), 'filled'); axis equal;

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

step = dx * 2;

opt_scale = 2;

% coarse
num_p = 2; % 2 control pnts for every segment
num_b = (num_p + 1) * (num_add + 1) - 1; % added nodes on one edge, including bridge nodes & bezier pnts
vdofs_num = ((num_b+2)^2*6 - num_b*12 - 8*2) * 3; % macro_dofs_num

ref_nodes = mesh3D(nelx, nely, nelz, microx, num_b, lx, opt_scale);
ref_nodes(:,1) = ref_nodes(:,1) - maxx;
ref_nodes(:,2) = ref_nodes(:,2) - maxy;
ref_nodes(:,3) = ref_nodes(:,3) - maxz;

[freedofs, fixeddofs, fext, alldofs] = design_domain(optDesign, maxx, minx, maxy, miny, maxz, minz, vF, ref_nodes, 0);

[iK_ma, jK_ma, edofMat_ma] = assemble_bezier(nelx, nely, nelz, num_b);

% fine
[~, ~, D1] = material_paras(1, nu0);
ke0 = intKE(dx/2, dx/2, dx/2, D1);
[iK_mi, jK_mi, edofMat_mi] = prepare_assemble(microx, microx, microx);

[boundarys, inner] = micro_boundary(microx);
newdofs = [boundarys(:); inner(:)];
dofid   = reorder_dofs(newdofs); 
sort_u  = prepare_index_micro_u(microx, edofMat_mi, dofid);
bdofs_num = length(boundarys);

bezier_B = parametric_B(microx, num_add, num_b);

[c_polyMesh, c_hex, c_nodesFine] = generateMesh_hex(nelx, nely, nelz, microx, lx, edofMat_mi);

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

%% iteration
datapath = "data/mat/cvt_cmpPBC_"+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, lx, 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, dx, loop);
        end

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

        %% simulation
        [U, c_psi, c_regulated] = simulation_CBN_hex(e1, e0, penal, freedofs, fext, iK_ma, jK_ma, ...
            microx, ke0, iK_mi, jK_mi, edofMat_mi, dofid, bdofs_num, bezier_B, c_rho);        
       
        c = fext' * U;
        v = volume_cvt(c_rho, dx, nelx, nely, nelz, lx);

        %% optimization   
        [dc, dv] = sensitivity_cvt_hex(c_dH, c_interEles, c_rho, c_psi, ...
            edofMat_ma, U, e0, e1, penal, ke0, nelx, nely, nelz, lx);

        % 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 
          
    def_nodes = extract_def_nodes_macro(U, nelx, nely, nelz, lx, num_b, edofMat_ma);
    figure(7); clf; scatter3(def_nodes(:,1), def_nodes(:,2), def_nodes(:,3), 'filled'); axis equal; 
    
    save(datapath, 'seeds', 'T', 'MMCs','c','E','vars_loop','obj_loop', 'U', 'c_dist', 'c_regulated');
    output_mmc(MMCs, maxx, minx, maxy, miny, maxz, minz, dx, loop);
else
    load(datapath);
end

%% recover deformation
% U = U / 10;
recover_def_mmc(c_dist, U, c_regulated, edofMat_ma, c_nodesFine, microx, nelx, nely, nelz, dofid, MMCs);


%% utils
function v = volume_cvt(rho_cell, dx, nelx, nely, nelz, lx)
    
    ncvt = size(rho_cell, 1);
    v_global = 0;
    
    parfor ele = 1:ncvt
        v = dx^3 .* rho_cell{ele};        
        v_global = v_global + sum(v);
    end

    v = v_global / (nelx * nely * nelz * lx^3);
end

function [c_polyMesh, c_hex, c_nodesFine] = generateMesh_hex(nelx, nely, nelz, microx, lx, edofMat_mi)
    nele = nelx * nely * nelz;

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

    [cx, cy, cz] = meshgrid(minx : maxx, miny : maxy, minz : maxz);
    cy = flipud(cy);
    ref_nodes = [cx(:), cy(:), cz(:)];
    ref = reshape(ref_nodes', [], 1);

    dx = lx / microx;
    [cx, cy, cz] = meshgrid(0 :dx: lx, 0 :dx: lx, 0 :dx: lx);
    cy = flipud(cy);
    nodes0 = [cx(:), cy(:), cz(:)];

    [~, ~, edofMat] = prepare_assemble(nelx, nely, nelz);

%     figure;

    c_polyMesh = cell(nele, 1);
    c_hex = cell(nele, 1);
    c_nodesFine = cell(nele, 1);
    for i = 1:nelx
        for j = 1:nely
            for k = 1:nelz
                ele = j + (i-1) * nely + (k-1) * nely * nelx;
                edof = edofMat(ele, :);
                re = ref(edof);
                pntsFiltered = reshape(re, 3, [])';

                faces = int32(convhull(pntsFiltered(:,1), pntsFiltered(:,2), pntsFiltered(:,3),'Simplify',true));
                vert1 = pntsFiltered(faces(:,1),:);
                vert2 = pntsFiltered(faces(:,2),:);
                vert3 = pntsFiltered(faces(:,3),:);
                c_polyMesh{ele} = [vert1, vert2, vert3]; % for the following intersection

                c_hex{ele} = edofMat_mi(:, 3:3:end) / 3;

                step = [lx * (i-1) + minx, lx * (nely - (j-1)) - lx - maxy, lx * (k-1) + minz];
                nodes = nodes0 + repmat(step, size(nodes0, 1), 1);
                c_nodesFine{ele} = nodes;

%                 scatter3(vert1(:,1), vert1(:,2), vert1(:,3), 'filled'); axis equal; hold on; 
%                 scatter3(vert2(:,1), vert2(:,2), vert2(:,3), 'filled'); axis equal; hold on; 
%                 scatter3(vert3(:,1), vert3(:,2), vert3(:,3), 'filled'); axis equal; hold on; 
%                 scatter3(nodes(:,1), nodes(:,2), nodes(:,3)); hold on;
            end
        end
    end
end

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