function [X,Y,varargout] = validateInputsForRegressionFitFunction(X,Y,trailingArgs,argsToEnsureOnDeviceIfPossible)
%VALIDATEINPUTSFORREGRESSIONFITFUNCTION Validate inputs for regression fit functions
%   This function validates input arguments to regression fitting functions
%   (for example, fitrtree, fitrensemble).
% 
%   The fit should be performed using the GPU if X is a gpuArray or is a
%   table containing gpuArrays, or if Y is a gpuArray. Otherwise all inputs
%   are gathered.
%
%   If the GPU is to be used for the fit, the inputs X, Y, "Weights",
%   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")));

useGpu = xContainsGpuArray || isgpuarray(Y);
if useGpu
    X = sendToDeviceIfPossible(X);
    Y = sendToDeviceIfPossible(Y);

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