function varargout = checkLayer(layer, validInputSize, opts)
% checkLayer   Check layer validity
%
%   checkLayer(layer, layout1, ..., layoutN) checks the validity of a layer
%   using the specified networkDataLayout objects, where N is the number of
%   layer inputs and layoutK corresponds to the input layer.InputNames(K).
%   This syntax is required for layers that inherit nnet.layer.Formattable.
%   
%   If the layer has learnable or state parameters, then specify 
%   a networkDataLayout object or an initialized layer.
%
%   checkLayer(layer, validInputSize) checks the validity of a layer using
%   generated data of size validInputSize. If the layer has multiple inputs
%   then validInputSize must be a cell array containing a size vector for
%   each input to the layer. 
%
%   The layer must be of type nnet.layer.Layer,
%   nnet.layer.RegressionLayer, or nnet.layer.ClassificationLayer. 
%
%   If the layer has learnable or state parameters, then first initialize 
%   the layer or specify networkDataLayout objects as input.
%
%   checkLayer(__,Name=Value) specifies additional options using one or
%   more name-value arguments.
%
%       'ObservationDimension'      - Dimension of data that indexes 
%                                     observations in layer input to check 
%                                     support for multiple observations, 
%                                     specified as a vector of positive 
%                                     integers. The default is the position 
%                                     of the "B" dimension in the layout 
%                                     objects. If the layouts do not
%                                     contain "B" or you specify valid 
%                                     input sizes as vectors, then by 
%                                     default, the function checks that the
%                                     layer supports a single observation.
%
%       'CheckCodegenCompatibility' - Flag to enable code generation tests,
%                                     specified as true or false. This 
%                                     option does not support layers that 
%                                     inherit from nnet.layer.Formattable.
%                                     Use analyzeNetworkForCodegen instead.
%                                     Default is false. 
%
%   See also: networkDataLayout

%   Copyright 2017-2023 The MathWorks, Inc.

arguments
    layer (:,1) {iAssertValidLayer}
end

arguments (Repeating)
    validInputSize {iAssertValidInputSize(layer,validInputSize)}
end

arguments
    opts.ObservationDimension {iAssertValidObsDimension} = []
    opts.CheckCodegenCompatibility (:,1) {iAssertBinary} = false
end

validInputSize = iCanonicalizeCellArrayInput(validInputSize);
format = iFormat(validInputSize);
[inputSize, isLayoutInput] = iSize(validInputSize, opts.ObservationDimension);

try
    opts.ObservationDimension = iDetermineObservationDimension(format,opts.ObservationDimension);
    opts.SequenceDimension = iDetermineSequenceDimension(opts.ObservationDimension, inputSize);
    iValidateInputOptions(layer,inputSize,opts);
catch MEParseArgumentsException
    throwAsCaller(MEParseArgumentsException);
end

% Store data needed for TestCase
config = nnet.checklayer.util.UseConfiguration;
config.Data.Layer = layer;
config.Data.InputSize = inputSize;
config.Data.Format = format;
config.Data.ObservationDimension = opts.ObservationDimension;
config.Data.HasGPU = canUseGPU;
config.Data.CheckCodegenCompatibility = opts.CheckCodegenCompatibility;
config.Data.SequenceDimension = opts.SequenceDimension;

% If the validInputSize is specified as a networkDataLayout and the layer
% has initialize defined, call initialize.
if isLayoutInput && iIsMethodDefined(layer,'initialize')
    config.Data.Layer = iInitialize(config.Data.Layer,config.Data.InputSize,config.Data.Format);
end

% Create TestSuite object
if isa(config.Data.Layer,'nnet.layer.Layer')
    if iIsMethodDefined(layer,'backward')
        suite = matlab.unittest.TestSuite.fromClass(?nnet.checklayer.TestLayerWithBackward);
    else
        suite = matlab.unittest.TestSuite.fromClass(?nnet.checklayer.TestLayerWithoutBackward);
    end
else
    if iIsMethodDefined(layer,'backwardLoss')
        suite = matlab.unittest.TestSuite.fromClass(?nnet.checklayer.TestOutputLayerWithBackward);
    else
        suite = matlab.unittest.TestSuite.fromClass(?nnet.checklayer.TestOutputLayerWithoutBackward);
    end
