function [varargout] = fitrkernel(X,Y,varargin)
%FITRKERNEL Fit a kernel regression model by explicit feature expansion.
%
%   Supported syntaxes for tall arrays:
%
%   OBJ = FITRKERNEL(X,Y)
%
%   [OBJ,FITINFO] = FITRKERNEL(X,Y) also returns a struct FITINFO
%   containing fit information.
%
%   [OBJ,FITINFO,HYPERPARAMETEROPTIMIZATIONRESULTS] = FITRKERNEL(X,Y) also
%   returns an object describing the results of hyperparameter optimization
%   when you specify the 'OptimizeHyperparameters' name-value pair
%   argument.
%
%   Supported parameter name/value pairs:
%
%       'BetaTolerance'      - default is relaxed to 1e-3
%       'BlockSize'
%       'BoxConstraint'
%       'NumExpansionDimensions'
%       'Epsilon'
%       'GradientTolerance'  - default is relaxed to 1e-5
%       'HessianHistorySize'
%       'HyperparameterOptimizationOptions'
%       'IterationLimit'     - default is relaxed to 20
%       'KernelScale'
%       'Lambda'
%       'Learner'
%       'OptimizeHyperparameters'
%       'RandomStream'
%       'Verbose'            - default is 1
%       'Weights'
%
%   X and Y must both be tall arrays. Weights when supplied must also be
%   a tall array.
%
%   For tall arrays FITRKERNEL implements LBFGS by distributing the
%   calculation of the loss and the gradient among different parts of the
%   tall array at each iteration. Before LBFGS, FITRKERNEL first
%   refines the initial estimates of the parameters by fitting the model
%   locally to parts of the data and combining the coefficients by
%   averaging.
%
%   See also fitrkernel, fitrlinear.

%   Copyright 2017-2023 The MathWorks, Inc.

%   Supported undocumented parameter name/value pairs:
%
%       'ADMMIterationLimit' 
%       'ADMMUpdateIterationLimit' 
%       'Consensus'      
%       'InitialStepSize'
%       'TallPassLimit' 
%       'Transformation' 
%       'WarmStartIterationLimit'
%
%   TIP:
%    - Predictor matrices that are larger than the cluster distributed memory
%      available for caching tall arrays may benefit from having more ADMM
%      iterations (ADMMIterationLimit).

if nargin > 2
    [varargin{:}] = convertStringsToChars(varargin{:});
end

narginchk(2,inf);
assert(istall(X),message('stats:tall:fitrkernel:ArgIn1NotTall')) 
assert(istall(Y),message('stats:tall:fitrkernel:ArgIn2NotTall'))

% Check if tall.fitclinear was called with an option to optimize
% hyperparameters, if so branch code here.
[IsOptimizing, argsWithoutOptim] = classreg.learning.paramoptim.parseOptimizationArgs(varargin);
if IsOptimizing
    nargoutchk(0,3) % error out before lengthy computations
    [X, Y, argsCached] = cacheXYWArgs(X, Y, varargin);
    [varargout{1:nargout}] = classreg.learning.paramoptim.fitoptimizing('fitrkernel',X,Y,argsCached{:});
    return
end

nargoutchk(0,2) % error out before lengthy computations

% Check if tall.fitrkernel was called either from RegressionKernel.resume. 
% If that is the case, a NVP with name 'Template' must have been passed, no
% partial matching is allowed (as it is an undocumented NVP), and the Value
% must be a FitTemplate object with 'Method' set to 'Kernel'. The NVP is
% removed from varargin. If conditions are not met, the NVP is not removed
% and normal parsing will error later with an unrecognized named parameter.
templateLoc = find(strcmpi('template',argsWithoutOptim));
if  numel(templateLoc)==1 && templateLoc<numel(argsWithoutOptim) && ...
        isa(argsWithoutOptim{templateLoc+1},'classreg.learning.FitTemplate') && ...
        strcmp(argsWithoutOptim{templateLoc+1}.Method,'Kernel')
    % If a template was passed then pull model parameters from template:
    temp = argsWithoutOptim{templateLoc+1};
    modelParams = temp.ModelParams;
    extraArgs = [temp.BaseFitObjectArgs argsWithoutOptim([1:templateLoc-1 templateLoc+2:end])];
