% COMPE 565 Homework 4
% 11/11/2018
% Name: Paul Jerrold Biglete
% REDID: 815430506
% Email: pbiglete@gmail.com
%
% Name: Van Zumarraga
% REDID: 818336877
% Email: vzumarraga@yahoo.com
%
% Name: Trisha Tolentino
% REDID: 818362695
% Email: trishadtolentino@gmail.com
%%
clc;
clear all;
close all;

%%
% Obtain video frames 6-10
video = VideoReader('walk_qcif.avi');
videoFrames = read(video, [6 10]);

% Convert RGB video frames to YCbCr
for i = 1:5
    video_frames(:,:,:,i) = rgb2ycbcr(videoFrames(:,:,:,i));
end

% Obtain video frame dimensions
numFrames = get(video, 'NumberOfFrames');
video_height = get(video, 'Height');
video_width = get(video, 'Width');
video_time = get(video, 'Duration');

% Variables & Objects
x_pos = 1; % x-axis index of frame
y_pos = 1; % y-axis index of frame
MB_size = 16;
P = 8; % displacement between MB and SW
SW_height = 32;
SW_width = 32;
sum = 0;
sw_count = 0; % debug var
MB_count = 0; % debug var
m = 1; % index SAD array
n = 1; % index SAD array
[frame_height,frame_width] = size(video_frames(:,:,1,1));
SAD_Count = 0;
TotalSADCount = 0;
Minimums = [];
Q28 = ones(8).*28;

%%
% Set I and P Frames (IPPPP); iterate in order of frames 6-10
for frameNum = 1:4

    %Compute the 8x8 Block DCT transform coefficients of the luminance and chrominance components of the image.
    Y_Component  = video_frames(:, :, 1, frameNum); % Seperate Y
    Cb_Component = video_frames(:, :, 2, frameNum); % Seperate Cb
    Cr_Component = video_frames(:, :, 3, frameNum); % Seperate Cr

    % 4:2:0, Subsample Cb and Cr components
    Cb_Subsample = Cb_Component(1:2:frame_height, 1:2:frame_width);
    Cr_Subsample = Cr_Component(1:2:frame_height, 1:2:frame_width);

    % Initialize blocks holding values of Y, Cb, and Cr blocks
    Y_Block = Y_Component;
    Cb_Block = Cb_Subsample;
    Cr_Block = Cr_Subsample;

    % Separates entire image into 8x8 blocks
    DCT_Block = @(block_struct) dct2(block_struct.data);

    DCT_Y_Block = blockproc(Y_Block, [8 8], DCT_Block);
    DCT_Cb_Block = blockproc(Cb_Block, [8 8], DCT_Block, 'PadPartialBlocks', true);
    DCT_Cr_Block = blockproc(Cr_Block, [8 8], DCT_Block, 'PadPartialBlocks', true);

%     % 1st Block 4th Row
%     Display_Block_1 = DCT_Y_Block(25:32, 1:8);
%     fprintf('\nDCT Coefficients Matrix of MB')
%     display(Display_Block_1);

    % Quantize blocks by dividing blocks by given matrices & rounding off
    Quantized_Y_Block  = @(block_struct) round(block_struct.data ./ Q28);
    Quantized_Cb_Block = @(block_struct) round(block_struct.data ./ Q28);
    Quantized_Cr_Block = @(block_struct) round(block_struct.data ./ Q28);

    % Apply values of processed quantized blocks to matrices
    Quantized_Luminance = blockproc(DCT_Y_Block, [8 8], Quantized_Y_Block);
    Quantized_Cb_Chrominance = blockproc(DCT_Cb_Block, [8 8], Quantized_Cb_Block);
    Quantized_Cr_Chrominance = blockproc(DCT_Cr_Block, [8 8], Quantized_Cr_Block);


    %%
    % Reference and current frames (Y_Component)
    ref_frame = double(video_frames(:,:,1,frameNum));
    current_frame = double(video_frames(:,:,1,frameNum+1));
    [ mV_X1, mV_Y1, mV_X2, mV_Y2 ] = motionVectors( ref_frame, current_frame, frame_height, frame_width, 8, 32);

