% demo_recon_time_cs_sbhe.m
%   video encoder and deocder using scrambled block Hadamard ensemble as
%   the measurement matrix
%
% Written: Hao Fang
% Created: Apr. 2013

%=======================================
clear all;
close all;
clc;

path(path, '../Videos');
path(path, '../Tools');
path(path, '../Tools/YUV');
path(path, '../ParallelCS');

% akiyo_qcif
% carphone_qcif
% claire_qcif
% coastguard_qcif
% foreman_qcif
% salesman_qcif
filename = 'akiyo_qcif';
yuvformat = 'qcif';

%======================================

switch yuvformat
    case 'cif'
        dims = [352 288];
        block_size = 16;
    case 'qcif'
        dims = [176 144];
        block_size = 8;
end
numfrm = 2;
startfrm = 0;
[Y, U, V] = yuv_import(sprintf('%s.yuv', filename), dims, numfrm, startfrm, yuvformat);

X_Ref = double(Y{1})-128; % range of Y-128: [-128, 127]
max_X_Ref = max(abs(X_Ref(:)));
X_Ref = X_Ref ./ max_X_Ref; % normalize since l1magic won't work for too large max(x)
X_Non_Ref = double(Y{1}-Y{2});
max_X_Non_Ref = max(abs(X_Non_Ref(:)));
X_Non_Ref = X_Non_Ref ./ max_X_Non_Ref;

% parameters
msr_ratio = 0.8;
N = dims(1)*dims(2);
ave_CR = 0.1:0.1:0.5;
iter = 200;
ave_rec_time = zeros(1, length(ave_CR));
ave_PSNR_Ref = zeros(1, length(ave_CR));
ave_PSNR_Non_Ref = zeros(1, length(ave_CR));
for jj = 1:length(ave_CR)
    CR_Ref = ave_CR(jj)*2*msr_ratio;
    CR_Non_Ref = ave_CR(jj)*2*(1-msr_ratio);
    K_Ref = ceil(N*CR_Ref);
    K_Non_Ref = ceil(N*CR_Non_Ref);
    
    PSNR_Ref = zeros(1, iter);
    PSNR_Non_Ref = zeros(1, iter);
    rec_time = zeros(1, iter);
    
    % sparsifying basis
    W_Ref = @(x) dct2(x);
    WT_Ref = @(x) idct2(x);
    W_Non_Ref = @(x) x;
    WT_Non_Ref = @(x) x;
    
    % measurement matrix block
    block_length = 32;
    H = hadamard(block_length);
    
    for ii = 1:iter
        % cs sampling
        perm_x_ref_ind = randperm(N);
        X_Ref_p = X_Ref(perm_x_ref_ind);
        X_Ref_p_2D = reshape(X_Ref_p, block_length, N/block_length);
        
        msr_ref_ind = randperm(N);
        Y_Ref_2D = H * X_Ref_p_2D;
        Y_Ref = Y_Ref_2D(msr_ref_ind(1:K_Ref));
        Y_Ref = reshape(Y_Ref, length(Y_Ref(:)), 1);
        
        perm_x_non_ref_ind = randperm(N);
        X_Non_Ref_p = X_Ref(perm_x_non_ref_ind);
        X_Non_Ref_p_2D = reshape(X_Non_Ref_p, block_length, N/block_length);
        
        msr_non_ref_ind = randperm(N);
        Y_Non_Ref_2D = H * X_Non_Ref_p_2D;
        Y_Non_Ref = Y_Non_Ref_2D(msr_non_ref_ind(1:K_Non_Ref));
        Y_Non_Ref = reshape(Y_Non_Ref, length(Y_Non_Ref(:)), 1);
        
        % sensing matrix
        A_Ref = @(theta) sensing_mtx(theta, H, WT_Ref,...
            dims(2), dims(1), block_length, perm_x_ref_ind, msr_ref_ind, K_Ref);
        AT_Ref = @(y) sensing_mtx_trans(y, H, W_Ref,...
            dims(2), dims(1), block_length, perm_x_ref_ind, msr_ref_ind, K_Ref);
        
        A_Non_Ref = @(theta) sensing_mtx(theta, H, WT_Non_Ref,...
            dims(2), dims(1), block_length, perm_x_non_ref_ind, msr_non_ref_ind, K_Non_Ref);
        AT_Non_Ref = @(y) sensing_mtx_trans(y, H, W_Non_Ref,...
            dims(2), dims(1), block_length, perm_x_non_ref_ind, msr_non_ref_ind, K_Non_Ref);
        
        % cs decoding
        t = cputime;
        s_Ref = l1eq_pd(AT_Ref(Y_Ref), A_Ref, AT_Ref, Y_Ref);
        rec_X_Ref = WT_Ref(reshape(s_Ref, dims(2), dims(1)));
        rec_Y_Ref = mod(rec_X_Ref.*max_X_Ref + 128, 255);
        
        s_Non_Ref = l1eq_pd(AT_Non_Ref(Y_Non_Ref), A_Non_Ref, AT_Non_Ref, Y_Non_Ref);
        rec_X_Non_Ref = reshape(WT_Non_Ref(s_Non_Ref), dims(2), dims(1));
        rec_Y_Non_Ref = rec_Y_Ref + rec_X_Non_Ref.*max_X_Non_Ref;
        rec_time(ii) = cputime - t;
        
        PSNR_Ref(ii) = psnr(double(Y{1}), double(rec_Y_Ref));
        PSNR_Non_Ref(ii) = psnr(double(Y{2}), double(rec_Y_Non_Ref));
    end
    ave_rec_time(jj) = mean(rec_time);
    ave_PSNR_Ref(jj) = mean(PSNR_Ref);
    ave_PSNR_Non_Ref(jj) = mean(PSNR_Non_Ref);
end

