%% =================================================================
% This script compares TRPCA methods by removing impulse noise
% listed as follows:
%     1. SNN-based TRPCA
%     2. TNN-based TRPCA
%     3. WSTNN-based TRPCA (the proposed method)
% You can:
%     1. Type 'Demo_TRPCA' to to run various methods and see the pre-computed results.
%     2. Select competing methods by turn on/off the enable-bits in Demo_TRPCA.m
%
% More detail can be found in [1]
% [1] Yu-Bang Zheng, Ting-Zhu Huang*, Xi-Le Zhao,Tai-Xiang Jiang,Teng-Yu Ji,and Tian-Hui Ma.
%     Tensor N-tubal rank and its convex relaxation for low-rank tensor recovery.
%
% Please make sure your data is in range [0, 1].
%
% Created by Yu-Bang Zheng £¨zhengyubang@163.com£©
% 11/19/2018

clear;
clc;
close all;
addpath(genpath('lib'));
addpath(genpath('data'));

EN_HaLRTC   = 0;
EN_TNN      = 0;
EN_WSTNN    = 1;
methodname  = { 'Observed','SNN', 'TNN',  'WSTNN'};
Mnum = length(methodname);
Re_tensor  =  cell(Mnum,1);
psnr       =  zeros(Mnum,1);
ssim       =  zeros(Mnum,1);
fsim       =  zeros(Mnum,1);
time       =  zeros(Mnum,1);

%% Load initial data
load('HSI_dc.mat')
if max(X(:))>1
    X = my_normalized(X);
end

Nway = size(X);
Ndim = ndims(X);

rhos = 0.4;  %% The ratio of corrupted pixels  
fprintf('=== The sample ratio is %4.2f ===\n', rhos);
%% index of the corrupted image
Xn = imnoise(X,'salt & pepper',rhos);   %% corrupted by uniform distributed values
   
i  = 1;
Re_tensor{i} = Xn;
[psnr(i), ssim(i), fsim(i)] = quality(X*255, Re_tensor{i}*255);
enList = 1;



%% Use HaLRTC

i = i+1;
if EN_HaLRTC
    %%%%%
    opts=[];
    alpha=ones(Ndim,1);
    opts.alpha=alpha/sum(alpha);
    opts.tol = 1e-4;
    opts.max_iter = 500;
    opts.rho = 1.1;
    opts.beta = 1e-2;
    opts.max_beta = 1e10;
    %opts.Xtrue=X;
    %%%%%
    fprintf('\n');
    disp(['performing ',methodname{i}, ' ... ']);
    t0= tic;
    [Re_tensor{i},~,~] = TRPCA_SNN(Xn,opts);    
    time(i)= toc(t0);
    [psnr(i), ssim(i), fsim(i)] = quality(X*255, Re_tensor{i}*255);
    enList = [enList,i]; 
end

%% Use TNN
i = i+1;
if EN_TNN
    %%%%%
    opts=[];   
    opts.tol = 1e-4;
    opts.max_iter = 500;
    opts.rho = 1.2;
    opts.beta = 1e-2;
    opts.max_beta = 1e10;
    %opts.Xtrue=X;
    opts.lambda = 1/sqrt(max(Nway(1),Nway(2))*Nway(3));
    %%%%%
    fprintf('\n');
    disp(['performing ',methodname{i}, ' ... ']);
    t0= tic;
    [Re_tensor{i},~,~] = TRPCA_TNN(Xn,opts);
    time(i)= toc(t0);
    [psnr(i), ssim(i), fsim(i)] = quality(X*255, Re_tensor{i}*255);
    enList = [enList,i];    
end

%% Use WSTNN
i = i+1;
if EN_WSTNN
    %%%%%
    % Please refer to our paper to set the parameters
    opts=[];
    alpha=[0,  0.001,  1;
           0,    0,    1;
           0,    0,    0]; 
    %%for Ndim-way tensors  (Ndim>3)
    % alpha = zeros(Ndim,Ndim);
    %         for ii=1:Ndim-1
    %             for jj=ii+1:Ndim
    %                 alpha(ii,jj)=1;
    %             end
    %         end 
    opts.alpha=alpha/sum(alpha(:));
    opts.tol = 1e-4;
    opts.max_iter = 500;
    opts.gamma = 1.2;
    omega = 100;
    opts.beta = opts.alpha/omega;
    opts.rho = 1/omega;
    opts.max_beta = 1e10*ones(Ndim,Ndim);
    opts.max_rho = 1e10;
    opts.lambda = Set_lambda(Nway,opts.alpha);
    %opts.Xtrue=X;
    %%%%%
    fprintf('\n');
    disp(['performing ',methodname{i}, ' ... ']);
    t0= tic;
    [Re_tensor{i},~,~]=TRPCA_WSTNN_ADMM(Xn,opts);
    time(i)= toc(t0);
    [psnr(i), ssim(i), fsim(i)] = quality(X*255, Re_tensor{i}*255);
    enList = [enList,i];    
end


%% Show result
fprintf('\n');
fprintf('================== Result =====================\n');
fprintf(' %8.8s    %5.4s    %5.4s    %5.4s   \n','method','PSNR', 'SSIM', 'FSIM');
for i = 1:length(enList)
    fprintf(' %8.8s    %5.3f    %5.3f    %5.3f   \n',...
        methodname{enList(i)},psnr(enList(i)), ssim(enList(i)), fsim(enList(i)));
end
fprintf('================== Result =====================\n');
figure,
showMSIResult(Re_tensor,X,min(X(:)),max(X(:)),methodname,enList,1,Nway(3))








