%% UPRAVENE DEMO
% Kod je vytoren na zakladě kodu DEMO JPEG: 
%       (c) 2010-2016 Pavel Rajmic, Brno University of Technology
% Autor: Jiri Bures 2022-2023
%%
% Priprava prostredi
close all; 
clc; 
clear all;
home; 

%% Parametry
quality = 50;       % ukladana kvalita obrazku: [0,100], vychozi je 50  
nEdges = 1;       % mnozství zachovanych pixelu z hran z [0,1] 
maxIterace = 1000;  % maximalni pocet iteraci proximálniho algoritmu

imgShowSizeX = 700; % sirka obrazku pri zobrazeni ve figure
imgShowSizeY = 700; % vyska obrazku pri zobrazeni ve figure

%% Nacteni obrazku ke komprimaci
disp(['Start: Nacteni obrazku...'])

% HS_raw.png snezenky_orez.png TEST.png 01 02
picture= 'ycbcr_orig.png';

imgLoad = imread(['obr\',picture]); % nacteni obrazku    

%% Kontrola velikosti obrazku 
%   - pripadne rozsirenim o okrajove hodnoty
disp(['Start: Kontrola velikosti...'])

x = size(imgLoad(:,1,1),1); % nacteni povodni velikosti x
y = size(imgLoad(1,:,1),2); % nacteni povodni velikosti y
dx = 0; dy = 0;             % pocet pridanych radku/sloupcu do nasobku 8

if x/8 ~= round(x/8) 
    dx = 8 - (x - 8*floor(x/8));
    for i = x+1:x+dx
        imgLoad(i,:,:) = imgLoad(x,:,:);
    end
end
if y/8 ~= round(y/8) 
    dy = 8 - (y - 8*floor(y/8));
    for i = y+1:y+dy
        imgLoad(:,i,:) = imgLoad(:,y,:);
    end
end

dim_ver = size(imgLoad(:,1,1),1);	% ulozeni nove velikosti x
dim_hor = size(imgLoad(1,:,1),2);  % ulozeni nove veliksoti y

%% Prevod do pozadovaneho brevneho prostoru a posun o -128
% Protoze pracujeme pouze s jasovou slozkou, prevadime do sedotonoveho
% obrazu, ktery oproti jasove slozce YCbCr pokryva cely rozsah [0, 255]
disp(['Start: Prevod barev...'])

% YCbCr schema:
% YCbCr = rgb2ycbcr(iLoad); %figure; imshow(YCbCr);         % prevod do YCbCr
% Y = YCbCr(:,:,1); %Cb = YCbCr(:,:,2); Cr = YCbCr(:,:,3);	% jasova slozky Y
% %figure; imshow(Y); figure; imshow(Cb); figure; imshow(Cr); % zobrazeni slozek
% image_orig = Y;

% Sedotonovy obraz:
if length(size(imgLoad))>2
    imGray = rgb2gray(imgLoad);	% prevedeni do odstinu sedi
    image_orig = imGray; 
else
    image_orig = imgLoad;
end;

% zobrazeni obrazku
figure, imshow(image_orig), 
truesize([imgShowSizeX imgShowSizeY]), 
title('Načtený obrázek');

image0 = cast(image_orig,'double'); % zmena datoveho typu na plovouci r. carku
image0 = image0 - 128;              % posun o -128

%% Podvzorkování
% Vynechavame, protoze pracujeme pouze s cernobilym obrazkem, tedy jasovou 
% slozkoua ta se obvykle nepodvozrkova
     
%% Detekce hran
disp(['Start: Detekce hran...'])

% nacteni hran
% EDs ... hranicni body vybrane (selected)
% EDp ... hranicni body (permutations)
% EDl ... ulozene hodnoty na hranicich (list)
EDl = zeros(dim_ver, dim_hor)-1000; % mit i zaporne hodnoty 0-128=-128

edges = edge(image0,'canny');	% vyber hran pomoci funkce MATLABu    
% figure, imshow(edges), title('Algoritmem nalezené hrany');  % zobrazení hran

% nahodny vyber bodu hran pro ulozeni
EDs = sparse(edges > 0);                % ulozeni hran do sparse matrix
n = sum(sum(EDs > 0));                  % pocet detekovanych bodu hran
EDp = randperm(n, round(n*nEdges));     % nahodny vyber

% vytvoreni matice s hodnotami pixelu pro vybrane body hran
[corX, corY, ~] = find(EDs);    % nacteni souradnic bodu hran
for i = 1:numel(EDp)
    % pro kazdy prvek EDv pouzije danou hodnotu EDv jako index poradi v [corX, corY, val],
    % coz jsou souradnice nahodne vybranych hran
    % pomoci techto souradnic a priradi hodnotu z puvodniho obrazku do nove matice
    EDl( corX(EDp(i)), corY(EDp(i)) ) = image0( corX(EDp(i)), corY(EDp(i)) );
end

% zobrazeni ulozenych bodu hran
figure, imagesc(EDl), colormap('gray'), colorbar,  
title('Zobrazení náhodně vybraných hran');

%% Nacteni kvantizacni tabulky
disp(['Start: Nacteni kvantizacni tabulky...'])

Qmtx = quant_matrix(quality); 

%% Kodovani
disp(['Start: Kodovani...'])

% inicializace promennych
no_blocks_ver = dim_ver/8;	% pocet bloku vertikalne
no_blocks_hor = dim_hor/8;  % pocet bloku horizontalne
B = zeros(8);               % matice 8x8 zpracovavaneho bloku
D = dctmtx(8);              % predpocitana DCT matice 8x8
                            %   - pouzivame pro vyssi rychlost
dctB = zeros(8);            % matice transformacnich koeficientu DCT
koefB = zeros(8);           % matice kvantizacnich koeficientu

imgSave = zeros(dim_ver, dim_hor); % matice simulujici JPEG soubor obrazku
x = zeros(1,8);             % rozsah bloku ve smeru X
y = zeros(1,8);             % rozsah bloku ve smeru Y

for v = 0:(no_blocks_ver-1)
    for h = 0:(no_blocks_hor-1)
        x = (v+1)*8-7:(v+1)*8;      % nastaveni rozsahu bloku
        y = (h+1)*8-7:(h+1)*8;      %     (v+1)*8-8+1:(v+1)*8-8+8
        B = image0(x, y);            % nacteni konkretniho bloku
        dctB = D*B*D';              % vypocet (dopredne) DCT(B)
        koefB = round(dctB./Qmtx);  % kvantizace
        % cik-cak razeni koeficientu% X nepovadime
        % Huffmanovo kodovani       % X neprovadime
        imgSave(x, y) = koefB;      % ulozeni bloku
    end
end

%% Dekodovani - "puvodni JPEG"
disp(['Start: Dekodovani - JPEG...'])

tic; % zacatek mereni casu
image = zeros(dim_ver, dim_hor); % inicializace pomocne obrazove matice

for v = 0:(no_blocks_ver-1)
    for h = 0:(no_blocks_hor-1) 
        x = (v+1)*8-7:(v+1)*8;  % nastaveni rozsahu bloku
        y = (h+1)*8-7:(h+1)*8;
        koefB = imgSave(x, y);  % nacteni bloku koeficientu
        % Huffmanovo dekodovani % X neprovadime
        % zpetne cik-cak razeni % X nepovadime    
        dctB = koefB.*Qmtx;     % dekvantizace
        B = D'*dctB*D;          % vypocet (zpetne) IDCT(dctB)
        image(x, y) = round(B); % prevod na celociselne hodntoy 
                      %uint8(round(B));
    end
end

image1 = image(1:dim_ver-dx, 1:dim_hor-dy); % uprava na puvodni velikost
image1 = image1 + 128;                      % posun zpet o 128
time1 = toc;    % konec mereni casu

% zobrazeni dekodovaneho obrazku
figure, imshow(uint8(image1)), 
truesize([imgShowSizeX imgShowSizeY]), 
title('dekodovany obrazek - puvodni algoritmus');

%% Dekodovani
disp(['Start: Dekodovani - upraveny algoritmus...'])

image = zeros(dim_ver, dim_hor); % vycisteni pomocne obrazove matice
tic;    % zacatek mereni casu

for v = 0:(no_blocks_ver-1)
    for h = 0:(no_blocks_hor-1)
        x = (v+1)*8-7:(v+1)*8;  % nastaveni rozsahu bloku
        y = (h+1)*8-7:(h+1)*8;
        koefB = imgSave(x, y);  % nacteni bloku koeficientu
        % Huffmanovo dekodovani % X neprovadime
        % zpetne cik-cak razeni % X nepovadime    
        dctB = koefB.*Qmtx;     % dekvantizace
        % aplikace proximalniho algotirmu   
            % "vysledny obr" = algorithm("zakodovane koeficienty", "ulozene hrany", 
            % 	"kvantizační matice", "pocet iteraci", "zobrazit každých i iterací" );
        EDlBlock = EDl(x, y);	% nacteni odpovidajiciho bloku matice s hranami    
        if max(max(EDlBlock)) > -1000
            % pokud se  nenajdou hrany nema smysl provadet algoritmus
            dctB = prox_algorithm(dctB, EDlBlock, Qmtx, maxIterace, 0);
        end
        B = D'*dctB*D;  % vypocet (zpetne) IDCT(dctB)
        image(x, y) = round( B );	% prevod na celociselne hodntoy
    end
end

image2 = image(1:dim_ver-dx, 1:dim_hor-dy); % uprava na puvodni velikost
image2 = image2 + 128;                      % posun zpet o 128

time2 = toc;    % konec mereni casu

% zobrazeni dekodovaneho obrazku
figure, imshow(uint8(image2)), 
truesize([imgShowSizeX imgShowSizeY]), 
title('dekodovany obrazek - novy algoritmus');

%% PSNR
% Vypocet hodnoty PSNR, pro dobre obrazky zhruba v intervalu [40,50]
disp(['Start: PSNR...'])

sizeX = dim_ver-dx; sizeY = dim_hor-dy;
image_orig = image_orig(1:sizeX, 1:sizeY);

MSEr1 = norm(double(image_orig)-double(image1),'fro')^2 / sizeX / sizeY;
PSNR1 = 10*log10(255^2 / MSEr1 );

MSEr2 = norm(double(image_orig)-double(image2),'fro')^2 / sizeX / sizeY;
PSNR2 = 10*log10(255^2 / MSEr2 );

%PSNR1 = psnr(uint8(image1), uint8(image_orig));
%PSNR2 = psnr(uint8(image2), uint8(image_orig));

%% SSIM
% Vypocet hodnoty SSIM, v rozmezi [0,1]; pro dobre obrazky je blize k 1
disp(['Start: SSIM...'])

SSIM1 = ssim(double(image1), double(image_orig));
SSIM2 = ssim(double(image2), double(image_orig));

%% Logovani do souboru
% 
% jemno obrazku & SXxSY & v1 & cas1 & MSNr1 & PSNR1 & SSIM1 \\
% jmeno obrazku & SXxSY & v2 & cas2 & MSNr2 & PSNR2 & SSIM2 \\@
% formatSpec1 = '%s & %4.0f x %4.0f & v1 & %4.4f & %4.4f & %2.4f & %1.4f \\\\\n';
% formatSpec2 = '%s & %4.0f x %4.0f & v2 & %4.4f & %4.4f & %2.4f & %1.4f \\\\@\n';
% fileID = fopen('vysledky.txt','w');
% 
% fprintf(fileID,formatSpec1',pic,sizeX,sizeY,time1,MSNr1,PSNR1,SSIM1);
% fprintf(fileID,formatSpec2',pic,sizeX,sizeY,time2,MSNr2,PSNR2,SSIM2);
% 
% fclose(fileID);

%% Zobrazeni vysledku a provnani obrazku
disp(['T1: ',num2str(time1)])
disp(['T2: ',num2str(time2)])
disp(['MSEr1: ',num2str(MSEr1)]);
disp(['MSEr2: ',num2str(MSEr2)]);
disp(['PSNR1: ',num2str(PSNR1)]);
disp(['PSNR2: ',num2str(PSNR2)]);
disp(['SSIM1: ',num2str(SSIM1)]);
disp(['SSIM2: ',num2str(SSIM2)]);

figure,
imagesc(abs(double(image_orig)-double(image1))), 
title(['Absolutni rozdily; PSNR1=' num2str(PSNR1)]), 
axis square tight off, colorbar;
figure,
imagesc(abs(double(image_orig)-double(image2))), 
title(['Absolutni rozdily; PSNR2=' num2str(PSNR2)]), 
axis square tight off, colorbar;




