function [Z,vtvZ] = imseg_hidd_vf(P, tau, mu, iters)
%
%   Energy segmentation using a hidden vector MRF. For details see 
%
%   J. Bioucas-Dias, F. Condessa, and J. Kovacevic, "Alternating direction 
%   optimization for image segmentation using hidden Markov measure field 
%   models", in IS&T/SPIE Electronic Imaging, International Society for 
%   Optics and Photonics,  pp. 90190P-90190P, February, 2014.
%
%
%  -- input parameters ---------------------------------
%  
%   P ->  [nc, nr*nc]  matrix containing the class conditional probalilities 
%                      in its columns
%
%  tau ->  regularization  parameter for of the vectot TV 
%
%  mu  ->  ADMM paramer
%
%  iters -> ADMM iters
%
%
%
%  -- Output parameter -------
%
%  Z    -> Extimated hidden fields
%  VTV  -> vector TV of the solution


%--------------------------------------------------------------------------
% BAYESIAN SEGMENTATION 
%--------------------------------------------------------------------------
%  P(x_i|y_i=k) = cte p_i(k)   -> computed by  the LORSAL
%  
%  p(y,z,x) = p(x|y)p(y|z)p(z)
%
%  Let us assume that  
%   
%        p(x|y) = prod p(x_i|y_i)
%        p(y|z) = prod p(y_i|z_i)    and   p(y_i = k|z_i)  = z_i(k) 
%                                        ==> z_i >=0 and sum(z_i) = 1
%
%         -log p(z)  = tau TV(z)
%       
%     
%     Computing the marginl on y, we have
%
%         p(z,x) = prod_i sum_k( p(x_i|y_i =k)z_i(k)) p(z)
%
%         -log p(z,x)  =  sum_i -log(f(z_i)) + lambda TV(z)
%
% 
%    where  f_i(z_i) = sum_k( p(x_i|y_i =k)z_i(k))
%
%    Segmentation
%
%    min   sum_i -log(f_1(z_i)) - tau TV(z)  s.t.:   z >= 0, sum(z_i) =1
%     z
%    
%     Equivalent problem
%
%    min   sum_i -log(f_i(z_i)) + tau TV(z) i_+(z) + i_1(z)
%     z
%
%     SPLITING
%
%    min  sum_i -log(f(v1_i)) + tau ||(V2,V3)|| i_+(V4) + i_1(V5)
%     z
%    
%    s.t.: v1 = Z
%          V2 = DHZ
%          V3 = DVZ
%          V4 = Z
%          V5 = Z
% 
%   ADMM
%   
%   L(V1,V2,V3,V4,D1,D2,D3,D4) =
%             sum_i -log (f(v1_i))   +  mu/2||Z-V1-D1||_F^2 +
%             tau ||(V2,V3)||       +  mu/2||ZDH-V2-D2||_F^2 +
%                                    +  mu/2||ZDV-V3-D3||_F^2 +
%             i_+ f(V4)              +  mu/2||Z-V4-D4||_F^2 +
%             i_1 f(V5)              +  mu/2||Z-V5-D5||_F^2 +
%
%
% Author: Jose Bioucas Dias, 2014 (bioucas@lx.it.pt)

[nr,nc,no_classes] = size(P);
n = nr*nc;
P = reshape(P, n,no_classes)';


%%
%------------------------------------------------------------------------
%  Define functions and operators
%------------------------------------------------------------------------
%
% ---------------------------------------------
% define a circular convolution (the same for all bands) accepting a 
% matrix  and returnig a matrix 