%%
    figure
    quiver(mV_X1,mV_Y1,mV_X2,mV_Y2);
    axis([0 176 0 144]);
    title(sprintf('Motion Vectors: Frames %d to %d',frameNum+5,frameNum+6));

    diff_frame = current_frame - ref_frame;
    % Display Reference
    figure
    subplot(1,2,1);
    imshow(uint8(ref_frame));
    title(sprintf('Reference Image for Frames %d-%d', frameNum+5, frameNum+6));

    %Displays the Error image
    subplot(1,2,2);
    imshow(diff_frame);
    title(sprintf('Error Image %d-%d', frameNum+5,frameNum+6));

    fprintf('\n\tFrames %d-%d', frameNum + 5, frameNum + 6);

     %CPU time calculation
    time(frameNum) = cputime-time(frameNum);
    fprintf('\n\t\tCPU time: %.5f seconds', time(frameNum));

    %PSNR / MSE
    sum = 0;
    for m = 1:144
        for n = 1:176
            sum = sum + (current_frame(m,n)-ref_frame(m,n))^2;
        end
    end
    MSE = sum/(144*176);
    PSNR(frameNum) = 10*log10(255*255/MSE);

    AvgSADValue = mean(Minimums);

    fprintf('\n\t\tMSE: %d', MSE);
    fprintf('\n\t\tPSNR: %d', PSNR(frameNum));
    fprintf('\n\t\tAverage SAD Value per MB: %d', AvgSADValue);
    fprintf('\n\t\tTotal SAD computations: %d\n', TotalSADCount);

    TotalSADCount = 0; % Reset for each frame
    AvgSADValue = 0;

    %%
    % Dequantize blocks by multiplying block data with given matrices for luminance and chrominance
    DeQuantized_Y_Block = @(block_struct) round(block_struct.data .* Q28);
    DeQuantized_Cb_Block = @(block_struct) round(block_struct.data .* Q28);
    DeQuantized_Cr_Block = @(block_struct) round(block_struct.data .* Q28);

    % Apply values of dequantized blocks into matrices
    DeQuantized_Luminance = blockproc(Quantized_Luminance, [8 8], DeQuantized_Y_Block);
    DeQuantized_Cb_Chrominance = blockproc(Quantized_Cb_Chrominance, [8 8], DeQuantized_Cb_Block);
    DeQuantized_Cr_Chrominance = blockproc(Quantized_Cr_Chrominance, [8 8], DeQuantized_Cr_Block);

%     Display_Block_3 = DeQuantized_Luminance(25:32, 1:8);
%     fprintf('Inverse Quantized DCT Matrix of 1st block in 4th row')
%     display(Display_Block_3);

    % Apply inverse DCT to blocks & apply values into matrices
    InverseDCT = @(block_struct) idct2 (block_struct.data);
    Inverse_Y_Block = blockproc(DeQuantized_Luminance, [8 8], InverseDCT);
    Inverse_Cb_Block = blockproc(DeQuantized_Cb_Chrominance, [8 8], InverseDCT);
    Inverse_Cr_Block = blockproc(DeQuantized_Cr_Chrominance, [8 8], InverseDCT);

    % Reconstructed Image
    % Perform Row Replication with the newly decoded block
    UpsampledRowRepImg(frame_height, frame_width, 3) = uint16(0);

    for rows = 1:2:frame_height
        for columns = 1:2:frame_width % Load Cb and Cr Subsamples into Upsampled
            UpsampledRowRepImg(rows, columns, 2) = Inverse_Cb_Block((rows / 2) + 0.5, (columns / 2) + 0.5);
            UpsampledRowRepImg(rows, columns, 3)= Inverse_Cr_Block((rows / 2) + 0.5, (columns / 2) + 0.5);
        end
    end

    for columns = 2:2:frame_width % Place odd column into even
        UpsampledRowRepImg(:, columns, 2) = UpsampledRowRepImg(:, columns - 1, 2);
        UpsampledRowRepImg(:, columns, 3) = UpsampledRowRepImg(:, columns - 1, 3);
    end

    for rows = 2:2:frame_height % Place odd row into even
        UpsampledRowRepImg(rows, :, 2) = UpsampledRowRepImg(rows - 1, :, 2);
        UpsampledRowRepImg(rows, :, 3) = UpsampledRowRepImg(rows - 1, :, 3);
    end

    % Reconstruct Image with Inverse_Y_Block
    UpsampledRowRepImg(:, :, 1) = Inverse_Y_Block;
    UpsampledRowRepRGBImg = uint8(UpsampledRowRepImg);
    UpsampledRowRepRGBImg = ycbcr2rgb(UpsampledRowRepRGBImg);

    figure(1);
    subplot(1,2,2), imshow(UpsampledRowRepRGBImg), title('Reconstructed Image');
    subplot(1,2,1), imshow(ycbcr2rgb(video_frames(:,:,:,frameNum))), title('Original Image');

    % Display Error Values for the Block, Reconstructed - Original
    Inverse_Y_Block = uint8(Inverse_Y_Block);
    Y_Block = uint8(Y_Block);
    Img_Error = Y_Block(25:32, 1:8) - Inverse_Y_Block(25:32, 1:8);

    fprintf('Error Matrix');
    display(Img_Error);

    % PSNR of Luminance Component
    % Tried calculating MSE using the equation but couldn't match the value
    % from using the immse function
    % Image_Error_Sum = (Inverse_Y_Block(25:32, 1:8) - Y_Block(25:32, 1:8)).^2
    % Image_Error_Sum = sum(sum(Image_Error_Sum)) / 64

    MSE = immse(Inverse_Y_Block(25:32, 1:8), Y_Block(25:32, 1:8));
    peak_signal = 255 * 255;
    PSNR = 10 * log10((peak_signal.^2) / MSE);

    fprintf('PSNR Value');
    display(PSNR);
    % PSNR = 87.4660
