# -*- coding: utf-8 -*-
# Copyright 2021 IRT Saint Exupéry, https://www.irt-saintexupery.com
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License version 3 as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with this program; if not, write to the Free Software Foundation,
# Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

# Contributors:
#    INITIAL AUTHORS - API and implementation and/or documentation
#        :author: Francois Gallard
#        :author: Damien Guenot
#        :author: Charlie Vanaret
#    OTHER AUTHORS   - MACROSCOPIC CHANGES
"""Basic display of optimization history: functions and x."""
from __future__ import division, unicode_literals

import logging
from typing import Iterable, List, Optional, Sequence, Tuple

import matplotlib.gridspec as gridspec
import pylab
from matplotlib.colors import SymLogNorm
from matplotlib.figure import Figure
from matplotlib.ticker import LogFormatter, MaxNLocator
from numpy import abs as np_abs
from numpy import append, arange, argmin, array, atleast_2d, concatenate
from numpy import e as npe
from numpy import hstack, isnan
from numpy import log10 as np_log10
from numpy import logspace
from numpy import max as np_max
from numpy import min as np_min
from numpy import ndarray, ones, ones_like
from numpy import sort as np_sort
from numpy import vstack, where
from numpy.linalg import norm

from gemseo.algos.opt_problem import OptimizationProblem
from gemseo.core.function import MDOFunction
from gemseo.post.core.colormaps import PARULA, RG_SEISMIC
from gemseo.post.core.hessians import SR1Approx
from gemseo.post.opt_post_processor import OptPostProcessor
from gemseo.utils.py23_compat import PY2

LOGGER = logging.getLogger(__name__)


