function [x countIter] = prox_algorithm(ckq, edges, Q, iterace, paramSp)
% ["vysledny obr"] = algorithm("zakodovane koeficienty", "ulozene hrany", "kvantizační matice", "pocet iteraci", "show picture");

%% Autor: Jiri Bures
% rok: 2023
% Vysoke uceni techicke v Brne

%% incializace:
[sx, sy] = size(ckq);	% velikost matice koeficientu (resp. zde take jako velikost matice obrazku)
x = zeros(sx,sy);       % hledane reseni (dekvantizovane koeficienty)
u1 = zeros(sx,sy,2);	% pocatecni bod duálni promenne; m = 1
u2 = zeros(sx,sy);      % pocatecni bod duálni promenne; m = 2

countIter = 0;	% pocita probehle iterace; zaroven zajistuje, ze algoritmus
                % se provede alespon 100x; experimentovani ukazalo, ze
                % vynuceni takoveho poctu iteraci dava lepsi vysledky
                
% podminky pro parametry: tau > 0, sigma > 0, rho > 0:
rho = 1.99;     % rho € ]0,2[, pokud  f = 0 (nas pripad) (doporuceno 1.99)
sigma= 0.28;    % tau*sigma*||sum(Lm*Lm)|| <= 1 ;  ||D*D|| <= 8, ||dct()*idct()|| >= 0
tau = 0.99/( sigma*(8+1) ); % tau= 1/( sigma*norm(9*idct(ckq)*idct2(ckq)));

epsilon = 1e-4;	% pozadovana presnost; ukoncovaci podminka

%% CYKLUS ALGORITMU
%tic     % zacatek mereni casu

for i = 0:iterace-1
    in = []; % hodnota vstupujici jako argument do proximalnich operatoru
    
    % xdual_(i+1)
    x_old = x;
    
    % reseni x
    in = x - tau .* ( dct2(operDdual(u1) + u2) );    % zkraceno z: dct2(operDdual(u1)) + dct2(u2)
    x_wave = prox_tau_g(in, ckq, Q);                 % Ck = prox_tau_g(Ck, CkQ, Qk)
    
    % x_(i+1)
    x = rho*x_wave + (1-rho)*x;
    
    % "vnoreny FOR"
        % dualni reseni u1
        in_u1 = [];
        in_u1 = u1 + sigma .* operD( idct2(2*x_wave - x_old) );
        u1_dual = prox_sigma_h1_adj(in_u1);
        % uh_(i+1)
        u1 = rho*u1_dual + (1-rho)*u1;
        
        % ualni reseni u2
        in_u2 = [];
        in_u2 = u2 + sigma .* ( idct2(2*x_wave - x_old) );
        u2_dual = prox_sigma_h2_adj(in_u2, sigma, edges);
        % uv_(i+1)
        u2 = rho*u2_dual + (1-rho)*u2;
        
        % zobrazeni prubehu  
        countIter = countIter + 1;
%         if(Spi >= paramSp && paramSp > 0)
%              figure; imshow(uint8(idct2(x))); title(['Cycle ' num2str(i+1)]);
%             Spi = 0;
%         end
        
        % ukonceni
        if norm( x(:) - x_old(:) ) < epsilon * norm( x_old(:) ) && countIter > 100
%             fprintf('Konečný počet iterací: %d\n', Spi);
            break;
        end
end

%toc     % konec mereni casu

end %KONEC FUNKCE


%% ========== OPERATORY ========== 
% operD ... D diference
function [Uout] = operD(Xin)
   % dopredne konecne diference
   % pozor! v teorii:
   %    kh = sloupcovy index - v matlabu x - radek, y - sloupec: A(:,kh) 
   %    kv = radkovy index - v matlabu A(kv,:)
   % lze pouzit: Uout1(:, :, 1) = Xin * Dh; Uout1(:, :, 2) = Dv * Xin;
   
   [sizex, sizey] = size(Xin);
   h_dif = [Xin(:,2:end)-Xin(:,1:end-1), zeros(sizex,1)]; % horizontalni
   v_dif = [Xin(2:end,:)-Xin(1:end-1,:); zeros(1,sizey)]; % vertikalni
   
   Uout(:, :, 1) = h_dif; 
   Uout(:, :, 2) = v_dif;
end
% =========================================================================
% operDdual ... adjungovany k D diferencim
function [Xout] = operDdual(Uin)
    % zpetne konecne diference
    % lze pouzit i: Uin1(:, :, 1)*Dh' + Dv'*Uin(:, :, 2)
    h_bdif = [-Uin(:,1,1), Uin(:,1:end-2,1)-Uin(:,2:end-1,1), Uin(:,end-1,1)];    % horizontal
    v_bdif = [-Uin(1,:,2); Uin(1:end-2,:,2)-Uin(2:end-1,:,2); Uin(end-1,:,2)];    % vertical

    Xout = h_bdif + v_bdif;
end

% =========================================================================
% PROX_(sigma.h1*)
function [Uout_w] = prox_sigma_h1_adj(Uin)
    % proximalni operator pro ||.||_1,2 normu
    % uh ... horizontalni hodnoty
    % uv ... vertikalni hodnoty
    
    u_h = Uin(:,:,1);
    u_v = Uin(:,:,2);

    M = max(sqrt(u_h.^2 + u_v.^2), 1);
    Uout_w(:,:,1) = u_h./M;
    Uout_w(:,:,2) = u_v./M;
end
% =========================================================================
% PROX_(sigma.h2*)
function result = prox_sigma_h2_adj(I, sgm, edg)
    % prox. operator indikatorove funkce, ktera nahrazuje hodnoty v upravovanem obrazku 
    % puvodnimi hodnotami ulozenych hran (Mx = Me)
    % I ..... vstupujici hodnoty
    % sgm ... parametr sigma
    % edg ... ulozene hodnoty na hranach
    
    M = I./sgm;
    M(edg > -1000) = edg(edg > -1000);    
    result = I - sgm.*M;
end
% =========================================================================
% PROX_(tau_g)
function Ck = prox_tau_g(Ck, CkQ, Qk)
    % prox. operator indikatorove funkce, ktera zajistuje, ze nove dekvantovane hodnoty
    % nejsou prilis daleko od puvodnich
    % Ck .... hledane hodnoty
    % CkQ ... matice znamych kvantizovaných koeficientu z koderu
    % Qk .... matice kvantizacnich koeficientu
    % dk .... delta_k limit pro rozdíl
    %            zvolena hodnota vychazi ze procesu kvantizace
    % |Ck - CkQ| <= Qk/2
    
    dk = Qk ./ 2;

    vBottom = CkQ - dk;   % dolni mez pripustneho intervalu pro hledany koeficient
    vTop = CkQ + dk;      % horni mez pripustneho intervalu pro hledany koeficient

    I = find(Ck < vBottom); Ck(I) = vBottom(I); % posun prilis nizkych hodnot na dolni mez
    I = find(Ck > vTop);    Ck(I) = vTop(I);    % posun prilis vysokych hodnot na horni mez
end