function [xe, x0s, reg_par_selected] = tvdenoise(x, t, varargin)
% TVDENOISE   Denoise 1-D sequence of data with TV regularisation.
%    XE = TVDENOISE(X) returns a denoised version of data contained in the
%    vector X.
% 
%    XE = TVDENOISE(X, T) takes the values in T as the abscissa. T must be
%    a vector of the same length as X or a scalar; if T is a scalar, it is
%    used as the constant spacing between the abscissa.
% 
%    XE = TVDENOISE(..., 'RegPar', A), where A is a scalar, uses A as
%    regularisation parameter. If A has more than one element, denoising
%    results are computed for each of them and TVDENOISE attempts to select
%    the optimum one. If RegPar is not specified, TVDENOISE considers the
%    following values: [1e-5, 1e-3, 1e-1]. It is recommended to provide a
%    single value or a set of values empirically selected for the given
%    problem under analysis.
% 
%    XE = TVDENOISE(..., 'DerivativeOrder', K) returns estimates of the
%    Kth derivative instead of the denoised data.
%    
%    XE = TVDENOISE(..., 'ConstraintOrder', J) performs total-variation
%    regularisation by constraining the L1 norm of the Jth derivative. By
%    default, the first derivative is constrained.
%    
%    XE = TVDENOISE(..., 'StdDev', S), where S is a vector of the same
%    length as X, uses elements of S as estimates of the standard
%    deviations of errors corrupting the data X. They are used for
%    weighting the penalty term and for optimising the regularisation
%    parameter (if the 'ncp' optimisation strategy is not selected). If S
%    is a scalar, TVDENOISE assumes standard deviation is equal to S for
%    each element of the sequence X, and S is only used for optimisation of
%    the regularisation parameter (if the 'ncp' optimisation strategy is
%    not selected).
% 
%    XE = TVDENOISE(..., 'OptimisationStrategy', STRATEGY) uses the
%    selected STRATEGY for optimisation of the regularisation parameter
%    (but only if the value of the optional argument 'RegPar' is not a
%    scalar - if it is a scalar, the regularisation parameter is not
%    optimised at all). The following strategies are available:
%    - 'dp' - a strategy based on the discrepancy principle; requires
%      estimates of the standard deviations of errors corrupting the data,
%      provided in the optional argument 'StdDev'.
%    - 'ncp' - a strategy based on the normalised cumulative periodogram.
%    If the parameter 'StdDev' is specified, the 'dp' strategy is used by
%    default; otherwise, the 'ncp' strategy is used.
%    
%    XE = TVDENOISE(..., 'MaxIterations', I) performs at most I iterations
%    for minimising the regularised penalty term.
% 
%    XE = TVDENOISE(..., 'RelTol', THR) stops iterating when the relative
%    change in the regularised penalty term is below THR.
% 
%    [XE, XS] = TVDENOISE(X, ...) returns the denoised version of the data
%    X except of the estimates of the derivative. XS = XE if the optional
%    parameter 'DerivativeOrder' is set to 0.
% 
%    [XE, XS, ALPHA] = TVDENOISE(X, ...) also returns the value of the
%    regularisation parameter selected by TVDENOISE.
% 
%    Jakub Wagner, June 30, 2021
%    Institute of Radioelectronics and Multimedia Technology
%    Warsaw University of Technology


% Check size of x
if ~isvector(x) || isscalar(x), error('x must be a vector.'); end
if isrow(x), x = x'; xrow = true; end
N = length(x);

% Abscissa not specified
if ~exist('t', 'var') || isempty(t), t = (1:N)'; end

% Spacing specified instead of abscissa
if isscalar(t), t = t*(0:N-1)'; end

% Check size of t
if ~isvector(t) || length(t) ~= N
    error('t must be a scalar or a vector of the same length as x.'); end
if isrow(t); t = t'; end

% Parse optional parameters
NArgs = numel(varargin);
if rem(NArgs, 2) ~= 0
    error('Wrong number of input arguments.')
end
for narg = 1:2:NArgs
    switch varargin{narg}
        case 'DerivativeOrder'
            derivative_order = varargin{narg + 1};
            if ~isscalar(derivative_order) || ...
                    floor(derivative_order) ~= derivative_order
                error('DerivativeOrder must be a scalar integer.')
            end
        case 'ConstraintOrder'
            constraint_order = varargin{narg + 1};
            if ~isscalar(constraint_order) || ...
                    floor(constraint_order) ~= constraint_order
                error('ConstraintOrder must be a scalar integer.')
            end
        case 'RegPar'
            reg_par = varargin{narg + 1};
            if ~isnumeric(reg_par) || any(reg_par <= 0)
                error('RegPar must only contain positive numbers.')
            end
        case 'StdDev'
            std_dev = varargin{narg + 1};
            if ~isscalar(std_dev) && ...
                    (~isvector(std_dev) || length(std_dev) ~= length(x))
                error(['StdDev must be a scalar or a vector of the same'...
                    'length as x.'])
            end
        case 'OptimisationStrategy'
            strategy = varargin{narg + 1};
            if ~ischar(strategy)
                error(['OptimisationStrategy must be ''dp'', ''lc'''...
                    'or ''ncp''.'])
            end
        case 'MaxIterations'
            i_max = varargin{narg + 1};
            if ~isscalar(i_max) || floor(i_max) ~= i_max
                error('MaxIterations must be a scalar integer.')
            end
        case 'RelTol'
            stop_crit_thr = varargin{narg + 1};
            if ~isscalar(stop_crit_thr) || stop_crit_thr <= 0
                error('MaxIterations must be a positive scalar number.')
            end
        otherwise
            error('Unknown optional parameter: %s', varargin{narg})
    end
