clear
close all

im = double ( imread('Pics/Lenna.png') ); 

Mb_size = 8;        % Macro block size
sample_ratio = 2;    % Chroma subsampling ratio
q = 1;               % Quantization depth

%% Calculate luma and chroma and downsample 
[Height,Width,Num_Components] = size(im);
im=im(1:floor(Height/(Mb_size*sample_ratio) )*(Mb_size*sample_ratio),1:floor(Width/(Mb_size*sample_ratio))*(Mb_size*sample_ratio),:);

if (Num_Components == 3)
R = im(:,:,1);  G = im(:,:,2);  B = im(:,:,3);
Y = 0.2126*R + 0.7152*G + 0.0722*B;
%Y = rgb2gray(im)
Cr = R - Y;
Cb = B - Y;
else
   Y = im;
   Cr = 0*Y;
   Cb = 0*Y;
end

filt = 1/sample_ratio^2*ones(sample_ratio,sample_ratio);

Y = Y - 128;
Cr_filt = conv2(Cr,filt,'same') - 128; 
Cb_filt = conv2(Cb,filt,'same') - 128;
Cr_ds = Cr_filt(1:sample_ratio:end,1:sample_ratio:end);
Cb_ds = Cb_filt(1:sample_ratio:end,1:sample_ratio:end);



%% Divide into macroblocks
Nmb = size(Y)/Mb_size;

Y_mb = mat2cell(Y, repmat(Mb_size,1,Nmb(1)), repmat(Mb_size,1,Nmb(2)));
Cr_mb = mat2cell(Cr_ds, repmat(Mb_size,1,Nmb(1)/sample_ratio), repmat(Mb_size,1,Nmb(2)/sample_ratio));
Cb_mb = mat2cell(Cb_ds, repmat(Mb_size,1,Nmb(1)/sample_ratio), repmat(Mb_size,1,Nmb(2)/sample_ratio));


%% Perform DCT for each block and quantize
DCT_mx = sqrt(2/Mb_size)*cos( pi/Mb_size*  (0:Mb_size-1)'*( 1/2 + (0:Mb_size-1)  ) );
DCT_mx(1,:) = DCT_mx(1,:)*1/sqrt(2);


Q = [16 11 10 16  24  40  51  61;
     12 12 14 19  26  58  60  55;
     14 13 16 24  40  57  69  56;
     14 17 22 29  51  87  80  62;
     18 22 37 56  68 109 103  77;
     24 35 55 64  81 104 113  92;
     49 64 78 87 103 121 120 101;
     72 92 95 98 112 100 103  99]*q;
[N,D]= rat(Mb_size/8);
Q = resample(resample(Q,N,D)',N,D)';


Y_dct = cell(Nmb(1),Nmb(2));
Cr_dct = cell(Nmb(1)/sample_ratio,Nmb(2)/sample_ratio);
Cb_dct = cell(Nmb(1)/sample_ratio,Nmb(2)/sample_ratio);

Y_zigzag = cell(Nmb(1),Nmb(2));
Cr_zigzag = cell(Nmb(1)/sample_ratio,Nmb(2)/sample_ratio);
Cb_zigzag = cell(Nmb(1)/sample_ratio,Nmb(2)/sample_ratio);
zz_i = zigzag(Mb_size);

for n = 1 : Nmb(2)
    for m = 1 : Nmb(1)
        
        % DCT transform
        Y_dct{m,n} = DCT_mx*Y_mb{m,n}*DCT_mx';
        % Quantize
        Y_dct{m,n} = round(Y_dct{m,n}./Q);
        % Zig-zag coding with deleting trailing zeros
        temp = Y_dct{m,n};
        temp2 = temp(zz_i);
        Y_zigzag{m,n} =  temp2(1:find(temp2,1,'last'));
    end
end

for n = 1 : Nmb(2)/sample_ratio
    for m = 1 : Nmb(1)/sample_ratio
        Cr_dct{m,n} = DCT_mx*Cr_mb{m,n}*DCT_mx';
        Cr_dct{m,n} = round(Cr_dct{m,n}./Q);  
       
        temp = Cr_dct{m,n};
        temp2 = temp(zz_i);
        Cr_zigzag{m,n} = temp2(1:find(temp2,1,'last'));
        
        Cb_dct{m,n} = DCT_mx*Cb_mb{m,n}*DCT_mx';
        Cb_dct{m,n} = round(Cb_dct{m,n}./Q);  
       
        temp = Cb_dct{m,n};
        temp2 = temp(zz_i);
        Cb_zigzag{m,n} = temp2(1:find(temp2,1,'last'));
    end
end

save('proba.mat','Y_zigzag','Cr_zigzag','Cb_zigzag','Mb_size','Q','sample_ratio')


%% Decoder
clear
close all

F = open('proba.mat');
Y_zigzag = F.Y_zigzag;
Cr_zigzag = F.Cr_zigzag;
Cb_zigzag = F.Cb_zigzag;
Mb_size = F.Mb_size;
Q = F.Q;
sample_ratio = F.sample_ratio;
filt = 1/sample_ratio^2*ones(sample_ratio,sample_ratio);

Nmb = size(Y_zigzag);
zz_i = zigzag(Mb_size);
DCT_mx = sqrt(2/Mb_size)*cos( pi/Mb_size*  (0:Mb_size-1)'*( 1/2 + (0:Mb_size-1)  ) );
DCT_mx(1,:) = DCT_mx(1,:)*1/sqrt(2);

for n = 1 : Nmb(2)
    for m = 1 : Nmb(1)
        temp = zeros(Mb_size);
        temp2 = zeros(Mb_size^2,1);
        temp2(1 : length(Y_zigzag{m,n}))  = Y_zigzag{m,n};
        temp(zz_i) =  temp2;
        Y_mb{m,n} = DCT_mx'*(temp.*Q)*DCT_mx;
    end
end
for n = 1 : Nmb(2)/sample_ratio
    for m = 1 : Nmb(1)/sample_ratio
       
        temp = zeros(Mb_size);
        temp2 = zeros((Mb_size)^2,1);
        temp2(1 : length(Cr_zigzag{m,n}))  = Cr_zigzag{m,n};
        temp(zz_i) =  temp2;
        
        Cr_mb{m,n} = DCT_mx'*(temp.*Q)*DCT_mx;        
        
        temp = zeros(Mb_size);
        temp2 = zeros((Mb_size)^2,1);
        temp2(1 : length(Cb_zigzag{m,n}))  = Cb_zigzag{m,n};
        temp(zz_i) =  temp2;
        Cb_mb{m,n} = DCT_mx'*(temp.*Q)*DCT_mx;
        
    end
end

Y = cell2mat( Y_mb )+128;
Cr_ds = cell2mat( Cr_mb )+128;
Cb_ds = cell2mat( Cb_mb )+128;

Cr_ = zeros(size(Cr_ds)*sample_ratio);
Cb_ = zeros(size(Cb_ds)*sample_ratio);
Cr_(1:sample_ratio:end,1:sample_ratio:end) = Cr_ds*sample_ratio^2;
Cb_(1:sample_ratio:end,1:sample_ratio:end) = Cb_ds*sample_ratio^2;

Cr_rec = conv2(Cr_,filt,'same');
Cb_rec = conv2(Cb_,filt,'same');

R = Cr_rec + Y;
B = Cb_rec + Y;
G = (Y - 0.2126*R - 0.0722*B)/0.7152;

out = zeros(size(R,1),size(R,2),3);
out(:,:,1) = R; out(:,:,2) = G; out(:,:,3) = B;

imshow(uint8( out ) )