end
numAvailableTests = numel(suite);

% Filter initialization tests.
if isa(config.Data.Layer,'nnet.layer.Layer')
    if ~iIsMethodDefined(layer,'initialize')
        disp( getString(message('nnet_cnn:checkLayer:InitializeNotDefined')) );
        suite = suite.selectIf(~iHasTag('CheckInitialization'));
        disp(" ")
    elseif ~isLayoutInput
        disp( getString(message('nnet_cnn:checkLayer:LayoutNotSpecified')) );
        suite = suite.selectIf(~iHasTag('CheckInitialization'));
        disp(" ")
    end
end
    
% The default case for ObservationDimension for a layer with N inputs is
% {NaN, NaN, ..., NaN} (of length N). If any are NaN, then we must filter
% multi-observation tests.
if iCellHasAnyNaN(config.Data.ObservationDimension)
    recommendationStr = getString(message('nnet_cnn:checkLayer:ObsDimRecommendation'));
    disp( getString(message('nnet_cnn:checkLayer:ObsDimNotSpecified',recommendationStr)) )
    suite = suite.selectIf(~iHasParameter('Property','Observations','Name','multiple'));
    disp(" ")
end

% Filter out GPU tests if GPU is not available
if ~config.Data.HasGPU
    disp( getString(message('nnet_cnn:checkLayer:NoGPUAvailable')) );
    suite = suite.selectIf(~iHasParameter('Property','Device','Name','gpu'));
    disp(" ")
end

% Filter out code-generation tests if CheckCodegenCompatibility not set
if ~config.Data.CheckCodegenCompatibility
    disp( getString(message('nnet_cnn:checkLayer:CheckCodegenCompatibilityNotSpecified')) );
    suite = suite.selectIf(~iHasTag('CheckCodegenCompatibility'));
    disp(" ")
end

%  Adjust codegen tests for FunctionLayer, which is not codegen compatible.
if config.Data.CheckCodegenCompatibility && isa(config.Data.Layer,'nnet.cnn.layer.FunctionLayer')
    suite = suite.selectIf(~iHasTag('CheckCodegenCompatibility'));
else
    suite = suite.selectIf(~iHasTag('CheckFunctionLayerCodegenCompatibility'));
end

%  Adjust tests for formattable layers, which require additional tests.
if ~isa(config.Data.Layer,'nnet.layer.Formattable')
    suite = suite.selectIf(~iHasTag('CheckFormattableLayer'));
end

numSkippedTests = numAvailableTests - numel(suite);

% Create TestRunner object and run test suite
runner = matlab.unittest.TestRunner.withTextOutput('Verbosity',2);
runner.addPlugin(matlab.unittest.plugins.DiagnosticsRecordingPlugin)
% Filter error during auto-initialization of dlnetwork Learnable properties
% using try...catch
try
    results = runner.run(suite);
catch originalError
    if originalError.identifier == "nnet_cnn:nnet:checklayer:TestCase:AutoInitializeFailure"
        err = MException(message("nnet_cnn:checkLayer:AutoInitializeFailure"));
        err = iAddCauseWithStackMessage(err, originalError.cause{1});
        throw(err);
    else
        rethrow(originalError);
    end
end

% Print results summary
iPrintTestSummary(results,numSkippedTests);

% Return TestResults array if requested
if nargout>=1
    varargout{1} = results;
end
end

function output = iCanonicalizeCellArrayInput(input)
% If the input is a cell array of cell arrays, return the first cell. This
% ensures that both possible syntaxes for calling this function, optional 
% positional arguments of validInputSize and a cell array of validInputSize
% is canonicalized like the optional positional arguments.
if iscell(input) && ~isempty(input) && iscell(input{1})
    output = input{1};
else
    output = input;
end
end

function [inputSize, isLayoutInput] = iSize(validInputSize, observationDimension)
inputSize = [];
isLayoutInput = false;
if isnumeric(validInputSize)
    % The batch dimension is implicit when validInputSize has one fewer
    % dimension than the value of observationDimension
    if length(validInputSize) == observationDimension - 1
        inputSize = [validInputSize NaN];
    else
        inputSize = validInputSize;
    end
elseif iIsNetworkDataLayout(validInputSize)
    inputSize = validInputSize.Size;
    isLayoutInput = true;
