function [m, v] = nbinstat(r,p)
%NBINSTAT Mean and variance of the negative binomial distribution, optimised for GPU.
%   [M, V] = NBINSTAT(R,P)
%
%   See also NBINSTAT, GPUARRAY.

%   See toolbox/stats/stats/nbinstat for implementation details.

%   Copyright 2021 The MathWorks, Inc.

if nargin < 2
    error(message('stats:nbinstat:TooFewInputs')); 
end

errorcode = distchck(2,r,p);
if errorcode > 0
    error(message('stats:nbinstat:InputSizeMismatch'));
end
outProto = ones("like", internal.stats.dominantType(r, p));

% Process each element independently using arrayfun. This causes the
% function to be compiled into an optimised GPU kernel.
[m,v] = arrayfun(@nbinstat_scalar, r, p, outProto);
end


function [m,v] = nbinstat_scalar(r, p, outProto)
if 0 < r && isfinite(r) && 0 < p && p <= 1
    q = 1 - p;
    m = r .* q ./ p;
    v = r .* q ./ (p .* p);
else
    % Everything else is NaN
    m = NaN("like",outProto);
    v = NaN("like",outProto);
end
end