function [X,Y,varargout] = validateInputsForClassificationFitFunction(X,Y,trailingArgs,argsToEnsureOnDeviceIfPossible)
%VALIDATEINPUTSFORCLASSIFICATIONFITFUNCTION Validate inputs for classification fit functions
%   This function validates input arguments to classification fitting
%   functions (for example, fitcknn, fitctree).
%
%   The fit should be performed using the GPU if X is a gpuArray or is a
%   table containing gpuArrays.  Otherwise all inputs are gathered. The
%   type of Y does not control whether the GPU is used or not because
%   Y represents categories, rather than numeric data.
%
%   If the GPU is to be used for the fit, the inputs X, Y, "Weights",
%   "Prior", and those specified in argsToEnsureOnDeviceIfPossible are sent
%   to the device if possible. The remaining inputs are gathered.

%   Copyright 2021-2022 The MathWorks, Inc.

% X should never be a dataset because dataset is not an inferior class to
% gpuArray, so the free function is called instead.
assert(~isa(X,"dataset"), "Input should not be dataset.");

xContainsGpuArray = isgpuarray(X) || (istable(X) && any(varfun(@isgpuarray, X, 'OutputFormat', "uniform")));
if xContainsGpuArray
    X = sendToDeviceIfPossible(X);
    Y = sendToDeviceIfPossible(Y);

    commonClassificationArgsToEnsureOnDevice = ["weights", "prior"];
    if nargin < 4
        argsToEnsureOnDeviceIfPossible = commonClassificationArgsToEnsureOnDevice;
    else
        argsToEnsureOnDeviceIfPossible = union(commonClassificationArgsToEnsureOnDevice, lower(argsToEnsureOnDeviceIfPossible));
    end
    varargout = ensureOnlySpecifiedPVPairsOnDevice(argsToEnsureOnDeviceIfPossible, trailingArgs{:});
else
    varargout = trailingArgs;
    [Y, varargout{:}] = gather(Y, varargout{:});
end
end