end


%%
% Function to obtain macroblock given frame and its coordinates
function [MB] = getMB(x_pos,y_pos,frame,frame_width,frame_height)
    x = 1;
    y = 1;
    MB_size = 16;
    P = 8; % displacement from edge of SW

    % Check for edges when extracting macroblock
    if((x_pos+MB_size) > frame_width)
        x_pos = frame_width - MB_size + 1; % Shift x position to fit MB
    end
    if((y_pos+MB_size) > frame_height)
        y_pos = frame_height - MB_size + 1; % Shift y position to fit MB
    end

    % Iterate through to instantiate MB
    for i = x_pos:x_pos+MB_size-1
        for j = y_pos:y_pos+MB_size-1
            MB(y,x) = frame(j,i); % Return specified macroblock
            y = y + 1;
        end
        y = 1;
        x = x + 1;
    end
    x = 1;
end

%%
% Function to obtain search window given frame and macroblock
function [SW] = getSW(x_pos,y_pos,frame,frame_width,frame_height)
    x = 1;
    y = 1;
    P = 8; % displacement between MB and SW
    SW_height = 32;
    SW_width = 32;

    % If first row or first column, truncate search window
    if(x_pos - P < 1)
        SW_width = SW_width - P;
    else
        x_pos = x_pos - P; % Shift search window to center MB
    end
    if(y_pos - P < 1)
        SW_height = SW_height - P;
    else
        y_pos = y_pos - P; % Shift search window to center MB
    end

    % Check for edge when extracting search window from frame & MB
    if((x_pos + SW_width) > frame_width)
       SW_width = frame_width - x_pos + 1; % Truncate search window
    end
    if((y_pos + SW_height) > frame_height)
        SW_height = frame_height - y_pos + 1; % Truncate search window
    end

    % Iterate through to instantiate SW
    for i = x_pos:x_pos+SW_width-1
       for j = y_pos:y_pos+SW_height-1
          SW(y,x) = frame(j,i); % Return specified search window
          y = y + 1;
       end
       y = 1;
       x = x + 1;
    end
    x = 1;
end

%%
% Function to obtain SAD value of macroblock (MB)
function [sadValue, TotalSADCount] = getSAD_val(ref_MB,current_MB,SAD_Count,TotalSADCount)
    MB_size = 16;
    sum = 0; % Sum of pixel value differences between ref and curr frames
    SAD_Count = TotalSADCount;

    % Iterate through MB & calculate SAD value
    for x = 1:MB_size
        for y = 1:MB_size
            sum = sum + abs((current_MB(y,x) - ref_MB(y,x)));
            SAD_Count = SAD_Count + 1;
        end
        sadValue = sum / (MB_size * MB_size);
        y = 1; % Reset to row 1 after each column
    end
    TotalSADCount = SAD_Count;
    sum = 0;

