function [ echo_rd_mtx, stat_RD, sigma_n_o ] = doppler_process_CS(echo_r_mtx, sigma_n, lambda, gamma, delta)
% Doppler processing for echo data using compressed sensing
%
% Usage: 
% [echo_rd_mtx, stat_RD, sigma_n_o] = doppler_process_CS(echo_r_mtx, sigma_n, lambda, gamma, delta)
%
% Inputs:
%   echo_r_mtx:     Range processing echo data
%   sigma_n:        Input noise standard deviation
%   lambda:         LASSO weight
%   gamma:          Compressed ratio(0.5)
%   delta:          Convergence normalized difference(1e-6)
%
% Outputs:
%   echo_rd_mtx:    Range-doppler processing echo data
%   stat_RD:        Range-doppler statistics
%   sigma_n_o:      Output noise standard deviation


    % parameters
    if nargin < 5
        delta = 1e-6;
    end
    if nargin < 4
        gamma = 0.5;
    end
    [lenA, lenR, M] = size(echo_r_mtx);
    N = round(M / gamma);
    iter_max_VAMP = 1000;  
    lambda_v = zeros(N, 1) + lambda;
    
    % generate mtx
    F_ori = dftmtx(N);
    F = F_ori(1:M,:);
    F_inv = conj(F) / N;
    
    % normalization
    A = (sqrt(N) * eye(M)) * F_inv;
    echo_r_mtx = sqrt(N) .* echo_r_mtx;
    sigma_n = sqrt(N) * sigma_n;
    
    % doppler processing
    echo_rd_mtx = zeros(lenA, lenR, N);
    stat_RD = zeros(lenA, lenR, N);
    sigma_n_o_cnt = zeros(lenA, lenR);
    for numA = 1: lenA
        parfor numR = 1: lenR
            sample = squeeze(echo_r_mtx(numA, numR, :));
            y = sample;
            x_LASSO = cVAMPro(y, A, lambda_v, delta, iter_max_VAMP);
            [x_d_CROD, sigma_CROD] = CROD(y, A, x_LASSO, lambda, sigma_n);
            sigma_n_o_cnt(numA, numR) = abs(sigma_CROD);
            stat_RD(numA, numR, :) = abs(fftshift(x_d_CROD) / sigma_CROD).^2;
            echo_rd_mtx(numA, numR, :) = fftshift(x_d_CROD);
        end
    end
%     sigma_n_o = mean(sigma_n_o_cnt, 2);
    sigma_n_o = mean(mean(sigma_n_o_cnt));
    
end


%% sub-functions
% algorithm for LASSO
% y: measurements
% A: measurement matrix
% lambda: LASSO weight
% tau: convergence normalized difference
% Kit: maximum number of iterations
% LASSO estimator
function x_hat_wl = cVAMPro(y, A, lambda, tau, Kit)

    % Initialization
    [M, N] = size(A);
    gamma = M / N;
    k = 0;
    p = ctranspose(A) * y;
    h_1 = p;
    Q_1 = gamma;
    tau_d = 1;

    % Iteration
    while ((k < Kit) && (tau_d > tau))
        % Factorized Part
        x_1 = ST(h_1, lambda, Q_1);
        chi_1 = F1(x_1, lambda, Q_1);
        % Message Passing
        h_2 = x_1 / chi_1 - h_1;
        Q_2 = 1 / chi_1 - Q_1;
        % Gaussian Part
        t1 = (p + h_2) / Q_2;
        t2 = ctranspose(A) * (A * (p + h_2)) / ((Q_2 + 1) * Q_2);
        x_2 = t1 - t2;
        chi_2 = gamma / (Q_2 + 1) + (1 - gamma) / Q_2;
        % Message Passing
        h_1_next = x_2 ./ chi_2 - h_2;
        Q_1_next = 1 / chi_2 - Q_2;
        tau_d = norm(h_1_next - h_1, Inf) / norm(h_1_next, Inf);
        k = k + 1;
        % output
        x_hat_wl = x_1;
        % next
        h_1 = h_1_next;
        Q_1 = Q_1_next;
    end

end


% soft threshold function
% x: processing object
% thd: threshold
% y: result
function x = ST(h_1, lambda, Q_1)
    [N, M] = size(h_1);
    x = zeros(N, M);

    for i = 1:N
        sign = h_1(i) ./ abs(h_1(i));
        diff = abs(h_1(i)) - lambda(i);
        x(i) = sign .* (diff ./ Q_1) .* SF(diff);
    end

end


% Heaviside's step function
function v = SF(a)

    if a > 0
        v = 1;
    elseif a == 0
        v = 0; % at zero points
    else
        v = 0;
    end

end


% Calculation of chi_1
function chi_1 = F1(x_1, lambda, Q_1)

    [N, M] = size(x_1);
    count = 0;
    
    for i = 1:N
        temp = Q_1 * abs(x_1(i)) + lambda(i);
        count = count + (2 - lambda(i) / temp) * SF(abs(x_1(i)));
%         count = count + (2-lambda(i)/temp) * (abs(x_1(i)) > 1e-4);
    end
    
    chi_1 = count / (2 * N * Q_1);
    
end

% calculate debiased LASSO estimator
% y: measurements
% A: measurement matrix
% x_LASSO: LASSO estimator
% lambda: LASSO weight
% sigma_n: input noise standard deviation
% x_d_CROD: debiased LASSO estimator
% sigma_CROD: equivalent noise standard deviation estimator
function [ x_d_CROD, sigma_CROD ] = CROD(y, A, x_LASSO, lambda, sigma_n)

    [m, n] = size(A);
    gamma = m / n;

    rho_active = sum(abs(x_LASSO) > 1e-3)/n;
    Q_hat = (gamma - rho_active)/(1 - rho_active);
    Rho = sum((abs(x_LASSO) > 1e-3).* (2 - lambda./(Q_hat*abs(x_LASSO) + lambda))) / 2 / n;
    diff = 1;
    while(diff > 1e-4)
        Rho_pre = Rho;
        Rho = sum((abs(x_LASSO) > 1e-3).* (2 - lambda./((gamma-Rho)/(1-Rho)*abs(x_LASSO) + lambda))) / 2 / n;
        diff = abs(Rho - Rho_pre);
    end
    Q_hat = (gamma-Rho)/(1-Rho);
    x_d_CROD = x_LASSO + A'*(y - A*x_LASSO)/Q_hat;

    RSS = sum(abs(y - A * x_LASSO).^2)/m;
    chi = Rho*(1 - Rho)/(gamma - Rho);
    if chi ~= 0
        chi_temp = sqrt((chi+1)*(chi+1)-4*gamma*chi);
        z = -(1 - chi + chi_temp) / (2*chi);
        z_prime = -(1 - 2*gamma*chi + chi + chi_temp) / (2*chi*chi*chi_temp);
        G_prime = (z + 1/chi);
        G_wprime = (z_prime + 1/chi/chi);
        chi_hat = gamma/2*G_wprime*RSS/(G_prime - chi*G_wprime)...
            + (G_prime*G_prime/2 - gamma/2*G_wprime)*sigma_n*sigma_n/(G_prime - chi*G_wprime);
    else
        G_prime = gamma;
        G_wprime = gamma*(1-gamma);
        chi_hat = gamma/2*G_wprime*RSS/(G_prime - chi*G_wprime)...
            + (G_prime*G_prime/2 - gamma/2*G_wprime)*sigma_n*sigma_n/(G_prime - chi*G_wprime);
    end
    sigma_CROD = sqrt(2*chi_hat) / Q_hat;

end