%% Ayodele Embry %% EE368B
%% PROJECT
%% - USE DISTORTION AS THE DECISION CRITERION
%%

I = imread('peppers.tif');
img = double(I);
dim = size(img,1);

step =16                     %% Step Size

figure;
subplot(3,3,1), imshow(I);
title(' step = 8, distortion');


%   *(1) = 16x16 blocks
%   *(2) = 8x8 blocks
%   *(3) = 4x4 blocks

M = [16, 8, 4]              %% Size of block
B_cnt = dim./M         %% # of blocks per image 
B_opt = size(M,2)      %% # of block sizes available

img_recon = zeros(dim,dim);
img_final = zeros(dim,dim);

lambda = [ 0, 5, 10, 30, 62, 70, 90];
%lambda = 0
lambda_sz = size(lambda, 2);

J = zeros(B_cnt(1),B_cnt(1));    %% Block Cost Function
D = zeros(B_cnt(1),B_cnt(1));    %% Block Distortion
R = zeros(B_cnt(1),B_cnt(1));    %% Block Rate

% QUADTREE CODE- 
%   | 0 | 0 0 0 0 | - use 16x16 block (0 bits per 16x16 pixels)
%   | 1 | 0 0 0 0 | - use 4 8x8 blocks (1 bit per 16x16 pixels)
%   | 1 | 1 0 0 0 | - use 3 8x8 blocks & 4 4x4 blocks (5 bits per 16x16 pixels)
%   | 1 | 1 0 1 0 | - use 3 8x8 blocks & 4 4x4 blocks (5 bits per 16x16 pixels)

QTcode = zeros(B_cnt(1),5,B_cnt(1),lambda_sz);

% per pixel QTrate
QTrate = [0,1,5]/(M(1)*M(1));     

B16 = zeros(M(1),M(1));         % 16x16 block
B8  = zeros(M(2), M(2),4);      % 8x8 block
B4  = zeros(M(3), M(3),4);      % 4x4 block

%B_dct16  = zeros(M(1),M(1));   % Block DCT
%B_dct8   = zeros(M(2),M(2),4);
%B_dct4   = zeros(M(3),M(3),4);

%B_q16  = zeros(M(1),M(1));     % Quantized block DCT
%B_q8   = zeros(M(2),M(2),4);
%B_q4   = zeros(M(3),M(3),4);

