function [r, tieadj] = tiedrank(x, tieflag, bidirectional, epsx, wts)
%TIEDRANK Compute the ranks of a sample, adjusting for ties.
%   [R, TIEADJ] = TIEDRANK(X) computes the ranks of the values in the
%   vector X.  If any X values are tied, TIEDRANK computes their average
%   rank.  The return value TIEADJ is an adjustment for ties required by
%   the nonparametric tests SIGNRANK and RANKSUM, and for the computation
%   of Spearman's rank correlation.
%
%   [R, TIEADJ] = TIEDRANK(X,1) computes the ranks of the values in the
%   vector X.  TIEADJ is a vector of three adjustments for ties required
%   in the computation of Kendall's tau.  TIEDRANK(X,0) is the same as
%   TIEDRANK(X).
%
%   [R, TIEADJ] = TIEDRANK(X,0,1) computes the ranks from each end, so
%   that the smallest and largest values get rank 1, the next smallest
%   and largest get rank 2, etc.  These ranks are used for the
%   Ansari-Bradley test.
%
%   [...] = TIEDRANK(X,0,0,EPSX) considers two X values within EPSX to be
%   tied. EPSX has the same size as X, and the comparison between X(J) and
%   X(K) is carried out with absolute tolerance EPSX(J)+EPSX(K).
%   TIEDRANK uses an absolute tolerance of zero by default.
%
%   [...] = TIEDRANK(X,0,0,0,W) computes weighted ranks.
%
%   See also ANSARIBRADLEY, CORR, PARTIALCORR, RANKSUM, SIGNRANK.

%   Copyright 2022-2024 The MathWorks, Inc.

if nargin < 2
    tieflag = false;
end
if nargin < 3
    bidirectional = 0;
end
if nargin<4 || isempty(epsx)
    epsx = zeros(size(x),"like",x);
end
if nargin<5 || isempty(wts)
    if isvector(x)
        wts = ones(1,length(x),"like",x);
    else
        wts = ones(1,size(x,1),"like",x);
    end
else
    wts = wts(:)';
end

if isvector(x)
   [r,tieadj] = tr(x,tieflag,bidirectional,epsx,wts);
else
   % Operate on each column vector of the input (possibly > 2 dimensional)
   sz = size(x);
   ncols = sz(2:end);  % for 2x3x4, ncols will be [3 4]
   r = zeros(sz,"like",x);
   if tieflag
      tieadj = zeros([3,ncols],"like",x);
   else
      tieadj = zeros([1,ncols],"like",x);
   end
   for j=1:prod(ncols)
      [r(:,j),tieadj(:,j)] = tr(x(:,j),tieflag,bidirectional,epsx(:,j),wts);
   end
end

% --------------------------------
function [r,tieadj] = tr(x,tieflag,bidirectional,epsx,wts)
%TR Local tiedrank function to compute results for one column

% Sort, then leave the NaNs (which are sorted to the end) alone
[sx, rowidx] = sort(x(:));
epsx = epsx(rowidx);
epsx = epsx(:);
numNaNs = sum(isnan(x));
xLen = numel(x) - numNaNs;
wts = wts(rowidx(1:xLen));

% Compute ranks
fwd = 1/2 + cumsum(wts) - wts/2;
if bidirectional
    % Use ranks counting from both low and high ends
    bac = 1/2 + cumsum(wts,"reverse") - wts/2;
    ranks = [min(fwd,bac), NaN(1,numNaNs)]';
else
    % Use ranks counting from low end
    ranks = [fwd, NaN(1,numNaNs)]';
end

if tieflag
    tieadj = zeros(3,1,"like",x);
else
    tieadj = zeros(1,1,"like",x);
end
if isUnderlyingType(x, "single")
   ranks = single(ranks);
end

% Adjust for ties.  Avoid using diff(sx) here in case there are infs.
ties = sx(1:xLen-1)+epsx(1:xLen-1) >= sx(2:xLen)-epsx(2:xLen);
tieloc = [find(ties); xLen+2];
maxTies = numel(tieloc);

tiecount = 1;
while (tiecount < maxTies)
    tiestart = tieloc(tiecount);
    ntied = 2;
    wtadj = wts(tiestart+1) .* wts(tiestart);
    while(tieloc(tiecount+1) == tieloc(tiecount)+1)
        wtadj = wtadj + sum(wts(tiestart+ntied) .* wts(tiestart+(0:ntied-1)));
        tiecount = tiecount+1;
        ntied = ntied+1;
    end

    if tieflag
        n2minusn = ntied*(ntied-1);
        % For Kendall, we need the first adjustment factor to take weights
        % into account. The others are used only for p-values. For constant
        % weights we would have
        %     tieadj = tieadj + [n2minusn/2; n2minusn*(ntied-2); n2minusn*(2*ntied+5)];
        tieadj = tieadj + [wtadj; n2minusn*(ntied-2); n2minusn*(2*ntied+5)];

    else
        % Spearman's correlation does not use the weights for computing the
        % coefficient value, and the hypothesis test functions don't accept
        % weights, so the tie adjustment is just based on counts.
        tieadj = tieadj + ntied*(ntied-1)*(ntied+1)/2;
    end
    
    % Compute weighted mean of tied ranks
    rvec = ranks(tiestart:tiestart+ntied-1);
    wvec = wts(tiestart:tiestart+ntied-1);
    ranks(tiestart:tiestart+ntied-1) = sum(wvec(:).*rvec(:)) / sum(wvec);
    tiecount = tiecount + 1;
end

% Broadcast the ranks back out, including NaN where required.
r(rowidx) = ranks;
r = reshape(r,size(x));