class OptHistoryView(OptPostProcessor):
    """The **OptHistoryView** post processing performs separated plots: the design
    variables history, the objective function history, the history of hessian
    approximation of the objective, the inequality constraint history, the equality
    constraint history, and constraints histories.

    By default, all design variables are considered. A sublist of design variables can
    be passed as options. Minimum and maximum values for the plot can be passed as
    options. The objective function can also be represented in terms of difference
    w.r.t. the initial value It is possible either to save the plot, to show the plot or
    both.
    """

    def __init__(
        self,
        opt_problem,  # type: OptimizationProblem
    ):  # type: (...) -> None
        super(OptHistoryView, self).__init__(opt_problem)
        self.cmap = PARULA  # "viridis"  # "jet"
        self.ineq_cstr_cmap = RG_SEISMIC  # "seismic" "PRGn_r"
        self.eq_cstr_cmap = "seismic"  # "seismic" "PRGn_r"

    def _plot(
        self,
        variables_names=None,  # type: Optional[Sequence[str]]
        obj_min=None,  # type: Optional[float]
        obj_max=None,  # type: Optional[float]
        obj_relative=False,  # type: bool
    ):  # type: (...) -> None
        """
        Args:
            variables_names: The names of the variables to display.
                If None, use all design variables.
            obj_max: The maximum value for the objective in the plot.
                If None, use the maximum value of the objective history.
            obj_min: The minimum value for the objective in the plot.
                If None, use the minimum value of the objective history.
            obj_relative: If True, plot the objective value difference
                with the initial value.
        """
        # compute the names of the inequality and equality constraints
        ineq_cstr = self.opt_problem.get_ineq_constraints()
        ineq_cstr_names = [c.name for c in ineq_cstr]
        eq_cstr = self.opt_problem.get_eq_constraints()
        eq_cstr_names = [c.name for c in eq_cstr]

        obj_name = self.opt_problem.get_objective_name()
        obj_history, x_history, n_iter = self._get_history(obj_name, variables_names)

        # design variables
        self._create_variables_plot(x_history, variables_names)

        # objective function
        self._create_obj_plot(
            obj_history,
            n_iter,
            obj_min,
            obj_max,
            obj_relative=obj_relative,
        )

        self._create_x_star_plot(x_history, n_iter)

        # Hessian plot
        plot_hessian = not self.database.is_func_grad_history_empty(obj_name)
        if plot_hessian:
            self._create_hessian_approx_plot(self.database, obj_name)

        # inequality and equality constraints
        self._plot_cstr_history(ineq_cstr_names, MDOFunction.TYPE_INEQ)
        self._plot_cstr_history(eq_cstr_names, MDOFunction.TYPE_EQ)

    def _plot_cstr_history(
        self,
        cstr_names,  # type: Sequence[str]
        cstr_type,  # type: str
    ):  # type: (...) -> None
        """Create the plot for (in)equality constraints.

        Args:
            cstr_names: The names of the constraints.
            cstr_type: The type of the constraints, either 'eq' or 'ineq'.
        """
        if cstr_names is not None:
            _, constraints_history = self._get_constraints(cstr_names)
            self._create_cstr_plot(
                constraints_history,
                cstr_type,
                cstr_names,
            )

    def _get_history(
        self,
        function_name,  # type: str
        variables_names,  # type: Sequence[str]
    ):  # type: (...) -> Tuple[ndarray,ndarray,int]
        """Access the optimization history of a function and the design variables at
        which it was computed.

        Args:
            function_name: The name of the function.
            variables_names: The names of the variables to display.

        Returns:
            The function values,
            the design variables values
            and the number of iterations.
        """
        f_hist, x_hist = self.database.get_func_history(function_name, x_hist=True)
        f_hist = array(f_hist).real
        x_hist = array(x_hist).real

        if variables_names is not None:
            # select only the interesting columns
            block_list = []
            column = 0
            for var in self.opt_problem.design_space.variables_names:
                if var in variables_names:
                    size = self.opt_problem.design_space.variables_sizes[var]
                    block_list.append(x_hist[:, column : column + size])
                    column += size
            # concatenate the blocks
            x_hist = hstack(block_list)
        n_iter = x_hist.shape[0]
        return f_hist, x_hist, n_iter

    def _get_constraints(
        self,
        cstr_names,  # type: Iterable[str]
    ):  # type: (...) -> Tuple[ndarray,List[ndarray]]
        """Extract a history of constraints.

        Args:
            cstr_names: The names of the constraints.

        Returns:
            The bounds of the constraints and history array.
        """
        available_data_names = self.database.get_all_data_names()
        for cstr_name in cstr_names:
            if cstr_name not in available_data_names:
                cstr_names.remove(cstr_name)

        constraints_history = []
        bnd_list = -1e300 * ones(len(cstr_names))

        for cstr_i, cstr_name in enumerate(cstr_names):

            cstr_history = array(self.database.get_func_history(cstr_name)).real

            constraints_history.append(cstr_history)
            bnd_list[cstr_i] = max(
                bnd_list[cstr_i],
                max(abs(np_min(cstr_history)), abs(np_max(cstr_history))),
            )

        return bnd_list, constraints_history

    def _normalize_x_hist(
        self,
        x_history,  # type: ndarray
        variables_names,  # type: Sequence[str]
    ):  # type: (...) -> ndarray
        """Normalize the design variables history.

        Args:
            x_history: The history for the design variables.
            variables_names: The names of the variables to display.

        Returns:
            The normalized design variables array.
        """
        x_hist_n = x_history.copy()
        lower_bounds = self.opt_problem.design_space.get_lower_bounds(variables_names)
        upper_bounds = self.opt_problem.design_space.get_upper_bounds(variables_names)
        norm_coeff = 1 / (np_abs(upper_bounds - lower_bounds))
        for i in range(x_history.shape[0]):
            x_hist_n[i, :] = (x_hist_n[i, :] - lower_bounds) * norm_coeff
        return x_hist_n

    def _create_variables_plot(
        self,
        x_history,  # type: ndarray
        variables_names,  # type: Sequence[str]
    ):  # type: (...) -> None
        """Create the design variables plot.

        Args:
             x_history: The history for the design variables.
             variables_names: The names of the variables to display.
        """
        if len(x_history) < 2:
            return
        n_variables = x_history.shape[1]
        norm_x_history = self._normalize_x_hist(x_history, variables_names)

        fig = pylab.plt.figure(figsize=(11, 6))
        grid = gridspec.GridSpec(1, 2, width_ratios=[15, 1], wspace=0.04, hspace=0.6)

        # design variables
        ax1 = fig.add_subplot(grid[0, 0])
        im1 = ax1.imshow(
            norm_x_history.T,
            cmap=self.cmap,
            interpolation="nearest",
            vmin=0.0,
            vmax=1.0,
            aspect="auto",
        )
        # y ticks: names of the variables
        design_space = self.opt_problem.design_space
        y_labels = []
        if variables_names is None:
            variables_names = design_space.variables_names
        for variable_name in variables_names:
            size = design_space.variables_sizes[variable_name]
            name = variable_name
            if size > 1:
                name += " (0)"
            y_labels.append(name)
            for i in range(1, size):
                y_labels.append("({})".format(i))

        ax1.set_yticks(arange(n_variables))
        ax1.set_yticklabels(y_labels)
        # ax1.invert_yaxis()

        ax1.set_title("Evolution of the optimization variables")
        ax1.xaxis.set_major_locator(MaxNLocator(integer=True))

        # colorbar
        ax2 = fig.add_subplot(grid[0, 1])
        fig.colorbar(im1, cax=ax2)

        # Set window size
        mng = pylab.plt.get_current_fig_manager()
        mng.resize(700, 1000)

        self._add_figure(fig, "variables")

    def _create_obj_plot(
        self,
        obj_history,  # type: ndarray
        n_iter,  # type: int
        obj_min=None,  # type: Optional[float]
        obj_max=None,  # type: Optional[float]
        obj_relative=False,  # type: bool
    ):  # type: (...) -> None
        """Creates the design variables plot.

        Args:
            obj_history: The history of the objective function.
            n_iter: The number of iterations.
            obj_max: The maximum value for the objective in the plot.
                If None, use the maximum value of the objective history.
            obj_min: The minimum value for the objective in the plot.
                If None, use the minimum value of the objective history.
            obj_relative: If True, plot the objective value difference
                with the initial value.
        """
        # if max problem, plot -objective value
        if not self.opt_problem.minimize_objective:
            obj_history = -obj_history
        if obj_relative:
            LOGGER.info(
                "Plot of optimization history "
                "with relative variation compared to "
                "initial point objective value = %s",
                obj_history[0],
            )
            obj_history -= obj_history[0]
        # Remove nans
        n_x = len(obj_history)
        x_absc = arange(n_x)
        x_absc_nan = None
        idx_nan = isnan(obj_history)

        if idx_nan.size > 0:
            obj_history = obj_history[~idx_nan]
            x_absc_nan = x_absc[idx_nan]
            x_absc = x_absc[~idx_nan]

        fmin = np_min(obj_history)
        fmax = np_max(obj_history)

        fig = pylab.plt.figure(figsize=(11, 6))
        # objective function
        pylab.plt.xlabel("Iterations", fontsize=12)
        pylab.plt.ylabel("Objective value", fontsize=12)

        pylab.plt.plot(x_absc, obj_history)

        if idx_nan.size > 0:
            for x_i in x_absc_nan:
                pylab.plt.axvline(x_i, color="purple")

        if obj_min is not None and obj_min < fmin:
            fmin = obj_min
        if obj_max is not None and obj_max > fmax:
            fmax = obj_max
        pylab.plt.ylim([fmin, fmax])
        pylab.plt.xlim([0, n_iter])
        ax1 = fig.gca()
        ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
        pylab.plt.grid(True)
        pylab.plt.title("Evolution of the objective value")

        # Set window size
        mng = pylab.plt.get_current_fig_manager()
        mng.resize(700, 1000)

        self._add_figure(fig, "objective")

    def _create_x_star_plot(
        self,
        x_history,  # type: ndarray
        n_iter,  # type:int
    ):
        """Create the design variables plot.

        Args:
            x_history: The history of the design variables.
            n_iter: The number of iterations.
        """
        fig = pylab.plt.figure(figsize=(11, 6))
        # objective function
        pylab.plt.xlabel("Iterations", fontsize=12)
        pylab.plt.ylabel("||x-x*||", fontsize=12)

        n_i = x_history.shape[0]
        _, x_opt, _, _, _ = self.opt_problem.get_optimum()

        normalize = self.opt_problem.design_space.normalize_vect

        x_xstar = [
            norm(normalize(x_history[i, :]) - normalize(x_opt)) for i in range(n_i)
        ]
        # Draw a vertical line at the optimum
        ind_opt = argmin(x_xstar)
        pylab.plt.axvline(x=ind_opt, color="r")
        pylab.plt.semilogy(arange(len(x_xstar)), x_xstar)
        # ======================================================================
        # try:
        #     pylab.plt.semilogy(np.arange(len(x_xstar)), x_xstar)
        # except ValueError:
        #     LOGGER.warning("Cannot use log scale for x_star plot since" +
        #                    "all values are not positive !")
        # ======================================================================
        ax1 = fig.gca()
        ax1.xaxis.set_major_locator(MaxNLocator(integer=True))
        pylab.plt.grid(True)
        pylab.plt.title("Distance to the optimum")
        pylab.plt.xlim([0, n_iter])

        # Set window size
        mng = pylab.plt.get_current_fig_manager()
        mng.resize(700, 1000)

        self._add_figure(fig, "x_xstar")

    @staticmethod
    def _cstr_number(
        cstr_names,  # type: Sequence[str]
        cstr_history,  # type: ndarray
    ):  # type: (...) -> int
        """Compute the total scalar constraints number.

        Args:
            cstr_names: The names of the constraints.
            cstr_history: The history of the constraints.

        Returns:
            The number of constraints.
        """
        n_cstr = 0
        for cstr_i in range(len(cstr_names)):
            c_hist_loc = atleast_2d(cstr_history[cstr_i]).T
            if c_hist_loc.shape[1] == 1:
                c_hist_loc = c_hist_loc.T
            n_cstr += c_hist_loc.shape[0]
        LOGGER.debug("Total constraints number =%s", n_cstr)
        return n_cstr

    def _create_cstr_plot(
        self,
        cstr_history,  # type: ndarray
        cstr_type,  # type: str
        cstr_names,  # type: Sequence[str]
    ):  # type: (...) -> None
        """Create the constraints plot: 1 line per constraint component.

        Args:
            cstr_history: The history of the constraints.
            cstr_type: The type of the constraints, either 'eq' or 'ineq'.
            cstr_names: The names of the constraints.
        """
        n_cstr = self._cstr_number(cstr_names, cstr_history)
        if n_cstr == 0:
            return

        # matrix of all constraints' values
        cstr_matrix = None
        vmax = 0.0
        cstr_labels = []

        max_iter = 0
        for (_, cstr_history_i) in enumerate(cstr_history):
            history_i = atleast_2d(cstr_history_i).T
            if history_i.shape[1] == 1:
                history_i = history_i.T

            if max_iter < history_i.shape[1]:
                max_iter = history_i.shape[1]

        for (i, cstr_history_i) in enumerate(cstr_history):
            history_i = atleast_2d(cstr_history_i).T
            if history_i.shape[1] == 1:
                history_i = history_i.T

            nb_components = history_i.shape[0]

            if history_i.shape[1] == max_iter:  # TEST
                for component_j in range(nb_components):
                    # compute the label of the constraint
                    if component_j == 0:
                        cstr_name = cstr_names[i]
                        if nb_components >= 2:
                            cstr_name += " (" + str(component_j + 1) + ")"
                        cstr_labels.append(cstr_name)
                    else:
                        cstr_labels.append("(" + str(component_j + 1) + ")")

                    history_i_j = atleast_2d(history_i[component_j, :])

                    # max value
                    notnans = ~isnan(history_i_j)
                    vmax = max(vmax, np_max(np_abs(history_i_j[notnans])))

                    # build the constraint matrix
                    if cstr_matrix is None:
                        cstr_matrix = history_i_j
                    else:
                        cstr_matrix = vstack((cstr_matrix, history_i_j))

        fig = self._build_cstr_fig(cstr_matrix, cstr_type, vmax, n_cstr, cstr_labels)

        self._add_figure(fig, "{}_constraints".format(cstr_type))

    def _build_cstr_fig(
        self,
        cstr_matrix,  # type: ndarray
        cstr_type,  # type: str
        vmax,  # type: float
        n_cstr,  # type: int
        cstr_labels,  # type: Sequence[str]
    ):  # type: (...) -> Figure
        """Build the constraints figure.

        Args:
            cstr_matrix: The matrix of constraints values.
            cstr_type: The type of the constraints, either 'eq' or 'ineq'.
            cstr_labels: The labels for the constraints.
            vmax: The maximum constraint absolute value.
            n_cstr: The number of constraints.
            cstr_labels: The labels of constraints names.

        Returns:
            The constraints figure.
        """
        # cmap of the constraints
        fullname = "equality"
        if cstr_type == MDOFunction.TYPE_EQ:
            cmap = self.eq_cstr_cmap
        else:
            cmap = self.ineq_cstr_cmap
            fullname = "in" + fullname

        idx_nan = isnan(cstr_matrix)
        hasnan = idx_nan.any()
        if hasnan > 0:
            cstr_matrix[idx_nan] = 0.0

        # generation of the image
        fig = pylab.plt.figure(figsize=(11, 6))
        grid = gridspec.GridSpec(1, 2, width_ratios=[15, 1], wspace=0.04, hspace=0.6)
        ax1 = fig.add_subplot(grid[0, 0])
        if PY2:
            im1 = ax1.imshow(
                cstr_matrix,
                cmap=cmap,
                interpolation="nearest",
                aspect="auto",
                norm=SymLogNorm(linthresh=1.0, vmin=-vmax, vmax=vmax),
            )
        else:
            im1 = ax1.imshow(
                cstr_matrix,
                cmap=cmap,
                interpolation="nearest",
                aspect="auto",
                norm=SymLogNorm(linthresh=1.0, vmin=-vmax, vmax=vmax, base=npe),
            )
        if hasnan > 0:
            x_absc_nan = where(idx_nan.any(axis=0))[0]
            for x_i in x_absc_nan:
                pylab.plt.axvline(x_i, color="purple")

        ax1.tick_params(labelsize=9)
        ax1.set_yticks(list(range(n_cstr)))
        ax1.set_yticklabels(cstr_labels)

        ax1.set_xlabel("evaluations")
        ax1.set_title("Evolution of the " + fullname + " constraints")
        ax1.xaxis.set_major_locator(MaxNLocator(integer=True))

        ax1.hlines(
            list(range(len(cstr_matrix))),
            [-0.5],
            [len(cstr_matrix[0]) - 0.5],
            alpha=0.1,
            lw=0.5,
        )

        # color map
        cax = fig.add_subplot(grid[0, 1])
        if 0.0 < vmax < 1.0:
            thick_min = int(np_log10(vmax))
        else:
            thick_min = 0
        if vmax > 1.0:
            thick_max = int(np_log10(vmax))
        else:
            thick_max = 0
        thick_num = thick_max - thick_min + 1
        levels_pos = logspace(thick_min, thick_max, num=thick_num, base=10.0)
        if vmax != 0.0:
            levels_pos = np_sort(append(levels_pos, vmax))
        levels_neg = np_sort(-levels_pos)
        levels_neg = append(levels_neg, 0)
        levels = concatenate((levels_neg, levels_pos))
        col_bar = fig.colorbar(im1, cax=cax, ticks=levels)
        col_bar.ax.tick_params(labelsize=9)
        cax.set_xlabel("symlog")

        mng = pylab.plt.get_current_fig_manager()
        # mng.full_screen_toggle()
        mng.resize(700, 1000)
        return fig

    def _create_hessian_approx_plot(
        self,
        history,  # type: ndarray
        obj_name,  # type: str
    ):  # type: (...) -> None
        """Create the plot of the Hessian approximation.

        Args:
            history: The optimization history.
            obj_name: The objective function name.
        """
        try:
            approximator = SR1Approx(history)
            _, diag, _, _ = approximator.build_approximation(
                funcname=obj_name, save_diag=True
            )
            diag = [ones_like(diag[0])] + diag  # Add first iteration blank
            diag = array(diag).T
        except ValueError:
            LOGGER.warning("Failed to create Hessian approximation", exc_info=True)
            return

        # if max problem, plot -Hessian
        if not self.opt_problem.minimize_objective:
            diag = -diag

        fig = pylab.plt.figure(figsize=(11, 6))
        grid = gridspec.GridSpec(1, 2, width_ratios=[15, 1], wspace=0.04, hspace=0.6)
        # matrix
        axe = fig.add_subplot(grid[0, 0])

        axe.set_title("Hessian diagonal approximation")
        axe.set_xlabel("Iterations", fontsize=12)
        axe.set_ylabel("Variable id", fontsize=12)
        axe.xaxis.set_major_locator(MaxNLocator(integer=True))
        vmax = max(abs(np_max(diag)), abs(np_min(diag)))
        linthresh = 10 ** (np_log10(vmax) - 5.0)
        if PY2:
            img = axe.imshow(
                diag.real,
                cmap=self.cmap,
                interpolation="nearest",
                aspect="auto",
                norm=SymLogNorm(linthresh=linthresh, vmin=-vmax, vmax=vmax),
            )
        else:
            img = axe.imshow(
                diag.real,
                cmap=self.cmap,
                interpolation="nearest",
                aspect="auto",
                norm=SymLogNorm(linthresh=linthresh, vmin=-vmax, vmax=vmax, base=npe),
            )
        axe.invert_yaxis()

        # colorbar
        vmax = max(abs(np_max(diag)), abs(np_min(diag)))
        thick_min = int(np_log10(linthresh))
        thick_max = int(np_log10(vmax))
        thick_num = thick_max - thick_min + 1
        levels_pos = logspace(thick_min, thick_max, num=thick_num, base=10.0)
        levels_pos = append(levels_pos, vmax)
        levels_neg = np_sort(-levels_pos)
        levels_neg = append(levels_neg, 0)
        levels = concatenate((levels_neg, levels_pos))

        l_f = LogFormatter(base=10, labelOnlyBase=False)
        cax = fig.add_subplot(grid[0, 1])
        cax.invert_yaxis()
        fig.colorbar(img, cax=cax, ticks=levels, format=l_f)

        mng = pylab.plt.get_current_fig_manager()
        # mng.full_screen_toggle()
        mng.resize(700, 1000)
        self._add_figure(fig, "hessian_approximation")
