% ModelManager: Class to handle model evaluation 

%   Copyright 2018-2020 The MathWorks, Inc.

classdef ModelManager < handle
    properties(Access = public)
        pool
        expensiveModel
        options
        EvaluationCount = 0
        Verbosity = 0
        % Have some evaluation in the queue; approximately 30% more than
        % the nWorkers (which should roughly account for the solver's
        % expense ratio)
        queue_size = 1.3
    end
    
    properties(Access=private)
        numWorkers = 1
        numWorkersAvailable = 1
        NumPendingTrials = 0
        expensiveModel_response = {}
        workerError
        pct_const_handle
    end
    
    properties(Hidden, Transient, Constant)
        varName = 'X';
        response_fields = {'Fval','Ineq','Stop'}
        use_pct_const = true;
    end
    
    properties(Hidden, Transient, Access = public)
        
        pendingTrials = struct('tid',{0}, ...
            'output', [], ...
            'pending', {true}, ...
            'task',{[]}, ...
            'ID',{0}, ...
            'clocktime',[])
        combinedModel
        output
    end
    
    methods(Access = public)
        
        function self = ModelManager(expensive, options, Verbosity_)
            
            self.expensiveModel = expensive.model;
            self.expensiveModel_response = expensive.response;
            
            % Create an empty output structure.
            output = struct();
            response = self.response_fields;
            for i = 1:length(response)
                output.(response{i}) = [];
            end
            output.Stop = false;
            self.pendingTrials.output = output;

            self.combinedModel = @(trial) conditionalModelEvaluator(trial, output, ...
                self.varName, expensive.model, expensive.response);
            
            self.output = output;
            self.options = options;
            self.Verbosity = Verbosity_;
            
            if options.UseParallel
                self.openPool();
                if ~isempty(self.pool)
                    self.numWorkersAvailable = self.pool.NumWorkers;
                    self.NumPendingTrials = 0;
                end
            else
                self.queue_size = 1;
            end
            
        end
        
        %==================================================================
        function evaluate(self, nextTrial)
           % nextTrial is (an array of) struct.            
            self.pendingTrials(end+1).trial = nextTrial;
            
            self.pendingTrials(end).pending = true;
            
            self.NumPendingTrials = self.NumPendingTrials + 1;
            self.numWorkersAvailable = self.numWorkersAvailable - 1;
            
            if self.Verbosity > 3
                fprintf('                      Sending trial with ID:%d for evaluation.\n', nextTrial.tid);
            end
            
            % Evaluate expensive model
            if self.options.UseParallel
                % Evaluate parallel now
                self.evaluateRemote(nextTrial);
            else
                self.evaluateLocal(nextTrial);
            end
            
        end
        
        %==================================================================
        function response = getNext(self)
            
            if self.NumPendingTrials == 0
                response = struct('tid',0);
                return;
            end
            
            if self.options.UseParallel
                
                % This will become a blocking call if no results are
                % available.
                ID_status = ~[self.pendingTrials.pending];
                index = find(ID_status);
                
                if ~isempty(index)
                    fetchID = index(1);
                else
                    fetchID = 2;
                end
                
                if ~isempty(self.pendingTrials(fetchID).task.Error)
                    self.workerError = self.pendingTrials(fetchID).task.Error;
                    %@todo: How to deal with error in function execution?
                    % Throw error by executing fetchNext
                    [~, self.pendingTrials(fetchID).output] = fetchNext(self.pendingTrials(fetchID).task);
                    self.pendingTrials(fetchID).output = struct([]);
                else
                    [~, self.pendingTrials(fetchID).output] = fetchNext(self.pendingTrials(fetchID).task);
                end

            else
                fetchID = 2; %length(self.pendingTrials);
            end
            
            response = self.prepareOutput(fetchID);
            
            self.NumPendingTrials = self.NumPendingTrials - 1;
            self.numWorkersAvailable = self.numWorkersAvailable + 1;
            self.EvaluationCount = self.EvaluationCount +  1;
            
            self.pendingTrials(fetchID) = [];
        end
        
        %==================================================================
        function TF = noMoreEvaluations(self)
            TF = self.options.MaxFunctionEvaluations <= ...
                self.EvaluationCount;
        end
        
        %==================================================================
        function nWorkers = availableWorkers(self)
            nWorkers = self.numWorkersAvailable*self.queue_size;
        end
        
        %==================================================================
        function TF = responseAvailable(self)
            if self.NumPendingTrials == 0
                TF = false;
                return;
            end
            
            if self.options.UseParallel
                TF = any(self.findFinishedTasks());
            else
                TF = true;
            end
        end
        
        %==================================================================
        function TF = isTrialInQueue(self)
            TF = self.NumPendingTrials > 0;
        end
        
        %==================================================================
        function cancelAll(self)
            
            if self.options.UseParallel && self.isTrialInQueue()
                 alltask = [self.pendingTrials.task];
                 cancel(alltask);
                 ID_status = ~[self.pendingTrials.pending];
                 self.pendingTrials(ID_status) = [];
                 
                 if self.Verbosity > 2 && numel(alltask) > 0
                     fprintf('\n%d remaining evaluation(s) cancelled.\n', numel(alltask));
                 end                 
                
                self.NumPendingTrials = 0;
            end
        end
        
        %==================================================================
        function cancelPending(self)
            
            if self.options.UseParallel && self.isTrialInQueue()
                ID_status = false(length(self.pendingTrials),1);
                % Cancel each task as soon as possible and then remove it
                % from the pendingTrials structure.
                
                for ii = 2:length(self.pendingTrials)                                        
                    tStart = self.pendingTrials(ii).task.StartDateTime;

                    if isempty(tStart)
                        cancel(self.pendingTrials(ii).task)
                        ID_status(ii) = true;
                        if self.Verbosity > 3
                            fprintf('Trial ID %d cancelled.\n', ii);
                        end
                    end
                end
                 
                 if self.Verbosity > 2 && nnz(ID_status) > 0
                     fprintf('\n%d pending evaluation(s) cancelled.\n', nnz(ID_status));
                 end
                 self.pendingTrials(ID_status) = [];                
                 self.NumPendingTrials = self.NumPendingTrials - nnz(ID_status);
                 
            end
        end
        

    end
    
    
    methods(Access = private)
        
        %==================================================================
        function self = evaluateLocal(self, nextTrial)
            t0 = tic;
            self.pendingTrials(end).output = self.combinedModel(nextTrial);
            self.pendingTrials(end).clocktime = toc(t0);
            
            % Set the pending flag to false now.
            self.pendingTrials(end).pending = false;
        end
        
        %==================================================================
        function evaluateRemote(self, nextTrial)

            self.openPool();
            
            if self.use_pct_const                
                self.pendingTrials(end).task = parfeval(self.pool, ... 
                    self.pct_const_handle.Value, 1, nextTrial);
            else
                self.pendingTrials(end).task = parfeval(self.pool, self.combinedModel, ...
                1,nextTrial);
            end
            self.pendingTrials(end).ID = self.pendingTrials(end).task.ID;
            
        end
        
        %==================================================================
        function response = prepareOutput(self, index)
           
            response = self.pendingTrials(index).trial;

            % Get model response from the Trials array
            results = self.pendingTrials(index).output;
            if isempty(results) % Error in PCT eval (MATLAB error)
                results = struct('error',true);
                fnames = fieldnames(self.output);
                for ii = 1:length(fnames)
                    results.(fnames{ii}) = NaN;
                end
            end
                    
            fnames = fieldnames(results);
            for ii = 1:length(fnames)
               response.(fnames{ii}) = results.(fnames{ii});
            end
            
            % Add 'update' field required to indicate that response is
            % updated.
            response.update = true;
            
            if self.Verbosity > 3
                fprintf('                      Trial ID %d evaluated.\n', response.tid);
            end
            
        end
        
        %==================================================================
        function ID_status = findFinishedTasks(self)            
            % Weird way to flush the PCT queue when there is a function
            % error on the worker.
            pause(eps); 
            ID_status = zeros(length(self.pendingTrials),1);
            for ii = 2:length(self.pendingTrials)
                tFinish = self.pendingTrials(ii).task.FinishDateTime;
                if ~isempty(tFinish)
                    ID_status(ii) = self.pendingTrials(ii).ID;
                    % Record time
                    tStart  = self.pendingTrials(ii).task.StartDateTime;                    
                    self.pendingTrials(ii).clocktime = seconds(tFinish - tStart);
                    
                    % Mark as finished
                    self.pendingTrials(ii).pending = false;
                end
            end
        end
        
        %==================================================================
        function openPool(self)
            
            try
                % Get a handle to the pool
                self.pool = gcp;
            catch % exception when there is no license
                % PCT not installed or pool creation threw an error
                self.pool = [];
            end
            
            if isempty(self.pool)
                % Warn to say that we are not using parallel processing
                warning(message('optimlib:commonMsgs:NoPCTLicense'));
                self.options.UseParallel = false;
                self.queue_size = 1;
                return
            else
                self.numWorkers = self.pool.NumWorkers;
            end
            
            % Send the function handle to all the workers
            if ~isempty(self.pool) && self.use_pct_const && isempty(self.pct_const_handle)
                %self.pct_const_handle = parallel.pool.Constant(@() self.combinedModel);
                self.pct_const_handle = copyFunctionHandleToWorkers(self.combinedModel);
            end
        end
        
    end
    
