function Stop = plotObjectiveModel(bo, State)
%plotObjectiveModel     PlotFcn for plotting the Objective
%function model surface, along with the estimated location of
%the minimum.
%   Stop = plotObjectiveModel(bo, State) plots the Objective
%   function model surface, along with the estimated location
%   of the minimum.
%
% See also: BAYESOPT, BAYESIANOPTIMIZATION
%

%   Copyright 2016-2022 The MathWorks, Inc.

persistent Axes Curves;
narginchk(2,2);
Stop = false;
VarSpec = bo.PrivOptions.VarSpec;
if ~bo.PrivOptions.FitModels || bo.NumVars > 2
    return;
end
set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
switch State
    case 'initial'
        Axes = setupFPlot(bo);
        Curves = {};
    case {'iteration', 'done'}
        if isvalid(Axes)
            [Axes, Curves] = updateFPlot(bo, Axes, Curves);
        end
    case 'standalone'
        StandaloneAxes = setupFPlot(bo);
        updateFPlot(bo, StandaloneAxes, {});
end

    function Axes = setupFPlot(Results)
        screenSize = get(groot,'ScreenSize');
        L = screenSize(1);
        B = screenSize(2);
        W = screenSize(3);
        H = screenSize(4);
        % (1) Plot True Objective on grid
        left = L;
        bottom = H/2 - 100;
        width = W/3 - 50;
        height = H/2;
        f = figure('Position',[left, bottom, width, height], 'Tag','bayesopt.ObjectiveModel');
        Axes = axes(f);
        title(Axes, bayesoptim.infoString('ObjModel'));
        zlabel(Axes, bayesoptim.infoString('EstObj'));
        if Results.NumVars == 2
            view(Axes, -45, 15);
        end
    end

    function [Axes, Curves] = updateFPlot(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 isempty(FModel)
            return;
        end
        XTrace = Results.XTrace;
        ObjectiveTrace = Results.ObjectiveTrace;
        ModelMinXTable = untransformPoints(Results, Results.IncumbentX, true);
        ModelMinF = Results.IncumbentF;

        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 = untransformPoints(Results, EvalGrid, false);
                PlotGrid = PlotGrid{:,1};
                % (1) Eval GPRs on eval grid
                % Objective
                [FPred, YSD, ~] = predict(FModel, LegalizedEvalGrid);
                FSD = bayesoptim.funStd(YSD, FModel.Sigma);
                % (2) Plot training points on all axes
                XPlottable = makeXTablePlottable(Results, XTrace);
                Curves{end+1} = plot(Axes, XPlottable, ObjectiveTrace, 'bo', 'MarkerFaceColor', 'blue');
                % (2.5) Plot pending "fantasy" points if any
                if ~isempty(Results.Parallel) && anyPendingX(Results.Parallel)
                    PendingXTable   = conditionalizeX(Results, Results.Parallel.PendingX);
                    PendingXToPlot  = makeXTablePlottable(Results, PendingXTable);
                    PendingObj      = Results.PendingObj;
                    Curves{end+1}   = plot(Axes, PendingXToPlot, PendingObj, 'go', 'MarkerFaceColor', 'green');
                end
                % (3) Plot FModel's mean, envelope
                Curves{end+1} = plot(Axes, PlotGrid, FPred, 'r');
                Curves{end+1} = plot(Axes, PlotGrid, FPred + FSD, 'b');
                Curves{end+1} = plot(Axes, PlotGrid, FPred - FSD, 'b');
                Curves{end+1} = plot(Axes, PlotGrid, FPred + FModel.Sigma, 'c');
                Curves{end+1} = plot(Axes, PlotGrid, FPred - FModel.Sigma, 'c');
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot(Axes, XNextToPlot, predictObjective(Results, Results.NextPoint), ...
                    MarkerSpec{:});
                % (6.1) Plot minimum if it exists
                if ~isempty(ModelMinXTable)
                    XMinPlottable = makeXTablePlottable(Results, ModelMinXTable);
                    Curves{end+1} = plot(Axes, XMinPlottable, ModelMinF, 'r*');
                end
                % (6.2) Plot legend
                if ~isempty(Results.Parallel) && anyPendingX(Results.Parallel)
                    LegendStrings = {...
                        bayesoptim.infoString('Observed'),...
                        bayesoptim.infoString('Evaluating'),...
                        bayesoptim.infoString('ModelMean'),...
                        bayesoptim.infoString('ModelErrorbars'),...
                        bayesoptim.infoString('NoiseErrorbars'),...
                        bayesoptim.infoString('NextPoint')};
                    CurveNums = [1,2,3,4,6,8];
                else
                    LegendStrings = {...
                        bayesoptim.infoString('Observed'),...
                        bayesoptim.infoString('ModelMean'),...
                        bayesoptim.infoString('ModelErrorbars'),...
                        bayesoptim.infoString('NoiseErrorbars'),...
                        bayesoptim.infoString('NextPoint')};
                    CurveNums = [1,2,3,5,7];
                end
                if ~isempty(ModelMinXTable)
                    LegendStrings{end+1} = bayesoptim.infoString('ModelMin');
                    CurveNums(end+1) = CurveNums(end)+1;
                end
                legend(Axes, [Curves{CurveNums}], LegendStrings{:}, 'Location', 'best');
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, bayesoptim.infoString('EstObj'));
            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 = untransformPoints(Results, EvalGrid, false);
                PlotGrid = PlotGrid{:,:};
                % Make a version of the PlotGrid in meshgrid shape:
                X1Plot = reshape(PlotGrid(:,1), reverse(Options.NumPlotGrid));
                X2Plot = reshape(PlotGrid(:,2), reverse(Options.NumPlotGrid));
                % (1) Eval GPRs on eval grid
                % Objective
                FPred = predict(FModel, LegalizedEvalGrid);
                % (2) Plot training points
                XPlottable = makeXTablePlottable(Results, XTrace);
                Curves{end+1} = plot3(Axes, XPlottable(:,1), XPlottable(:,2), ObjectiveTrace, ...
                    'bo', 'MarkerFaceColor', 'blue');
                % (2.5) Plot pending "fantasy" points if any
                if ~isempty(Results.Parallel) && anyPendingX(Results.Parallel)
                    PendingXTable   = conditionalizeX(Results, Results.Parallel.PendingX);
                    PendingXToPlot  = makeXTablePlottable(Results, PendingXTable);
                    PendingObj      = Results.PendingObj;
                    Curves{end+1} = plot3(Axes, PendingXToPlot(:,1), PendingXToPlot(:,2), PendingObj, ...
                        'go', 'MarkerFaceColor', 'green');
                end
                % (3) Plot FModel's mean
                SurfAndContour = surfc(Axes, X1Plot, X2Plot, reshape(FPred, reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'r', 'FaceAlpha', .5, 'LineStyle', '-');
                Curves = [Curves, {SurfAndContour(1), SurfAndContour(2)}];
                % (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), predictObjective(Results, Results.NextPoint), ...
                    MarkerSpec{:});
                % (6.1) Plot model min if exists
                if ~isempty(ModelMinXTable)
                    XMinPlottable = makeXTablePlottable(Results, ModelMinXTable);
                    Curves{end+1} = plot3(Axes, XMinPlottable(1), XMinPlottable(2), ModelMinF, 'r*');
                end

                % (6.2) Plot legend
                if ~isempty(Results.Parallel) &&  anyPendingX(Results.Parallel)
                    LegendStrings = {...
                        bayesoptim.infoString('Observed'),...
                        bayesoptim.infoString('Evaluating'),...
                        bayesoptim.infoString('ModelMean'),...
                        bayesoptim.infoString('NextPoint')};
                    CurveNums = [1,2,3,5];
                else
                    LegendStrings = {...
                        bayesoptim.infoString('Observed'),...
                        bayesoptim.infoString('ModelMean'),...
                        bayesoptim.infoString('NextPoint')};
                    CurveNums = [1,2,4];
                end
                if ~isempty(ModelMinXTable)
                    LegendStrings{end+1} = bayesoptim.infoString('ModelMin');
                    CurveNums(end+1) = CurveNums(end)+1;
                end
                legend(Axes, [Curves{CurveNums}], LegendStrings{:}, 'Location', 'best');

                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, VarSpec.Names{2});
        end
        % set axes scale and limits
        Axes = bayesoptim.setAxesScaleAndLimits(Axes, VarSpec);
    end
end
