function Stop = plotAcquisitionFunction(bo, State)
%plotAcquisitionFunction     PlotFcn for plotting the
%Acquisition Function surface.
%   Stop = plotAcquisitionFunction(bo, State) plots the
%   Acquisition function surface. The value is set to NaN for
%   points that violate the xConstraintFcn.
%
% See also: BAYESOPT, BAYESIANOPTIMIZATION
%

%   Copyright 2016-2022 The MathWorks, Inc.

persistent Axes Curves;
narginchk(2,2);
Stop = false;
if bo.NumVars > 2 || acqFcnIsNonBayes(bo)
    return;
end
% Plot AF using fantasy models if UseParallel is true
if bo.PrivOptions.UseParallel && ~isempty(bo.FantasyModels)
    bo = bo.FantasyModels;
end
set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
switch State
    case 'initial'
        Axes = setupPlot(bo);
        Curves = {};
    case {'iteration', 'done'}
        if isvalid(Axes)
            [Axes, Curves] = updatePlot(bo, Axes, Curves);
        end
    case 'standalone'
        StandaloneAxes = setupPlot(bo);
        updatePlot(bo, StandaloneAxes, {});
end

    function Axes = setupPlot(Results)
        screenSize = get(groot,'ScreenSize');
        L = screenSize(1);
        B = screenSize(2);
        W = screenSize(3);
        H = screenSize(4);
        left = W/3 + 1;
        bottom = H/2 - 100;
        width = W/3 - 50;
        height = H/2;
        f = figure('Position',[left, bottom, width, height],'Tag','bayesopt.AcqFcn');
        Axes = axes(f);
        title(Axes, bayesoptim.infoString('AcqFcn', Results.PrivOptions.AcquisitionFunctionName));
        if Results.NumVars == 2
            view(Axes, -45, 15);
        end
    end

    function [Axes, Curves] = updatePlot(Results, Axes, Curves)
        import bayesoptim.*
        set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
        HoldOff = bayesoptim.holdOn(Axes);
        Options = Results.PrivOptions;
        VarSpec = Options.VarSpec;
        FModel = Results.ObjectiveFcnModel;
        if acqFcnIsNonBayes(Results)
            return;
        end
        if isempty(FModel)
            return;
        end
        ObjectiveEvaluationTimeModel = Results.ObjectiveEvaluationTimeModel;
        if isempty(ObjectiveEvaluationTimeModel) && ismember(Options.AcquisitionFunctionName, ...
                {'expectedimprovementpersecond', 'expectedimprovementpersecondplus'})
            return;
        end
        cellfun(@delete, Curves);
        Curves = {};
        switch Results.NumVars
            case 1
                % Create grid to evaluate ObjectiveFcn on (in design matrix shape)
                EvalGrid = NDGrid(VarSpec.LBTrans, VarSpec.UBTrans, Options.NumPlotGrid);
                LegalizedEvalGrid = legalizePoints(Results, EvalGrid);
                % Transform eval grid to native space for plotting
                PlotGrid = table2array(untransformPoints(Results, EvalGrid, false));
                % Compute AF at gridpoints and at next point
                switch Options.AcquisitionFunctionName
                    case 'probabilityofimprovement'
                        AFGrid = probabilityOfImprovement(LegalizedEvalGrid, ...
                            FModel, Results.IncumbentF, Results.ObjectiveFcnModel.Sigma);
                        AFNext = probabilityOfImprovement(legalizePoints(Results, Results.XNext), ...
                            FModel, Results.IncumbentF, Results.ObjectiveFcnModel.Sigma);
                    case 'lowerconfidencebound'
                        AFGrid = lowerConfidenceBound(LegalizedEvalGrid, ...
                            FModel, 2);
                        AFNext = lowerConfidenceBound(legalizePoints(Results, Results.XNext), ...
                            FModel, 2);
                    case {'expectedimprovement', 'expectedimprovementplus'}
                        AFGrid = expectedImprovement(LegalizedEvalGrid, ...
                            FModel, ...
                            Results.IncumbentF);
                        AFNext = expectedImprovement(legalizePoints(Results, Results.XNext), ...
                            FModel, ...
                            Results.IncumbentF);
                    case {'expectedimprovementpersecond', 'expectedimprovementpersecondplus'}
                        AFGrid = expectedImprovementPerSecond(LegalizedEvalGrid, ...
                            FModel, ObjectiveEvaluationTimeModel, Results.IncumbentF);
                        AFNext = expectedImprovementPerSecond(legalizePoints(Results, Results.XNext), ...
                            FModel, ObjectiveEvaluationTimeModel, Results.IncumbentF);
                end
                % NaN-out points that violate xConstraint
                mask = double(satisfiesXConstraint(Results, EvalGrid));
                mask(mask==0) = NaN;
                AFGrid = AFGrid .* mask;
                % Apply constraint weighting
                AFGrid = AFGrid .* ProbAllConstraintsSatisfied(Results, EvalGrid);
                % Plot surface
                Curves{end+1} = plot(Axes, PlotGrid, AFGrid, 'r');
                % Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot(Axes, XNextToPlot, AFNext, MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, Options.AcquisitionFunctionName);
            case 2
                % Create grid to evaluate ObjectiveFcn on (in design matrix shape)
                EvalGrid = NDGrid(VarSpec.LBTrans, VarSpec.UBTrans, Options.NumPlotGrid);
                LegalizedEvalGrid = legalizePoints(Results, EvalGrid);
                % Transform eval grid to native space for plotting
                PlotGrid = table2array(untransformPoints(Results, EvalGrid, false));
                % Make a version of the PlotGrid in meshgrid shape:
                X1Plot = reshape(PlotGrid(:,1), reverse(Options.NumPlotGrid));
                X2Plot = reshape(PlotGrid(:,2), reverse(Options.NumPlotGrid));
                % Compute AF at gridpoints and at next point
                switch Options.AcquisitionFunctionName
                    case 'probabilityofimprovement'
                        AFGrid = probabilityOfImprovement(LegalizedEvalGrid, ...
                            FModel, Results.IncumbentF, Results.ObjectiveFcnModel.Sigma);
                        AFNext = probabilityOfImprovement(legalizePoints(Results, Results.XNext), ...
                            FModel, Results.IncumbentF, Results.ObjectiveFcnModel.Sigma);
                    case 'lowerconfidencebound'
                        AFGrid = lowerConfidenceBound(LegalizedEvalGrid, ...
                            FModel, 2);
                        AFNext = lowerConfidenceBound(legalizePoints(Results, Results.XNext), ...
                            FModel, 2);
                    case {'expectedimprovement', 'expectedimprovementplus'}
                        AFGrid = expectedImprovement(LegalizedEvalGrid, ...
                            FModel, Results.IncumbentF);
                        AFNext = expectedImprovement(legalizePoints(Results, Results.XNext), ...
                            FModel, Results.IncumbentF);
                    case {'expectedimprovementpersecond', 'expectedimprovementpersecondplus'}
                        AFGrid = expectedImprovementPerSecond(LegalizedEvalGrid, ...
                            FModel, ObjectiveEvaluationTimeModel, Results.IncumbentF);
                        AFNext = expectedImprovementPerSecond(legalizePoints(Results, Results.XNext), ...
                            FModel, ObjectiveEvaluationTimeModel, Results.IncumbentF);
                end
                % NaN-out points that violate xConstraint
                mask = double(satisfiesXConstraint(Results, LegalizedEvalGrid));
                mask(mask==0) = NaN;
                AFGrid = AFGrid .* mask;
                % Apply constraint weighting
                AFGrid = AFGrid .* ProbAllConstraintsSatisfied(Results, LegalizedEvalGrid);

                % Plot surface
                Curves{end+1} = surfc(Axes, X1Plot, X2Plot, reshape(AFGrid, reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'b', 'FaceAlpha', .6);
                % Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot3(Axes, XNextToPlot(1), XNextToPlot(2), AFNext, ...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, VarSpec.Names{2});
                zlabel(Axes, Options.AcquisitionFunctionName);
        end
        % set axes scale and limits
        Axes = bayesoptim.setAxesScaleAndLimits(Axes, VarSpec);
    end
end