elseif iscell(validInputSize)
    if isempty(observationDimension)
        observationDimension = NaN;
    end
    if ~iscell(observationDimension)
        observationDimension = num2cell(observationDimension);
    end
    if length(observationDimension) == 1
        observationDimension = repmat(observationDimension, 1, length(validInputSize));
    end
    [inputSize, isLayoutInput] = cellfun(@iSize, validInputSize, observationDimension, UniformOutput=false);
    isLayoutInput = all(cell2mat(isLayoutInput));
end
end

function format = iFormat(validInputSize)
format = '';

if iIsNetworkDataLayout(validInputSize)
    format = validInputSize.Format;
elseif isempty(validInputSize)
    format = {};
elseif iscell(validInputSize)
    format = cellfun(@iFormat, validInputSize, UniformOutput=false);
end
end

function observationDimension = iParseBatchDimensionAndObservationDimension(formatObservationDimension, inputObservationDimension)
% If neither format with a batch dimension nor ObservationDimension is
% given, then this defaults to a cell array of NaN. If only one is given
% this takes the value of that one. If both are specified this verifies the
% two dimensions match.
if isnan(formatObservationDimension)
    observationDimension = inputObservationDimension;
elseif isnan(inputObservationDimension)
    observationDimension = formatObservationDimension;
elseif formatObservationDimension ~= inputObservationDimension
    error(message('nnet_cnn:checkLayer:BatchAndObservationDimensionMismatch'));
else
    observationDimension = inputObservationDimension;
end
end

function iValidateInputOptions(layer, size, opts)

if iIsFormattableLayer(layer)
    iAssertNoCheckCodegenCompatibilityGivenFormattableLayer(opts.CheckCodegenCompatibility)
end

if iIsOutputLayer(layer)
    iAssertObsDimIsTwoFourOrFive(opts.ObservationDimension);
else
    iAssertValidSizesMatchLayerInputs(layer, size);
end

if opts.CheckCodegenCompatibility
    iAssertObsDimIsTwoOrFourForCodegen(opts.ObservationDimension);
    iAssertInputSizeMatchesWithObsDimForCodegen(size, opts.ObservationDimension);
end

end

function batchDimension = iGetBatchDimension(format)
if iscell(format)
    batchDimension = cellfun(@iGetBatchDimension, format, UniformOutput=false);
else
    batchDimension = strfind(format, 'B');
    if isempty(batchDimension)
        batchDimension = NaN;
    end
end
end

function tf = iCellHasAnyNaN(observationDimension)
tf = any(isnan(cell2mat(observationDimension)));
end

function tf = iIsOutputLayer(layer)
tf = isa(layer,'nnet.layer.RegressionLayer') || isa(layer,'nnet.layer.ClassificationLayer');
end

function tf = iIsFormattableLayer(layer)
tf = (isa(layer, 'nnet.layer.Layer') || isa(layer, 'nnet.layer.RegressionLayer') ...
    || isa(layer, 'nnet.layer.ClassificationLayer')) && isa(layer, 'nnet.layer.Formattable');
end

function tf = iIsNetworkDataLayout(inputSize)
tf = isa(inputSize, 'networkDataLayout');
end

function iAssertBinary(value)
nnet.internal.cnn.options.OptionsValidator.assertBinary(value);
end

function iAssertValidLayer(layer)
validLayers = {'nnet.layer.Layer','nnet.layer.RegressionLayer','nnet.layer.ClassificationLayer'};
validateattributes(layer, validLayers, {'scalar'});
end

function iAssertValidInputSize(layer,validInputSize)
validateattributes(validInputSize, {'numeric', 'cell', 'networkDataLayout'},{})
if iIsFormattableLayer(layer)
    if iscell(validInputSize) && ~isempty(validInputSize)
        cellfun(@(inputSize) iAssertValidInputSize(layer, inputSize), validInputSize)
    elseif isa(validInputSize, 'networkDataLayout')
        % No further checks required
    else
        error(message('nnet_cnn:checkLayer:FormattableRequiresNetworkDataLayout'))
    end
elseif iscell(validInputSize)
    if isempty(validInputSize)
        validateattributes(validInputSize,{'cell'}, {'nonempty'})
    else
        cellfun(@(inputSize) iAssertValidInputSize(layer, inputSize), validInputSize)
    end
