% This function takes the struct created when using the vdpgm toolbox,
% creates the GMM for all clusters. This object is used to determine the 
% which cluster each test sample belongs to. Finally, the cluster labels are
% used to determine the final hypothesized label for each test sample.

function [output_labels] = gen_test_labels(dpm_results, test_data, cluster_labels)
    %Find the number of clusters from dpm_results. If the algorithm used is
    %AVDPM, then the number of clusters is based on K-1 instead of k. This
    %discrepancy is a from the vdpgm code (See readme)
    if isequal(dpm_results.algorithm, 'vdp')
        k = dpm_results.K - 1;  
    else
        k = dpm_results.K;
    end
    
    mu = [];
    sigma = [];
    % Collect all of the clusters'means and variances to build one large GMM 
    for i = 1:k
        if cluster_labels(i) ~= 0
            mu =  [mu; dpm_results.hp_posterior.m(:,i)'];
            tmp = dpm_results.hp_posterior.B{i} / dpm_results.hp_posterior.eta(i);
            tmp = triu(tmp) +triu(tmp)';
            sigma = cat(3, sigma, tmp);
        end
    end
    % Create the GMM
    obj = gmdistribution(mu,sigma);
    % Find which cluster each sample is assigned to
    test_idx = cluster(obj,test_data);
    % Determine the final output label for each sample
    actual_cluster_labels = cluster_labels(cluster_labels ~=0);
    output_labels = actual_cluster_labels(test_idx);

end