end


%==================================================================
function output = conditionalModelEvaluator(nextTrial, output, ...
    varName, expensiveModel, expensiveModel_response)
% trial is a struct with field
% .(self.varName) : decision vector

out = feval(expensiveModel,nextTrial.(varName));
for ii = 1:length(expensiveModel_response)
    % Response is always a row vector for surrogateopt
    output.(expensiveModel_response{ii}) = out.(expensiveModel_response{ii})(:)';
end

end

function C = copyFunctionHandleToWorkers(fHandle)
%copyFunctionHandleToWorkers Copy objective function handle to parallel
%workers, sending any required code.
%
%   C = copyFunctionHandleToWorkers(fHandle) copies the function handle
%   fHandle to all workers in the current parallel pool, and sends any code
%   required to execute the function on the workers. C is an object of
%   class parallel.pool.Constant.

% IMPORTANT: This function must attach the files BEFORE sending the function handle.
attachMissingFunctionHandleFilesToPool(fHandle, gcp);
% Use an anonymous function that will be evaluated on workers to set the
% function handle into the constant.
C = parallel.pool.Constant(@() fHandle);
    function attachMissingFunctionHandleFilesToPool(fHandle, pool)
        % The command
        % parallel.internal.apishared.AttachedFiles.convertFunctionHandleForDependencyAnalysis(fHandle)
        % is internal and may not be supported in the future.
        NeededFunctionsAndFiles = parallel.internal.apishared.AttachedFiles.convertFunctionHandleForDependencyAnalysis(fHandle);
        Paths = cellfun(@which, NeededFunctionsAndFiles, 'UniformOutput',false);
        Builtin = cellfun(@isBuiltin, Paths);
        ToAttach = setdiff(Paths(~Builtin), pool.AttachedFiles);
        if ~isempty(ToAttach)
            pool.addAttachedFiles(ToAttach);
        end
    end

    function tf = isBuiltin(path)
        tf = strncmp(path, 'built-in', 8);
    end
end


