# 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.
# Contributors:
#    INITIAL AUTHORS - API and implementation and/or documentation
#        :author: Gilberto Ruiz
#    OTHER AUTHORS   - MACROSCOPIC CHANGES
"""
Solve a mixed optimization problem.
===================================
.. _mixed_optimization_scenario:
"""

# %%
# Introduction
# ------------
# In this example,
# we will solve a very simple mixed optimization problem with a full factorial approach.
# To do so, we will split the problem into a discrete optimization problem
# to enumerate all the combinations
# and a continuous one to solve the associated sub-problem.
# Thus, we will use the
# :class:`.MDOScenarioAdapter` to wrap a :class:`.BaseScenario` and treat it as a
# discipline whose inputs are some or all of its design space variables and whose
# outputs are some or all of its functions (objective, constraints or observables).
# Keep in mind that this approach may be very time-consuming.
# It works well when the dimension of the integers to
# explore is not too large. Otherwise, a dedicated algorithm may be better suited, such
# as the ones available in the ``gemseo-pymoo`` plugin.

# %%
# Imports
# -------
# All the imports needed for the tutorial are performed here.
from __future__ import annotations

from numpy import ndarray  # noqa: TC002
from numpy.linalg import norm

from gemseo import configure_logger
from gemseo import create_design_space
from gemseo import create_discipline
from gemseo import create_scenario
from gemseo.disciplines.scenario_adapters.mdo_scenario_adapter import MDOScenarioAdapter
from gemseo.settings.doe import PYDOE_FULLFACT_Settings
from gemseo.settings.formulations import DisciplinaryOpt_Settings
from gemseo.settings.opt import NLOPT_COBYLA_Settings
from gemseo.settings.post import OptHistoryView_Settings

configure_logger()


# %%
# Optimization problem definition.
# --------------------------------
# We define the following optimization problem:
#
# .. math::
#
#    \begin{aligned}
#    \text{minimize the objective function }&\text{f(x,y)}=|x| + |y| \\
#    \text{with respect to the design variables }&x,\,y \\
#    \text{subject to the general constraint }
#    & g(x,y) \geq 2\\
#    \text{subject to the bound constraints }
#    & 0.0 \leq x \leq 1.0\\
#    & 0 \leq y_0 \leq 1\\
#    & 0 \leq y_2 \leq 2
#    \end{aligned}
#
# and where the general constraint is:
#
# .. math::
#
#    g(x,y) = x + y
#
# Where :math:`y` is an integer vector with two components and :math:`x` is a float
# vector with two components.

# %%
# Optimization problem reformulation.
# -----------------------------------
# The problem can be split using the :class:`.MDOScenarioAdapter`. To do this we will
# divide the design space in two, a continuous one and a discrete one.
# The :class:`.MDOScenarioAdapter` will wrap the continuous inner scenario as a
# discipline to be executed taking the inputs from the discrete design space.
# These inputs are generated by the outer :class:`.DOEScenario` using a full factorial
# method.
# It is possible of course to use different DOE algorithms or even generate the samples
# first, filter them, and then launch a class:`.CustomDOE` with the filtered samples.
#
# The reformulated optimization problem would read as follows:
# For the outer DOE Scenario:
#
# .. math::
#
#    \begin{aligned}
#    \text{minimize the objective function }&\text{f(x,y)}=|x| + |y| \\
#    \text{with respect to the design variables }&y \\
#    \text{subject to the general constraint }
#    & g(x,y) \geq 2\\
#    \text{subject to the bound constraints }
#    & 0 \leq y_0 \leq 1\\
#    & 0 \leq y_2 \leq 2
#    \end{aligned}
#
# For the inner MDO Scenario:
#
# .. math::
#
#    \begin{aligned}
#    \text{minimize the objective function }&\text{f(x,y)}=|x| + |y| \\
#    \text{with respect to the design variables }&x \\
#    \text{subject to the general constraint }
#    & g(x,y) \geq 2\\
#    \text{subject to the bound constraints }
#    & 0.0 \leq x \leq 1.0
#    \end{aligned}
#
# In the next steps, we will build both scenarios and connect them to solve the full
# problem.


# %%
# Create an :class:`.AutoPyDiscipline` to compute the objective and the constraints.
# ----------------------------------------------------------------------------------
# Since the expressions of our toy problem are very simple, we can use an
# :class:`.AutoPyDiscipline` to compute the objective and constraints.
# Note that there are no strong couplings in our expressions, which means we could also
# compute both the objective and constraints with a single discipline if we wished to.
def obj(x: ndarray, y: ndarray) -> float:
    """A simple Python function to compute f(x,y)."""
    f = norm(x) + norm(y)
    return f  # noqa: RET504


def const(x: ndarray, y: ndarray) -> ndarray:
    """A simple Python function to compute g(x,y)."""
    g = x + y
    return g  # noqa: RET504


objective = create_discipline("AutoPyDiscipline", name="f(x,y)", py_func=obj)
constraint = create_discipline("AutoPyDiscipline", name="g(x,y)", py_func=const)

