function [m,v] = gevstat(k,sigma,mu)
%GEVSTAT Mean and variance of the generalized extreme value distribution, optimised for GPU.
%   [M,V] = GEVSTAT(K,SIGMA,MU)
%
%   See also GEVSTAT, GPUARRAY.

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

%   Copyright 2021 The MathWorks, Inc.

if nargin < 3
    error(message('stats:gevstat:TooFewInputs'));
end

err = internal.stats.statsizechk(3,k,sigma,mu);
if err > 0
    error(message('stats:gevstat:InputSizeMismatch'));
end

% Return NaN for out of range parameters.
sigma(sigma <= 0) = NaN;

outProto = ones("like", internal.stats.dominantType(k,sigma,mu));

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


function [m,v] = gevstat_scalar(k, sigma, mu, outProto)

if k<1
    lngam1 = gammaln(1 - k);
else
    lngam1 = Inf("like",k);
end
if k<0.5
    lngam2 = gammaln(1 - 2*k);
else
    lngam2 = Inf("like",k);
end

% Calculate mean
if abs(k) < 1e-8
    tm = cast(-psi(1),"like",k);  % Euler's constant
elseif k < 1
    tm = expm1(lngam1) ./ k; % (gam1 - 1) ./ k
elseif k >= 1
    tm = inf("like", k);
else % k < 0
    tm = nan("like", k);
end

% Calculate variance
if abs(k) < 5e-6
    tv = cast(pi,"like",k)^2 / 6;
elseif k < 1/2
    tv = (expm1(lngam2) - expm1(2*lngam1)) ./ k.^2; % (gam2 - gam1.^2) ./ k.^2
elseif k >= 1/2
    tv = inf("like",k);
else % k < 0
    tv = nan("like",k);
end

m = mu + sigma .* cast(tm,"like",outProto);
v = sigma.^2 .* cast(tv,"like",outProto);
end