function Stop = plotConstraintModels(bo, State)
%plotConstraintModels     PlotFcn for plotting Constraint-model
%surfaces.
%   Stop = plotConstraintModels(bo, State) plots each
%   Constraint-model surface, including the built-in Error
%   constraint if there were any errors. It also plots a
%   Prob(Feasible) surface.
%
% See also: BAYESOPT, BAYESIANOPTIMIZATION
%

%   Copyright 2016-2022 The MathWorks, Inc.

persistent Axes Curves NoErrorsYet;
narginchk(2,2);
Stop = false;
if ~bo.PrivOptions.FitModels || bo.NumVars > 2
    return;
end
set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
NumConstraints = bo.PrivOptions.NumCoupledConstraints;
% Order of plots is: {C1, ..., Cn, ProbFeas, Error}. Note that
% the error plot will not exist until there is an error.
switch State
    case 'initial'
        % User constraints
        for i = 1:NumConstraints
            Axes{i} = setupPlot(bo, bayesoptim.infoString('ConstraintNum', i), ...
                bayesoptim.infoString('DegreeOfViolation'), ...
                sprintf('bayesopt.Constraint%d', i));
            Curves{i} = {};
        end
        % Prob feasible
        Axes{NumConstraints+1} = setupPlot(bo, bayesoptim.infoString('ProbFeas'), ...
            bayesoptim.infoString('ProbFeas'), 'bayesopt.ProbFeas');
        Curves{NumConstraints+1} = {};
        % Error Constraint
        NoErrorsYet = true;
        if ~isempty(bo.ErrorModel)
            NoErrorsYet = false;
            Axes{NumConstraints+2} = setupPlot(bo, bayesoptim.infoString('ErrorConstraint'), ...
                bayesoptim.infoString('DegreeOfViolation'), 'bayesopt.ErrorConstraint');
            Curves{NumConstraints+2} = {};
        end
    case {'iteration', 'done'}
        % User constraints
        for i = 1:NumConstraints
            if isvalid(Axes{i})
                [Axes{i}, Curves{i}] = updateConstraintPlot(bo, i, Axes{i}, Curves{i});
            end
        end
        % Prob feasible
        if isvalid(Axes{NumConstraints+1})
            [Axes{NumConstraints+1}, Curves{NumConstraints+1}] = updateProbFeasPlot(bo, Axes{NumConstraints+1}, Curves{NumConstraints+1});
        end
        % Error Constraint
        if ~isempty(bo.ErrorModel)
            if NoErrorsYet
                NoErrorsYet = false;
                Axes{NumConstraints+2} = setupPlot(bo, bayesoptim.infoString('ErrorConstraint'), ...
                    bayesoptim.infoString('DegreeOfViolation'), 'bayesopt.ErrorConstraint');
                Curves{NumConstraints+2} = {};
            end
            if isvalid(Axes{NumConstraints+2})
                [Axes{NumConstraints+2}, Curves{NumConstraints+2}] = updateErrorPlot(...
                    bo, Axes{NumConstraints+2}, Curves{NumConstraints+2});
            end
        end
    case 'standalone'
        % User constraints
        for i = 1:NumConstraints
            StandaloneAxes{i} = setupPlot(bo, bayesoptim.infoString('ConstraintNum', i), ...
                bayesoptim.infoString('DegreeOfViolation'), ...
                sprintf('bayesopt.Constraint%d', i));
            updateConstraintPlot(bo, i, StandaloneAxes{i}, {});
        end
        % Prob feasible
        StandaloneAxes{NumConstraints+1} = setupPlot(bo, bayesoptim.infoString('ProbFeas'), ...
            bayesoptim.infoString('ProbFeas'), 'bayesopt.ProbFeas');
        updateProbFeasPlot(bo, StandaloneAxes{NumConstraints+1}, {});
        % Error Constraint
        if ~isempty(bo.ErrorModel)
            StandaloneAxes{NumConstraints+2} = setupPlot(bo, bayesoptim.infoString('ErrorConstraint'), ...
                bayesoptim.infoString('DegreeOfViolation'), 'bayesopt.ErrorConstraint');
            updateErrorPlot(bo, StandaloneAxes{NumConstraints+2}, {});
        end
