function layer = imageInputLayer(varargin)
% imageInputLayer   Image input layer
%
%   layer = imageInputLayer(inputSize) defines an image input layer.
%   inputSize is the size of the input images for the layer. It must be a
%   row vector of two or three numbers.
%
%   layer = imageInputLayer(inputSize, 'PARAM1', VAL1, 'PARAM2', VAL2, ...)
%   specifies optional parameter name/value pairs for creating the layer:
%
%    'DataAugmentation'  The 'DataAugmentation' parameter is not recommended.
%                        Use augmentedImageDatastore instead.
%
%    'Normalization'     Data normalization applied when data is forward
%                        propagated through the input layer. Valid values
%                        are:
%                           'zerocenter'        - zero-center normalization
%                           'zscore'            - z-score normalization
%                           'rescale-symmetric' - rescale to [-1 1]
%                           'rescale-zero-one'  - rescale to [0 1]
%                           'none'              - do not normalize
%                           function handle     - custom normalization
%
%                        Default: 'zerocenter'
%
%    'NormalizationDimension'
%                        Dimension over which the same normalization is
%                        applied, specified as one of the following:
%                          - 'auto'    - If the ResetInputNormalization
%                                        training option is false and you
%                                        specify normalization statistics,
%                                        then normalize over dimensions
%                                        matching the statistics.
%                                        Otherwise, recompute the
%                                        statistics and apply channel-wise
%                                        normalization.
%                          - 'channel' - Channel-wise normalization
%                          - 'element' - Element-wise normalization
%                          - 'all'     - Normalize all values using
%                                        scalar statistics
%
%                        Default: 'auto'
%
%    'Mean'              The mean image used for zero center and z-score
%                        normalization. This can be a H-by-W-by-C array, a
%                        1-by-1-by-C array of means per channel, a scalar
%                        or [], where H is the height, W the width, and C
%                        the number of channels.
%
%                        Default: []
%
%    'StandardDeviation' The standard deviation used for z-score normalization.
%                        This can be a H-by-W-by-C array, a 1-by-1-by-C
%                        array of standard deviations per channel, a
%                        scalar, or [], where H is the height, W the width,
%                        and C the number of channels.
%
%                        Default: []
%
%    'Min'               Minimum value used for rescaling.
%                        This can be a H-by-W-by-C array, a 1-by-1-by-C
%                        array of minima per channel, a scalar, or [],
%                        where H is the height, W the width, and C the
%                        number of channels.
%
%                        Default: []
%
%    'Max'               Maximum value used for rescaling.
%                        This can be a H-by-W-by-C array, a 1-by-1-by-C
%                        array of maxima per channel, a scalar, or [],
%                        where H is the height, W the width, and C the
%                        number of channels.
%
%                        Default: []
%
%    'Name'              A name for the layer.
%
%                        Default: ''
%
%   Example 1:
%       % Create an image input layer for 28-by-28 color images.
%
%       layer = imageInputLayer([28 28 3]);
%
%   Example 2:
%       % Create an image input layer for 28-by-28 grayscale images. Set
%       % the mean image for zero-center normalization to matrix ones.
%
%       layer = imageInputLayer([28 28], 'Mean', ones([28 28]));
%
%       % Reset the mean image to empty
%
%       layer.Mean = [];
%
%   See also nnet.cnn.layer.ImageInputLayer, convolution2dLayer,
%   fullyConnectedLayer, maxPooling2dLayer, augmentedImageDatastore.
%
%   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>

%   Copyright 2015-2019 The MathWorks, Inc.

% Parse the input arguments.
inputArguments = iParseInputArguments(varargin{:});

normalization = iCreateTransform(...
    inputArguments.Normalization, inputArguments.InputSize);

augmentations = iCreateTransforms(...
    inputArguments.DataAugmentation, inputArguments.InputSize);

