function [cleanupObj, cleanupWorkersObj] = setAugLagFcnHandleOnWorkers(useparallel,objfun,confun)

% This function stores the objective and and constraint function locally on
% main process and sends them to workers so that the AugLagFun function can
% be evaluated later locally or in parallel.
%
% 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.
%
% This function also creates onCleanup object and returns to the caller.
% This will make sure that the persistent storage on the workers will be
% removed.
%
% If there are any recursive calls to any of the three solvers, the older
% values are stored as part of cleanup object. When the caller returns, 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 that will ensure that 
%             the cleanup work on the main process is completed.
% cleanupWorkersObj : This is a onCleanup object that will ensure that the
%                     cleanup work on the workers is completed.

%   Copyright 2021-2022 The MathWorks, Inc.

cleanupWorkersObj = [];

% Set the persistent storage on the local host.
[old_objfun, old_confun] = augLagFun([], objfun, confun);
cleanupObj = onCleanup(@() CleanupLocal(old_objfun, old_confun));

if ~validateopts_UseParallel(useparallel,true,true)
    % When user does not set UseParallel option to true, return now.
    return;
end

% At this point, we know that solvers will execute parfor body. 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.
    cleanupWorkersObj = parallel.pool.Constant(@SetupWorker);
    wait(parfevalOnAll(pool, @(~) [], 0, cleanupWorkersObj));
end


function cleanupObj = SetupWorker()
% Setup persistent storage on this current worker and return a cleanup
% object to restore it afterwards.
[objfun, confun] = getSetOptimFcns();
[old_objfun, old_confun] = augLagFun([], objfun, confun);
cleanupObj = onCleanup(@() augLagFun([], 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)
augLagFun([], old_objfun, old_confun);