else 
    % Use KernelParams.make to grab all input parameters that fitrkernel
    % would take from varargin:
    [modelParams,extraArgs] = classreg.learning.modelparams.KernelParams.make('regression',argsWithoutOptim{:});
end

modelParams = fillDefaultParamsForTall(modelParams);

[extraParams,extraArgs] = makeExtraForTall(modelParams,extraArgs{:});
tallPassLimit = extraParams.TallPassLimit;

% Parameters that can be extracted from KernelParams
epsilon = modelParams.Epsilon;
lambda = modelParams.Lambda;
boxConstraint = modelParams.BoxConstraint;
betaTol = modelParams.BetaTolerance;
doBias = modelParams.FitBias;
gradTol = modelParams.GradientTolerance;
lossfun = modelParams.LossFunction;
regularizer = modelParams.Regularization;
doridge = strcmp(regularizer,'ridge');
verbose = modelParams.VerbosityLevel;
rho = modelParams.Consensus;
admmIterationLimit = modelParams.ADMMIterationLimit;
initialBeta = modelParams.InitialBeta;
initialBias = modelParams.InitialBias;
featureMapper = modelParams.FeatureMapper;
kernelScale = modelParams.KernelScale;
transformation = modelParams.Transformation;
numExpansionDimensions = modelParams.NumExpansionDimensions;
blockSize = modelParams.BlockSize;
stream = modelParams.Stream;

% Parameters used only in the centralized LBFGS
iterationlimit = modelParams.IterationLimit;
hessianHistorySize = modelParams.HessianHistorySize;
lineSearch = modelParams.LineSearch;
initialStepSize = modelParams.InitialStepSize;

% Parameters for the LBFGS inside ADMM iterations:
iterationlimit_ADMMLBFGS = [modelParams.WarmStartIterationLimit modelParams.ADMMUpdateIterationLimit];
hessianHistorySize_ADMMLBFGS = 15;  % fixed, cannot be controlled by user
dowolfe_ADMMLBFGS = true;           % fixed, cannot be controlled by user
doBias_ADMMLBFGS  = doBias;         % already restricted to always true above
gradTol_ADMMLBFGS = gradTol;        % same as outer loop
betaTol_ADMMLBFGS = betaTol;        % same as outer loop
doridge_ADMMLBFGS = doridge;        % already restricted to always true above
lossfun_ADMMLBFGS = lossfun;        % already set to 'mse' or 'hinge' above

%-% Resolve numExpansionDimensions if needed now (if is set to 'auto'), 
% numExpansionDimensions is needed for maxChunkSize and therefore it cannot be calculated 
% after the first pass over the data. Other parameters that may still be
% set to 'auto' or empty will be processed  later once we have more info
% about the input data.
numExpansionDimensions =  resolveAutoNumExpansionDimensionsTall(numExpansionDimensions,X);

% Figure out maxChunkSize from BlockSize
maxChunkSize = RegressionKernel.estimateMaxChunkSize(numExpansionDimensions,blockSize,true);
maxChunkSize = min(maxChunkSize,20000); % No more than datastore default

%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%
%-% START  RegressionLinear.PrepareData  %-%

% Extract BaseFitObjectArgs related to regression:
args = {'responsetransform'};
defs = {                 []};
[transformer,~,BaseFitObjectArgsExtra] = internal.stats.parseArgs(args,defs,extraArgs{:});

%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%
%-% START   classreg.learning.Linear.prepareDataCR