% Create an internal representation of an image input layer.
internalLayer = nnet.internal.cnn.layer.ImageInput(...
    inputArguments.Name, ...
    inputArguments.InputSize, ...
    normalization, ...
    augmentations);

% Assign Statistics
internalLayer.Mean = inputArguments.Mean;
internalLayer.Std = inputArguments.StandardDeviation;
internalLayer.Min = inputArguments.Min;
internalLayer.Max = inputArguments.Max;
internalLayer.NormalizationDimension = inputArguments.NormalizationDimension;

% Pass the internal layer to a function to construct a user visible image
% input layer.
layer = nnet.cnn.layer.ImageInputLayer(internalLayer);
end

function inputArguments = iParseInputArguments(varargin)
varargin = nnet.internal.cnn.layer.util.gatherParametersToCPU(varargin);
parser = iCreateParser();
parser.parse(varargin{:});
inputArguments = iConvertToCanonicalForm(parser.Results);
end

function p = iCreateParser(varargin)
p = inputParser;

defaultName = '';
defaultTransform = 'zerocenter';
defaultTrainTransform = 'none';
defaultDimension = 'auto';

addRequired(p,  'InputSize', @iAssertValidInputSize);
addParameter(p, 'Normalization', defaultTransform, @(x)~isempty(iCheckAndReturnValidNormalization(x)));
addParameter(p, 'NormalizationDimension', defaultDimension);
addParameter(p, 'DataAugmentation', defaultTrainTransform, @(x)~isempty(iCheckAndReturnValidAugmentation(x)));
addParameter(p, 'Name', defaultName, @iAssertValidLayerName);
addParameter(p, 'Mean', []);
addParameter(p, 'StandardDeviation', []);
addParameter(p, 'Min', []);
addParameter(p, 'Max', []);
addParameter(p, 'AverageImage', []);
end

function iAssertValidLayerName(name)
nnet.internal.cnn.layer.paramvalidation.validateLayerName(name)
end

function iAssertValidInputSize(sz)
validateattributes(sz,{'numeric'},{'nonempty','row','real','finite','integer','positive'})
isValidSize = iIsRowVectorOfTwoOrThree(sz) && ...
    (iIsValidGrayscaleImageSize(sz) || iIsValidRGBImageSize(sz) || iIsValidMultiChannelImageSize(sz));
if ~isValidSize
    error(message('nnet_cnn:layer:ImageInputLayer:InvalidImageSize'));
end
end

function args = iConvertToCanonicalForm(params)
try    
    args = struct;
    % Make sure integral values are converted to double
    args.InputSize = double(iMakeIntoRowVectorOfThree(params.InputSize));
    args.Normalization = iCheckAndReturnValidNormalization(params.Normalization);
    args.NormalizationDimension = iCheckAndReturnValidDimension( ...
        params.NormalizationDimension, args.Normalization );
    args.Mean = iCheckAndReturnSingleMean( params, ...
        args.Normalization, args.NormalizationDimension, args.InputSize);
    args.StandardDeviation = iCheckAndReturnSingleStatistics( params.StandardDeviation, 'StandardDeviation', ...
        args.Normalization, args.NormalizationDimension, args.InputSize);
    args.Min = iCheckAndReturnSingleStatistics( params.Min, 'Min', ...
        args.Normalization, args.NormalizationDimension, args.InputSize);
    args.Max = iCheckAndReturnSingleStatistics( params.Max, 'Max', ...
        args.Normalization, args.NormalizationDimension, args.InputSize);
    args.DataAugmentation = iCheckAndReturnValidAugmentation(params.DataAugmentation);
    args.Name = char(params.Name); % make sure strings get converted to char vectors
catch e
    % Reduce the stack trace of the error message by throwing as caller
    throwAsCaller(e)
end
end

function meanImage = iCheckAndReturnSingleMean( params, ...
    normalization, normalizationDimension, inputSize)
if isempty(params.Mean)
    % Use value of AverageImage NVP as the Mean if Mean not specified
    meanImage = params.AverageImage;
