function [phi, dphi, wfun, ds] = penalty_funs(pen, a)
% [phi, dphi, wfun, ds] = penalty_funs(pen)
%
% Penalty and associated functions
%
% Input
%   pen : 'L1', 'log', 'atan', 'mcp'
%
% Ouput
%   phi : penalty function
%   dphi : derivative of phi
%   wfun : x/dphi(x)
%   ds : derivative of s where s is defined as s(x) = phi(x) - abs(x)
%
% Ivan Selesnick, selesi@nyu.edu

switch pen
    case 'L1'        
        phi = @(x) abs(x);
        dphi = @(x) sign(x);
        wfun = @(x) abs(x);
        ds = @(x) zeros(size(x));
        
    case 'log'        
        phi = @(x) 1/a * log(1 + a*abs(x));
        dphi = @(x) 1 ./(1 + a*abs(x)) .* sign(x);
        wfun = @(x) abs(x) .* (1 + a*abs(x));
        % ds = @(x) dphi(x) - sign(x);
        ds = @(x) - a * x ./ (1 + a*abs(x));
        
    case 'atan'        
        phi = @(x) 2./(a*sqrt(3)) .* (atan((2*a.*abs(x)+1)/sqrt(3)) - pi/6);
        wfun = @(x) abs(x) .* (1 + a.*abs(x) + a.^2.*abs(x).^2);
        dphi = @(x) 1 ./(1 + a*abs(x) + a.^2.*abs(x).^2) .* sign(x);
        ds = @(x) dphi(x) - sign(x);
        
    case 'mcp'
        phi = @(t) (abs(t) - a/2 * t.^2).*(abs(t) <= 1/a) + 1/(2*a)*(abs(t) > 1/a);
        dphi = @(t) max(1 - a*abs(t), 0) .* sign(t);
        wfun = @(x) abs(x)./max(1-abs(a*x), 0);   % Note: divide by zero is infinity
        ds = @(x) min(max(-a*x, -1), 1);        
end

