"""
This module contains classes and methods to construct plots with data
coming from PELE reports.
"""
import os

from pele_platform.Utilities.Helpers.helpers import check_make_folder
from pele_platform.analysis import DataHandler
from pele_platform.constants import constants


class Plotter(object):
    """
    It handles the plots.
    """

    def __init__(self, dataframe, logger=None):
        """
        It initializes the plotter with the dataframe that contains the
        data to plot.

        Parameters
        ----------
        dataframe : a pandas.DataFrame object
            The dataframe containing the information from PELE reports
        logger : a Logger object
            The logger manager to stream out the log messages. Default is
            None
        """
        self._dataframe = dataframe
        self._logger = logger

    def plot_two_metrics(self, metric_to_x, metric_to_y, metric_to_z=None,
                         output_name=None, output_folder="."):
        """
        Given 2 or 3 metrics, it generates the scatter plot. In case that
        a 3rd metric is supplied, it will be represented as the color bar.

        Parameters
        ----------
        metric_to_x : str or int
            The metric id to plot in the X axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_y : str or int
            The metric id to plot in the Y axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_z : str or int
            The metric id to plot in the color bar. It can be either a string
            with the name of the metric or an integer with the column index.
            Default is None
        output_name : str
            The name that will be given to the resulting plot. Default is None
        output_folder : str
            The path where the plot will be saved. Default is '.', so it
            will be stored in the local directory
        """
        from pele_platform.Utilities.Helpers.helpers import backup_logger

        check_make_folder(output_folder)

        # Initialize a data handler from the current dataframe and get column names
        metric_to_x, metric_to_y, metric_to_z = self._get_column_names(
            metric_to_x, metric_to_y, metric_to_z)

        # Prepare plot name
        if output_name is None:
            if metric_to_z is not None:
                output_name = "{}_{}_{}_plot.png".format(metric_to_x,
                                                         metric_to_y,
                                                         metric_to_z)
            else:
                output_name = "{}_{}_plot.png".format(metric_to_x, metric_to_y)

        # Replace whitespaces in the output name
        output_name = output_name.replace(" ", "_")
        output_name = os.path.join(output_folder, output_name)

        # Generate plot with matplotlib
        import matplotlib

        matplotlib.use("Agg")
        from matplotlib import pyplot as plt

        fig, ax = plt.subplots()

        if metric_to_z is not None:
            scatter = ax.scatter(self._dataframe[metric_to_x],
                                 self._dataframe[metric_to_y],
                                 c=self._dataframe[metric_to_z],
                                 s=20)
            cbar = plt.colorbar(scatter)
            cbar.ax.set_ylabel(metric_to_z)
            ax.set_xlabel(metric_to_x)
            ax.set_ylabel(metric_to_y)
            plt.savefig(output_name)
            backup_logger(self._logger,
                          "Plotted {} vs {} vs {}".format(metric_to_x,
                                                          metric_to_y,
                                                          metric_to_z))

        else:
            ax.scatter(self._dataframe[metric_to_x],
                       self._dataframe[metric_to_y])
            ax.set_xlabel(metric_to_x)
            ax.set_ylabel(metric_to_y)
            plt.savefig(output_name)
            backup_logger(self._logger,
                          "Plotted {} vs {}".format(metric_to_x,
                                                    metric_to_y))

        plt.close("all")
        return output_name

    def plot_kde(self, metric_to_x, metric_to_y, output_folder, kde_structs):
        """
        Given 2 metrics, it generates the kde plot.

        Parameters
        ----------
        metric_to_x : str or int
            The metric id to plot in the X axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_y : str or int
            The metric id to plot in the Y axis. It can be either a string
            with the name of the metric or an integer with the column index
        output_folder : str
            The path where the plot will be saved
        kde_structs : int
            The number of structures to represent in the plot
        """
        import seaborn as sb

        check_make_folder(output_folder)
        metric_to_x, metric_to_y, metric_to_z = \
            self._get_column_names(metric_to_x, metric_to_y)

        # Define output path
        output_name = "{}_{}_kde.png".format(metric_to_x, metric_to_y)
        output_name = output_name.replace(" ", "_")
        output_name = os.path.join(output_folder, output_name)

        # Filter out the number of structures from dataframe to plot
        structures_to_keep = min(int(kde_structs), len(self._dataframe) - 1)
        sorted_df = self._dataframe.sort_values(metric_to_y, ascending=True)
        top = sorted_df[0:structures_to_keep]

        x_values = top[metric_to_x]
        y_values = top[metric_to_y]

        color1 = "lightskyblue"
        color2 = "royalblue"

        ax = sb.JointGrid(x=x_values, y=y_values)
        ax.plot_joint(sb.scatterplot, color=color1, edgecolor=color2, marker='o', alpha=0.7, s=20)
        sb.kdeplot(x=x_values, ax=ax.ax_marg_x, color=color1, shade=True, alpha=0.5, edgecolor=color2)
        sb.kdeplot(y=y_values, ax=ax.ax_marg_y, color=color1, shade=True, alpha=0.5, edgecolor=color2)
        ax.ax_joint.set_xlabel(metric_to_x, fontweight='bold')
        ax.ax_joint.set_ylabel(metric_to_y, fontweight='bold')

        ax.savefig(output_name)

        return output_name

    def plot_clusters(self, metric_to_x, metric_to_y, output_folder,
                      clusters, representative_structures=None):
        """
        It creates a scatter plot with the two metrics that are supplied
        and displays the points belonging to each top cluster with a
        different color.

        Parameters
        ----------
        metric_to_x : str or int
            The metric id to plot in the X axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_y : str or int
            The metric id to plot in the Y axis. It can be either a string
            with the name of the metric or an integer with the column index
        output_folder : str
            The path where the plot will be saved
        clusters : a numpy.array object
            The array of cluster labels that were obtained
        representative_structures : dict[str, tuple[str, int]]
            Dictionary containing the representative structures that
            were selected. Cluster label is the key and value is a list
            with [trajectory, step] of each cluster. If supplied, points
            belonging to representative structures will be represented
        """
        import copy
        from matplotlib.colors import LinearSegmentedColormap

        from pele_platform.analysis.clustering import get_cluster_label
        from pele_platform.Utilities.Helpers.helpers import backup_logger

        check_make_folder(output_folder)

        # Initialize a data handler from the current dataframe
        metric_to_x, metric_to_y, metric_to_z = \
            self._get_column_names(metric_to_x, metric_to_y)

        import matplotlib

        matplotlib.use("Agg")
        from matplotlib import pyplot as plt
        from matplotlib import colors, cm

        # Initialize figure
        fig, ax = plt.subplots(figsize=(6, 6), dpi=100, facecolor="w",
                               edgecolor="k")
        fig.subplots_adjust(right=0.8)  # To make room for the legend

        # Set axis labels
        plt.xlabel(metric_to_x)
        plt.ylabel(metric_to_y)

        # Configurate grid
        ax.set_axisbelow(True)
        ax.grid(True)
        ax.xaxis.grid(color="#AEB6BF", linestyle="dashed")
        ax.yaxis.grid(color="#AEB6BF", linestyle="dashed")
        ax.spines["right"].set_visible(False)
        ax.spines["top"].set_visible(False)
        ax.spines["left"].set_color("black")
        ax.spines["bottom"].set_color("black")
        ax.set_facecolor("#E6E9EB")

        # Extract the list of cluster labels
        cluster_labels = sorted(list(set(clusters)))

        # Configurate colormap
        if len(cluster_labels) > 18:
            cmap = copy.copy(cm.get_cmap("jet"))
        elif 9 < len(cluster_labels) <= 18:
            cmap = LinearSegmentedColormap.from_list('custom_tab20',
                                                     constants.custom_colors)
        else:
            cmap = copy.copy(cm.get_cmap("Set1"))

        norm = colors.Normalize(vmin=0, vmax=len(cluster_labels))
        cmap.set_under("grey")

        # Values to plot
        all_xs = self._dataframe[metric_to_x]
        all_ys = self._dataframe[metric_to_y]

        colors_used = []
        rep_struct_marker = []
        if representative_structures is not None:
            marker_cm = cm.get_cmap('binary')
            marker_norm = colors.Normalize(vmin=0, vmax=1)
            # Draw points with representative structures
            trajectories = self._dataframe['trajectory']
            steps = self._dataframe['numberOfAcceptedPeleSteps']

            rep_trajs = [traj for (traj, step)
                         in representative_structures.values()]
            rep_steps = [step for (traj, step)
                         in representative_structures.values()]

            for current_cluster in cluster_labels:
                xs = []
                ys = []
                for x, y, cluster, traj, step in zip(all_xs, all_ys,
                                                     clusters,
                                                     trajectories,
                                                     steps):
                    if cluster == current_cluster:
                        traj_idxs = set([i for i, x in enumerate(rep_trajs)
                                         if x == traj])
                        step_idxs = set([i for i, x in enumerate(rep_steps)
                                         if x == step])
                        if len(traj_idxs.intersection(step_idxs)) == 1:
                            sc = ax.scatter([x, ], [y, ], c=[1, ],
                                            zorder=3, marker='x',
                                            cmap=marker_cm,
                                            norm=marker_norm)
                            rep_struct_marker = sc.legend_elements()[0]
                        else:
                            xs.append(x)
                            ys.append(y)

                # In case there is only one point and it is the
                # representative structure
                if len(xs) == 0:
                    continue

                if current_cluster == -1:
                    zorder = 1
                else:
                    zorder = 2

                sc = ax.scatter(xs, ys, c=[current_cluster, ] * len(xs),
                                cmap=cmap, norm=norm, alpha=0.7,
                                zorder=zorder)
                colors_used += sc.legend_elements()[0]

        else:
            # Draw points without representative structures
            for current_cluster in cluster_labels:
                xs = []
                ys = []
                for x, y, cluster in zip(all_xs, all_ys, clusters):
                    if cluster == current_cluster:
                        xs.append(x)
                        ys.append(y)
                if current_cluster == -1:
                    zorder = 1
                else:
                    zorder = 2

                sc = ax.scatter(xs, ys, c=[current_cluster, ] * len(xs),
                                cmap=cmap, norm=norm, alpha=0.7,
                                zorder=zorder)
                colors_used += sc.legend_elements()[0]

        # Configure legend
        cluster_names = []
        for cluster_id in cluster_labels:
            if cluster_id == -1:
                cluster_names.append("Others")
            else:
                cluster_names.append(get_cluster_label(cluster_id))

        if cluster_names[0] == "Others":
            n = cluster_names.pop(0)
            c = colors_used.pop(0)
            cluster_names.append(n)
            colors_used.append(c)

        if len(rep_struct_marker) == 1:
            cluster_names.append("Representative\nstructure")

        ax.legend(colors_used + rep_struct_marker, cluster_names,
                  title="Clusters", loc='center left',
                  bbox_to_anchor=(1, 0.5))

        # Set output name
        if representative_structures is not None:
            output_name = "{}_{}_representatives_plot.png".format(metric_to_x,
                                                                  metric_to_y)
        else:
            output_name = "{}_{}_plot.png".format(metric_to_x, metric_to_y)
        output_name = output_name.replace(" ", "_")
        output_name = os.path.join(output_folder, output_name)

        plt.savefig(output_name, dpi=200, edgecolor="k",
                    orientation="portrait", transparent=True,
                    bbox_inches="tight")

        backup_logger(self._logger,
                      "Plotted {} vs {}".format(metric_to_x, metric_to_y))
        plt.close("all")

    def _get_column_names(self, metric_to_x, metric_to_y, metric_to_z=None):
        """
        Gets column names based on indices.
        Parameters
        ----------
        metric_to_x : int
            Index of column to be plotted on x-axis
        metric_to_y : int
            Index of column to be plotted on y-axis
        metric_to_z : int
            Index of column to be used for the colour bar.
        Returns
        -------
            Column names of x, y and z axes.
        """

        data_handler = DataHandler.from_dataframe(self._dataframe)
        # Ensure that metrics are strings pointing to dataframe columns
        if str(metric_to_x).isdigit():
            metric_to_x = data_handler.get_column_name(metric_to_x)
        if str(metric_to_y).isdigit():
            metric_to_y = data_handler.get_column_name(metric_to_y)
        if metric_to_z is not None and str(metric_to_z).isdigit():
            metric_to_z = data_handler.get_column_name(metric_to_z)
        return metric_to_x, metric_to_y, metric_to_z