% Extract remaining BaseFitObjectArgs related to classification:
args = {'weights' 'predictornames' 'responsename' 'categoricalpredictors' 'observationsin' 'ordinalpredictors'};
defs = {       []               []             []                      []           'rows'                  []};
[W,predictorNames,responseName,categoricalPredictors,obsIn,~] = internal.stats.parseArgs(args,defs,BaseFitObjectArgsExtra{:});
%  Note: the only AllowedBaseFitObjectArgs not extracted can be scoretransform', 
%        if present in BaseFitObjectArgs parseArgs will error

% Check input type.
X = tall.validateType(X,upper(mfilename),{'float'},1);
X = lazyValidate(X,{@(x) ismatrix(x),'stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:BadXType'});
 
% Check input orientation
obsIn = validatestring(obsIn,{'rows' 'columns'},'classreg.learning.FullClassificationRegressionModel:prepareDataCR','ObservationsIn');
assert(strcmp(obsIn,'rows'),message('stats:tall:fitrkernel:BadObservationsIn'))
obsIn = false; % the Impl of the only binary learners supported for tall expect obs in columns

%-% Validate/set default weights
if istall(W)
    if any(strcmp({'categorical','datetime','duration','calendarDuration','cell','char'},W.Adaptor.Class))
        error(message('stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:BadW'));
    end
    W = tall.validateColumn(W,'stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:BadW');
    W = lazyValidate(W,{@(w) isfloat(w),'stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:BadW'});
    W = lazyValidate(W,{@(w) ~any(w<0,1),'stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:NegativeWeights'});
    % Note: Checking for Inf's departs from In-memory behavior, this may change in
    % a future release.    
    W = lazyValidate(W,{@(w) ~any(isinf(w),1),'stats:tall:fitrkernel:InfiniteWeights'});
else
    assert(isempty(W),message('stats:tall:fitrkernel:BadWeights'))
    W = slicefun(@(x) ones(size(x,1),1),Y);
end

% Check input size
[X,Y,W] = validateSameTallSize(X,Y,W);

% Remove obs with NaNs
lidx_a = ~any(isnan(X),2);
% Remove obs with NaNs or zeros in W
lidx_b = ~isnan(W) & W~=0;

% Process predictor names: 
%    test on D is deferred
if ~isempty(predictorNames) && ~isnumeric(predictorNames) && ~iscellstr(predictorNames)
    if ~ischar(predictorNames)
        error(message('stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:BadPredictorType'));
    end
    predictorNames = cellstr(predictorNames);
end

% Process response name
if isempty(responseName)
     responseName = 'Y';
elseif ~internal.stats.isString(responseName)
    error(message('stats:classreg:learning:internal:utils:BadResponseName'));
end

% Find categorical predictors
if ~isempty(categoricalPredictors)
    error(message('stats:tall:fitrkernel:InvalidCategoricalPredictors'));
end

%-% END   classreg.learning.Linear.prepareDataCR
%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%
%-% CONTINUE  RegressionLinear.PrepareData 

% Check Y type
Y = tall.validateColumn(Y,'stats:classreg:learning:regr:FullRegressionModel:prepareData:BadYType');
Y = tall.validateType(Y,upper(mfilename),{'float'},2);

% Remove NaN's in Y
lidx_c = ~isnan(Y);

rowsused = lidx_a&lidx_b&lidx_c;
Ne_ = sum(rowsused,1);
[N_,D_] = size(X);

%-% Filter...
sidx = substruct('()', {rowsused, ':'});
X = subsref(X,sidx);
Y = subsref(Y,sidx);
W = subsref(W,sidx);

%-% Chunk data
[X_,Y_,W_,Wk_,Ck_] = fixedchunkfun(maxChunkSize,@firstPassFun,X,Y,W);
% Wk_                     Weight sums per chunk
% Ck_                     Chunk counter
sw0_ = sum(Wk_,1);      % Total sum of weights EXCLUDING chunks with only one class
N_chunks_ = sum(Ck_,1); % Number of chunks

% TDigest to compute mean and median
T = TDigest(Y);

