%--------------------------
 % @Author: Jingqiao Hu
 % @Date: 2021-02-25 17:56:07
 % @LastEditTime: 2022-01-24 23:55:27

 % control pnts of one face of this macro-ele
 % p00-p01-p02-p03
 % p10-p11-p12-p13
 % p20-p21-p22-p23
 % p30-p31-p32-p33
 % for one face, the order of c-pnts is col by col
 % NOTE: assume num_add = 0
%--------------------------
function B = parametric_B(microx, num_add, num_b)   

    num_nodes = (microx+1)^2;
    
    Be = parametric_B_one_face(microx, num_add, num_b); % [3, num_nodes^2*3, nnodes]
    Be1 = permute(Be, [2,1,3]);  % [~, 3, nnodes]
    Be2 = reshape(Be1, [], 3 * num_nodes)'; % [n, ~]    
 
    [up_i, front_i, down_i, back_i, left_i, right_i, dnum_i] = surface_dofs_id(microx);
    [up_a, front_a, down_a, back_a, left_a, right_a, dnum_a] = surface_dofs_id(num_b + 1);

    B = zeros(dnum_i, dnum_a);

    % NOTE: these six faces have shared edges, which should be compute only by one face 
    % B assumes the control pnts are ordered as col by col, but nodes in this surface is row by row
    Be_row = row_order(Be2, num_b);
    B(up_i, up_a) = Be_row;

    fid_saved = remove_shared_edges(up_i, front_i);
    B(front_i(fid_saved), front_a) = Be_row(fid_saved, :);
    
    last_id = [up_i(:); front_i(:)];
    did_saved = remove_shared_edges(last_id, down_i);
    B(down_i(did_saved), down_a) = Be_row(did_saved, :);

    last_id = [up_i(:); front_i(:); down_i(:)];
    bid_saved = remove_shared_edges(last_id, back_i);
    B(back_i(bid_saved), back_a) = Be_row(bid_saved, :);

    last_id = [up_i(:); front_i(:); down_i(:); back_i(:)];
    lid_saved = remove_shared_edges(last_id, left_i);
    B(left_i(lid_saved), left_a) = Be_row(lid_saved, :);
    
    last_id = [up_i(:); front_i(:); down_i(:); back_i(:); left_i(:)];
    rid_saved = remove_shared_edges(last_id, right_i);
    B(right_i(rid_saved), right_a) = Be_row(rid_saved, :);
end

function B3 = row_order(Be, num_b)
    num_nodes = size(Be, 2);
    B1 = reshape(Be, [], 3, num_b+2, num_b+2);
    B2 = permute(B1, [1,2,4,3]);
    B3 = reshape(B2, [], num_nodes);
end

function id_saved = remove_shared_edges(last_id, this_id)
    [~, id_shared, ~] = intersect(this_id, last_id,'stable');
    id_saved = setdiff(1:length(this_id), id_shared);    
end

function B = parametric_B_one_face(microx, num_add, num_b)

    segment_length = microx / (num_add+1);
    Be = parametric_B_one_divided_face(segment_length); % 3,48,s^2
    
    idx0_Be = reshape(1 : 4^2, 4, 4);
    jdx0_Be = reshape(1 : (segment_length+1)^2, segment_length+1, segment_length+1);
    
    % for the dimension correspondance
    i1_e = 1:4;    
    j1_e = 1:4;        
    idx1_Be = idx0_Be(i1_e, j1_e);
    idx2_Be = [3*idx1_Be(:)-2, 3*idx1_Be(:)-1, 3*idx1_Be(:)];    
            
    B = zeros(3, 3*(num_b+2)^2, (microx+1)^2);    

    idx0 = reshape(1 : (num_b+2)^2, num_b+2, num_b+2);
    jdx0 = reshape(1 : (microx+1)^2, microx+1, microx+1);

    for i = 1 : num_add+1

        i1 = (i-1)*3 + 1 : i*3 + 1;
        
        if i < num_add+1
            i2 = (i-1)*segment_length + 1 : i*segment_length ;
            i2_e = 1:segment_length;        
        else
            i2 = (i-1)*segment_length + 1 : i*segment_length + 1;
            i2_e = 1:segment_length + 1;        
        end
           
        for j = 1 : num_add+1
            
            j1 = (j-1)*3 + 1 : j*3 + 1;
            
            if j < num_add+1
                j2 = (j-1)*segment_length + 1 : j*segment_length ;            
                j2_e = 1:segment_length;
            else
                j2 = (j-1)*segment_length + 1 : j*segment_length + 1;            
                j2_e = 1:segment_length + 1;
            end
            
            jdx2_Be = jdx0_Be(j2_e, i2_e);

            idx1 = idx0(i1, j1);
            idx2 = [3*idx1(:)-2, 3*idx1(:)-1, 3*idx1(:)];            
            jdx1 = jdx0(j2, i2);

            B(:, idx2(:), jdx1(:)) = Be(:, idx2_Be(:), jdx2_Be(:));
        end
    end
end

% B assume the control pnts is col by col
function B = parametric_B_one_divided_face(microx)
    B = zeros(3, 48, (microx+1)^2);
    I = eye(4);
    for i = 0 : microx
        for j = 0 : microx
            u = i / microx; % project to [0,1]
            v = j / microx; 
            Bu = Bezier_matrix(u); % [3,12]
            Bv = Bezier_matrix(v); % [3,12]
            
            B(:,:,j+1+i*(microx+1)) = Bu * kron(Bv, I); % [3,12]*[12,48]=[3,48]
        end
    end
end

function B = Bezier_matrix(t)
    b = [1, t, t^2, t^3] * [1, 0,0,0;
                            -3,3,0,0;
                            3,-6,3,0;
                            -1,3,-3,1]; % 1,4
    I = eye(3);
    B = kron(b, I); % 3,12
end

% test
% [up_i, front_i, down_i, back_i, left_i, right_i, dnum_i] = surface_dofs_id(microx);
% fid_saved = remove_shared_edges(up_i, front_i);
% 
% num_b = microx-1;
% [~, ~, edofMat_full] = assemble_bezier(nelx, nely, nelz, num_b);
% [~, ~, ~, alldofs] = design_domain(opt_scale, optDesign, nelx, nely, nelz, microx, vF, num_b, ref_nodes, lx, 1);
% ufull = zeros(alldofs, 1);
% tmp = ref_nodes';
% tmp = tmp(:);
% for ele = 1: nelx*nely*nelz
%     edof = edofMat_ma(ele, :);
%     ue = tmp(edof);
%     ub = bezier_B * ue;
% 
%     edof = edofMat_full(ele, :);
%     ufull(edof) = ub;
% end
% deform_nodes = reshape(ufull, 3, [])';
% figure; scatter3(deform_nodes(:,1), deform_nodes(:,2), deform_nodes(:,3), 'b', 'filled'); axis equal;
% hold on; scatter3(ref_nodes(:,1), ref_nodes(:,2), ref_nodes(:,3), 'k'); axis equal;