end

%%
% Function to obtain position of minimum SAD value within search window
function [x_diff,y_diff, TotalSADCount, minSAD] = getMin_pos(x_pos,y_pos,current_MB,SW,SAD_Count,TotalSADCount)
    [SW_height, SW_width] = size(SW);

    x = 1;
    y = 1;
    MB_size = 16;
    SAD_array = ones(SW_height-MB_size+1, SW_width-MB_size+1) * (2^16);

    for x = 1:SW_width-MB_size+1
        for y = 1:SW_height-MB_size+1
            ref_MB = getMB(x,y,SW,SW_width,SW_height); % Obtain ref macroblock from SW
            [SAD_val, TotalSADCount] = getSAD_val(ref_MB,current_MB,SAD_Count,TotalSADCount);
            SAD_array(y, x) = SAD_val;
        end
    end

    % Obtain minimum SAD value (best matching MB) from SAD_array
    minSAD = min(min(SAD_array));

    for u = 1:SW_width-MB_size+1
        for v = 1:SW_height-MB_size+1
            if(minSAD == SAD_array(v,u))
                x_diff = u;
                y_diff = v;
            end
        end
    end

    x_diff = x_pos + x_diff - 1;
    y_diff = y_pos + y_diff - 1;

end

%%
function [ mV_X1, mV_Y1, mV_X2, mV_Y2 ] = motionVectors( refFrame, curFrame, FRAME_HEIGHT, FRAME_WIDTH, MB_SIZE, SEARCH_WIN )

    %% Motion vectors
    motionVectorsY = zeros(FRAME_HEIGHT/MB_SIZE,FRAME_WIDTH/MB_SIZE,2);
    motionVectorsX = zeros(FRAME_HEIGHT/MB_SIZE,FRAME_WIDTH/MB_SIZE,2);

    %% Keep track of the sum of MAD for a frame
    frMadTotal = 0;

    %% Error
    errImages = zeros(FRAME_HEIGHT,FRAME_WIDTH);

    %% Motion vectors index
    rowIdx = 1;
    colIdx = 1;

    %% Loop for searching block
    for blkRow=1:MB_SIZE:FRAME_HEIGHT
        for blkCol=1:MB_SIZE:FRAME_WIDTH

            %% Original location of MB
            motionVectorsY(rowIdx,colIdx,1) = blkRow;
            motionVectorsX(rowIdx,colIdx,1) = blkCol;

            %% Exact current MB from current frame
            curMB = curFrame(blkRow:blkRow+MB_SIZE-1, ...
                blkCol:blkCol+MB_SIZE-1);

            %& Exact search window from current MB
            [top, left, right, bottom] = windowSearchEdges(blkRow, ...
                blkCol, MB_SIZE, SEARCH_WIN, FRAME_HEIGHT, FRAME_WIDTH);

            %% Motion vector with Exhaustive Search
            [x1, y1, minMad] = exhaustiveSearch(...
                refFrame, curMB, top, left, right, bottom, MB_SIZE);

            %% sum of MAD for a specific frame
            frMadTotal = frMadTotal + minMad;
            if (frMadTotal < 128)
                x1 = 0;
                y1 = 0;
            end

            %% Best match location of MB
            motionVectorsY(rowIdx,colIdx,2) = y1;
            motionVectorsX(rowIdx,colIdx,2) = x1;

            %% pick MB with minimum MAD
            targetMB = refFrame(y1:y1+MB_SIZE-1, x1:x1+MB_SIZE-1);
            errImages(blkRow:blkRow+MB_SIZE-1,blkCol:blkCol+MB_SIZE-1)...
                = int16(curMB) - int16(targetMB);

            colIdx = colIdx+1; %index +1 after each iteration
        end
        colIdx = 1; %Reset after last column searched
        rowIdx = rowIdx+1; % index +1 after each iteration
    end


    %% Original
    mV_X1 = motionVectorsX(:, :, 2);
    mV_Y1 = motionVectorsY(:, :, 2);
    %% Best Match
    mV_X2 = motionVectorsX(:, :, 2);
    mV_Y2 = motionVectorsY(:, :, 2);

end