elseif isa(validInputSize, 'networkDataLayout')
    % No further checks required
else
    validateattributes(validInputSize,{'numeric'}, {'nonempty','vector'})
    iAssertValidSizeVector(validInputSize)
end
end

function iAssertNoCheckCodegenCompatibilityGivenFormattableLayer(checkCodegenCompatibility)
if checkCodegenCompatibility
    error(message('nnet_cnn:checkLayer:NoCheckCodgenCompatibilityGivenFormattableLayer'))
end
end

function iAssertValidSizeVector(sz,varargin)
validateattributes(sz, {'numeric'}, {'vector', 'nonempty', 'integer', 'positive', 'row'},varargin{:});
end

function iAssertValidObsDimension(obsDimension)
validateattributes(obsDimension, {'numeric'}, {'integer', 'positive'});
end

function iAssertObsDimIsTwoFourOrFive(obsDimension)
if iscell(obsDimension)
    cellfun(@iAssertObsDimIsTwoFourOrFive, obsDimension);
elseif isnan(obsDimension) || obsDimension==2 || obsDimension==4 || obsDimension==5
    % ObservationDimension is valid
else
    recommendationStr = getString(message('nnet_cnn:checkLayer:ObsDimRecommendation'));
    error(message('nnet_cnn:checkLayer:WrongObsDimForOutputLayer',recommendationStr))
end
end

function iAssertObsDimIsTwoOrFourForCodegen(obsDimension)
if iscell(obsDimension)
    cellfun(@iAssertObsDimIsTwoOrFourForCodegen, obsDimension);
else
    recommendationStr = getString(message('nnet_cnn:checkLayer:ObsDimRecommendationForCodegen'));
    if isnan(obsDimension)
        error(message('nnet_cnn:checkLayer:ObsDimNotSpecifiedForCodegen',recommendationStr))
    elseif ~(obsDimension==2 || obsDimension==4)
        error(message('nnet_cnn:checkLayer:WrongObsDimForCodegen',recommendationStr))
    end
end
end

function iAssertInputSizeMatchesWithObsDimForCodegen(inputSize, obsDimension)
if iscell(inputSize)
    % Allow for multiple inputs
    cellfun(@iAssertInputSizeMatchesWithObsDimForCodegen, inputSize, obsDimension);
else
    sequenceDimension = iDetermineSequenceDimension(obsDimension, inputSize);
    if isnumeric(inputSize)
        % Verify the sequence dimension is supported for Codegen
        iAssertInputSizeVectorMatchesWithObsDimForCodegen(inputSize, obsDimension, sequenceDimension)
    else
        for i = 1:numel(inputSize)
            iAssertInputSizeVectorMatchesWithObsDimForCodegen(inputSize{i}, obsDimension, sequenceDimension)
        end
    end
end
end

function iAssertInputSizeVectorMatchesWithObsDimForCodegen( inputSize, obsDimension, sequenceDimension )
numInputDims = numel(inputSize);
if isempty(sequenceDimension)
    if ~any(numInputDims == [(obsDimension-1) obsDimension])
        % if obsDimension is 4 and no sequenceDimension, then expected input is in h-by-w-by-c-(by-N) format
        % if obsDimension is 2 and no sequenceDimension, then expected input is in c-(by-N) format
        error(message('nnet_cnn:checkLayer:WrongInputSizeForCodegen'))
    end
else
    % if sequenceDimension exists, then expected input for obsDimension 2 is in c-by-N-by-S format 
    % input in h-by-w-by-c-by-N-by-S format (obsDim = 4) is not supported for codegen
    if obsDimension ~= 2 || numInputDims ~= sequenceDimension
        error(message('nnet_cnn:checkLayer:WrongInputSizeForCodegen'))
    end
end
end

function iAssertValidSizesMatchLayerInputs( layer, validInputSize )
if iscell(validInputSize)
    numSizes = numel(validInputSize);
else
    numSizes = 1;
end
if numSizes ~= layer.NumInputs
    error(message('nnet_cnn:checkLayer:NumSizesNumInputsMismatch',numSizes,layer.NumInputs));
end
end

