clc, close all, clear all
addpath('./tools/')
addpath('./partial/')
addpath(genpath('./manopt/manopt/'))

options = {};

options.n_eigen = 50;           % no. eigenpairs used
options.tv_sigma = 0.2;         % variance of indicator func. in TV term
options.tv_mean = 0.5;          % mean =
options.descr_max_radius = 5;   % max support radius for rigid descriptors (in % of diameter)
options.intinv_volstep = 0.3;   % voxelization step for integral invariants (in % of diameter)
options.shot_bins = 5;          % no. bins for SHOT descriptor
options.shot_min_neighs = 3;    % min no. neighbors for SHOT descriptor

% -------------------------------------------------------------------------
% Load the shapes
% -------------------------------------------------------------------------
result_folder = dir('init_point');
% delete first two rows
result_folder(1:2,:) = [];

for j = 1:length(result_folder)
    file = result_folder(j).name;
    temp = strsplit(file, {'_','.'});
    shape_name = string(temp(2));
    partial_name = append(shape_name, '_', string(temp(3)));
    full_path = append('SGP_dataset/null/', shape_name,'.mat');
    partial_path = append('SGP_dataset/selected/cuts/',partial_name,'.mat');
    full = load(full_path);
    partial = load(partial_path);

    M = full.N;
    N = partial.N;

    M.TRIV = M.tri;
    M.VERT = M.xyz;
    M.n = size(M.xyz,1);
    M.S_tri = calc_tri_areas(M);

    N.TRIV = N.tri;
    N.VERT = N.xyz;
    N.n = size(N.xyz,1);
    N.S_tri = calc_tri_areas(N);

    [M.evecs, M.evals, M.S, M.W] = calc_LB(M, options.n_eigen);
    [N.evecs, N.evals, N.S, N.W] = calc_LB(N, options.n_eigen);

    % -------------------------------------------------------------------------
    % Compute dense descriptors
    % -------------------------------------------------------------------------

    diam = sqrt(sum(M.S_tri));
    [F1, F2] = calc_dense_descriptors(N, diam, options);
    [G1, G2] = calc_dense_descriptors(M, diam, options);
    F = [F1 F2];
    G = [G1 G2];
    clear F1 F2 G1 G2
    
    % -------------------------------------------------------------------------
    % Match the partial shape N to the full shape M. This code reproduces the
    % results shown in Figure 9 of the paper.
    % -------------------------------------------------------------------------
    v_init = ones(M.n,1);
    C_init = [];
    
    mu1 = 1;   % slanted-diagonal mask
    mu2 = 1e3; % sub-orthogonality
    
    
    for i=1:3
        
        [C, est_rank] = optimize_C(M, N, G, F, v_init, C_init, mu1, mu2);
        
        fprintf('Estimated rank: %d\n', est_rank);
        
        [~, matches] = run_icp(M, N, est_rank, C, 0);
        
        [Co, matches] = run_icp(M, N, est_rank, C, 1);
        
        C_init = [Co zeros(size(C,1),size(C,1)-size(Co,2))];
        
        v = optimize_v(M, N, G, F, C_init, 1, 1e2, options);
        v = (0.5*tanh(6*(v-0.5))+0.5);
        
        v_init = v;        
    end
    save_path = append('matches_pfm_50/matches_pfm_',partial_name);
    save(save_path,'matches');
end
