function X = softmax(X, varargin)
%SOFTMAX   Apply softmax activation
%
%   Y = SOFTMAX(X) computes the softmax activation over the channel 'C'
%   dimension of input data X. X must be a formatted dlarray. Output Y is a
%   formatted dlarray with the same dimension labels as X.
%
%   Y = SOFTMAX(X,'DataFormat',FMT) computes the softmax activation over
%   the channel 'C' dimension of the input X, where X is an unformatted
%   dlarray. FMT specifies the dimension labels of the input data X. FMT
%   must be a char array or string. Output Y is an unformatted dlarray with
%   the same dimension order as X.
%
%   Example:
%       % Create input data as a formatted dlarray and compute softmax
%       % output
%       x = dlarray(randn(128,10),'BC');
%       y = softmax(x);
%
%   See also FULLYCONNECT, DLARRAY, BATCHNORM, RELU

%   Copyright 2019-2020 The MathWorks, Inc.

narginchk(1,inf);

% Parse Name-Value pairs
args = iParseInputArguments(varargin{:});
% Handle input format
[X,perm] = deep.internal.dlarray.validateDataFormatArg(X, args.DataFormat);

% Extract the data
fd = X.FormattedData;
X.FormattedData = [];
[fd,xdata] = extractData(fd);

% Retrieve the channel dimension and error if not a unique C dimension
channelDim = channelIndex(fd);
if numel(channelDim) > 1 || channelDim > length(fd.DimensionMetadata)
    error(message('deep:dlarray:MustHaveOneChannelDimension'))
end

% Call internal API
xdata = internal_softmax(xdata, channelDim);

% Manage format if input data was numeric or unformatted
if ~isempty(args.DataFormat) 
    % Replace new empty FormattedData object
    fd = deep.internal.format.FormattedData([]);
end

% Insert the FormattedData back in the dlarray
fd = insertData(fd, xdata);
X.FormattedData = fd;

% For unformatted, revert the internal permutation vector to its original
% state
if ~isempty(args.DataFormat) 
    X = ipermute(X,perm);
end
end


function inputArguments = iParseInputArguments(varargin)
% Parse the optional input arguments

% Default values
defaultFormat = [];

% Create parser
parser = inputParser;
parser.FunctionName = 'softmax';
parser.addParameter('DataFormat', defaultFormat, @(x)validateattributes(x, {'char', 'string'}, {'nonempty'}));

% Parse the arguments
parser.parse(varargin{:});
inputArguments = parser.Results;
end
