%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%% Clear previous work
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
clear;
close all;
clc;

addpath('utils')
addpath('../tensorlab/')
if isempty(gcp('nocreate'))
    parpool('threads');  % Start with default settings
end
% fix random seed for reproducibility
rng(1);

%% generate synthetic data for testing
I   = 2;    % the number of factor matrices
AVG = 200;
M_candidate = 48:4:72;
N = 80*ones(I,1);Nbar = prod(N);
% sparsity pattern
K = 15*ones(I,1); % hierarchical sparse, #nonzeros of each dimension
% SNR condition
SNR = 20; % Define the range of SNR values for testing
error = zeros(AVG,length(M_candidate),7);
time  = zeros(AVG,length(M_candidate),7);
srr   = zeros(AVG,length(M_candidate),7);


for avg = 1:AVG

    avg


    for m = 1:length(M_candidate)

        % initialize dimensions
        M = M_candidate(m)*ones(I,1);Mbar = prod(M);

        % generate measuring dictionaries
        A = cell(1,I);
        for i = 1:I
            A{i} = randn(M(i),N(i));
            col_norm = vecnorm(A{i});
        end

        % generate sparse vector x
        tX = generate_hierarchical_sparse_tensor(N, K);
        % generate noisy measurements
        tY_ori         = tmprod(tX,A,1:I);
        control.tX     = tX; % for debugging

        signal_power   = norm(vec(tY_ori),'fro')^2/numel(tY_ori); %average signal power per asymbol
        n_var          = (signal_power)/(10^(SNR/10));
        noise          = sqrt(n_var)*randn(size(tY_ori));
        tY             = tY_ori + noise;
        A_full         = kron(A{2},A{1});

        tol_coef = 0.05;
        %% sparse vector estimation
        %% case 1: without the knowledge of sparsity level
        %% using SBL: sequential
        result_mtsbl = tenMulReSBL_hi(tY,A,N,M);
        disp('Multi-Stage SBL-seq is done!')
        %% using SBL: parallel
        result_mtsbl_pl = tenMulReSBL_hi_pl(tY,A,N,M);
        disp('Multi-Stage SBL-pl is done!')
        %% case 2: with the knowledge of sparsity level
        %% using HTP
        result_mthtp = tenMulReHTP_hi(tY,A,N,M,K);
        disp('Multi-Stage HTP is done!')
        %% using IHT
        result_mtiht = tenMulReIHT(tY,A,N,M,K);
        disp('Multi-Stage IHT is done!')
        %% Benchmark 1: HiHTP
        result_hihtp = HiHTP(A_full, vec(tY), N(2), N(1), K(2), K(1), 100);
        disp('HiHTP is done!')
        %% Benchmark 2: HTP
        result_htp = htp(A_full, vec(tY), prod(K));
        disp('HTP is done!')
        %% Benchmark 3: IHT
        result_iht = iht(A_full, vec(tY), prod(K));
        disp('iht is done!')
        %%

        error(avg,m,1) = (norm(vec(result_mtsbl{1} - tX), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,1) = result_mtsbl{2};
        srr(avg,m,1) = recover_rate(vec(result_mtsbl{1}),vec(tX));

        error(avg,m,2) = (norm(vec(result_mtsbl_pl{1} - tX), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,2) = result_mtsbl_pl{2};
        srr(avg,m,2) = recover_rate(vec(result_mtsbl_pl{1}),vec(tX));

        error(avg,m,3) = (norm(vec(result_mthtp{1} - tX), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,3) = result_mthtp{2};
        srr(avg,m,3) = recover_rate(vec(result_mthtp{1}),vec(tX));

        error(avg,m,4) = (norm(vec(result_hihtp{1} - vec(tX)), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,4) = result_hihtp{2};
        srr(avg,m,4) = recover_rate(vec(result_hihtp{1}),vec(tX));

        error(avg,m,5) = (norm(vec(result_htp{1} - vec(tX)), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,5) = result_htp{2};
        srr(avg,m,5) = recover_rate(vec(result_htp{1}),vec(tX));

        error(avg,m,6) = (norm(vec(result_mtiht{1} - (tX)), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,6) = result_mtiht{2};
        srr(avg,m,6) = recover_rate(vec(result_mtiht{1}),vec(tX));

        error(avg,m,7) = (norm((result_iht{1} - vec(tX)), 'fro')/norm(vec(tX),'fro'))^2;
        time(avg,m,7) = result_iht{2};
        srr(avg,m,7) = recover_rate(vec(result_iht{1}),vec(tX));
    end
end

save hie_s_measurement.mat

