function X = fullyconnect(X, W, B, varargin)
%FULLYCONNECT    Sum all weighted input data and apply a bias
%
%   Y = FULLYCONNECT(X,WEIGHTS,BIAS) computes the weighted sum of input
%   data X using the operation Y = WEIGHTS*X + BIAS. Input X must be a
%   formatted dlarray. Inputs WEIGHTS and BIAS can be formatted dlarrays,
%   unformatted dlarrays, or numeric arrays. Output Y is a dlarray with one
%   'C' dimension, and the same number of 'B' and 'T' dimensions as there
%   are in input X. When X has no 'B' or 'T' dimensions, the default format
%   of Y is 'CB'. The size of the 'C' dimension of Y is equal to the number
%   of output features.
%
%   WEIGHTS must have at least two dimensions. If WEIGHTS is an unformatted
%   dlarray or numeric array, the size of the first dimension must match
%   the number of output features. If WEIGHTS is a formatted dlarray, the
%   size of the 'C' dimension must match the number of output features. The
%   number of weights for each output feature must match the number of
%   elements in each observation of X along the 'B' or 'T' dimensions. If
%   necessary, X and WEIGHTS are reshaped to compute the fullyconnect
%   operation.
%
%   BIAS must be a vector containing as many elements as the number of
%   output features. If BIAS is a formatted dlarray, the non-singleton
%   dimension must be labeled with 'C'.
%
%   Y = FULLYCONNECT(X,WEIGHTS,BIAS,'DataFormat',FMT) computes the weighted
%   sum of an unformatted dlarray or numeric array X. FMT specifies the
%   dimension labels for X. FMT must be a char array or string. If the
%   input X is a numeric array, at least one of WEIGHTS or BIAS must be a
%   dlarray. The output Y is an unformatted dlarray where the size of the
%   first dimension equals the number of output features. The size of the
%   other dimensions of Y corresponds to the size of the 'B' and 'T'
%   dimensions of X in the same order as they appear in X, where present.
%
%   Example 1:
%       % Create the input data X as a formatted dlarray, and learnable
%       % parameters weights and bias as unformatted dlarrays
%       X = dlarray(randn(7,7,32),'SSCB');
%       weights = dlarray(ones(10,7,7,32)); bias = dlarray(ones(10,1));
%       % Compute fullyconnect
%       Z = fullyconnect(X,weights,bias);
%       % Output: Z = 10(C) x 1(B) dlarray double
%
%   Example 2:
%       % Create the input data X as an unformatted dlarray, and learnable
%       % parameters weights and bias as numeric arrays
%       X = dlarray(randn(7,7,32));
%       weights = ones(10,1568); bias = ones(10,1);
%       % Compute fullyconnect
%       Z = fullyconnect(X,weights,bias,'DataFormat','SSCB');
%       % Output: Z = 10x1 dlarray double
%
%   Example 3:
%       % Create the input data X as a formatted dlarray, and learnable
%       % parameters weights and bias as numeric arrays
%       X = dlarray(ones(4,2,3),'CBT');
%       weights = ones(10,4); bias = ones(10,1);
%       % Compute fullyconnect
%       Z = fullyconnect(X,weights,bias);
%       % Output: Z = 10(C) x 2(B) x 3(T) dlarray double
%
%   See also DLARRAY, DLCONV, BATCHNORM, RELU

%   Copyright 2019-2020 The MathWorks, Inc.

narginchk(3,inf);

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

% Extract X data in the permuted order
X = permuteExplicitly(X);
xsize = size(X);

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

% Retrieve the batch dimensions of the input
batchDims = indexOfLabels(fd, 'BT');

% Prepare the output label. Remove dimensions from the dlarray (unformatted).
newfmt = ''; 
% If X was formatted, assign the new format on the basis of the label's batch
% dimensions
if isempty(args.DataFormat)
    if isempty(batchDims)
        newfmt = 'CB';
    else
        newfmt = ['C', fd.DimensionMetadata(batchDims)];
    end
end
fd.DimensionMetadata = newfmt;

[xdata, xsize, batchDims] = iPermuteIfRequired(xdata, xsize, batchDims);

% Extract the weights and bias, and verify the sizes match the input
wdata = iValidateWeights(W, xsize, batchDims);
bdata = iValidateBias(B, size(wdata, 1));

% The dlarray methods should not accept logical inputs 
if islogical(wdata) || islogical(bdata)
    error(message('deep:dlarray:LogicalsNotSupported'));
end

% Internal API
xdata = internal_fullyconnect(xdata, wdata, bdata, batchDims); 

% Reshape the fullyconnect output to have both 'B' and 'T' observation
% dimensions, where present
if numel(batchDims) > 1
    % batchDims elements are, in order, the index of 'B' and 'T' dims. For
    % unformatted input X, the observation dimensions must appear in output
    % in the same order as they are in 'DataFormat'
    if ~isempty(args.DataFormat) && perm(batchDims(1)) > perm(batchDims(2))
        % Check that there are exactly two batchDims
        assert(numel(batchDims)==2)
        % perm contains an entry for both observation dimensions. If 'T'
        % appears before 'B' in 'DataFormat', perm(batchDims) is not
        % ordered and batchDims must be reversed
        batchDims = batchDims([2 1]);
    end
    % Reshape X back to the user format
    xdata = iReshapeXBackToUserFormat(xdata, xsize, batchDims);