end

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

    function [Axes, Curves] = updateConstraintPlot(Results, ConstraintNum, Axes, Curves)
        import bayesoptim.*
        if isempty(Results.ConstraintModels{ConstraintNum})
            return;
        end

        set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
        HoldOff = bayesoptim.holdOn(Axes);
        Options = Results.PrivOptions;
        VarSpec = Options.VarSpec;
        Model = Results.ConstraintModels{ConstraintNum};
        XTrace = Results.XTrace;
        ValTrace = Results.ConstraintsTrace(:,ConstraintNum);

        cellfun(@delete, Curves);
        Curves = {};
        switch Results.NumVars
            case 1
                XPlottable = makeXTablePlottable(Results, XTrace);
                % 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
                YPred = predict(Model, LegalizedEvalGrid);
                % (2) Plot training points on all axes
                plot(Axes, XPlottable, ValTrace, 'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot Model
                Curves{end+1} = plot(Axes, PlotGrid, YPred, 'r');
                Curves{end+1} = plot(Axes, [PlotGrid(1); PlotGrid(end)], [0;0], 'g');     % Plot line at Y=0
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                ConstraintVals = predictConstraints(Results, Results.NextPoint);
                Curves{end+1} = plot(Axes, XNextToPlot, ConstraintVals(ConstraintNum), ...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, bayesoptim.infoString('ConstraintNum', ConstraintNum));
            case 2
                XPlottable = makeXTablePlottable(Results, XTrace);
                % 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
                YPred = predict(Model, LegalizedEvalGrid);
                % (2) Plot training points
                Curves{end+1} = plot3(Axes, XPlottable(:,1), XPlottable(:,2), ValTrace, ...
                    'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot model
                Curves{end+1} = surfc(Axes, X1Plot, X2Plot, reshape(YPred, reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'r', 'FaceAlpha', .5, 'LineStyle', '-');
                Curves{end+1} = surf(Axes, X1Plot, X2Plot, zeros(reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'g', 'FaceAlpha', .5, 'LineStyle','-');     % Plot surface at Y=0
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                ConstraintVals = predictConstraints(Results, Results.NextPoint);
                Curves{end+1} = plot3(Axes, XNextToPlot(1), XNextToPlot(2), ConstraintVals(ConstraintNum),...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, VarSpec.Names{2});
                zlabel(Axes, bayesoptim.infoString('DegreeOfViolation'));
        end
        % set axes scale and limits
        Axes = bayesoptim.setAxesScaleAndLimits(Axes, VarSpec);
    end

    function [Axes, Curves] = updateErrorPlot(Results, Axes, Curves)
        import bayesoptim.*
        if isempty(Results.ErrorModel)
            return;
        end

        set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
        HoldOff = bayesoptim.holdOn(Axes);
        Options = Results.PrivOptions;
        VarSpec = Options.VarSpec;
        Model = Results.ErrorModel;
        XTrace = Results.XTrace;
        ValTrace = Results.ErrorTrace;

        cellfun(@delete, Curves);
        Curves = {};
        switch Results.NumVars
            case 1
                XPlottable = makeXTablePlottable(Results, XTrace);
                % 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
                YPred = predict(Model, LegalizedEvalGrid);
                % (2) Plot training points
                plot(Axes, XPlottable, ValTrace, 'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot Model
                Curves{end+1} = plot(Axes, PlotGrid, YPred, 'r');
                Curves{end+1} = plot(Axes, [PlotGrid(1); PlotGrid(end)], [0;0], 'g');     % Plot line at Y=0
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot(Axes, XNextToPlot, predictError(Results, Results.NextPoint), ...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, bayesoptim.infoString('DegreeOfViolation'));
            case 2
                XPlottable = makeXTablePlottable(Results, XTrace);
                % 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
                YPred = predict(Model, LegalizedEvalGrid);
                % (2) Plot training points
                Curves{end+1} = plot3(Axes, XPlottable(:,1), XPlottable(:,2), ValTrace, ...
                    'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot model
                Curves{end+1} = surfc(Axes, X1Plot, X2Plot, reshape(YPred, reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'r', 'FaceAlpha', .5, 'LineStyle', '-');
                Curves{end+1} = surf(Axes, X1Plot, X2Plot, zeros(reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'g', 'FaceAlpha', .5, 'LineStyle','-');     % Plot surface at Y=0
                % (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), ...
                    predictError(Results, Results.NextPoint),...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, VarSpec.Names{2});
                zlabel(Axes, bayesoptim.infoString('DegreeOfViolation'));
        end
        % set axes scale and limits
        Axes = bayesoptim.setAxesScaleAndLimits(Axes, VarSpec);
    end

    function [Axes, Curves] = updateProbFeasPlot(Results, Axes, Curves)
        import bayesoptim.*
        set(0, 'defaultfigureunits', 'pixels');     % Make sure units are pixels
        HoldOff = bayesoptim.holdOn(Axes);
        Options = Results.PrivOptions;
        VarSpec = Options.VarSpec;
        XTrace = Results.XTrace;
        ValTrace = Results.FeasibilityProbabilityTrace;

        cellfun(@delete, Curves);
        Curves = {};
        switch Results.NumVars
            case 1
                XPlottable = makeXTablePlottable(Results, XTrace);
                % 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
                YPred = ProbAllConstraintsSatisfied(Results, LegalizedEvalGrid);
                % (2) Plot training points on all axes
                Curves{end+1} = plot(Axes, XPlottable, ValTrace, 'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot Model
                Curves{end+1} = plot(Axes, PlotGrid, YPred, 'r');
                % (6) Plot next point
                XNextToPlot = makeXTablePlottable(Results, Results.NextPoint);
                MarkerSpec = {'o', 'MarkerSize', 8, 'Color', 'black', ...
                    'MarkerFaceColor', 'black'};
                Curves{end+1} = plot(Axes, XNextToPlot, ...
                    ProbAllConstraintsSatisfied(Results, Results.XNext), ...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, bayesoptim.infoString('ProbFeas'));
            case 2
                XPlottable = makeXTablePlottable(Results, XTrace);
                % 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
                YPred = ProbAllConstraintsSatisfied(Results, LegalizedEvalGrid);
                % (2) Plot training points
                Curves{end+1} = plot3(Axes, XPlottable(:,1), XPlottable(:,2), ValTrace, ...
                    'bo', 'MarkerFaceColor', 'blue');
                % (3) Plot model
                Curves{end+1} = surfc(Axes, X1Plot, X2Plot, reshape(YPred, reverse(Options.NumPlotGrid)), ...
                    'FaceColor', 'r', 'FaceAlpha', .5, 'LineStyle', '-');
                % (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), ...
                    ProbAllConstraintsSatisfied(Results, Results.XNext),...
                    MarkerSpec{:});
                % (7) Label axes
                xlabel(Axes, VarSpec.Names{1});
                ylabel(Axes, VarSpec.Names{2});
                zlabel(Axes, bayesoptim.infoString('ProbFeas'));
        end
        % set axes scale and limits
        Axes = bayesoptim.setAxesScaleAndLimits(Axes, VarSpec);
    end
end
