function [X, Out] = LRTC_WSTNN_ADMM(B, Omega, opts)
%%% Solve the WSTNN-based LRTC problem by ADMM
%
% min_X ||X||_WSTNN, s.t. P_Omega(X) = P_Omega(B)
% where ||X||_WSTNN=sum_{1=<k1<k2=<N} alpha_(k1,k2)||X_(k1,k2)||_TNN
% --------------------------------------------
% Input:
%       B       -    the observed tensor with missing elements, 
%                    please make sure B is size of d1*d2*d3 and in range [0, 1].
%       Omega   -    index of the known elements.
%       opts    -    structure value in Matlab. The fields are
%           opts.alpha      -   weights, 
%           opts.tol        -   termination tolerance,
%           opts.maxit      -   maximum number of iterations,
%           opts.beta       -   stepsize for dual variable updating in ADMM,
%           opts.max_beta   -   maximum stepsize,
%           opts.rho        -   rho>=1, ratio used to increase beta.
%         
% Output:
%       X       -    the recovery tensor.
%       Out     -    structure value in Matlab. The fields are
%           Out.PSNR        -   PSNR of each step,
%           Out.Res         -   Res of each step, 
%               Res: the relative square error of two successive recovered tensors,

% Date 11/19/2018
% Written by Yu-Bang Zheng (zhengyubang@163.com)
%%

if ~exist('opts', 'var')
    error('Not enough inputs!');
end
if isfield(opts, 'tol');         tol = opts.tol;              end
if isfield(opts, 'maxit');       maxit = opts.maxit;    end
if isfield(opts, 'rho');         rho = opts.rho;              end
if isfield(opts, 'beta');        beta = opts.beta;                end
if isfield(opts, 'max_beta');    max_beta = opts.max_beta;        end
if isfield(opts, 'alpha');       alpha = opts.alpha;        end

X = B;
Nway = size(B);
N = ndims(B);

Y = cell(N,N);
for i=1:N
    for j=1:N
    Y{i,j} = zeros(Nway); %% the auxiliary tensor
    end
end
M = Y; %% the Lagrange multiplier
temp = Y;
Out.Res=[];Out.PSNR=[];

for k = 1:maxit
    Xold = X;
    
    %% solve Y-subproblem
    tau = alpha./beta;
    for i=1:N-1
        for j=i+1:N
            tempX   = tensor_permute(X,Nway,i,j);
            tempM   = tensor_permute(M{i,j},Nway,i,j);
            Y{i,j}  = tensor_ipermute(prox_tnn_my(tempX + tempM/beta(i,j),tau(i,j)),Nway,i,j);
            temp{i,j} = Y{i,j}-M{i,j}/beta(i,j);
        end
    end
    
    %% solve X-subproblem
    tempsum = zeros(Nway);
    for i=1:N-1
        for j=1+i:N
        tempsum = tempsum+ beta(i,j)*temp{i,j};
        end
    end
    X = tempsum/sum(beta(:));
    X(Omega) = B(Omega);
    
    %% check the convergence
    if isfield(opts, 'Xtrue')
        XT=opts.Xtrue;
        psnr=PSNR(X,XT);
        Out.PSNR = [Out.PSNR,psnr];
    end
    res=norm(X(:)-Xold(:))/norm(Xold(:));
    Out.Res = [Out.Res,res];
    
    if k==1 || mod(k, 10) == 0
        if isfield(opts, 'Xtrue')
            fprintf('WSTNN: iter = %d   PSNR=%f   res=%f   \n', k, psnr, res);
        else
            fprintf('WSTNN: iter = %d   res=%f   \n', k, res);
        end
    end
    if res < tol 
        break;
    end
    
    %% update Lagrange multiplier
    for i=1:N-1
        for j=i:N 
           M{i,j} = M{i,j}+beta(i,j) * (X-Y{i,j});
        end
    end
    beta = min(rho * beta, max_beta);
end
end