% Turn-off display from tall framework...
if verbose<3
    session = matlab.bigdata.internal.startMultiExecution(...
        'OutputFunction',@(x,y,z)fprintf(''),...
        'PrintBasicInformation',false); %#ok<NASGU>
end

%-% First pass over the data...
[N,Ne,D,sw0,N_chunks,Tg] = gather(N_,Ne_,D_,sw0_,N_chunks_,T);

assert(sw0>0,message('stats:tall:fitrkernel:TooFewObs'))

% (deferred) Check input size
if N==0
    error(message('stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:NoData'));
end

% (deferred) Process predictor names
if isempty(predictorNames)
    predictorNames = D;
elseif isnumeric(predictorNames)
    if ~(isscalar(predictorNames) && predictorNames==D)
        error(message('stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:BadNumericPredictor', D));
    end
else
    if length(predictorNames)~=D || length(unique(predictorNames))~=D
        error(message('stats:classreg:learning:FullClassificationRegressionModel:prepareDataCR:PredictorMismatch', D));
    end
end
predictorNames = predictorNames(:)';

Wk_ = Wk_./sw0;          % Weight per chunk
W_ = fixedchunkfun(1,@(w) {w{1}./sw0},W_);

defaultResponse.Median = icdf(Tg,0.5);
defaultResponse.Mean   = trimmedMean(Tg,0);

% Summarize data:
dataSummary.PredictorNames = predictorNames;
dataSummary.ResponseName = responseName;
dataSummary.CategoricalPredictors = [];
dataSummary.VariableRange = {};
dataSummary.TableInput = false;
dataSummary.ObservationsInRows = false;
dataSummary.OrdinalPredictors = false(1,D);

transformer = classreg.learning.regr.FullRegressionModel.processResponseTransform(transformer);

% This function may be called from bayes optimization, make sure we do not
% double cache the tall arrays
if ~matlab.bigdata.internal.util.hasMarkForReuse(X_,Y_,W_)
    [X_,Y_,W_] = markforreuse(X_,Y_,W_);
end

if verbose>0
    fprintf(getString(message('stats:tall:fitrkernel:FoundNChunks',N_chunks)))
end

%-% END  RegressionLinear.PrepareData  %-%
%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%%-%

keepHist = true;

% resolve parameters that may still be set to 'auto' or empty
% (Order matters!). All these modifications are not placed back
% into modelParams.
featureMapper = RegressionKernel.resolveEmptyFeatureMapper(featureMapper,D,numExpansionDimensions,transformation,stream);
[lambda,boxConstraint] = RegressionKernel.resolveAutoLambdaOrEmptyBoxConstraint(lambda,boxConstraint,Ne);
kernelScale = resolveAutoKernelScaleTall(kernelScale,X,Y,Ne);
epsilon = resolveAutoEpsilonTall(epsilon,lossfun,Y);
          
initialBeta = RegressionKernel.resolveEmptyInitialBeta(initialBeta,numExpansionDimensions);
initialBias = resolveEmptyInitialBiasTall(initialBias,lossfun,Y,W,epsilon,Ne);
Beta = [initialBias;initialBeta]; %including Bias
            
% Make the objective function
objgraF_ = makeobjgradF(X_,Y_,W_,lossfun,lambda./N_chunks,featureMapper,kernelScale,epsilon);

% Wrap objective function with verbosity and progress tracking behavior
progressF = classreg.learning.linearutils.linearSolverProgressFunction(objgraF_,false,admmIterationLimit>1,verbose>0,keepHist,@clientfun);

