classdef MultiGPStrategy < bayesoptim.GPStrategy
    % MULTIGPSTRATEGY This class contains methods in which a Multiple
    % Gaussian Process model is used as the estimator

    %   Copyright 2020-2021 The MathWorks, Inc.

    properties
        GPActiveSetSize
    end

    methods
        function this = MultiGPStrategy(GPActiveSetSize)
            this.GPActiveSetSize = GPActiveSetSize;
        end

        function Args = overrideGPFitMethodArgs(this, X, Y, group)
            isg = X(:,1)==group;
            NumFiniteYInGroup = sum(isfinite(Y(isg)));
            if this.GPActiveSetSize <= NumFiniteYInGroup
                Args = {'FitMethod', 'sd', 'ActiveSetSize', this.GPActiveSetSize};
            else
                Args = {'FitMethod', 'exact', 'ActiveSetSize', []};
            end
        end


        function BayesOptModel = iFitRobust(this, BayesOptModel, X, F, SigmaLowerBound, varargin)
            % Fit a Multi Gaussian Process model to the given data.

            % default hyperparameters set in original SMAC paper

            if all(isnan(F))
                bayesoptim.err('AllModelTargetsNaN');
            end

            % Fit a MultiGP model
            if isempty(BayesOptModel)
                multiGP = bayesoptim.fitrmultigp(X, F, SigmaLowerBound, varargin{:});
                GPModel = multiGP.Models;

                % There should be at least one (non-empty) cell in GPModel,
                % given that BayesOptModel is empty. Typically this is the
                % first model, however, if the first objective function was
                % a NaN, this could be a later cell value.
                idx = ~(cellfun('isempty',GPModel));
                if sum(idx) < 1 % Ensure that at least one model is non-empty
                    BayesOptModel = [];
                else
                    idx = find(idx,1,'first');
                    GPModel = GPModel{idx};
                    BayesOptModel = bayesoptim.BayesOptGPModel(GPModel.Sigma, multiGP, ...
                        GPModel.KernelInformation, GPModel.Beta);
                end
            else
                newGroup = X(end, 1);
                fitMethodArgs = overrideGPFitMethodArgs(this, X, F, newGroup);
                BayesOptModel.Model = refit(BayesOptModel.Model, newGroup, X, F, SigmaLowerBound, varargin{:}, fitMethodArgs{:});
            end
        end
    end
end