end

fd = insertData(fd, xdata);
X.FormattedData = fd;
% Remove any permutation from the output dlarray
X.DataToViewMap = [];
X.ViewToDataMap = [];

end

function xdata = iReshapeXBackToUserFormat(xdata, xsize, batchDims)
% Reshape back to original X number of batch dims
xdata = reshape(xdata, [size(xdata,1), xsize(batchDims)]);
end

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

% Default values
defaultFormat = [];

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

% Parse the arguments
parser.parse(varargin{:});
% Ensure the arguments are host arrays
inputArguments = structfun(@gather, parser.Results, 'Uniform', false);
end

function [xdata, sizeOfX, batchDims] = iPermuteIfRequired(xdata, sizeOfX, batchDims)

if ~isempty(batchDims)
    % Find the number of dimensions of the original input X (only singleton
    % 'U' dimensions are discarded).
    maxDim = max([batchDims length(sizeOfX)]);
    
    % Bring dims to connect at front followed by the ones treated as batch.
    targetBatchDims = maxDim-numel(batchDims)+1:maxDim;
    
    if ~all(ismember(targetBatchDims, batchDims))
        % Permute data if batch dims are not already at the end of X.
        connectDims = setdiff(1:maxDim, batchDims);
        xdata = permute(xdata, [connectDims, batchDims]);
        % Update sizeOfX to reflect xdata permutation (N.B. if one of the
        % batch dims is singleton, calling xsize = size(X) here will not
        % return the correct length for xsize)
        sizeOfX = sizeOfX([connectDims, batchDims]);
    end
     
    % batchDims are now at the end
    batchDims = targetBatchDims;
else
    % If the input has no batch dims, it is assumed to be after the last
    % dimension
    batchDims = length(sizeOfX)+1;
    sizeOfX = [sizeOfX 1];
end
end

function weights = iValidateWeights(weights, dataSize, batchDims)
% This function gathers the weight values, verifies if the size matches
% the data size, and return a nbOutputsxnbInputs matrix.

% If weights are unformatted dlarray or numeric array, the size of the first
% dimension is the number of output features.
nbClassesDim = 1;

if isa(weights, 'dlarray')
    
    weights = permuteExplicitly(weights);
    
    % Extract the values
    fdw = weights.FormattedData;
    [fd, weights] = extractData(fdw);
    
    if ~isempty(fd.DimensionMetadata) % If formatted dlarray
        % Retrieve the dimension related to the classes
        nbClassesDim = channelIndex(fdw);
    end
else
   deep.internal.dlarray.validateNonDlarray(weights);
end

% Verify the number of connections for the weights and for the input match
weightSize = size(weights);
nbConnectionsWeight = prod(weightSize((1:numel(weightSize)) ~= nbClassesDim));
nbConnectionsInput = prod(dataSize(1:end-numel(batchDims)));
if nbConnectionsWeight ~= nbConnectionsInput
    error(message('deep:dlarray:FullyconnectInputMismatchWeights', nbConnectionsWeight, nbConnectionsInput))
end

% Reshape the weights as nbOutputs x nbInputs
weights = iReformatWeights(weights, nbClassesDim);

end

function weights = iReformatWeights(weights, nbOutputsDim)
% This function reformats the weights as a nbOutputsxnbInputs matrix

% Get the number of outputs
nbOutputs = size(weights, nbOutputsDim);
% A permutation is required if the output dimensions is not the first one
if nbOutputsDim ~= 1
    pVector = unique([nbOutputsDim 1: ndims(weights)], 'stable');
    weights = permute(weights, pVector);
end
% Reshape as a nbOutputsxnbInputs matrix
weights = reshape(weights, nbOutputs, []);
end

function bias = iValidateBias(bias, nbClasses)
% This function gather the bias values and verify if the size matches the
% expected size.

% Check that bias is a vector
if ~isvector(bias)
    error(message('deep:dlarray:FullyconnectBiasWithManyNonSingletonDims', nbClasses));
end

if isa(bias, 'dlarray')
    
    bias = permuteExplicitly(bias);
        
    % Extract the values
    fdw = bias.FormattedData;
    [fd, bias] = extractData(fdw);
    if ~isempty(fd.DimensionMetadata) % If formatted dlarray
        channelIdx = channelIndex(fd);
        % Bias must be either a scalar labeled with 'C' or a vector in
        % which the only non-singleton dimension must be labeled with 'C'.
        % The 'C' dimension contains as many elements as the number of
        % output features.
        if ~contains(fd.DimensionMetadata,'C') || sum(size(bias, channelIdx)) ~= nbClasses 
            error(message('deep:dlarray:FullyconnectChannelSizeMismatchBiasSize', nbClasses));
        end
    end
else
    deep.internal.dlarray.validateNonDlarray(bias);
end

if numel(bias) ~= nbClasses
    error(message('deep:dlarray:FullyconnectInputMismatchBias', nbClasses));
end


%  Reshape bias as a column vector
bias = reshape(bias,[],1);

end