# %%
# Create the design space for the entire problem.
# -----------------------------------------------
# We can define a :class:`.DesignSpace` for the whole problem and then filter
# either the continuous variables or the discrete ones.
design_space = create_design_space()
design_space.add_variable("x", lower_bound=0, upper_bound=1, value=1.0, size=2)
design_space.add_variable(
    "y", lower_bound=[0, 0], upper_bound=[1, 2], value=1, size=2, type_="integer"
)

# %%
# Create the design space for the inner scenario.
# -----------------------------------------------
# The inner scenario is the one that solves the continuous optimization problem, and as
# such, it only needs to include the continuous design variables. We use the
# :meth:`.DesignSpace.filter` method to keep ``x`` and we set ``copy`` to ``True`` to
# keep the original ``design_space`` unchanged, as we will use it later.
design_space_inner_scenario = design_space.filter(keep_variables=["x"], copy=True)

# %%
# Create the inner MDO scenario.
# ------------------------------
inner_scenario = create_scenario(
    [objective, constraint],
    "f",
    design_space_inner_scenario,
    formulation_settings_model=DisciplinaryOpt_Settings(),
)
inner_scenario.set_algorithm(NLOPT_COBYLA_Settings(max_iter=100))
inner_scenario.add_constraint("g", constraint_type="ineq", value=2)

# %%
# Create the scenario adapter.
# ----------------------------
# An :class:`.MDOScenarioAdapter` wraps an entire :class:`.BaseScenario` as a
# :class:`.Discipline`, its inputs are all or
# part of the design space variables and its outputs are all or part of the objective
# values, constraints or observables.
# Here we select the variables of the inner scenario that we wish to set as
# inputs/outputs for the adapter.
input_names = ["y"]
output_names = ["f", "g"]

# %%
# The argument ``set_x0_before_opt`` allows us to set the starting point of the adapted
# scenario from the outer DOE scenario values.
adapted_inner_scenario = MDOScenarioAdapter(
    inner_scenario,
    input_names,
    output_names,
    set_x0_before_opt=True,
)
# %%
#
# .. tip::
#
#    You may be interested in keeping the optimization history of the inner scenario for
#    each of the executions launched by the outer scenario. To do this, set the argument
#    ``keep_opt_history`` to ``True``, this option will store the databases in memory
#    and make them accessible via the :attr:`.MDOScenarioAdapter.databases` attribute.
#    Keep in mind that depending on the size of the database, storing it in memory may
#    lead to a significant increase in memory usage.
#    If you prefer to store the databases on disk instead, set the argument
#    ``save_opt_history`` to ``True``. An ``hdf5`` file will be saved on the disk at
#    each new execution. You may also choose a prefix for the name of these files with
#    the argument ``opt_history_file_prefix``. If no prefix is given, the default
#    prefix is ``"database"``.
#    Both ``keep_opt_history`` and ``save_opt_history`` are independent of
#    each-other.

# %%
# Create the design space for the outer DOE scenario.
# ---------------------------------------------------
# The outer scenario is the one that solves the discrete optimization problem, and as
# such, it only needs to include the integer design variables. Once again, we use the
# :meth:`.DesignSpace.filter` method to keep ``y``, the ``copy`` argument ensures that
# the original ``design_space`` remains unchanged in case you need it for other
# purposes.
design_space_outer_scenario = design_space.filter(keep_variables="y", copy=True)

# %%
# Create the outer DOE scenario.
# ------------------------------
outer_scenario = create_scenario(
    adapted_inner_scenario,
    "f",
    design_space_outer_scenario,
    formulation_settings_model=DisciplinaryOpt_Settings(),
    scenario_type="DOE",
)

# %%
# Here, we add the constraints on the outer scenario in order to be able to know if a
# given set of integers returns a feasible solution once the inner scenario has been
# executed.
outer_scenario.add_constraint("g", constraint_type="ineq", value=2)

# %%
# Show an xDSM of the process, including the outer and inner scenarios.
outer_scenario.xdsmize(save_html=False)

# %%
# Execute the outer scenario (which contains the inner scenario) and solve the whole
# problem.
# The console will show the progress of the optimization. For each DOE point it will
# show the optimization of the continuous problem and its optimal result.
outer_scenario.execute(PYDOE_FULLFACT_Settings(n_samples=9))

# %%
#
# .. tip::
#
#    In a :class:`.DOEScenario`, we know a priori the samples that will be evaluated.
#    This means we can run the outer scenario in parallel if we set the setting
#    ``n_processes`` to at least 2. Note that if you are running the outer scenario in
#    parallel and requesting the databases of the continuous optimizations on the disk,
#    you will need to instantiate the :class:`.MDOScenarioAdapter` with the argument
#    ``naming="UUID"``, which is multiprocessing-safe.
#    Running in parallel also means that the option ``keep_opt_history`` will not work
#    because we are unable to copy the databases from the sub-processes to the main
#    process.

# %%
# Plot the objective and constraint history for the scenario.
# -----------------------------------------------------------
# At the end of the optimization we see the results of the problem. The optimal
# solution would be in this case the iteration of the DOE that gave the minimum for
# :math:`f(x,y)` while respecting the constraint :math;`g(x,y)`. The full problem
# optimal result will contain the values for :math:`x`, :math:`y`, :math:`f`, and
# :math:`g`.
outer_scenario.post_process(OptHistoryView_Settings(save=False, show=True))