function A = onehotdecode(B, classes, featdim, typename)
% ONEHOTDECODE Decode probability vectors into class labels.
%
%   A = ONEHOTDECODE(B,CLASSES,FEATDIM) reduces a dlarray of probability
%   vectors into the most probable class label for each observation. Each
%   vector in B along the dimension specified by FEATDIM is replaced by the
%   element of CLASSES in the same position as the highest value along the
%   vector. Input B is an n-dimensional dlarray containing values between 0
%   and 1. CLASSES must be a cell array of char vectors, a string or
%   numeric vector, or a 2-dimensional char array that specifies the
%   classes used for decoding. FEATDIM is a positive integer such that
%   size(B,featdim) is equal to length(CLASSES). If a probability vector in
%   B contains NaN values, then that observation is decoded to the class
%   that has the largest probility that is not NaN. If an observation
%   contains only NaN values, then that observation is decoded to the first
%   class label.
% 
%   A = ONEHOTDECODE(...,TYPENAME) specifies the data type of the decoded
%   labels. Valid options are 'categorical', 'string', and numeric types
%   such as 'double'. If you specify a numeric type, CLASSES must be a
%   numeric array, or a string array or a cellstr array containing strings
%   or char vectors of numbers. The default value is 'categorical'.
%
%   Example: One-hot decode the output of a softmax operation
% 
%     classes = ["dog","cat","tiger"];
%     networkOutput = dlarray(randi(10,3,5),'CB');
% 
%     probabilities = softmax(networkOutput);
%     featureDim = finddim(probabilities, 'C');
%     classesPredicted = onehotdecode(probabilities,classes,featureDim);
% 
%   See also ONEHOTDECODE, ONEHOTENCODE, DLARRAY.

%   Copyright 2020 The MathWorks, Inc.

arguments
    B {mustBeDlarray}
    classes
    featdim double
    typename {iAssertValidTypename} = "categorical"
end

A = onehotdecode(extractdata(B),classes,featdim,typename);

end

function iAssertValidTypename(typename)
supportedTypes=["categorical", "string", "double", "single", ...
    "int8" "int16" "int32" "int64" "uint8" "uint16" "uint32" "uint64"];
validatestring(typename, supportedTypes, 'onehotdecode', 'typename');
end

function mustBeDlarray(x)
if ~isdlarray(x)
   error(message('deep:dlarray:FirstInputDlarray'));
end
end