else
    % Ignore AverageImage NVP if Mean was specified
    meanImage = params.Mean;
end
meanImage = iCheckAndReturnSingleStatistics( meanImage, 'Mean', ...
    normalization, normalizationDimension, inputSize);
end

function statsValue = iCheckAndReturnSingleStatistics( userValue, name, ...
    normalization, normalizationDimension,inputSize)
% Validate user-provided statistics with respect to normalization method 
% and input size. Return the statistic in single precision.
if ~isempty(userValue)
    nnet.internal.cnn.layer.paramvalidation.validateNormalizationStatistics(...
        userValue, name, normalization, normalizationDimension, inputSize);
end
statsValue = single(userValue);
end

function validSize = iMakeIntoRowVectorOfThree(inputSize)
if(iIsRowVectorOfTwo(inputSize))
    validSize = [inputSize 1];
else
    validSize = inputSize;
end
end

function x = iCheckAndReturnValidNormalization(x)
validateattributes(x,{'string','char','function_handle'},{})
if isa(x,'function_handle')
    % Checks are performed at train/inference time when normalization is
    % applied
else
    validTransforms = {'zerocenter', 'zscore', 'rescale-symmetric', ...
        'rescale-zero-one', 'none'};
    x = validatestring(x, validTransforms);
end
end

function x = iCheckAndReturnValidAugmentation(x)
validTransforms = {'randcrop', 'randfliplr', 'none'};
x = iIsValidCellStringOrStringArg(x, validTransforms);

if iscellstr(x)
    if numel(x) > 1 && iAreNotUniqueStrings(x)
        error(message('nnet_cnn:layer:ImageInputLayer:AugmentationsMustBeUnique'));
    end
        
    if numel(x) > 1 && ismember('none', x)
        error(message('nnet_cnn:layer:ImageInputLayer:NoneNotAllowedWithOthers'));
    end

    if numel(x) == 1
        % return string if only single cell element.
        x = x{1};
    end
end
end

function dimValue = iCheckAndReturnValidDimension(dimValue,normalization)
dimValue = nnet.internal.cnn.layer.paramvalidation.validateNormalizationDimension(dimValue,normalization);
end

function x = iIsValidCellStringOrStringArg(x, validstr)
if iscellstr(x) && numel(x) > 0
    x = cellfun(@(str)validatestring(str, validstr), ...
        x, 'UniformOutput', false);
else
    x = validatestring(x, validstr);
end
end

function tf = iIsValidGrayscaleImageSize(sz)
% Return true if size is [M N] or [M N 1]. Assumes input sz is already
% validated as a 2 or 3 element vector.
if numel(sz) == 2
    tf = true;
else
    tf = sz(end) == 1;
end
end

function tf = iIsValidRGBImageSize(sz)
tf = numel(sz) == 3 && sz(end) == 3;
end

function tf = iIsValidMultiChannelImageSize(sz)
% Placeholder, if more constraints need to be added
tf = numel(sz) == 3;
end

function tf = iIsRowVectorOfTwoOrThree(x)
tf = (iIsRowVectorOfTwo(x) || iIsRowVectorOfThree(x));
end

function tf = iAreNotUniqueStrings(x)
tf = numel(unique(x)) ~= numel(x);
end

function tf = iIsRowVectorOfTwo(x)
tf = isrow(x) && numel(x)==2;
end

function tf = iIsRowVectorOfThree(x)
tf = isvector(x) && all(size(x) == [1 3]);
end

function tformarray = iCreateTransforms(type, imageSize)
type = cellstr(type);
tformarray = nnet.internal.cnn.layer.InputTransform.empty();
for i = 1:numel(type)
    tnew = iCreateTransform(type{i}, imageSize);
    tformarray = [tformarray tnew]; %#ok<AGROW>
end
end

function transform = iCreateTransform(type, dataSize)
transform = nnet.internal.cnn.layer.InputTransformFactory.create(type, dataSize);
end