%%%%%%%%%%%%%%%%%%%%%%%%%%% ADMM FIT %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
Beta = internal.stats.bigdata.ADMMimplTall(X_,Y_,W_,Beta,rho,lambda,doridge,...
    N_chunks,Wk_,@markforreuse,@updateforreuse,@fixedchunkfun,...
    betaTol,gradTol,admmIterationLimit,tallPassLimit,progressF,verbose,...
    {lossfun_ADMMLBFGS,epsilon},betaTol_ADMMLBFGS,gradTol_ADMMLBFGS,...
    iterationlimit_ADMMLBFGS, doridge_ADMMLBFGS,...
    hessianHistorySize_ADMMLBFGS,dowolfe_ADMMLBFGS,doBias_ADMMLBFGS,...
    featureMapper,transformation,kernelScale);

%%%%%%%%%%%%%%%%%%%%%%%%%%% CENTRALIZED LBFGS %%%%%%%%%%%%%%%%%%%%%%%%%%%%%
Beta = internal.stats.bigdata.LBFGSimplTall(Beta,progressF,verbose,...
    betaTol,gradTol,iterationlimit,tallPassLimit,...
    hessianHistorySize,lineSearch,initialStepSize);

beta = Beta(2:end);
bias = Beta(1);

if verbose>0
    printLastLine(progressF);
end


if verbose>0
    % Adapt the history of linearSolverProgressFunction to
    % RegressionKernel's history
    fiHisNew = struct('ObjectiveValue',progressF.History.ObjectiveValue(:),...
        'GradientMagnitude',progressF.History.GradientMagnitude(:),...
        'Solver',categorical(repmat({'LBFGS-tall'},numel(progressF.History.Solver),1)),...
        'IterationNumber',progressF.History.IterationNumber(:),...
        'DataPass',progressF.History.DataPass(:),...
        'RelativeChangeInBeta',progressF.History.RelativeChangeBeta(:),...
        'ElapsedTime',progressF.History.ElapsedTime(:));
else
    fiHisNew = [];
end
                
fitInfo = struct('Solver','LBFGS-tall',...
    'LossFunction',lossfun,...
    'Lambda',lambda,...
    'BetaTolerance', betaTol,...
    'GradientTolerance',gradTol,...
    'ObjectiveValue',progressF.History.ObjectiveValue(max(1,end):end),...
    'GradientMagnitude',progressF.History.GradientMagnitude(max(1,end):end),...
    'RelativeChangeInBeta',progressF.History.RelativeChangeBeta(max(1,end):end),...
    'FitTime',progressF.History.ElapsedTime(max(1,end):end),...
    'History',fiHisNew);
    
model = RegressionKernel.makeNoFit(beta,bias,lambda,featureMapper,...
    kernelScale,boxConstraint,epsilon,modelParams,dataSummary,fitInfo,transformer,...
    [],[],defaultResponse);

varargout= {model,fitInfo};

end

function objgraF = makeobjgradF(X,Y,W,lossfun,lambdaK,FM,sigma,epsilon)
objgraF = @fcn;
    function [obj,gra] = fcn(Beta)
        [objs,grads] = fixedchunkfun(1, @(x,y,w) ...
        chunkObjGraFun(Beta,x,y,w,lossfun,lambdaK,FM,sigma,epsilon), ...
            X,Y,W);
        obj = sum(objs,1);
        gra = sum(grads,1);
    end
end

