%%%% Inter-frame (inter-wavelength) motion estiomation and compensation
%%%% of real-time PAUS imaging system using PatchMatch
%%%% An in-vivo mouse demonstration where 10 different laser wavelengths
%%%% were used. 
%%%% Created by Geng-Shi Jeng
%%%% All copyright reserved 2020 by Bioengineering, University of Washington, Seattle, USA.                           

clear all
load PAUS_10frame_motion

laser_wav = [700 715:20:875]; %[735 875 795 700 815 755  855  715  775  835];  
Nv = length(laser_wav);

%%% Original B mode size
[lat_N,axi_N,~] = size(bmode_all);

%%% Parameters for PatchMatch
krn_r=3;	    %Kernel size limit (rows)
krn_c=8;		%Kernel size limit (columns)

srch_ru=3; 		%Search area limit (rows, UP)
srch_rd=3;		%Search area limit (rows, DOWN)
srch_cl=15;  	%Search area limit (columns, LEFT)
srch_cr=15;		%Search area limit (columns, RIGHT)

itera = 4;          % Iteration for PatchMatch 
N_rs = 1;           % Random search


%%%%%% Imaging parameter  
TH = -50;  % Thresholding for the lowest pixel value

%%%% Define ROI where motion estimation and compensation are done. 
ROI_lat_s = 1;   % start index of lateral ROI
ROI_lat_e = 89; % ending index of lateral ROI
ROI_lat_i = 2;   % interval of lateral ROI 

ROI_axi_s = 140; % start index of axial ROI
ROI_axi_e = 200; % ending index of axial ROI
ROI_axi_i = 1;   % interval of axial ROI 


%%%% Size and indexing of the ROI
ssz = [((ROI_lat_e-ROI_lat_s)/ROI_lat_i+1) ((ROI_axi_e-ROI_axi_s)/ROI_axi_i+1)]- 2*[krn_r krn_c];
indxx_all = (1:ssz(1))'*ones(1,ssz(2)); % lateral indexing
indxz_all = ones(ssz(1),1)*(1:ssz(2));  % axial indexing

%%%% interpolation of indexing in the ROI
indxx_intp =  1:1/ROI_lat_i:ssz(1);
indxz_intp =  1:1/ROI_axi_i:ssz(2);
indxx_all_intp = indxx_intp'*ones(1,length(indxz_intp)); % lateral indexing
indxz_all_intp = ones(length(indxx_intp),1)*indxz_intp;  % axial indexing
[tmp1, tmp2] = size (indxz_all_intp);

%%% Let indexing after interpolation to be intergral 
[indx_all_intp_int,indz_all_intp_int]= meshgrid(1:tmp1,1:tmp2);

%%%% Initializatoin
Axi_dis_med_all = zeros(tmp2,tmp1);
Lat_dis_med_all = zeros(tmp2,tmp1);

%%% Main program
nn = 2; 
for jj = 2:10 

refimg = bmode_all(ROI_lat_s:ROI_lat_i:ROI_lat_e,ROI_axi_s:ROI_axi_i:ROI_axi_e,jj-1);
tarimg = bmode_all(ROI_lat_s:ROI_lat_i:ROI_lat_e,ROI_axi_s:ROI_axi_i:ROI_axi_e,jj);

Def_db = 20*log10(abs(tarimg));
flag_val  = Def_db > TH;
   
tic
%%% PatchMatch core function
NNF = PM2DabsRT_mex(tarimg,refimg, krn_r,krn_c,srch_ru,srch_rd,srch_cl,srch_cr, N_rs, itera);

%% Displacement 
Displ(:,:,1) = NNF(:,:,1) - (krn_r+indxx_all); % lateral
Displ(:,:,2) = NNF(:,:,2) - (krn_c+indxz_all); % axial 

toc

%%% Thresholding and post-processing (median filtering)
Lat_dis = (Displ(:,:,1).*flag_val(krn_r+1:end-krn_r,krn_c+1:end-krn_c))';
Axi_dis = (Displ(:,:,2).*flag_val(krn_r+1:end-krn_r,krn_c+1:end-krn_c))';
Lat_dis_med = medfilt2(Lat_dis,[17 7]);
Axi_dis_med = medfilt2(Axi_dis,[17 7]);

%%% Interpolation & Accumulation
Axi_dis_med_all  = Axi_dis_med_all + interp2(indxx_all',indxz_all',Axi_dis_med,indxx_all_intp',indxz_all_intp');
Lat_dis_med_all  = Lat_dis_med_all + interp2(indxx_all',indxz_all',Lat_dis_med,indxx_all_intp',indxz_all_intp');

%%% Alignment and indexing 
indx_z_def = (ROI_axi_s + krn_c*ROI_axi_i-1) + indz_all_intp_int + round(Axi_dis_med_all*ROI_axi_i);
indx_x_def = (ROI_lat_s + krn_r*ROI_lat_i-1) + indx_all_intp_int + round(Lat_dis_med_all*ROI_lat_i);

%%% Load PA data 
pamode = pamode_all(:,:,jj);
%%% Load US data
DefImag = bmode_all(:,:,jj)';
     
%%% Compensating motion
DefImag_cor = DefImag(indx_z_def + (indx_x_def-1)*axi_N);
pamode_cor = pamode(indx_z_def + (indx_x_def-1)*axi_N);


%%% Plot  
SL_plot = SL_pos(ROI_lat_s + krn_r*ROI_lat_i:ROI_lat_s + krn_r*ROI_lat_i+tmp1-1)*1000; %
range_plot = range(ROI_axi_s + krn_c*ROI_axi_i:ROI_axi_s + krn_c*ROI_axi_i+tmp2-1)*1000;

figure

ax2 = subplot(221);
imagesc(SL_pos*10^3,range*10^3,20*log10(abs(pamode))) %SL_pos*1000,range*1000,
title(['PA(original) ' num2str(laser_wav(nn)) 'nm'])
colormap(ax2,hot(256))
axis([SL_pos(ROI_lat_s + krn_r*ROI_lat_i)*10^3 SL_pos(ROI_lat_s + krn_r*ROI_lat_i+tmp1-1)*10^3 range(ROI_axi_s + krn_c*ROI_axi_i)*10^3 range(ROI_axi_s + krn_c*ROI_axi_i+tmp2-1)*10^3]) 
caxis([-30 0])
colorbar
xlabel('mm')
ylabel('mm')

ax5 = subplot(223);
imagesc(SL_plot,range_plot,20*log10(abs(pamode_cor))) %SL_pos*1000,range*1000,
title(['PA(corrected) ' num2str(laser_wav(nn)) 'nm'])
colormap(ax5,hot(256))
caxis([-30 0])
colorbar
xlabel('mm')
ylabel('mm')

ax6 = subplot(224);
imagesc(SL_plot,range_plot,Axi_dis_med_all*ROI_axi_i)
colormap(ax6,parula(256))
colorbar
caxis([-15 15])
title('Axial Displ(pixel)')
xlabel('mm')
ylabel('mm')


ax3 = subplot(222);
imagesc(SL_plot,range_plot,Lat_dis_med_all*ROI_lat_i)
colormap(ax3,parula(256))
colorbar
caxis([-6 6])
title('Lateral Displ(pixel)')
xlabel('mm')
ylabel('mm')
     
 nn = nn +1;

end