clear
close all

%[mov,imgRgb] = loadFileYuv('foreman_qcif.yuv', 176, 144, (1:300));
%frames = mov;
load('foreman.mat');
load('Q.mat');
%% Parameters
GOP_structure = 'IBBBP';        % Group of pictures sequence
B_size = 8;                         % Block size ( DCT and quantizer )
MB_size = 16;                       % Macro-block size ( Motion prediction )
Wq = 4;                             % Quantizer scale
search_width = 8;                   % Width of search area in pixels ( search area: MB_size + 2*search_width square )

%% Encoder
Qi = interp_Q(Qi,B_size);
Qp = interp_Q(Qp,B_size);

GOP_length = length(GOP_structure);
order = get_enc_order( GOP_structure );
encode_order =  GOP_structure( order );

encoded_stream = struct('frame_data',0,'frame_type',0,'motion_vector_x',0,'motion_vector_y',0);
past_ref = 0;
future_ref = 0;

wb = waitbar(0,'Encoding frames');
for n = 1 : length(frames)
    waitbar(n/length(frames),wb);
    m = mod(n-1,GOP_length)+1;
    frame_num_to_encode =  floor((n-1)/GOP_length)*GOP_length+order(m);
    frame_in = crop_pic( frames(frame_num_to_encode).cdata , MB_size );
    
    switch encode_order(m)
        case 'I'
            past_ref = future_ref;
            frame_encoded = encode_I_frame( frame_in, B_size, MB_size, Wq*Qi );
            future_ref  = decode_I_frame( frame_encoded, MB_size, B_size, Wq*Qi  );
            %ref_backward = frame_in;
            encoded_stream = add_frame(encoded_stream, frame_encoded, 'I', 0, 0,n);
        case 'P'
            past_ref = future_ref;
            [frame_encoded, u, v]  = encode_P_frame( frame_in, past_ref,  B_size, MB_size, Wq*Qp, search_width );
            future_ref   = decode_P_frame( frame_encoded, past_ref ,MB_size, B_size, Wq*Qp, u, v  );
           % ref_backward = frame_in;
            encoded_stream = add_frame(encoded_stream, frame_encoded, 'P', u, v,n);
        case 'B'
            [frame_encoded, u_f, v_f, u_b, v_b]  = encode_B_frame( frame_in, past_ref, future_ref , B_size, MB_size, Wq*Qp, search_width );
            encoded_stream = add_frame(encoded_stream, frame_encoded, 'B', [u_f;u_b], [v_f;v_b],n);
    end
    
end
close(wb);

%% Decoder
output_stream = struct('frame_data',0,'frame_type',0,'motion_vector_x',0,'motion_vector_y',0);
past_ref = 0;
future_ref = 0;

wb = waitbar(0,'Decoding frames');
for n = 1 : length(encoded_stream)
    waitbar(n/length(encoded_stream),wb);
    
    m = mod(n-1,length(order))+1;
    orig_num =  floor((n-1)/GOP_length)*GOP_length+order(m);
    switch encoded_stream(n).frame_type
        case 'I'
            past_ref = future_ref;
            decoded_frame = decode_I_frame( encoded_stream(n).frame_data, MB_size, B_size, Wq*Qi  );
            future_ref  = decoded_frame;
            output_stream = add_frame( output_stream, decoded_frame, 'I', 0, 0, orig_num );
        case 'P'
            past_ref = future_ref;
            decoded_frame = decode_P_frame( encoded_stream(n).frame_data, past_ref, MB_size, ...
                B_size, Wq*Qp, encoded_stream(n).motion_vector_x,encoded_stream(n).motion_vector_y );
            future_ref  = decoded_frame;
            output_stream = add_frame( output_stream, decoded_frame, 'P', 0, 0, orig_num );
        case 'B'
            decoded_frame = decode_B_frame( encoded_stream(n).frame_data, past_ref, future_ref, MB_size, B_size, Wq*Qp,...
                encoded_stream(n).motion_vector_x(1:end/2,:),encoded_stream(n).motion_vector_y(1:end/2,:),...
                encoded_stream(n).motion_vector_x(end/2+1:end,:),encoded_stream(n).motion_vector_y(end/2+1:end,:));
            output_stream = add_frame( output_stream, decoded_frame, 'B', 0, 0, orig_num );
    end

end
close(wb);

%% Playback
figure('units','normalized','outerposition',[0 0 1 1])
for n = 1:length(output_stream)
    subplot(1,2,1)
    imshow( imresize(  frames(n).cdata , 2) );
    subplot(1,2,2)
    imshow( imresize( output_stream(n).frame_data ,2 ) );
    drawnow
end
%%
v = VideoWriter('proba.avi');
open(v);
for n = 1:length(output_stream)
    writeVideo(v, output_stream(n).frame_data );
end
close(v);