% TrialQueue: Queue to store design points generated by solver. Update
% responses when model is evaluated.
% 

%   Copyright 2018-2019 The MathWorks, Inc.

classdef TrialQueue < handle
    
    properties(Access=public)
        TrialsArray
    end
    
    properties(SetAccess = private, GetAccess = public)
        TrialStruct
        response_fields = {}
        all_responses
        response_field_ok = false;
        tid = 0;
        evalCount = 0
        Verbosity = 0
    end
    
    properties(Hidden, Constant)
        varName = 'X';
        rejectFailed = false;
    end
    
    methods(Access=public)
        function self = TrialQueue(responses_,verbosity_)
            if nargin < 2
                verbosity_ = 0;
            end
            self.Verbosity = verbosity_;
            
            self.all_responses = responses_;
            % Some fields are essential. Must have a self.varName field.
            self.TrialStruct = struct(self.varName,[], 'tid',0, 'flag','', ...
                'waiting',false,'pending',false,'available',false, ... 
                'obsolete',false,'failed',false);

            self.TrialsArray = self.TrialStruct;
        end
        
        %==================================================================
        function tids = addTrial(self, trials, flag)
            % Add a new trial(s) in the queue waiting for evaluation
            tids = zeros(size(trials,1),1);
            for i = 1:size(trials,1)
                id_ = ~isreal(trials(i,:)) || ~all(isfinite(trials(i,:))) || ... 
                      ~isnumeric(trials(i,:)); % 
                if any(id_)
                   error(message('globaloptim:bmo:TrialsNotReal')); 
                end
                temp = self.TrialStruct;
                temp.(self.varName) = trials(i,:); % trial is row vector
                
                self.tid = self.tid + 1; % Increment the tid counter
                % Modify trial attributes to indicate their id and status
                temp.tid = self.tid;
                temp.waiting = true;
                temp.flag = flag;
                % Add the trial in the queue
                self.TrialsArray(end+1) = temp;                
                tids(i) = self.tid; % Return the id of the trials
            end
            
        end
        
        %==================================================================
        function trial = getWaitingTrial(self)
            % Next trial waiting for evaluation.
            id = find([self.TrialsArray.waiting],1);
            TF = ~isempty(id);
            if TF
                fnames = self.response_fields;
                temp = self.TrialsArray(id);
                trial.tid = temp.tid;
                trial.(self.varName) = temp.(self.varName);
                trial.flag = temp.flag;
                for i = 1:length(fnames)
                    trial.(fnames{i}) = temp.(fnames{i});
                end
                self.TrialsArray(id).pending = true;
                self.TrialsArray(id).waiting = false;
            else
                trial = struct([]);
            end
        end
        
        %==================================================================
        function TF = updateTrial(self,trials,pending,complete)
            % Update trial values
            % We get the intersection of fields; do not assume all
            % response_fields are present.
            fnamesIn = fieldnames(trials);
            
            if ~ismember('tid',fnamesIn) % Internal only error.
                error(message('globaloptim:bmo:TrialUpdateWithNoTID'));
            end
            
            if ~self.response_field_ok
                % Response fields are not finalized yet. Check for fields.
                
                fnames = intersect(self.all_responses, fnamesIn);
                % Check if we have the field already
                newFields = setdiff(fnames,self.response_fields);
                
                if ~isempty(newFields)
                    % We don't have all the response fields, add now.
                    for ii = 1:length(newFields)
                        self.response_fields{end+1} = newFields{ii};
                        self.TrialStruct.(newFields{ii}) = [];
                        self.TrialsArray(1).(newFields{ii}) = [];
                    end
                end
            else
                % Get all the response fields
                fnames = intersect(fnamesIn, self.response_fields);
            end
            
            tids = arrayfun(@(x) x.tid,  trials);
            to_remove = false(numel(tids),1);
            for ii = 1:length(tids)
                id = arrayfun( @(x) x.tid == tids(ii), self.TrialsArray);
                
                if ~any(id) % Internal only error.
                    error(message('globaloptim:bmo:TrialQueueUpdateNoTID',tids(ii)));
                end
                
                if complete
                    self.TrialsArray(id).waiting = false;
                end
                    
                self.TrialsArray(id).pending = pending;
                self.TrialsArray(id).available = complete;
                
                for jj = 1:length(fnames)
                    value = trials(ii).(fnames{jj})(:)';
                    id_ = ~isreal(value) || ~all(isfinite(value)) || ... 
                          ~isnumeric(value); %
                    if any(id_)
                        to_remove(ii) = true;
                        self.TrialsArray(id).failed = true;
                    end
                    self.TrialsArray(id).(fnames{jj}) = value;
                end
            end
            
            if complete
                % Update evalCount
                self.evalCount = self.evalCount + length(tids);
                self.response_field_ok = true;
            end
            if self.rejectFailed && any(to_remove)
                badTids = tids(to_remove);
                self.removeTrial(badTids)
                if self.Verbosity > 2 % testing
                    fprintf(getString(message('globaloptim:bmo:DiscardingNonFiniteComplex', ...
                        num2str(nnz(badTids)))));
                end
            end
            TF = true;
        end
        
        %==================================================================
        function trials = getEvaluated(self,tids, partial)
            % Trials that are evaluated. The trial is removed from the
            % queue after this operation. Only if ALL the tids are
            % available then TF is successful and trials is returned, else
            % not. If caller can work with one tid at a time then call with
            % one tid and not several tids.
            
            % Select trials corresponding to tid
            id = find(arrayfun( @(x) ismember(x.tid, tids), self.TrialsArray));
            
            % Filter the ones that are avaialble
            if nargin > 2 && ~partial
                avail_id = find([self.TrialsArray(id).available]);            
                % Ensure we have all the tids availble
                TF = length(avail_id) == length(id);
            else
                TF = ~isempty(id);
            end
            
            if TF
                response_fields_ = self.response_fields;
                X = vertcat(self.TrialsArray(id).(self.varName));
                flags = {self.TrialsArray(id).flag}';
                obsolete_flag = vertcat(self.TrialsArray(id).obsolete);
                failed_flag = vertcat(self.TrialsArray(id).failed);
                trials = struct('tid',num2cell(vertcat(self.TrialsArray(id).tid)), ...
                    self.varName,mat2cell(X, ones(size(X,1),1),size(X,2)), ...
                    'obsolete',num2cell(obsolete_flag), ...
                    'flag',flags,'failed',num2cell(failed_flag));
                
                for ii = 1:length(id)
                    for i = 1:length(response_fields_)
                        if isfield(self.TrialsArray(id(ii)), response_fields_{i})
                            trials(ii).(response_fields_{i}) = ...
                                self.TrialsArray(id(ii)).(response_fields_{i});
                        end
                    end
                end
                
            else
                trials = struct([]);
            end
        end
        
        %==================================================================
        function removeTrial(self,tids)
            % Select trials corresponding to tid
            id = arrayfun( @(x) ismember(x.tid, tids), self.TrialsArray);
            self.TrialsArray(id) = [];
        end
        
        %==================================================================
        function makeActiveTrialsObsolete(self, remove)
            % Remove 'waiting' trials from the queue and mark 'pending' as
            % obsolete. 
            
            if nargin < 2
                % Optionally do not remove waiting trials
                remove = true;
            end
            
            id = [self.TrialsArray.waiting];
            
            if remove
                self.TrialsArray(id) = [];
            elseif any(id)
                % turn the obsolete flag only.
                for ii = 1:length(id)
                    if id(ii)
                        self.TrialsArray(ii).obsolete = true;
                    end
                end
            end
            
            id = [self.TrialsArray.pending];
            % Turn the obsolete flag only.
            if any(id)
                for ii = 1:length(id)
                    if id(ii)
                        self.TrialsArray(ii).obsolete = true;
                    end
                end
            end
            
        end

        %==================================================================
        function makePendingTrialsWaiting(self)
            % Trials pending expensive model evaluation are to be marked as
            % waiting (possibly because model evaluation failed and we want
            % to try again)
            
            id = [self.TrialsArray.pending];            
            if any(id)
                for ii = 1:length(id)
                    if id(ii)                        
                        self.TrialsArray(ii).waiting = true;
                        self.TrialsArray(ii).pending = false;
                    end
                end
            end
            
        end
        %==================================================================
        function num = getNumEvaluatedTrials(self)
            num = nnz([self.TrialsArray.available]);
        end

        %==================================================================
        function tids = getEvaluatedTrialTIDs(self)
            tids = [self.TrialsArray([self.TrialsArray.available]).tid];
        end
        
        %==================================================================
        function num = getNumWaitingTrials(self)
            num = nnz([self.TrialsArray.waiting]);
        end
        
        %==================================================================
        function num = getNumPendingTrials(self)
            num = nnz([self.TrialsArray.pending]);
        end
        
        %==================================================================
        function trials = getPendingTrials(self)
            id = find([self.TrialsArray.pending]);
            if ~isempty(id)
                X = vertcat(self.TrialsArray(id).(self.varName));
                trials = struct('tid',num2cell(vertcat(self.TrialsArray(id).tid)), ...
                        self.varName,mat2cell(X, ones(size(X,1),1),size(X,2)));
            else
                trials = struct([]);
            end
        end        
        
    end

end