"""
This module contains classes and methods to construct plots with data
coming from PELE reports.
"""
import os

from pele_platform.Utilities.Helpers.helpers import check_make_folder
from pele_platform.analysis import DataHandler
from pele_platform.constants import constants


class Plotter(object):
    """
    It handles the plots.
    """

    def __init__(self, dataframe, logger=None):
        """
        It initializes the plotter with the dataframe that contains the
        data to plot.

        Parameters
        ----------
        dataframe : a pandas.DataFrame object
            The dataframe containing the information from PELE reports
        logger : a Logger object
            The logger manager to stream out the log messages. Default is
            None
        """
        self._dataframe = dataframe
        self._logger = logger

    def plot_two_metrics(self, metric_to_x, metric_to_y, metric_to_z=None,
                         output_name=None, output_folder="."):
        """
        Given 2 or 3 metrics, it generates the scatter plot. In case that
        a 3rd metric is supplied, it will be represented as the color bar.

        Parameters
        ----------
        metric_to_x : str or int
            The metric id to plot in the X axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_y : str or int
            The metric id to plot in the Y axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_z : str or int
            The metric id to plot in the color bar. It can be either a string
            with the name of the metric or an integer with the column index.
            Default is None
        output_name : str
            The name that will be given to the resulting plot. Default is None
        output_folder : str
            The path where the plot will be saved. Default is '.', so it
            will be stored in the local directory
        """
        from pele_platform.Utilities.Helpers.helpers import backup_logger

        check_make_folder(output_folder)

        # Initialize a data handler from the current dataframe and get column names
        metric_to_x, metric_to_y, metric_to_z = self._get_column_names(
            metric_to_x, metric_to_y, metric_to_z)

        # Prepare plot name
        if output_name is None:
            if metric_to_z is not None:
                output_name = "{}_{}_{}_plot.png".format(metric_to_x,
                                                         metric_to_y,
                                                         metric_to_z)
            else:
                output_name = "{}_{}_plot.png".format(metric_to_x, metric_to_y)

        # Replace whitespaces in the output name
        output_name = output_name.replace(" ", "_")
        output_name = os.path.join(output_folder, output_name)

        # Generate plot with matplotlib
        import matplotlib

        matplotlib.use("Agg")
        from matplotlib import pyplot as plt

        fig, ax = plt.subplots()

        if metric_to_z is not None:
            scatter = ax.scatter(self._dataframe[metric_to_x],
                                 self._dataframe[metric_to_y],
                                 c=self._dataframe[metric_to_z],
                                 s=20)
            cbar = plt.colorbar(scatter)
            cbar.ax.set_ylabel(metric_to_z)
            ax.set_xlabel(metric_to_x)
            ax.set_ylabel(metric_to_y)
            plt.savefig(output_name)
            backup_logger(self._logger,
                          "Plotted {} vs {} vs {}".format(metric_to_x,
                                                          metric_to_y,
                                                          metric_to_z))

        else:
            ax.scatter(self._dataframe[metric_to_x],
                       self._dataframe[metric_to_y])
            ax.set_xlabel(metric_to_x)
            ax.set_ylabel(metric_to_y)
            plt.savefig(output_name)
            backup_logger(self._logger,
                          "Plotted {} vs {}".format(metric_to_x,
                                                    metric_to_y))

        plt.close("all")
        return output_name

    def plot_kde(self, metric_to_x, metric_to_y, output_folder, kde_structs):
        """
        Given 2 metrics, it generates the kde plot.

        Parameters
        ----------
        metric_to_x : str or int
            The metric id to plot in the X axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_y : str or int
            The metric id to plot in the Y axis. It can be either a string
            with the name of the metric or an integer with the column index
        output_folder : str
            The path where the plot will be saved
        kde_structs : int
            The number of structures to represent in the plot
        """
        import seaborn as sb

        check_make_folder(output_folder)
        metric_to_x, metric_to_y, metric_to_z = \
            self._get_column_names(metric_to_x, metric_to_y)

        # Define output path
        output_name = "{}_{}_kde.png".format(metric_to_x, metric_to_y)
        output_name = output_name.replace(" ", "_")
        output_name = os.path.join(output_folder, output_name)

        # Filter out the number of structures from dataframe to plot
        structures_to_keep = min(int(kde_structs), len(self._dataframe) - 1)
        sorted_df = self._dataframe.sort_values(metric_to_y, ascending=True)
        top = sorted_df[0:structures_to_keep]

        x_values = top[metric_to_x]
        y_values = top[metric_to_y]

        color1 = "lightskyblue"
        color2 = "royalblue"

        ax = sb.JointGrid(x=x_values, y=y_values)
        ax.plot_joint(sb.scatterplot, color=color1, edgecolor=color2, marker='o', alpha=0.7, s=20)
        sb.kdeplot(x=x_values, ax=ax.ax_marg_x, color=color1, shade=True, alpha=0.5, edgecolor=color2)
        sb.kdeplot(y=y_values, ax=ax.ax_marg_y, color=color1, shade=True, alpha=0.5, edgecolor=color2)
        ax.ax_joint.set_xlabel(metric_to_x, fontweight='bold')
        ax.ax_joint.set_ylabel(metric_to_y, fontweight='bold')

        ax.savefig(output_name)

        return output_name

    def plot_clusters(self, metric_to_x, metric_to_y, output_folder,
                      clusters, representative_structures=None):
        """
        It creates a scatter plot with the two metrics that are supplied
        and displays the points belonging to each top cluster with a
        different color.

        Parameters
        ----------
        metric_to_x : str or int
            The metric id to plot in the X axis. It can be either a string
            with the name of the metric or an integer with the column index
        metric_to_y : str or int
            The metric id to plot in the Y axis. It can be either a string
            with the name of the metric or an integer with the column index
        output_folder : str
            The path where the plot will be saved
        clusters : a numpy.array object
            The array of cluster labels that were obtained
        representative_structures : dict[str, tuple[str, int]]
            Dictionary containing the representative structures that
            were selected. Cluster label is the key and value is a list
            with [trajectory, step] of each cluster. If supplied, points
            belonging to representative structures will be represented
        """
        import copy
        from matplotlib.colors import LinearSegmentedColormap

        from pele_platform.analysis.clustering import get_cluster_label
        from pele_platform.Utilities.Helpers.helpers import backup_logger

        check_make_folder(output_folder)

        # Initialize a data handler from the current dataframe
        metric_to_x, metric_to_y, metric_to_z = \
            self._get_column_names(metric_to_x, metric_to_y)

        import matplotlib

        matplotlib.use("Agg")
        from matplotlib import pyplot as plt
        from matplotlib import colors, cm

        # Initialize figure
        fig, ax = plt.subplots(figsize=(6, 6), dpi=100, facecolor="w",
                               edgecolor="k")
        fig.subplots_adjust(right=0.8)  # To make room for the legend

        # Set axis labels
        plt.xlabel(metric_to_x)
        plt.ylabel(metric_to_y)

        # Configurate grid
        ax.set_axisbelow(True)
        ax.grid(True)
        ax.xaxis.grid(color="#AEB6BF", linestyle="dashed")
        ax.yaxis.grid(color="#AEB6BF", linestyle="dashed")
        ax.spines["right"].set_visible(False)
        ax.spines["top"].set_visible(False)
        ax.spines["left"].set_color("black")
        ax.spines["bottom"].set_color("black")
        ax.set_facecolor("#E6E9EB")

        # Extract the list of cluster labels
        cluster_labels = sorted(list(set(clusters)))

        # Configurate colormap
        if len(cluster_labels) > 18:
            cmap = copy.copy(cm.get_cmap("jet"))
        elif 9 < len(cluster_labels) <= 18:
            cmap = LinearSegmentedColormap.from_list('custom_tab20',
                                                     constants.custom_colors)
        else:
            cmap = copy.copy(cm.get_cmap("Set1"))

        norm = colors.Normalize(vmin=0, vmax=len(cluster_labels))
        cmap.set_under("grey")

        # Values to plot
        all_xs = self._dataframe[metric_to_x]
        all_ys = self._dataframe[metric_to_y]

        colors_used = []
        rep_struct_marker = []
        if representative_structures is not None:
            marker_cm = cm.get_cmap('binary')
            marker_norm = colors.Normalize(vmin=0, vmax=1)
            # Draw points with representative structures
            trajectories = self._dataframe['trajectory']
            steps = self._dataframe['numberOfAcceptedPeleSteps']

            rep_trajs = [traj for (traj, step)
                         in representative_structures.values()]
            rep_steps = [step for (traj, step)
                         in representative_structures.values()]

            for current_cluster in cluster_labels:
                xs = []
                ys = []
                for x, y, cluster, traj, step in zip(all_xs, all_ys,
                                                     clusters,
                                                     trajectories,
                                                     steps):
                    if cluster == current_cluster:
                        traj_idxs = set([i for i, x in enumerate(rep_trajs)
                                         if x == traj])
                        step_idxs = set([i for i, x in enumerate(rep_steps)
                                         if x == step])
                        if len(traj_idxs.intersection(step_idxs)) == 1:
                            sc = ax.scatter([x, ], [y, ], c=[1, ],
                                            zorder=3, marker='x',
                                            cmap=marker_cm,
                                            norm=marker_norm)
                            rep_struct_marker = sc.legend_elements()[0]
                        else:
                            xs.append(x)
                            ys.append(y)

                # In case there is only one point and it is the
                # representative structure
                if len(xs) == 0:
                    continue

                if current_cluster == -1:
                    zorder = 1
                else:
                    zorder = 2

                sc = ax.scatter(xs, ys, c=[current_cluster, ] * len(xs),
                                cmap=cmap, norm=norm, alpha=0.7,
                                zorder=zorder)
                colors_used += sc.legend_elements()[0]

        else:
            # Draw points without representative structures
            for current_cluster in cluster_labels:
                xs = []
                ys = []
                for x, y, cluster in zip(all_xs, all_ys, clusters):
                    if cluster == current_cluster:
                        xs.append(x)
                        ys.append(y)
                if current_cluster == -1:
                    zorder = 1
                else:
                    zorder = 2

                sc = ax.scatter(xs, ys, c=[current_cluster, ] * len(xs),
                                cmap=cmap, norm=norm, alpha=0.7,
                                zorder=zorder)
                colors_used += sc.legend_elements()[0]

        # Configure legend
        cluster_names = []
        for cluster_id in cluster_labels:
            if cluster_id == -1:
                cluster_names.append("Others")
            else:
                cluster_names.append(get_cluster_label(cluster_id))

        if cluster_names[0] == "Others":
            n = cluster_names.pop(0)
            c = colors_used.pop(0)
            cluster_names.append(n)
            colors_used.append(c)

        if len(rep_struct_marker) == 1:
            cluster_names.append("Representative\nstructure")

        ax.legend(colors_used + rep_struct_marker, cluster_names,
                  title="Clusters", loc='center left',
                  bbox_to_anchor=(1, 0.5))

        # Set output name
        if representative_structures is not None:
            output_name = "{}_{}_representatives_plot.png".format(metric_to_x,
                                                                  metric_to_y)
        else:
            output_name = "{}_{}_plot.png".format(metric_to_x, metric_to_y)
        output_name = output_name.replace(" ", "_")
        output_name = os.path.join(output_folder, output_name)

        plt.savefig(output_name, dpi=200, edgecolor="k",
                    orientation="portrait", transparent=True,
                    bbox_inches="tight")

        backup_logger(self._logger,
                      "Plotted {} vs {}".format(metric_to_x, metric_to_y))
        plt.close("all")

    def _get_column_names(self, metric_to_x, metric_to_y, metric_to_z=None):
        """
        Gets column names based on indices.
        Parameters
        ----------
        metric_to_x : int
            Index of column to be plotted on x-axis
        metric_to_y : int
            Index of column to be plotted on y-axis
        metric_to_z : int
            Index of column to be used for the colour bar.
        Returns
        -------
            Column names of x, y and z axes.
        """

        data_handler = DataHandler.from_dataframe(self._dataframe)
        # Ensure that metrics are strings pointing to dataframe columns
        if str(metric_to_x).isdigit():
            metric_to_x = data_handler.get_column_name(metric_to_x)
        if str(metric_to_y).isdigit():
            metric_to_y = data_handler.get_column_name(metric_to_y)
        if metric_to_z is not None and str(metric_to_z).isdigit():
            metric_to_z = data_handler.get_column_name(metric_to_z)
        return metric_to_x, metric_to_y, metric_to_z