for w=1:lambda_sz

   %% Loop through each 16x16 block and make the decision whether or
   %%  not to break down the block into 8x8 and/or 4x4 blocks

   for i=1:B_cnt(1)
      for j=1:B_cnt(1)
         row = M(1)*(i-1);
         col = M(1)*(j-1);

         % Create  16x16 block
         B16(:,:) = img(row+1:row+M(1),col+1:col+M(1));

         %Subdivide the 16x16 block into 4 8x8 blocks
         B8(:,:,1) = img(row+1:row+M(2), col+1:col+M(2));
         B8(:,:,3) = img(row+1:row+M(2), col+1+M(2):col+2*M(2));
         B8(:,:,2) = img(row+1+M(2):row+2*M(2), col+1:col+M(2));
         B8(:,:,4) = img(row+1+M(2):row+2*M(2), col+1+M(2):col+2*M(2));

         %% For 16x16 block, find the DCT, distortion, bitrate, IDCT
         %%   and cost function
         [D16, R16, B_idct16 ] = blk_calc(B16, step);
         J16 = D16 + lambda(w).*R16;

         %% For each 8x8 block, find the DCT, distortion, bitrate, IDCT
         %%   and cost function
         for p=1:4
            [D8(p), R8(p), B_idct8(:,:,p)] = blk_calc(B8(:,:,p), step);
            J8(p)  = D8(p) + lambda(w).*R8(p);
         end

         % Find the average distortion,bitrate, & cost for the 4 8x8 blocks 
         % Note that for the cost, you must add the additional bits for
         %   the quadtree code
         D8_QT = sum(D8)/4;
         J8_QT = sum(J8)/4 + lambda(w)*QTrate(2);

         % Compare costs; if the cost of 4 8x8 blocks is less than the
         %  the cost of a single 16x16 block, break up the block
         if D16 > D8_QT

            % Update the quadtree code
            QTcode(i,1,j,w) = 1;

            % Now, compare each 8x8 block with its corresponding 4x4 blocks
            for q=1:4

               %Subdivide the 8x8 block into 4 4x4 blocks
               B4(:,:,1) = B8(1:4,1:4,q);
               B4(:,:,2) = B8(1:4,5:8,q);
               B4(:,:,3) = B8(5:8,1:4,q);
               B4(:,:,4) = B8(5:8,5:8,q);

               %% For each 4x4 block, find the DCT, distortion, bitrate, IDCT
               %%   and cost function
               for p=1:4
                  [D4(p,q), R4(p,q), B_idct4(:,:,p)] = blk_calc(B4(:,:,p), step);
                  J4(p,q) = D4(p,q) + lambda(w).*R4(p,q);
               end

               J4_QT =  sum(J4(:,q))/4 + lambda(w)*QTrate(3);
               D4_QT =  sum(D4(:,q))/4;

               % Compare costs; if the cost of 4 4x4 blocks is less than the
               %  the cost of a single 8x8 block, break up the block
               if D8(q) > D4_QT

                  % Update the quadtree code
                  QTcode(i,1+q,j,w) = 1;

                  % Create a "new" 8x8 idct block made up of 4x4 blocks 
                  B8_new(1:4,1:4,q) = B_idct4(:,:,1);
                  B8_new(1:4,5:8,q) = B_idct4(:,:,2);
                  B8_new(5:8,1:4,q) = B_idct4(:,:,3);
                  B8_new(5:8,5:8,q) = B_idct4(:,:,4);
                  J8_new(q) = J4_QT;
                  D8_new(q) = sum(sum(D4(:,q)))/4;
                  R8_new(q) = sum(sum(R4(:,q)))/4 + QTrate(3);

               % If it costs less to use this 8x8 block, don't split it
               else
                  B8_new(:,:,q) = B_idct8(:,:,q);
                  J8_new(q) = J8(q);
                  D8_new(q) = D8(q);
                  R8_new(q) = R8(q) + QTrate(2);
               end
            end % for q

            
            % Create a "new" 16x16 idct block made up of 8x8 blocks 
            B16_new(1:8,1:8) = B8_new(:,:,1);
            B16_new(1:8,9:16) = B8_new(:,:,2);
            B16_new(9:16,1:8) = B8_new(:,:,3);
            B16_new(9:16,9:16) = B8_new(:,:,4);
 
            % Store the cost, distortion and rate for the new image
            J(i,j) = sum(J8_new)/4;
            D(i,j) = sum(D8_new)/4;
            R(i,j) = sum(R8_new)/4;

            % Add the created 16x16 block to the reconstructed image
            img_recon(row+1:row+M(1),col+1:col+M(1)) = B16_new;

         % If it costs less to use the 16x16 block, don't split it
         else

            % Add the 16x16 block to the reconstructed image
            img_recon(row+1:row+M(1),col+1:col+M(1)) = B_idct16;

            J(i,j) = J16;
            D(i,j) = D16;
            R(i,j) = R16;
         end

      end  % for j
   end  % for i

   % find the total cost function, rate and distortion for the entire image
   %   (4x4, 8x8, 16x16, best) 
   JT(w) = sum(sum(J(:,:)))/(B_cnt(1)*B_cnt(1))
   DT(w) = sum(sum(D(:,:)))/(B_cnt(1)*B_cnt(1))
   RT(w) = sum(sum(R(:,:)))/(B_cnt(1)*B_cnt(1))

%   subplot(3,3,w+1), imagesc(img_recon);
%   colormap=('gray');
%   axis('equal');

end % for w

PSNR = 10*log10(255*255./DT);
%figure;
%plot(RT, PSNR(:,w));
%xlabel('Rate[bits/pixel]');
%ylabel('PSNR [dB]');
%title('QT Distortion Decision');

%figure;
%plot(lambda, JT);
%xlabel('Lambda ');
%ylabel('Lagrangian Cost Function');
%title('QT Distortion Decision');


%figure;
%imagesc(img_recon);
%colormap=('gray');
%axis('equal');
%title(' peppers.tif, step = 1, distortion');






