function cleanupObj = setOptimFcnHandleOnWorkers(useparallel,funfcn,confcn,alreadySet)

% Used for parallel finite difference by fmincon, fminimax, and fgoalattain
% solvers.
% This function sends objective and constraint functions to workers which
% later are used by parfinitedifference function. parfinitedifference is
% used by fmincon, fminimax, and fgoalattain solvers. These solvers must
% call this function to setup the workers.
%
% The idea is that we want to avoid sending function handles in every
% optimization iteration. This helps when function handle is large.
%
% This function sends to the workers objective and constraint function
% handles. These handles are stored as persistent on the workers.
% parfinitedifference function gets the function handles from the workers. 
%
% This function also creates onCleanup object and returns to the caller.
% This will make sure that when the fmincon, fminimax, or fgoalattain
% solvers are done, the persistent storage on the workers will be removed. 
%
% If there are any recursive calls to any of the three solvers, the
% persistent storage is limited to the current stack entry. The older
% values are stored as part of cleanup object. When the stack unwinds, the
% old values are restored.
%
% Inputs:
% useparallel: must be true or false indicates if parallel finite
%              difference is to be used or not
% funfcn: (Objective function information) A cell array used by
%         optimization nonlinear solvers. The first element indicates if
%         finite difference is to be used or not. The third element
%         contains the function handle.
% confcn: Similar syntax as funfcn but for constraint function.
%
% Output:
% cleanupObj: The caller must create a variable as the output of this
%             function. This is a onCleanup object (see help for this
%             class) that will ensure that the cleanup work on the workers
%             is completed

%   Copyright 2014-2022 The MathWorks, Inc.

cleanupObj = [];
if ~useparallel
    % When user does not set UseParallel option to true, return now.
    return;
end

% If the memory on the workers has already been set, return now.
if nargin > 3 && alreadySet
    return;
end

if strcmp(funfcn{1},'fun')
    % Solvers will use finite difference for objective; get the function
    % handle.
    objfun = funfcn{3};
else
    % No finite difference; set the function handle to empty.
    objfun = '';
end

if strcmp(confcn{1},'fun')
    % Solvers will use finite difference for objective; get the function
    % handle.
    confun = confcn{3};
else
    % No finite difference; set the function handle to empty.
    confun = '';
end

% Check for both objfun and confun empty: return immediately if no finite
% difference is needed (irrespective of useparallel value).
if isempty(objfun) && isempty(confun)
    return;
end

% At this point, we know that solvers will execute parfor body for finite
% difference. We need to access the workers now in order to send function
% handles. It will start the pool if user preference is set.
try
    % Get a handle to the pool.
    pool = gcp;
catch
    % PCT not installed or pool creation threw an error
    pool = [];
end

if ~isempty(pool)
    % Update the worker persistent storage. This usage of Constant will run
    % SetupWorker per worker when first used with a pool, and then hold the 
    % returning object (I.E. one cleanup object per worker) to restore 
    % persistent storage across the pool on cleanup.
    cleanupObj = parallel.pool.Constant(@() SetupWorker(objfun, confun));
    wait(parfevalOnAll(pool,@(~) [], 0, cleanupObj));
else
    % Warn to say that we are not using parallel processing
    warning(message('optimlib:commonMsgs:NoPCTLicense'));
    
    % The pool is not available to use. The solver is still going to use the
    % parfor loop (because finite difference is 'on' and UseParallel is true)
    % This means, we must set the persistent storage on the local host.        
    [old_objfun, old_confun] = getSetOptimFcns(objfun, confun);
    cleanupObj = onCleanup(@() CleanupLocal(old_objfun, old_confun));
end

function cleanupObj = SetupWorker(objfun, confun)
% Setup persistent storage on this current worker and return a cleanup
% object to restore it afterwards.
[old_objfun, old_confun] = getSetOptimFcns(objfun, confun);
cleanupObj = onCleanup(@() getSetOptimFcns(old_objfun, old_confun));

function CleanupLocal(old_objfun, old_confun)
% Cleanup function to restore persistent storage (in getSetOptimFcns file
% on host) to old values.

% Restore the persistent storage (could be [] and it's ok)
getSetOptimFcns(old_objfun, old_confun);