This module contains classes and methods to construct plots with data coming from PELE reports. It handles the plots. It initializes the plotter with the dataframe that contains the         data to plot.          Parameters         ----------         dataframe : a pandas.DataFrame object             The dataframe containing the information from PELE reports         logger : a Logger object             The logger manager to stream out the log messages. Default is             None Given 2 or 3 metrics, it generates the scatter plot. In case that         a 3rd metric is supplied, it will be represented as the color bar.          Parameters         ----------         metric_to_x : str or int             The metric id to plot in the X axis. It can be either a string             with the name of the metric or an integer with the column index         metric_to_y : str or int             The metric id to plot in the Y axis. It can be either a string             with the name of the metric or an integer with the column index         metric_to_z : str or int             The metric id to plot in the color bar. It can be either a string             with the name of the metric or an integer with the column index.             Default is None         output_name : str             The name that will be given to the resulting plot. Default is None         output_folder : str             The path where the plot will be saved. Default is '.', so it             will be stored in the local directory # Initialize a data handler from the current dataframe and get column names # Prepare plot name # Replace whitespaces in the output name # Generate plot with matplotlib Given 2 metrics, it generates the kde plot.          Parameters         ----------         metric_to_x : str or int             The metric id to plot in the X axis. It can be either a string             with the name of the metric or an integer with the column index         metric_to_y : str or int             The metric id to plot in the Y axis. It can be either a string             with the name of the metric or an integer with the column index         output_folder : str             The path where the plot will be saved         kde_structs : int             The number of structures to represent in the plot # Define output path # Filter out the number of structures from dataframe to plot It creates a scatter plot with the two metrics that are supplied         and displays the points belonging to each top cluster with a         different color.          Parameters         ----------         metric_to_x : str or int             The metric id to plot in the X axis. It can be either a string             with the name of the metric or an integer with the column index         metric_to_y : str or int             The metric id to plot in the Y axis. It can be either a string             with the name of the metric or an integer with the column index         output_folder : str             The path where the plot will be saved         clusters : a numpy.array object             The array of cluster labels that were obtained         representative_structures : dict[str, tuple[str, int]]             Dictionary containing the representative structures that             were selected. Cluster label is the key and value is a list             with [trajectory, step] of each cluster. If supplied, points             belonging to representative structures will be represented # Initialize a data handler from the current dataframe # Initialize figure # To make room for the legend # Set axis labels # Configurate grid # Extract the list of cluster labels # Configurate colormap # Values to plot # Draw points with representative structures # In case there is only one point and it is the # representative structure # Draw points without representative structures # Configure legend # Set output name Gets column names based on indices.         Parameters         ----------         metric_to_x : int             Index of column to be plotted on x-axis         metric_to_y : int             Index of column to be plotted on y-axis         metric_to_z : int             Index of column to be used for the colour bar.         Returns         -------             Column names of x, y and z axes. # Ensure that metrics are strings pointing to dataframe columns