%--------------------------
 % @Author: Jingqiao Hu
 % @Date: 2021-09-13 20:11:46
 % @LastEditTime: 2021-09-18 10:06:45

 % input:
 % harm_u is the harmonic solutions
 % microx is the size of the fine mesh in one coarse element
 % lx is the length of one coarse element

 % output: n_ij
 % TODO: run
%--------------------------
function n_ma = optimize_NH(harm_u, microx, lx, addNum)

    nele_ma = size(harm_u, 1);
    dx = lx / microx;
    num_ndofs = 2*(microx+1)^2;
    
    n_ma = cell(nele_ma, 1);

    [dNh_ele, ~, ~] = shape_gradient_ele(dx/2, dx/2, 4); % dNdx with 4 gauss points, 4*8 & 3*8
    [~, ~, edofMat_mi, ~] = forAssemble(microx, microx);
    
    [vdofs, ~, ~, ~, ~] = multi_microBC(microx, addNum, 0); 

    [A1, b1, Q_obj] = prepare_optNH(edofMat_mi, microx, dNh_ele, vdofs);

    parfor ele = 1:nele_ma
        hue = harm_u{ele};
        [A2, b2] = harmonic_constraint(num_ndofs, hue, vdofs);
        
        % assemble all constraints together
        A0 = [A1; A2];
        A = sparse(A0);
        b = sparse([b1; b2]);

        % optimization
        n = quadratic_programming(A, b, Q_obj);
        n_opt = reshape(n, [], num_ndofs); % [ndofs, vdofs]
        n_ma{ele} = n_opt; 
    end  

end

function [A, b, Q] = prepare_optNH(edofMat_mi, microx, dNh, vdofs)
        
    num_vdofs  = length(vdofs);
    alldofs_mi = 2*(microx+1)^2;
    
    Q = obj_Q(dNh, edofMat_mi, microx, num_vdofs);

    % constraints
    [A, b] = geometric_constraint(vdofs, alldofs_mi);        
end

function [A3, b3] = harmonic_constraint(num_ndofs, harm_u, vdofs)
    num_vdofs = length(vdofs);
    testNum = size(harm_u,2);

    % cons3: harmonic cons
    b3 = harm_u(:); % [2n*3, 1]

    A3 = zeros(num_ndofs * testNum, num_ndofs * num_vdofs);
    I3 = eye(num_ndofs);
    for i = 1:testNum
        R3 = harm_u(vdofs, i); % [2v, 1]
        A3i = kron(I3, R3'); % [1,2v] -> [2n,2n*2v]
        
        A3(num_ndofs*(i-1)+1:num_ndofs*i, :) = A3i;
    end    
end

% rewrite constraints to a big Ax = b, i.e. [A1;A2;A3]x = [b1;b2;b3]
function [A, b] = geometric_constraint(vdofs, num_ndofs)
    num_vdofs = length(vdofs);

    n1 = (1:num_vdofs*num_ndofs)';
    n2 = reshape(1:num_vdofs*num_ndofs, 8, []);

    % cons1: \sum_i n_ij = I
    I = eye(2);    
    R1 = repmat(I, 1, num_vdofs/2); % [2, 2v] 
    
    I = eye(num_ndofs);
    A1 = kron(I, R1); % [2,2v] -> [2*2n, 2v*2n]

    C1 = repmat(eye(2), 1, num_ndofs/2); % [2, 2n]
    b1 = C1(:);
    
    % cons2: n_ij = \delta_ij * I, for all X_i^H, X_j^H
    R2 = zeros(num_vdofs, num_ndofs); % [2v, 2n]
    R2(:, vdofs) = eye(num_vdofs);
    
    I2 = eye(num_vdofs);
    A2 = kron(R2, I2); % [2v*2v, 2v*2n]

    C2 = eye(num_vdofs);
    b2 = C2(:);

    A = [A1; A2];
    b = [b1; b2];    
end

% expand orignal n:[2n,2v] to [2n*2v,1]
% min \sum tr(n^T * dNh^T * dNh * n) = \sum tr(n^T * Q * n)
function Q = obj_Q(dNh, edofMat_mi, microx, num_vdofs)
    eleNum_mi = microx^2;
    alldofs_mi = 2*(microx+1)^2;

    % for obj, dNh: [4,8] -> [4m,8m] -> [4m,2n] ->[4m*2v,2n*2v]
    Q = 0;
    I1 = eye(size(edofMat_mi, 1)); % [m,m]
    I2 = eye(num_vdofs);           % [2v,2v]

    for gp = 1:4
        dNhg = dNh{gp};
        dNh1 = kron(I1, dNhg); % [4m,8m]
        
        dNh2 = zeros(4*eleNum_mi, alldofs_mi); % [4m,2n]
        for ele = 1:eleNum_mi
            edof = edofMat_mi(ele, :);
            dNh2(:,edof) = dNh2(:,edof) + dNh1(:,8*(ele-1)+1:8*ele);
        end

        % post multiplcation need change kron sequence
        dNh_opt = sparse(kron(dNh2, I2)); % [4m*2v, 2n*2v]

        Q = Q + dNh_opt' * dNh_opt;
    end
end