clear
close all

load('foreman.mat');

MB_size = 8;
search_width = 8; 

I = crop_pic(  grayscale( frames(1).cdata ), MB_size );
B = crop_pic(  grayscale( frames(3).cdata ), MB_size  );
P = crop_pic(  grayscale( frames(4).cdata ), MB_size  );
%%
Nmb = size(I)/MB_size;
I_mb = mat2cell(I, repmat(MB_size,1,Nmb(1)), repmat(MB_size,1,Nmb(2)));
P_mb = mat2cell(P, repmat(MB_size,1,Nmb(1)), repmat(MB_size,1,Nmb(2)));
B_mb = mat2cell(B, repmat(MB_size,1,Nmb(1)), repmat(MB_size,1,Nmb(2)));

P_pred_mb_f = cell(size(I_mb));
P_pred_mb_b = cell(size(I_mb));
P_pred_mb = cell(size(I_mb));

%[ u_f , v_f ]       = get_frame_motion_vectors( B_mb, I_mb, MB_size, search_width );
%[ u_b , v_b ]       = get_frame_motion_vectors( B_mb, P_mb, MB_size, search_width );

pel_size = 'half';
for n = 1 : size(P_mb,1)
   for m = 1 : size(P_mb,2)
       
        [u_f(n,m), v_f(n,m)]  = get_motion_vector( B_mb{n,m}, n,m, I, MB_size, search_width,'ext', pel_size );
        [u_b(n,m), v_b(n,m)]  = get_motion_vector( B_mb{n,m}, n,m, P, MB_size, search_width,'ext', pel_size );
        
        P_pred_mb_f{n,m} = motion_prediction( B_mb{n,m}, I, u_f(n,m), v_f(n,m), n, m ,pel_size);
        P_pred_mb_b{n,m} = motion_prediction( B_mb{n,m}, P, u_b(n,m), v_b(n,m), n, m , pel_size);
        P_pred_mb{n,m} = (P_pred_mb_f{n,m}+P_pred_mb_b{n,m})/2;
   end
   
end

P_pred_F = cell2mat(P_pred_mb_f);
P_pred_B = cell2mat(P_pred_mb);
%%
figure('units','normalized','outerposition',[0 0 1 1])

max_I = max(max(I));
PSNR_orig = 10*log10( max_I^2/(sum( sum(  (P-I).^2 ) ) /numel(P).^2 ));
PSNR_Ppred = 10*log10( max_I^2/(sum( sum(  P_pred_F.^2 ) ) /numel(P).^2 ));
PSNR_Bpred = 10*log10( max_I^2/(sum( sum(  P_pred_B.^2 ) ) /numel(P).^2 ));

subplot(1,3,1);
pcolor(flipud(P-I));shading interp;colormap gray
caxis([-1e1,1e1]*10)
axis equal tight
title( sprintf( 'PSNR = %f', PSNR_orig  ) );
subplot(1,3,2);
pcolor(flipud(P_pred_F));shading interp;colormap gray
caxis([-1e1,1e1]*10)
axis equal tight
title( sprintf( 'PSNR = %f', PSNR_Ppred) );
subplot(1,3,3);
pcolor(flipud(P_pred_B));shading interp;colormap gray
caxis([-1e1,1e1]*10)
axis equal tight
title( sprintf( 'PSNR = %f', PSNR_Bpred) );
%%
figure('units','normalized','outerposition',[0 0 1 1])
[X,Y] = meshgrid((MB_size/2:MB_size:size(I,2)),(MB_size/2:MB_size:size(I,1)));

q = 1/2;
quiver(X,Y,flipud(u_f)*q,flipud(-v_f)*q,0,'-k','LineWidth',1.5);
hold on
quiver(X,Y,flipud(u_b)*q,flipud(-v_b)*q,0,'-k','LineWidth',1.5);
plot(X,Y,'ok','MarkerSize',3)
hold on
for n = 1:size(P,1)/MB_size
plot( linspace(1,size(P,2),100),(size(P,1)-n*MB_size)*ones(100,1) ,'-k') 
end
for n = 1:size(P,2)/MB_size
plot( n*MB_size*ones(100,1), linspace(1,size(P,1),100) ,'-k')
end
xlim([1,size(I,2)]);
ylim([1,size(I,1)]);
axis equal tight
title('B picture motion vectors')