function regularized_fn = penalize_square(cost_fn, alpha)
%PENALIZE_SQUARE regularize a cost function for use in MINIMIZE_ARGS
%
% regularized_fn = penalize_square(cost_fn, alpha)
%
% cost_fn has the following signature:
%     [f, df_dx1, df_dx2, ..., df_dxK] = cost_fn(x1, x2, ..., xK)
%
% alpha is a single scalar, a vector of length K, or a cell array of length K.
% For cell arrays, each element alpha{k} must be a scalar or match x_k in size.
%
% The cost function has the following penalty added to it:
%     \sum_k alpha_k.*x_k(:)'*x_k(:)
% The derivatives of the new function are corrected accordingly.
%
% See also: minimize_args, utility_to_cost, var_checkgrad

if (~iscell(alpha)) && (~any(alpha))
    % By spotting this optimization here, external code doesn't need to be
    % cluttered with checks for alpha==0. Possible future optimization: could
    % also check inside cell arrays.
    regularized_fn = cost_fn;
else
    regularized_fn = @(varargin) penalize_square_helper(cost_fn, alpha, varargin);
end

function varargout = penalize_square_helper(fn, alpha, args)

varargout = cell(1, max(1, nargout));
[varargout{:}] = fn(args{:});

alpha_arg = args;
if length(alpha) == 1
    if iscell(alpha)
        alpha = alpha{1};
    end
    for ii = 1:length(args)
        alpha_arg{ii} = alpha_arg{ii}*alpha;
    end
else
    assert(length(alpha) == length(args));
    assert(isvector(alpha));
    if iscell(alpha)
        for ii = 1:length(args)
            alpha_arg{ii} = alpha_arg{ii}.*alpha{ii};
        end
    else
        for ii = 1:length(args)
            alpha_arg{ii} = alpha_arg{ii}*alpha(ii);
        end
    end
end

varargout{1} = varargout{1} + sum(cellfun(@(x,y) x(:)'*y(:), args, alpha_arg));
if nargout > 1
    varargout(2:end) = cellmap(@(x, y) x + 2*y, ...
            varargout(2:end), alpha_arg(1:(nargout-1)));
end