function Xhat = rfimpute(X,varargin)
%RFIMPUTE Impute missing values by random forest.
%   XHAT = RFIMPUTE(X) replaces NaNs in X with values fitted by random
%   forest and returns XHAT of the same size and type as X. Pass X as an
%   N-by-P matrix for N observations and P predictors.
%
%   XHAT = RFIMPUTE(X,Name,Value) specifies optional parameter name/value
%   pairs:
%       'CategoricalPredictors' - Logical vector with P elements. If J-th
%                                 element is true, RFIMPUTE will treat J-th
%                                 predictor as categorical. By default, all
%                                 predictors are assumed to be numeric.
%
%   Copyright 2018 The MathWorks, Inc.

D = size(X,2);

if D<2
    error('Unable to carry out imputation for less than 2 predictors.');
end

% Process input args
args = {'categoricalpredictors'};
defs = {                     []};
catpred = internal.stats.parseArgs(args,defs,varargin{:});

if isempty(catpred)
    catpred = false(1,D);
else
    if ~islogical(catpred) || ~isvector(catpred) || numel(catpred)~=D
        error('''CategoricalPredictors'' value must be a logical vector with %i elements.',D);
    end
end

idxnan = isnan(X);

Xhat = X;

for d=1:D
    tstart = tic;
    
    % Missing observations for this predictor
    missing = idxnan(:,d);
    
    % Indices for all predictors but this one
    allpreds_but_this_one = true(1,D);
    allpreds_but_this_one(d) = false;
    
    % Categorical predictors for this subproblem
    cat_allpreds_but_this_one = catpred;
    cat_allpreds_but_this_one(d) = [];
    
    % If this predictor is categorical, do classification
    if catpred(d)
        fun = @fitcensemble;
    else
        fun = @fitrensemble;
    end
    
    % If no missing values, proceed
    if ~any(missing)
        continue;
    end
    
    % Form data for learning a regression model
    y  = X(~missing,d);
    X1 = X(~missing,allpreds_but_this_one);
    
    % If there are missing values in other predictors, it is best to use
    % surrogate splits. Using surrogate splits can slow down training
    % considerably. If speed is of prime concern, don't use them.
    if any(any(isnan(X1)))
        tmp = templateTree('Surrogate','on');
    else
        tmp = templateTree('Surrogate','off');
    end
    
    % Train random forest model with predictor d excluded
    rf = fun(X1,y,'Method','Bag','NumLearningCycles',200,...
        'Learners',tmp,'CategoricalPredictors',cat_allpreds_but_this_one);
    
    % Fill missing values for predictor d
    Xhat(missing,d) = predict(rf,X(missing,allpreds_but_this_one));
    
    telapsed = toc(tstart);
    
    fprintf('Predictor %6i filled in %7.4g s.\n',d,telapsed);
end

end