function iPrintTestSummary(results,numSkippedTests)
passed = sum([results.Passed]);
failed = sum([results.Failed]);
incomplete = sum([results.Incomplete]);
testTime = sum([results.Duration]);

disp( getString(message('nnet_cnn:checkLayer:TestSummaryHeading')) )
disp( getString(message('nnet_cnn:checkLayer:TestSummaryFirstLine',passed,failed,incomplete,numSkippedTests)) )
disp( getString(message('nnet_cnn:checkLayer:TestSummarySecondLine',num2str(testTime))) )
end

function layer = iInitialize(layer, inputSize, format)
layouts = cellfun(@networkDataLayout, inputSize, format, UniformOutput=false);
msg = message("nnet_cnn:internal:cnn:layer:CustomLayer:InitializeErrored", layer.Name);
try 
    layer = nnet.internal.cnn.util.UserCodeException.fevalUserCode(msg,@()initialize(layer, layouts{:}));
catch originalError
    throwAsCaller(originalError.buildMExceptionForDisplay())
end
end

function selector = iHasParameter(varargin)
selector = matlab.unittest.selectors.HasParameter(varargin{:});
end

function tf = iIsMethodDefined(layer,methodName)
tf = nnet.internal.cnn.layer.util.isMethodDefined(layer,methodName);
end

function selector = iHasTag(varargin)
selector = matlab.unittest.selectors.HasTag(varargin{:});
end

function observationDimension = iDetermineObservationDimension(format, inputObservationDimension)
formatObservationDimension = iGetBatchDimension(format);
if isscalar(inputObservationDimension) && ~isscalar(formatObservationDimension)
    inputObservationDimension = repmat(inputObservationDimension, 1, numel(formatObservationDimension));
end
if isempty(inputObservationDimension) && ~isempty(formatObservationDimension)
    inputObservationDimension = NaN(1, numel(formatObservationDimension));
end
inputObservationDimension = num2cell(inputObservationDimension);
observationDimension = cellfun(@iParseBatchDimensionAndObservationDimension, ...
    formatObservationDimension, inputObservationDimension, UniformOutput=false);
end

function sequenceDim = iDetermineSequenceDimension(obsDim, validInputSize)

if isempty(validInputSize)
    % This is only possible if the no validInputSize is specified
    sequenceDim = [];
elseif iscell(obsDim)
    % For multiple inputs, determine the sequence dimension of each
    sequenceDim = cellfun(@iDetermineSequenceDimension, obsDim, validInputSize, UniformOutput=false);
else
    numInputDims = numel(validInputSize);

    if numInputDims == (obsDim + 1)
        % sequenceDim is the dimension after obsDim if data is sequence data
        sequenceDim = numInputDims;
    else
        % if no sequence input, then sequenceDim is empty
        sequenceDim = []; 
    end
end
end

function err = iAddCauseWithStackMessage(err, cause)
% Trim the stack from the autoInitialize call. Avoid ambiguity by using the
% full file path
pattern = fullfile("+nnet","+checklayer","+util","autoInitializeLayer.m");
stackFiles = {cause.stack.file};
pruneFrom = find(cellfun(@(s)endsWith(s,pattern),stackFiles),1,'first');
if ~isempty(pruneFrom) && numel(stackFiles)>0
    report = iCropErrorReportBelowFilePreserveCause(cause, stackFiles{pruneFrom});
    cause = MException(cause.identifier, "%s", report);
end
err = addCause(err, cause);
end

function report = iCropErrorReportBelowFilePreserveCause(exception, file)
% Find the start of the stack for the file, using the full report with
% hyperlinks to ensure full file paths are present
fullReport = getReport(exception, 'extended', 'hyperlinks', 'on');
idx = strfind(fullReport, file);
% Count lines in this report up to the newline before this file
cropLine = numel(regexp(fullReport(1:idx(1)), '\n'));
% Find newlines in the default report
defaultReport = getReport(exception);
newLines = regexp(defaultReport, '\n');
% Remove all the unwanted stack
report = strtrim(defaultReport(1:newLines(cropLine)));
% Put back cause's cause if it existed.
idx = strfind(defaultReport, 'Caused by:');
if ~isempty(idx)
    cropLine = numel(regexp(defaultReport(1:idx(1)),'\n'));
    report = [report, defaultReport(newLines(max(1,cropLine-1)):end)];
end
end