function [obj,gra] = chunkObjGraFun(Beta,x,y,w,lossfun,lambda,FM,sigma,epsilon)
assert(numel(x)==1); % sanity check
x = x{1};
y = y{1};
w = w{1};
useBias = true;
doridge = true;
if isempty(x)
    obj = (Beta(2:end)'*Beta(2:end))*lambda/2;
    gra = [0;Beta(2:end)*lambda]';
else
    xm = map(FM,x,sigma)';
    if isempty(epsilon)
        % When lossfun is 'mse' epsilon may be empty, but 
        % classreg.learning.linearutils.objgrad expects a scalar
        epsilon = 0;
    end
    [obj,gra] = classreg.learning.linearutils.objgrad( Beta(2:end),Beta(1),xm,y,w,...
        lossfun,doridge,lambda,epsilon,useBias, ...
        zeros(numel(Beta)-1,1,'like',Beta),0,0 );
    gra = [gra(end);gra(1:end-1)]';
    % obj = sum(max(0,1-(x*Beta(2:end)+Beta(1)).*y).*w)+(Beta(2:end)'*Beta(2:end))*lambda/2;
    % gra = (sum((y.*(x*Beta(2:end)+Beta(1))<1) .* (-y.*[ones(numel(w),1) x].*w))' + [0;Beta(2:end)*lambda])';
end
end

function [X_,Y_,W_,sw,ck] = firstPassFun(X,Y,W)
% Outputs:
%     X_,Y_,W_    cell-ified data
%     sw          chunk weight
%     ck          chunk counter
ck = 1;
sw = sum(W);
X_ = {X};
Y_ = {Y};
W_ = {W};
end

function initialBias = resolveEmptyInitialBiasTall(initialBias,lossFunction,Y,W,epsilon,N)
% Mimics RegressionKernel.resolveEmptyInitialBias for tall inputs
if isempty(initialBias)
    session1 = matlab.bigdata.internal.startMultiExecution(... % No display for next gather
        'OutputFunction',@(x,y,z)fprintf(''),...
        'PrintBasicInformation',false);
    switch lossFunction
        case 'epsiloninsensitive'  
            YW = gather(datasample([Y W],min(200,N),'replace',false));
            % Epsilon is not used for initial Bias to avoid cases that could
            % lead to zero objective gradient...
            epsilon = 0;
            initialBias =  classreg.learning.linearutils.fitbias(...
                'epsiloninsensitive',YW(:,1),zeros(size(YW,1),1,'like',YW),YW(:,2),epsilon);
        case 'mse'
            initialBias = gather(sum(Y.*W)./sum(W));
    end
    endMultiExecution(session1);
end
end

function epsilon = resolveAutoEpsilonTall(epsilon,lossFunction,Y)
% Mimics RegressionKernel.resolveAutoEpsilon for tall inputs
if strcmp(epsilon,'auto')
    switch lossFunction
        case 'epsiloninsensitive'
            session1 = matlab.bigdata.internal.startMultiExecution(... % No display for next gather
                'OutputFunction',@(x,y,z)fprintf(''),...
                'PrintBasicInformation',false);
            epsilon = diff(gather(quantile(Y,[.25 .75],'Method','approximate')))./13.49;
            endMultiExecution(session1);
            if epsilon == 0
                epsilon = 1;
            end
        case 'mse'
            epsilon = [];
    end
end
end

function kernelScale = resolveAutoKernelScaleTall(kernelScale,X,Y,N)
% Mimics RegressionKernel.resolveAutoKernelScale for tall inputs
if strcmp(kernelScale,'auto')
    session1 = matlab.bigdata.internal.startMultiExecution(... % No display for next gather
        'OutputFunction',@(x,y,z)fprintf(''),...
        'PrintBasicInformation',false);
    Xs = gather(datasample(X,min(200,N),'replace',false));
    Ys = []; % optimalKernelScale ignores Y when type is 0 or 1
    endMultiExecution(session1);
    kernelScale = classreg.learning.svmutils.optimalKernelScale(Xs,Ys,0); % type=0 for regression
end
end

function numExpansionDimensions = resolveAutoNumExpansionDimensionsTall(numExpansionDimensions,X)
% Mimics RegressionKernel.resolveAutoNumExpansionDimensions for tall inputs
if strcmp(numExpansionDimensions,'auto')
    session1 = matlab.bigdata.internal.startMultiExecution(... % No display for next gather
        'OutputFunction',@(x,y,z)fprintf(''),...
        'PrintBasicInformation',false);
    D = gather(size(X,2));
    endMultiExecution(session1);
    numExpansionDimensions = RegressionKernel.resolveAutoNumExpansionDimensions(numExpansionDimensions,D);
end
end

% LocalWords:  NVP bayes NChunks objgrad gra ified sw
