# Copyright 2021 IRT Saint Exupéry, https://www.irt-saintexupery.com
#
# This work is licensed under a BSD 0-Clause License.
#
# Permission to use, copy, modify, and/or distribute this software
# for any purpose with or without fee is hereby granted.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL
# WARRANTIES WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL
# THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, INDIRECT,
# OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING
# FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
# NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
# WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
"""
Customize with matplotlib
=========================

In this example,
we will see how to modify the matplotlib figures
generated by an :class:`.BasePost`.
This can be useful to finely tune a graph for a presentation or a paper.
"""

from __future__ import annotations

from matplotlib import pyplot as plt

from gemseo import create_design_space
from gemseo import create_discipline
from gemseo import create_scenario

# %%
# We consider a minimization problem over the interval :math:`[0,1]`
# of the :math:`f(x)=x^2` objective function:
discipline = create_discipline("AnalyticDiscipline", expressions={"y": "x**2"})

design_space = create_design_space()
design_space.add_variable("x", lower_bound=0.0, upper_bound=1.0)

scenario = create_scenario(
    [discipline], "y", design_space, formulation_name="DisciplinaryOpt"
)

# %%
# We solve this optimization problem with the gradient-free algorithm COBYLA:
scenario.execute(algo_name="NLOPT_COBYLA", max_iter=10)

# %%
# Then,
# we can post-process this :class:`.MDOScenario` with an :class:`.OptHistoryView`.
opt_post_processor = scenario.post_process(
    post_name="OptHistoryView", show=True, save=False
)

# %%
# Instead of saving or showing this :class:`.OptHistoryView`,
# and so the associated :attr:`.OptHistoryView.figures`,
# we will slightly modify the latter which are matplotlib figures by default.
figures = opt_post_processor.figures
print(figures.keys())

# %%
# By default,
# the color bar representing the evolution of the optimization variables
# does not use labels.
# To add a custom label indicating
# that the optimization variables are scaled in :math:`[0,1]`,
# we can get the matplotlib ``Axes`` of the matplotlib figures:
figure = figures["variables"]
axes = figure.axes

# %%
# and change the y-label of the color bar:
axes[1].set_ylabel("Optimization variables scales in [0,1]")

# %%
# We can also change the y-label of the graph:
axes[0].set_ylabel("Optimization variables")

# %%
# Lastly,
# we can plot all the figures:
plt.show()

# %%
# or save the modified figure:
figure.savefig("variables.png")