end

% Default: denoising, no differentiation
if ~exist('derivative_order', 'var'), derivative_order = 0; end

% Default: constraint on L1-norm of first derivative
if ~exist('constraint_order', 'var'), constraint_order = 1; end

% Default set of values of regularisation parameter
if ~exist('reg_par', 'var')
    reg_par = 10.^(-5:2:-1);
end

% Default maximum number of iterations
if ~exist('i_max', 'var'), i_max = 1000; end

% Check if estimates of standard deviations are provided
if ~exist('std_dev', 'var')
    std_dev_provided = false;
    std_dev = ones(size(x));
    default_strategy = 'ncp';
else
    std_dev_provided = true;
    if isscalar(std_dev), std_dev = std_dev * ones(size(x)); end
    default_strategy = 'dp';
end

% Default strategy for optimisation of regularisation parameter
if ~exist('strategy', 'var'), strategy = default_strategy; end

% Default threshold on relative change of objective function
if ~exist('stop_crit_thr', 'var'), stop_crit_thr = 1e-12; end

% Parameter preventing division by zero
e = 1e-16;

% Shift data so that x(1) = 0
xp = x - x(1);

% Matrix of coefficients of progressive-difference formula
D = pdmatrix(t, constraint_order-derivative_order);

% Matrix of coefficients of trapezoid rule
Q = qmatrix(t, derivative_order);

% Weighting matrix based on estimates of standard deviations of data
W = diag(std_dev.^-2) * min(std_dev)^2;

% Parameters used for optimisation of regularisation parameter.
switch lower(strategy)
    case 'dp'
        if ~std_dev_provided, error(['Cannot use ''dp'' optimisation'...
                'strategy without estimates of standard deviations of'...
                'data. Specify ''StdDev'' parameter or use ''ncp'''...
                'strategy.']); end
        uDP = 2;
        trS = sum(std_dev.^2);
    case 'ncp'
        Nl = floor(N/2);
        l = (1:Nl)' / Nl;
    otherwise
        error('Unknown OptimisationStrategy: %s', strategy)
end

% Number of values of regularisation parameter to be checked
NRP = numel(reg_par);

% Criterion minimised in order to select regularisation parameter
reg_par_crit = zeros(NRP, 1);

% Matrix of denoised data or estimates of derivative (depending on
% derivative_order) for different values of regularisation parameter
Xe = zeros(N,NRP);

% For each considered value of regularisation parameter
for nrp = 1:NRP
    
    alpha = reg_par(nrp);
    
    % Vector of denoised data or estimates of derivative (depending on
    % derivative_order)
    xe = zeros(N, 1);
    
    % Iterative minimisation of norm(Q*xe-xp, 2) + alpha*norm(D*xe, 1)
    stop_crit_previous = Inf;
    for i = 1:i_max
        
        % Compute change in xe
        E = diag(1./sqrt((D*xe).^2 + e));
        L = D'*E*D;
        H = Q'*W*Q + alpha*L;
        g = Q'*W*(Q*xe - xp) + alpha*L*xe;
        s = H \ -g;
        xe = xe + s;
        
        % Check stopping criterion
        stop_crit_current = norm(Q*xe-xp, 2) + alpha*norm(D*xe, 1);
        if (stop_crit_previous - stop_crit_current) ...
                / stop_crit_previous <= stop_crit_thr
            break
        end
        stop_crit_previous = stop_crit_current;
    end
    
    Xe(:, nrp) = xe;
    
    % Vector of approximation residuals
    residual = Q*xe - xp;
    
    % Compute criteria for optimisation of regularisation parameter
    switch lower(strategy)
        case 'dp'
            % Criterion based on discrepancy principle
            reg_par_crit(nrp) = abs(sum(residual.^2, 1) - uDP*trS);
        case 'ncp'
            % Criterion based on normalised cumulative periodogram
            rf = fft(residual);
            rp = abs(rf(1:Nl)).^2;
            rncp = cumsum(rp) / sum(rp);
            reg_par_crit(nrp) = sqrt(sum((rncp - l).^2));
    end
end


% Select optimum value of regularisation parameter
nrp_sel = find(reg_par_crit == min(reg_par_crit), 1, 'last');
xe = Xe(:, nrp_sel);
reg_par_selected = reg_par(nrp_sel);

% Shift back denoised data
if derivative_order == 0, xe = xe + x(1); end

% Return denoised data if needed
if nargout >= 2, x0s = Q*xe + x(1); end

% Return row vectors if x was row
if exist('xrow', 'var') && xrow
    xe = xe';
    if nargout >= 2, x0s = x0s'; end
end