ConvC = @(X,FK,nb)  reshape(real(ifft2(fft2(reshape(X', nr,nc,nb)).*repmat(FK,[1,1,nb]))), nr*nc,nb)';

% convert matrix to image
conv2im  = @(X,nb)  reshape(X',nr,nc,nb);
% convert image to matrix
conv2mat = @(X,nb)  reshape(X,nr*nc,nb)';


% define horiontal difference operator kernel
dh = zeros(nr,nc);
dh(1,1) = 1;
%dh(1,nc) = -1;
% dh(1,1) = 2;
 dh(1,2) = -1;
% dh(1,nc) = -1;
% 
dv = zeros(nr,nc);
dv(1,1) = 1;
%dv(nr,1) = -1;
% dv(1,1) = 2;
 dv(2,1) = -1;
% dv(nr,1) = -1;
% 

% compute FFTs for filtering
FDH = fft2(dh);
FDHC = conj(FDH);
FDV = fft2(dv);
FDVC = conj(FDV);

%  more filters in the frequency domain
I_DH = FDHC ./(abs(FDH).^2+ abs(FDV).^2 + 3);
I_DV = FDVC ./(abs(FDH).^2+ abs(FDV).^2 + 3);
II     = 1  ./(abs(FDH).^2+ abs(FDV).^2 + 3);


% some constants
P2 =  sum(P.^2);

Iu = eye(no_classes)- 1/no_classes*ones(no_classes);
B = ones(no_classes,n)/no_classes;

% define and initialize variables
Z = P;
V1 = Z;
D1 = 0*Z;
V2 = Z;
D2 = 0*Z;
V3 = Z;
D3 = 0*Z;
V4 = Z;
D4 = 0*Z;
V5 = Z;
D5 = 0*Z;

% All variables (U,V1,V2,V3,V4,V5, D1,D2,D3,D4,D5)  are matrices

% SALSA  (ADMM) Algorithm
for i=1:iters
    
%            min     mu/2||Z-V1-D1||_F^2 +
%             Z      mu/2||ZDH-V2-D2||_F^2 +
%                    mu/2||ZDV-V3-D3||_F^2 +
%                    mu/2||Z-V4-D4||_F^2 +
%                    mu/2||Z-V5-D5||_F^2 +
%         
%  

Z =    ConvC(V1+D1, II,no_classes) + ConvC(V2+D2, I_DH,no_classes) + ...
       ConvC(V3+D3, I_DV,no_classes) + ConvC(V4+D4, II,no_classes)  + ...
       ConvC(V5+D5, II,no_classes); 


%  min sum_i -log (f(v1_i))  +  mu/2||Z-V1-D1||_F^2 +
%   V1
NU1 = Z-D1;
pnu = sum(P.*NU1);
aux = (pnu + sqrt(pnu.^2 +4*P2/mu))/2;
V1  = NU1 + P./(repmat(aux*mu,no_classes,1));


%   min tau ||(V2,V3)||_1      +  mu/2||ZDH-V2-D2||_F^2 +
%    V2,V3                     +  mu/2||ZDV-V3-D3||_F^2 +
NU2 =  ConvC(Z,FDH,no_classes) - D2;
NU3 =  ConvC(Z,FDV,no_classes) - D3;
[V2,V3] = vector_soft_col_iso(NU2,NU3,tau/mu);


%non isotropic VTV
% V2 = vector_soft_col(NU2,tau/mu);
% V3 = vector_soft_col(NU3,tau/mu);
% 

% quadratic regularizer
%    V2 = (mu/tau)/(1+mu/tau) * NU2;
%    V3 = (mu/tau)/(1+mu/tau) * NU3;

% i_+ f(V4)  +  mu/2||Z-V4-D4||_F^2 +
NU4 = Z-D4;
V4 = max(0,NU4);


% i_1 f(V5)  +  mu/2||Z-V5-D5||_F^2 +
NU5 = Z-D5;
V5 = Iu*NU5 + B;



fprintf('iter = %d, e1 = %2.2f, e2 = %2.2f, e3 = %2.2f, e4 = %2.2f, e5 = %2.2f\n', ...
        i,  norm(NU1+D1-V1, 'fro'), norm(NU2+D2-V2, 'fro'),  norm(NU3+D3-V3, 'fro'), ...
    norm(NU4+D4-V4, 'fro'), norm(NU5+D5-V5, 'fro') )
 
 % update Lagrange multipliers
 D1 = -NU1 + V1;    % D1 -(ZB-V1)
 D2 = -NU2 + V2;    % D2 -(ZDH-V2)
 
 D3 = -NU3 + V3;    % D3 -(ZDV-V3)
 D4 = -NU4 + V4;    % D4 -(Z-V4)
 D5 = -NU5 + V5;    % D5 -(Z-V5)
 
 
 Zim = conv2im(Z,no_classes);
 imagesc(Zim(:,:,2))

 
 figure(gcf)
 drawnow;
 
end

% project onto the probability simplex
Z = max(0,Z);
Z = Z./repmat(sum(Z),no_classes,1);




% compute TV Z
vtvZ = sum(sqrt(sum([ConvC(Z,FDH,no_classes).^2; ConvC(Z,FDV,no_classes).^2])));



