function choices = nchoosek_owr(nn, kk, force_set)
%NCHOOSEK_OWR like NCHOOSEK but k ordered choices with replacement
%
%     choices = nchoosek_owr(nn, kk[, force_set=0])
%
% If nn is a non-negative integer scalar, return the number of choices, for each
% setting of kk. That is, nn.^kk
%
% If nn is a vector, give all possible sets of length-kk choices from this
% vector. The choices matrix is (nn^kk x kk). kk must be a scalar.
%
% This style of function overloading (copied from Matlab's nchoosek) is
% dangerous. Note what happens if your 'set' is coincidentally of length 1. This
% problem is rife in Matlab code. You can avoid bugs caused by this "feature" by
% setting the optional third argument to be non-zero. This forces the first
% argument, nn, to be interpreted as a set, even if nn is a single non-negative
% integer.

% Iain Murray, October 2008, January 2009, February 2009

if ~exist('force_set', 'var')
    force_set = 0;
end

is_nonneg_integer = @(x) isscalar(x) && (x >= 0) && (round(x) == x);

if is_nonneg_integer(nn) && (~force_set)
    choices = nn.^kk; % Note: equal to 1 for kk=0
elseif isvector(nn)
    if ~is_nonneg_integer(kk)
        error('kk must be a non-negative integer when returning choices from set nn');
    end
    if kk == 0
        choices = zeros(1, 0); % There is one choice, which is the empty set
    else
        nn = nn(:);
        choices = nn;
        num = length(choices);
        vec = @(x) x(:);
        for ii = 1:(kk-1)
            prev_length = size(choices, 1);
            choices = [vec(repmat(nn', prev_length, 1)), repmat(choices, num, 1)];
        end
    end
else
    error('nn must be a non-negative integer or a vector');
end