# -*- coding: utf-8 -*-


# Global imports
from __future__ import unicode_literals
import os
import glob
import argparse as ap
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("tkagg")
import seaborn as sns
from math import isnan

# Local imports
from PELEParseReports import *


# Script information
__author__ = ["<NAME>", "<NAME>"]
__license__ = "MIT"
__version__ = "1.0.1"
__maintainer__ = ["<NAME>", "<NAME>"]
__email__ = ["<EMAIL>","<EMAIL>"]

# Possible font of the axis and the title
Dict_of_fonts={"title" : {'family': 'serif',
        'weight': 'bold',
        'size': 16,},
"axis" : {'family': 'serif',
        'weight': 'normal',
        'size': 14,}}


# Functions
def parseArgs():
    """Parse arguments from command-line

    RETURNS
    -------
    reports : string
              list of report files to look for data
    x_data : string
             data to parse and assign to the X axis
    y_data : string
             data to parse and assign to the Y axis
    z_data : string
             data to parse and assign to the colorbar or the 3rd axis
    z_data : string
             data to parse and assign to the colorbar of the 3D scatter plot
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    title : string
                  title of the generated figure
    font: list of strings
		  key representing the font style in the font dictionary
    color : string
                  optional color to use for the figure
    size: integer
          The font size of all the labels in the plot
    scatterplot: boolean
          Perform the 2D scatter plot of two PELE metrics and a 3rd metric as colorbar (hover function)
    twodensityplot: boolean
          Perform the 2D density plot of two PELE metrics (with the points represented)
    densityplot: boolean
          Perform the density plot of one PELE metric
    pointplot: boolean
          Perform the 2D scatter plot of two PELE metrics and a 3rd metric as colorbar
    threeDscatterplot: boolean
          Perform the 3D scatter plot of three PELE metrics and a 4th metric as colorbar (hover function)
    bestquantile : float
          it sets the best quantile out of the data to be plotted
    """

    parser = ap.ArgumentParser(description='Script used to generate plots from the metrics \
        of the reports files from a PELE simulation')
    optional = parser._action_groups.pop()
    required = parser.add_argument_group('required arguments')
    required.add_argument("-i", "--input", required=True, metavar="FILE",
                          type=str, nargs='*', help="path to report files")
    optional.add_argument("-X", "--xaxis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to plot on the X axis", default=None)
    optional.add_argument("-Y", "--yaxis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to plot on the Y axis", default=None)
    optional.add_argument("-Z", "--zaxis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to represent in the colorbar", default=None)
    optional.add_argument("-Z2","--z2axis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to represent in the colorbar of the 3D scatter plot", default=None)
    optional.add_argument("-o", "--output", metavar="PATH", type=str,
                          help="output path to save figure", default=None)
    optional.add_argument("-r", "--Zmin", metavar="FLOAT", type=float,
                          help="minimum Z value for the colorbar (or Z2 value in the 3D scatter plot)",
                          default=None)
    optional.add_argument("-R", "--Zmax", metavar="FLOAT", type=float,
                          help="maximum Z value for the colorbar (or Z2 value in the 3D scatter plot)",
                          default=None)
    optional.add_argument("-T","--title", metavar="STRING", type=str,
			  help = "title of the figure", default="")
    optional.add_argument("-F","--font", metavar="STRING [STRING]", type=str,nargs='*',
                          help = "a list of the name of the font of the axis and the title", default="")
    optional.add_argument("-CO","--color", metavar="STRING", type=str,
                          help = "The color that you want to apply to the plot (only use when no colorbar is used)", default="")
    optional.add_argument("-CM","--colormap", metavar="STRING", type=str,
                          help = "The color that you want to apply to the colorbar (only for SP and TP and must be matplotlib valid)", default="")
    optional.add_argument("-S","--size", metavar="INTEGER", type=int,
                          help = "the size of the font for all the plot (only for SP, DP, PP, and TP)", default=12)
    optional.add_argument("-SP","--scatterplot"
                          ,help = "Perform the archetypical PELEPlot", action = "store_true")
    optional.add_argument("-DDP","--twodensityplot"
                          ,help = "Perform a 2D density plot of the two selected variables", action = "store_true")
    optional.add_argument("-DP","--densityplot"
                          ,help = "Perform the densityPlot of the metric specified in X", action = "store_true")
    optional.add_argument("-PP","--pointplot"
                          ,help = "Perform the pointPlot of the metric specified in X and Y", action = "store_true")
    optional.add_argument("-TP","--threeDscatterplot"
                          ,help = "Perform the 3D scatter plot of the metric specified in X, Y, and Z [and Z 2]", action = "store_true")
    optional.add_argument("-BQ","--bestquantile", metavar="FLOAT", type=float
                          ,help = "Take only the best quantile to plot (in DDP, DP, and PP)", default=None)
    parser._action_groups.append(optional)
    args = parser.parse_args()

    reports = parseReports(args.input, parser)

    x_data = args.xaxis
    y_data = args.yaxis
    z_data = args.zaxis
    z2_data = args.z2axis

    if z_data is None:
        if args.Zmin is not None or args.Zmax is not None:
            print("No data to represent in the colorbar (Z axis). " +
                  "Custom ranges are ignored.")
        z_min = None
        z_max = None
    else:
        z_min = args.Zmin
        z_max = args.Zmax

    output_path = args.output
    title = args.title
    font=args.font
    if len(font)==0:
        font=["",""]
    color = args.color; colormap = args.colormap
    size = args.size
    SP = args.scatterplot
    DDP = args.twodensityplot
    DP = args.densityplot
    PP = args.pointplot
    TP = args.threeDscatterplot
    BQ = args.bestquantile

    return reports, x_data, y_data, z_data, z2_data, z_min, z_max, output_path, title, font, color, colormap, size, SP, DDP, DP, PP, TP, BQ


def addUnits(metric_name):
    """Add units according to the input metric

    PARAMETERS
    ----------
    metric_name : string
                  name of the metric to plot

    RETURNS
    -------
    label : string
            name of the metric to plot with the units that were added to it
    """

    if "energy" in metric_name.lower():
        label = metric_name + " ($kcal/mol$)"
    elif "energies" in metric_name.lower():
        label = metric_name + " ($kcal/mol$)"
    elif "distance" in metric_name.lower():
        label = metric_name + " ($\AA$)"
    elif "rmsd" in metric_name.lower():
        label = metric_name + " ($\AA$)"
    else:
        label = metric_name
    return label
            
def parseAxisData(axis_data):
    """It sets the columns and label of the data that wants to be plotted

    PARAMETERS
    ----------
    axis_data : string
                axis data to parse

    RETURNS
    -------
    parsed_data : tuple of a list and a string
                  the list specifies the report columns that want to be plotted
                  in the axis and the string sets the name of the axis
    """

    if axis_data is None:
        return ([None, ], None)
    else:
        try:
            rows = [int(axis_data[0]), ]
        except ValueError:
            print("Warning: axis data not recognized: {}".format(axis_data))
            return ([None, ], None)
        if len(axis_data) == 1:
            return (rows, None)
        elif len(axis_data) > 1:
            label_index = 1
            while axis_data[label_index] == "+":
                label_index += 2
                try:
                    rows.append(int(axis_data[2]))
                except (ValueError, IndexError):
                    print("Warning: axis data not recognized: " +
                          "{}".format(axis_data))
                    return ([None, ], None)
                if len(axis_data) == label_index:
                    return (rows, '?')
            label = addUnits(axis_data[label_index])
            return (rows, label)

        print("Warning: axis data not recognized: {}".format(axis_data))
        return ([None, ], None)


def scatterPlot(reports,
                x_rows = [None, ], y_rows = [None, ], z_rows = [None, ],
                x_name = None, y_name = None, z_name = None,
                output_path = None, z_max = None, z_min = None, title = "", font = ["",""], color = "", colormap = "", size=12):
    """Represent the scatter plot

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    z_rows : list of integers
             integers which specify the report columns to represent in the
             colorbar
    x_name : string
             label of the X axis
    y_name : string
             label of the Y axis
    z_name : string
             label of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    title: string
	   it sets the title name of the plot
    color : string
              color to be used for the figure (when only 2 variables are used)
    size: integer
          The font size of all the labels in the plot
    """

    # The different variables are created and the size of the labels is set
    x_values, y_values, z_values, labels, annotations = [], [], [], [], []
    plt.rcParams.update({'font.size': size})

    # Set the rows and their labels to perform the scatter plot if they are not specified
    with open(reports[0], 'r') as report_file:
        line = report_file.readline()
        if None in x_rows:
            x_rows = [7, ]
            x_name = "RMSD ($\AA$)"
        if None in y_rows:
            y_rows = [5, ]
            y_name = "Energy ($kcal/mol$)"
        if x_name is None:
            x_name = str(line.split("    ")[x_rows[0] - 1])
        if y_name is None:
            y_name = str(line.split("    ")[y_rows[0] - 1])
        if (None not in z_rows) and (z_name is None):
            z_name = str(line.split("    ")[z_rows[0] - 1])
            z_name = addUnits(z_name)

    # Get the report files and save the directory where the report is saved and their number
    for report in reports:
        report_directory = os.path.dirname(report)
        report_number = os.path.basename(report).split('_')[-1].split('.')[0]
        
        # Open the report file and save the valeus that will be represented in the 2D scatter plot
        with open(report, 'r') as report_file:
            next(report_file)
            for i, line in enumerate(report_file):
                x_total = 0.
                y_total = 0.
                z_total = 0.

                for x_row in x_rows:
                    x_total += float(line.split()[x_row - 1])

                for y_row in y_rows:
                    y_total += float(line.split()[y_row - 1])

                if None not in z_rows:
                    for z_row in z_rows:
                        z_total += float(line.split()[z_row - 1])

                if isnan(x_total) or isnan(y_total) or isnan(z_total):
                    continue

                x_values.append(x_total)
                y_values.append(y_total)
                z_values.append(z_total)

                epoch = report_directory.split('/')[-1]
                if not epoch.isdigit():
                    epoch = '0'

                annotations.append("Epoch: " + epoch + "\n" +
                                   "Trajectory: " + report_number + "\n" +
                                   "Model: " + str(i + 1))

                labels.append(0)

    if z_max is None:
        z_max = max(z_values)

    if z_min is None:
        z_min = min(z_values)

    if z_min == z_max:
        cmap = plt.cm.autumn
    else:
        if colormap!="":
            cmap = plt.cm.get_cmap("{}".format(colormap))
        else:
            cmap = plt.cm.plasma

    norm = plt.Normalize(z_min, z_max)

    fig, ax = plt.subplots()

    if output_path is not None:
        s = 20
    else:
        s = None

    if color!="" and (None in z_rows):
        sc = plt.scatter(x_values, y_values, c=color, s=s)
    else:
        sc = plt.scatter(x_values, y_values, c=z_values, cmap=cmap, s=s,
                            norm=norm)

    ax.margins(0.05)
    ax.set_facecolor('white')
    if font[0]!="":
        plt.ylabel(y_name,Dict_of_fonts[font[0]])
        plt.xlabel(x_name,Dict_of_fonts[font[0]])
    else:
        plt.ylabel(y_name)
        plt.xlabel(x_name)
    if font[1]!="":
        plt.title(title,Dict_of_fonts[font[1]])
    else:
        plt.title(title)

    annot = ax.annotate("", xy=(0, 0), xytext=(20, 20),
                        textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w"),
                        arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)

    # Activate the colorbar only if the Z axis contains data to plot
    if None not in z_rows:
        cbar = plt.colorbar(sc, drawedges=False)
        if font[0]!="":
            cbar.set_label(z_name,Dict_of_fonts[font[0]])
        else:
            cbar.set_label(z_name)

    def update_annot(ind):
        """Update the information box of the selected point"""
        pos = sc.get_offsets()[ind["ind"][0]]
        annot.xy = pos
        annot.set_text(annotations[int(ind["ind"][0])])
        if color!="" and (None in z_rows):
            annot.get_bbox_patch().set_facecolor(color)
        else:
            annot.get_bbox_patch().set_facecolor(cmap(norm(
                z_values[ind["ind"][0]])))
            

    def hover(event):
        """Action to perform when hovering the mouse on a point"""
        vis = annot.get_visible()
        if event.inaxes == ax:
            cont, ind = sc.contains(event)
            if cont:
                update_annot(ind)
                annot.set_visible(True)
                fig.canvas.draw_idle()
            else:
                if vis:
                    annot.set_visible(False)
                    fig.canvas.draw_idle()

    # Respond to mouse motion
    fig.canvas.mpl_connect("motion_notify_event", hover)

    # Save or display the plot depending on whether an output path was set or not
    if output_path is not None:
        plt.savefig(output_path)
    else:
        plt.show()

def TwoDensityPlot(reports,
                x_rows = [None, ], y_rows = [None, ],
                x_name = None, y_name = None,
                output_path = None, title = "", font = ["",""], color = "", size=12, bestquantile = None):
    """Represent the density plot between two variables

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    z_rows : list of integers
             integers which specify the report columns to represent in the
             colorbar
    x_name : string
             label of the X axis
    y_name : string
             label of the Y axis
    z_name : string
             label of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    title: string
       it sets the title name of the plot
    color : string
        color to be used for the figure
    size: integer
          The font size of all the labels in the plot
    """

    # The different variables are created and the size of the labels is 
    x_values, y_values = [], []
    plt.rcParams.update({'font.size': size})
    # Set the rows and their labels to perform the scatter plot if they 
    with open(reports[0], 'r') as report_file:
        line = report_file.readline()
        if None in x_rows:
            x_rows = [7, ]
            x_name = "RMSD ($\AA$)"
        if None in y_rows:
            y_rows = [5, ]
            y_name = "Energy ($kcal/mol$)"
        if x_name is None:
            x_name = str(line.split("    ")[x_rows[0] - 1])
        if y_name is None:
            y_name = str(line.split("    ")[y_rows[0] - 1])
    # Get the report files and save the directory where the report is saved and their number
    for report in reports:
        report_directory = os.path.dirname(report)
        report_number = os.path.basename(report).split('_')[-1].split('.')[0]
        
        # Open the report file and save the valeus that will be represented in the 2D scatter plot
        with open(report, 'r') as report_file:
            next(report_file)
            for i, line in enumerate(report_file):
                x_total = 0.
                y_total = 0.
                for x_row in x_rows:
                    x_total += float(line.split()[x_row - 1])
                for y_row in y_rows:
                    y_total += float(line.split()[y_row - 1])
                if isnan(x_total) or isnan(y_total):
                    continue
                x_values.append(x_total)
                y_values.append(y_total)

    if color=="":
        color = "blue"

    # Get the best quantile out of the data according to the X axis to be plotted
    if bestquantile is not None:
        df = pd.DataFrame(list(zip(x_values,y_values)),columns=["X axis","Y axis"])
        quantile_df = df[df["X axis"] < df["X axis"].quantile(bestquantile)]
        x_values, y_values = quantile_df["X axis"].values,quantile_df["Y axis"].values

    sns.set(rc={'figure.figsize':(11.7,8.27),"font.size":28,"axes.titlesize":20,"axes.labelsize":28, "xtick.labelsize":24,"ytick.labelsize":24},style="white")
    plot = sns.JointGrid(x_values, y_values,space=0)
    plot = plot.plot_joint(plt.scatter, edgecolor="k", c=color)
    sns.scatterplot(x_values, y_values, color=color) 
    sns.distplot(y_values, kde=True, ax=plot.ax_marg_y, vertical=True, color=color,axlabel=False)
    sns.distplot(x_values, kde=True, ax=plot.ax_marg_x, color=color,axlabel=False)

    if font[0]!="":
        plt.ylabel(y_name,Dict_of_fonts[font[0]])
        plt.xlabel(x_name,Dict_of_fonts[font[0]])
    else:
        plt.ylabel(y_name)
        plt.xlabel(x_name)
    if font[1]!="":
        plt.title(title,Dict_of_fonts[font[1]])
    else:
        plt.title(title)

    # Save or display the plot depending on whether an output path was set or not
    if output_path is not None:
        plt.savefig(output_path)
    else:
        plt.show()

def densityPlot(reports, x_rows = [None, ], x_name = None, title = "", color = "", size = 12, bestquantile = None):
    """Represent the density plot

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    x_name : string
             label of the X axis
    title: string
       it sets the title name of the plot
    size: integer
          The font size of all the labels in the plot
    color : string
              color to be used for the figure
    """
    
    x_values = []
    plt.rcParams.update({'font.size': size})

    for report in reports:
        report_file = open(report,'r')
        i=0
        for line in report_file:
            if i!=0:
                x_values.append(float(line.split()[x_rows[0] - 1]))
            else:
                pass
            i+=1

    if color=="":
        color = "blue"

        # Get the best quantile out of the data according to the X axis to be plotted
    if bestquantile is not None:
        df = pd.DataFrame(list(x_values),columns=["X axis"])
        quantile_df = df[df["X axis"] < df["X axis"].quantile(bestquantile)]
        x_values = quantile_df["X axis"].values

    sns.distplot(x_values, color=color)
    plt.title(title)
    plt.xlabel(x_name)
    if x_name !=None:
        plt.ylabel("Density (1/{})".format(x_name.split("(")[1][0:-1]))
    plt.show()

def pointPlot(reports, x_rows = [None, ], x_name = None, y_rows = [None, ], y_name = None, title = "", color = "", size = 12, bestquantile = None):
    """Represent the scatter point plot

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    x_name : string
             label of the X axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    y_name : string
             label of the Y axis
    title: string
       it sets the title name of the plot
    color : string
              color to be used for the figure
    size: integer
          The font size of all the labels in the plot
    """    

    x_values,y_values = [],[]
    plt.rcParams.update({'font.size': size})

    for report in reports:
        report_file = open(report,'r')
        i=0
        for line in report_file:
            if i!=0:
                x_values.append(float(line.split()[x_rows[0] - 1]))
                y_values.append(float(line.split()[y_rows[0] - 1]))
            else:
                pass
            i+=1

    # Get the best quantile out of the data according to the X axis to be plotted
    if bestquantile is not None:
        df = pd.DataFrame(list(zip(x_values,y_values)),columns=["X axis","Y axis"])
        quantile_df = df[df["X axis"] < df["X axis"].quantile(bestquantile)]
        x_values, y_values = quantile_df["X axis"].values,quantile_df["Y axis"].values    

    if color=="":
        color="blue"

    plt.scatter(x_values,y_values,c=color)
    plt.title(title)
    plt.xlabel(x_name)
    plt.ylabel(y_name)
    plt.show()

def ThreeDPlot(reports, x_rows = [None, ], x_name = None, y_rows = [None, ], y_name = None,
    z_rows = [None, ], z_name = None, z2_rows = [None, ], z2_name = None,
    z_max = None, z_min = None, output_path = None, title = "", font = ["",""], color = "", colormap = "", size = 12):
    """Represent the 3D scatter plot of some PELE metrics

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    x_name : string
             label of the X axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    y_name : string
             label of the Y axis
    z_rows : list of integers
             integers which specify the report columns to represent in the Z
             axis
    z_name : string
             label of the Z axis
    z2_rows : list of integers
             integers which specify the report columns to represent in the colorbar
    z2_name : string
             label of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    title: string
       it sets the title name of the plot
    color : string
              color to be used for the figure (when only 3 variables are used)
    size: integer
          The font size of all the labels in the plot
    """

    # The different variables are created and the size of the labels is set
    x_values, y_values, z_values, z2_values, labels, annotations = [], [], [], [], [], []
    plt.rcParams.update({'font.size': size})

    # Set the rows and their labels to perform the scatter plot if they are not specified
    with open(reports[0], 'r') as report_file:
        line = report_file.readline()
        if None in x_rows:
            x_rows = [7, ]
            x_name = "RMSD ($\AA$)"
        if None in y_rows:
            y_rows = [5, ]
            y_name = "Energy ($kcal/mol$)"
        if None in z_rows:
            z_rows = [6, ]
            z_name = "SASA of the ligand"
        if x_name is None:
            x_name = str(line.split("    ")[x_rows[0] - 1])
        if y_name is None:
            y_name = str(line.split("    ")[y_rows[0] - 1])
        if z_name is None:
            z_name = str(line.split("    ")[z_rows[0] - 1])
        if (None not in z2_rows) and (z2_name is None):
            z2_name = str(line.split("    ")[z2_rows[0] - 1])
            z2_name = addUnits(z2_name)

    # Get the report files and save the directory where the report is saved and their number
    for report in reports:
        report_directory = os.path.dirname(report)
        report_number = os.path.basename(report).split('_')[-1].split('.')[0]

        # Open the report file and save the valeus that will be represented in the 3D scatter plot
        with open(report, 'r') as report_file:
            next(report_file)
            for i, line in enumerate(report_file):
                x_total = 0.
                y_total = 0.
                z_total = 0.
                z2_total = 0.

                for x_row in x_rows:
                    x_total += float(line.split()[x_row - 1])

                for y_row in y_rows:
                    y_total += float(line.split()[y_row - 1])

                for z_row in z_rows:
                    z_total += float(line.split()[z_row - 1])

                if None not in z2_rows:
                    for z2_row in z2_rows:
                        z2_total += float(line.split()[z2_row - 1])

                if isnan(x_total) or isnan(y_total) or isnan(z_total) or isnan(z2_total):
                    continue

                x_values.append(x_total)
                y_values.append(y_total)
                z_values.append(z_total)
                z2_values.append(z2_total)

                epoch = report_directory.split('/')[-1]
                if not epoch.isdigit():
                    epoch = '0'

                annotations.append("Epoch: " + epoch + "\n" +
                                   "Trajectory: " + report_number + "\n" +
                                   "Model: " + str(i + 1))

                labels.append(0)

    if z_max is None:
        z_max = max(z2_values)

    if z_min is None:
        z_min = min(z2_values)

    if z_min == z_max:
        cmap = plt.cm.autumn
    else:
        if colormap!="":
            cmap = plt.cm.get_cmap("{}".format(colormap))
        else:
            cmap = plt.cm.plasma

    norm = plt.Normalize(z_min, z_max)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    if color!="" and (None in z2_rows):
        sc = ax.scatter(x_values, y_values, z_values, c=color)
    else:
        sc = ax.scatter(x_values, y_values, z_values, c=z2_values, cmap=cmap,
                            norm=norm)

    ax.margins(0.05)
    ax.set_facecolor('white')
    if font[0]!="":
        ax.set_ylabel(y_name,Dict_of_fonts[font[0]])
        ax.set_xlabel(x_name,Dict_of_fonts[font[0]])
        ax.set_zlabel(z_name,Dict_of_fonts[font[0]])
    else:
        ax.set_ylabel(y_name)
        ax.set_xlabel(x_name)
        ax.set_zlabel(z_name)
    if font[1]!="":
        ax.set_title(title,Dict_of_fonts[font[1]])
    else:
        ax.set_title(title)

    annot = ax.annotate("", xy=(0, 0), xytext=(20, 20),
                        textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w"),
                        arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)

    # Activate the colorbar only if the Z2 axis contains data to plot
    if None not in z2_rows:
        cbar = plt.colorbar(sc, drawedges=False)
        if font[0]!="":
            cbar.set_label(z2_name,Dict_of_fonts[font[0]])
        else:
            cbar.set_label(z2_name)

    def update_annot(ind):
        """Update the information box of the selected point"""
        pos = sc.get_offsets()[ind["ind"][0]]
        annot.xy = pos
        annot.set_text(annotations[int(ind["ind"][0])])
        if color!="" and (None in z2_rows):
            annot.get_bbox_patch().set_facecolor(color)
        else:
            annot.get_bbox_patch().set_facecolor(cmap(norm(
                z2_values[ind["ind"][0]])))

    def hover(event):
        """Action to perform when hovering the mouse on a point"""
        vis = annot.get_visible()
        if event.inaxes == ax:
            cont, ind = sc.contains(event)
            if cont:
                update_annot(ind)
                annot.set_visible(True)
                fig.canvas.draw_idle()
            else:
                if vis:
                    annot.set_visible(False)
                    fig.canvas.draw_idle()

    # Respond to mouse motion
    fig.canvas.mpl_connect("motion_notify_event", hover)

    # Save or display the plot depending on whether an output path was set or not
    if output_path is not None:
        plt.savefig(output_path)
    else:
        plt.show()

def main():
    """Main function

    It is called when this script is the main program called by the interpreter
    """

    # Parse command-line arguments
    reports, x_data, y_data, z_data, z2_data, z_min, z_max, output_path, title, font, color, colormap, size, SP, DDP, DP, PP, TP, BQ = parseArgs()

    # Parse axis data to label it properly
    x_rows, x_name = parseAxisData(x_data)
    y_rows, y_name = parseAxisData(y_data)
    z_rows, z_name = parseAxisData(z_data)
    z2_rows, z2_name = parseAxisData(z2_data)

    # Generate the plot
    if SP:
        scatterPlot(reports,
                    x_rows=x_rows, y_rows=y_rows, z_rows=z_rows,
                    x_name=x_name, y_name=y_name, z_name=z_name,
                    z_min=z_min, z_max=z_max,
                    output_path=output_path,title=title,font=font,color=color,colormap=colormap,size=size)
    if DDP:
        TwoDensityPlot(reports,
                    x_rows=x_rows, y_rows=y_rows,
                    x_name=x_name, y_name=y_name,
                    output_path=output_path,title=title,font=font,color=color,size=size,bestquantile=BQ)
    if DP:
        densityPlot(reports,
                    x_rows=x_rows,x_name=x_name,title=title,color=color,size=size,bestquantile=BQ)
    if PP:
        pointPlot(reports,
                  x_rows=x_rows, y_rows=y_rows,x_name=x_name,y_name=y_name,title=title,color=color,size=size,bestquantile=BQ)
    if TP:
        ThreeDPlot(reports,
                    x_rows=x_rows, y_rows=y_rows, z_rows=z_rows, z2_rows=z2_rows,
                    x_name=x_name, y_name=y_name, z_name=z_name, z2_name=z2_name,
                    z_min=z_min, z_max=z_max,
                    output_path=output_path,title=title,font=font,color=color,colormap=colormap,size=size)


if __name__ == "__main__":
    """Call the main function"""
    main()

# -*- coding: utf-8 -*-


# Global imports
from __future__ import unicode_literals
import os
import glob
import argparse as ap
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("tkagg")
import seaborn as sns
from math import isnan

# Local imports
from PELEParseReports import *


# Script information
__author__ = ["<NAME>", "<NAME>"]
__license__ = "MIT"
__version__ = "1.0.1"
__maintainer__ = ["<NAME>", "<NAME>"]
__email__ = ["<EMAIL>","<EMAIL>"]

# Possible font of the axis and the title
Dict_of_fonts={"title" : {'family': 'serif',
        'weight': 'bold',
        'size': 16,},
"axis" : {'family': 'serif',
        'weight': 'normal',
        'size': 14,}}


# Functions
def parseArgs():
    """Parse arguments from command-line

    RETURNS
    -------
    reports : string
              list of report files to look for data
    x_data : string
             data to parse and assign to the X axis
    y_data : string
             data to parse and assign to the Y axis
    z_data : string
             data to parse and assign to the colorbar or the 3rd axis
    z_data : string
             data to parse and assign to the colorbar of the 3D scatter plot
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    title : string
                  title of the generated figure
    font: list of strings
		  key representing the font style in the font dictionary
    color : string
                  optional color to use for the figure
    size: integer
          The font size of all the labels in the plot
    scatterplot: boolean
          Perform the 2D scatter plot of two PELE metrics and a 3rd metric as colorbar (hover function)
    twodensityplot: boolean
          Perform the 2D density plot of two PELE metrics (with the points represented)
    densityplot: boolean
          Perform the density plot of one PELE metric
    pointplot: boolean
          Perform the 2D scatter plot of two PELE metrics and a 3rd metric as colorbar
    threeDscatterplot: boolean
          Perform the 3D scatter plot of three PELE metrics and a 4th metric as colorbar (hover function)
    bestquantile : float
          it sets the best quantile out of the data to be plotted
    """

    parser = ap.ArgumentParser(description='Script used to generate plots from the metrics \
        of the reports files from a PELE simulation')
    optional = parser._action_groups.pop()
    required = parser.add_argument_group('required arguments')
    required.add_argument("-i", "--input", required=True, metavar="FILE",
                          type=str, nargs='*', help="path to report files")
    optional.add_argument("-X", "--xaxis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to plot on the X axis", default=None)
    optional.add_argument("-Y", "--yaxis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to plot on the Y axis", default=None)
    optional.add_argument("-Z", "--zaxis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to represent in the colorbar", default=None)
    optional.add_argument("-Z2","--z2axis", metavar="INTEGER [METRIC]",
                          type=str, nargs='*', help="column numbers and " +
                          "metric to represent in the colorbar of the 3D scatter plot", default=None)
    optional.add_argument("-o", "--output", metavar="PATH", type=str,
                          help="output path to save figure", default=None)
    optional.add_argument("-r", "--Zmin", metavar="FLOAT", type=float,
                          help="minimum Z value for the colorbar (or Z2 value in the 3D scatter plot)",
                          default=None)
    optional.add_argument("-R", "--Zmax", metavar="FLOAT", type=float,
                          help="maximum Z value for the colorbar (or Z2 value in the 3D scatter plot)",
                          default=None)
    optional.add_argument("-T","--title", metavar="STRING", type=str,
			  help = "title of the figure", default="")
    optional.add_argument("-F","--font", metavar="STRING [STRING]", type=str,nargs='*',
                          help = "a list of the name of the font of the axis and the title", default="")
    optional.add_argument("-CO","--color", metavar="STRING", type=str,
                          help = "The color that you want to apply to the plot (only use when no colorbar is used)", default="")
    optional.add_argument("-CM","--colormap", metavar="STRING", type=str,
                          help = "The color that you want to apply to the colorbar (only for SP and TP and must be matplotlib valid)", default="")
    optional.add_argument("-S","--size", metavar="INTEGER", type=int,
                          help = "the size of the font for all the plot (only for SP, DP, PP, and TP)", default=12)
    optional.add_argument("-SP","--scatterplot"
                          ,help = "Perform the archetypical PELEPlot", action = "store_true")
    optional.add_argument("-DDP","--twodensityplot"
                          ,help = "Perform a 2D density plot of the two selected variables", action = "store_true")
    optional.add_argument("-DP","--densityplot"
                          ,help = "Perform the densityPlot of the metric specified in X", action = "store_true")
    optional.add_argument("-PP","--pointplot"
                          ,help = "Perform the pointPlot of the metric specified in X and Y", action = "store_true")
    optional.add_argument("-TP","--threeDscatterplot"
                          ,help = "Perform the 3D scatter plot of the metric specified in X, Y, and Z [and Z 2]", action = "store_true")
    optional.add_argument("-BQ","--bestquantile", metavar="FLOAT", type=float
                          ,help = "Take only the best quantile to plot (in DDP, DP, and PP)", default=None)
    parser._action_groups.append(optional)
    args = parser.parse_args()

    reports = parseReports(args.input, parser)

    x_data = args.xaxis
    y_data = args.yaxis
    z_data = args.zaxis
    z2_data = args.z2axis

    if z_data is None:
        if args.Zmin is not None or args.Zmax is not None:
            print("No data to represent in the colorbar (Z axis). " +
                  "Custom ranges are ignored.")
        z_min = None
        z_max = None
    else:
        z_min = args.Zmin
        z_max = args.Zmax

    output_path = args.output
    title = args.title
    font=args.font
    if len(font)==0:
        font=["",""]
    color = args.color; colormap = args.colormap
    size = args.size
    SP = args.scatterplot
    DDP = args.twodensityplot
    DP = args.densityplot
    PP = args.pointplot
    TP = args.threeDscatterplot
    BQ = args.bestquantile

    return reports, x_data, y_data, z_data, z2_data, z_min, z_max, output_path, title, font, color, colormap, size, SP, DDP, DP, PP, TP, BQ


def addUnits(metric_name):
    """Add units according to the input metric

    PARAMETERS
    ----------
    metric_name : string
                  name of the metric to plot

    RETURNS
    -------
    label : string
            name of the metric to plot with the units that were added to it
    """

    if "energy" in metric_name.lower():
        label = metric_name + " ($kcal/mol$)"
    elif "energies" in metric_name.lower():
        label = metric_name + " ($kcal/mol$)"
    elif "distance" in metric_name.lower():
        label = metric_name + " ($\AA$)"
    elif "rmsd" in metric_name.lower():
        label = metric_name + " ($\AA$)"
    else:
        label = metric_name
    return label
            
def parseAxisData(axis_data):
    """It sets the columns and label of the data that wants to be plotted

    PARAMETERS
    ----------
    axis_data : string
                axis data to parse

    RETURNS
    -------
    parsed_data : tuple of a list and a string
                  the list specifies the report columns that want to be plotted
                  in the axis and the string sets the name of the axis
    """

    if axis_data is None:
        return ([None, ], None)
    else:
        try:
            rows = [int(axis_data[0]), ]
        except ValueError:
            print("Warning: axis data not recognized: {}".format(axis_data))
            return ([None, ], None)
        if len(axis_data) == 1:
            return (rows, None)
        elif len(axis_data) > 1:
            label_index = 1
            while axis_data[label_index] == "+":
                label_index += 2
                try:
                    rows.append(int(axis_data[2]))
                except (ValueError, IndexError):
                    print("Warning: axis data not recognized: " +
                          "{}".format(axis_data))
                    return ([None, ], None)
                if len(axis_data) == label_index:
                    return (rows, '?')
            label = addUnits(axis_data[label_index])
            return (rows, label)

        print("Warning: axis data not recognized: {}".format(axis_data))
        return ([None, ], None)


def scatterPlot(reports,
                x_rows = [None, ], y_rows = [None, ], z_rows = [None, ],
                x_name = None, y_name = None, z_name = None,
                output_path = None, z_max = None, z_min = None, title = "", font = ["",""], color = "", colormap = "", size=12):
    """Represent the scatter plot

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    z_rows : list of integers
             integers which specify the report columns to represent in the
             colorbar
    x_name : string
             label of the X axis
    y_name : string
             label of the Y axis
    z_name : string
             label of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    title: string
	   it sets the title name of the plot
    color : string
              color to be used for the figure (when only 2 variables are used)
    size: integer
          The font size of all the labels in the plot
    """

    # The different variables are created and the size of the labels is set
    x_values, y_values, z_values, labels, annotations = [], [], [], [], []
    plt.rcParams.update({'font.size': size})

    # Set the rows and their labels to perform the scatter plot if they are not specified
    with open(reports[0], 'r') as report_file:
        line = report_file.readline()
        if None in x_rows:
            x_rows = [7, ]
            x_name = "RMSD ($\AA$)"
        if None in y_rows:
            y_rows = [5, ]
            y_name = "Energy ($kcal/mol$)"
        if x_name is None:
            x_name = str(line.split("    ")[x_rows[0] - 1])
        if y_name is None:
            y_name = str(line.split("    ")[y_rows[0] - 1])
        if (None not in z_rows) and (z_name is None):
            z_name = str(line.split("    ")[z_rows[0] - 1])
            z_name = addUnits(z_name)

    # Get the report files and save the directory where the report is saved and their number
    for report in reports:
        report_directory = os.path.dirname(report)
        report_number = os.path.basename(report).split('_')[-1].split('.')[0]
        
        # Open the report file and save the valeus that will be represented in the 2D scatter plot
        with open(report, 'r') as report_file:
            next(report_file)
            for i, line in enumerate(report_file):
                x_total = 0.
                y_total = 0.
                z_total = 0.

                for x_row in x_rows:
                    x_total += float(line.split()[x_row - 1])

                for y_row in y_rows:
                    y_total += float(line.split()[y_row - 1])

                if None not in z_rows:
                    for z_row in z_rows:
                        z_total += float(line.split()[z_row - 1])

                if isnan(x_total) or isnan(y_total) or isnan(z_total):
                    continue

                x_values.append(x_total)
                y_values.append(y_total)
                z_values.append(z_total)

                epoch = report_directory.split('/')[-1]
                if not epoch.isdigit():
                    epoch = '0'

                annotations.append("Epoch: " + epoch + "\n" +
                                   "Trajectory: " + report_number + "\n" +
                                   "Model: " + str(i + 1))

                labels.append(0)

    if z_max is None:
        z_max = max(z_values)

    if z_min is None:
        z_min = min(z_values)

    if z_min == z_max:
        cmap = plt.cm.autumn
    else:
        if colormap!="":
            cmap = plt.cm.get_cmap("{}".format(colormap))
        else:
            cmap = plt.cm.plasma

    norm = plt.Normalize(z_min, z_max)

    fig, ax = plt.subplots()

    if output_path is not None:
        s = 20
    else:
        s = None

    if color!="" and (None in z_rows):
        sc = plt.scatter(x_values, y_values, c=color, s=s)
    else:
        sc = plt.scatter(x_values, y_values, c=z_values, cmap=cmap, s=s,
                            norm=norm)

    ax.margins(0.05)
    ax.set_facecolor('white')
    if font[0]!="":
        plt.ylabel(y_name,Dict_of_fonts[font[0]])
        plt.xlabel(x_name,Dict_of_fonts[font[0]])
    else:
        plt.ylabel(y_name)
        plt.xlabel(x_name)
    if font[1]!="":
        plt.title(title,Dict_of_fonts[font[1]])
    else:
        plt.title(title)

    annot = ax.annotate("", xy=(0, 0), xytext=(20, 20),
                        textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w"),
                        arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)

    # Activate the colorbar only if the Z axis contains data to plot
    if None not in z_rows:
        cbar = plt.colorbar(sc, drawedges=False)
        if font[0]!="":
            cbar.set_label(z_name,Dict_of_fonts[font[0]])
        else:
            cbar.set_label(z_name)

    def update_annot(ind):
        """Update the information box of the selected point"""
        pos = sc.get_offsets()[ind["ind"][0]]
        annot.xy = pos
        annot.set_text(annotations[int(ind["ind"][0])])
        if color!="" and (None in z_rows):
            annot.get_bbox_patch().set_facecolor(color)
        else:
            annot.get_bbox_patch().set_facecolor(cmap(norm(
                z_values[ind["ind"][0]])))
            

    def hover(event):
        """Action to perform when hovering the mouse on a point"""
        vis = annot.get_visible()
        if event.inaxes == ax:
            cont, ind = sc.contains(event)
            if cont:
                update_annot(ind)
                annot.set_visible(True)
                fig.canvas.draw_idle()
            else:
                if vis:
                    annot.set_visible(False)
                    fig.canvas.draw_idle()

    # Respond to mouse motion
    fig.canvas.mpl_connect("motion_notify_event", hover)

    # Save or display the plot depending on whether an output path was set or not
    if output_path is not None:
        plt.savefig(output_path)
    else:
        plt.show()

def TwoDensityPlot(reports,
                x_rows = [None, ], y_rows = [None, ],
                x_name = None, y_name = None,
                output_path = None, title = "", font = ["",""], color = "", size=12, bestquantile = None):
    """Represent the density plot between two variables

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    z_rows : list of integers
             integers which specify the report columns to represent in the
             colorbar
    x_name : string
             label of the X axis
    y_name : string
             label of the Y axis
    z_name : string
             label of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    title: string
       it sets the title name of the plot
    color : string
        color to be used for the figure
    size: integer
          The font size of all the labels in the plot
    """

    # The different variables are created and the size of the labels is 
    x_values, y_values = [], []
    plt.rcParams.update({'font.size': size})
    # Set the rows and their labels to perform the scatter plot if they 
    with open(reports[0], 'r') as report_file:
        line = report_file.readline()
        if None in x_rows:
            x_rows = [7, ]
            x_name = "RMSD ($\AA$)"
        if None in y_rows:
            y_rows = [5, ]
            y_name = "Energy ($kcal/mol$)"
        if x_name is None:
            x_name = str(line.split("    ")[x_rows[0] - 1])
        if y_name is None:
            y_name = str(line.split("    ")[y_rows[0] - 1])
    # Get the report files and save the directory where the report is saved and their number
    for report in reports:
        report_directory = os.path.dirname(report)
        report_number = os.path.basename(report).split('_')[-1].split('.')[0]
        
        # Open the report file and save the valeus that will be represented in the 2D scatter plot
        with open(report, 'r') as report_file:
            next(report_file)
            for i, line in enumerate(report_file):
                x_total = 0.
                y_total = 0.
                for x_row in x_rows:
                    x_total += float(line.split()[x_row - 1])
                for y_row in y_rows:
                    y_total += float(line.split()[y_row - 1])
                if isnan(x_total) or isnan(y_total):
                    continue
                x_values.append(x_total)
                y_values.append(y_total)

    if color=="":
        color = "blue"

    # Get the best quantile out of the data according to the X axis to be plotted
    if bestquantile is not None:
        df = pd.DataFrame(list(zip(x_values,y_values)),columns=["X axis","Y axis"])
        quantile_df = df[df["X axis"] < df["X axis"].quantile(bestquantile)]
        x_values, y_values = quantile_df["X axis"].values,quantile_df["Y axis"].values

    sns.set(rc={'figure.figsize':(11.7,8.27),"font.size":28,"axes.titlesize":20,"axes.labelsize":28, "xtick.labelsize":24,"ytick.labelsize":24},style="white")
    plot = sns.JointGrid(x_values, y_values,space=0)
    plot = plot.plot_joint(plt.scatter, edgecolor="k", c=color)
    sns.scatterplot(x_values, y_values, color=color) 
    sns.distplot(y_values, kde=True, ax=plot.ax_marg_y, vertical=True, color=color,axlabel=False)
    sns.distplot(x_values, kde=True, ax=plot.ax_marg_x, color=color,axlabel=False)

    if font[0]!="":
        plt.ylabel(y_name,Dict_of_fonts[font[0]])
        plt.xlabel(x_name,Dict_of_fonts[font[0]])
    else:
        plt.ylabel(y_name)
        plt.xlabel(x_name)
    if font[1]!="":
        plt.title(title,Dict_of_fonts[font[1]])
    else:
        plt.title(title)

    # Save or display the plot depending on whether an output path was set or not
    if output_path is not None:
        plt.savefig(output_path)
    else:
        plt.show()

def densityPlot(reports, x_rows = [None, ], x_name = None, title = "", color = "", size = 12, bestquantile = None):
    """Represent the density plot

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    x_name : string
             label of the X axis
    title: string
       it sets the title name of the plot
    size: integer
          The font size of all the labels in the plot
    color : string
              color to be used for the figure
    """
    
    x_values = []
    plt.rcParams.update({'font.size': size})

    for report in reports:
        report_file = open(report,'r')
        i=0
        for line in report_file:
            if i!=0:
                x_values.append(float(line.split()[x_rows[0] - 1]))
            else:
                pass
            i+=1

    if color=="":
        color = "blue"

        # Get the best quantile out of the data according to the X axis to be plotted
    if bestquantile is not None:
        df = pd.DataFrame(list(x_values),columns=["X axis"])
        quantile_df = df[df["X axis"] < df["X axis"].quantile(bestquantile)]
        x_values = quantile_df["X axis"].values

    sns.distplot(x_values, color=color)
    plt.title(title)
    plt.xlabel(x_name)
    if x_name !=None:
        plt.ylabel("Density (1/{})".format(x_name.split("(")[1][0:-1]))
    plt.show()

def pointPlot(reports, x_rows = [None, ], x_name = None, y_rows = [None, ], y_name = None, title = "", color = "", size = 12, bestquantile = None):
    """Represent the scatter point plot

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    x_name : string
             label of the X axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    y_name : string
             label of the Y axis
    title: string
       it sets the title name of the plot
    color : string
              color to be used for the figure
    size: integer
          The font size of all the labels in the plot
    """    

    x_values,y_values = [],[]
    plt.rcParams.update({'font.size': size})

    for report in reports:
        report_file = open(report,'r')
        i=0
        for line in report_file:
            if i!=0:
                x_values.append(float(line.split()[x_rows[0] - 1]))
                y_values.append(float(line.split()[y_rows[0] - 1]))
            else:
                pass
            i+=1

    # Get the best quantile out of the data according to the X axis to be plotted
    if bestquantile is not None:
        df = pd.DataFrame(list(zip(x_values,y_values)),columns=["X axis","Y axis"])
        quantile_df = df[df["X axis"] < df["X axis"].quantile(bestquantile)]
        x_values, y_values = quantile_df["X axis"].values,quantile_df["Y axis"].values    

    if color=="":
        color="blue"

    plt.scatter(x_values,y_values,c=color)
    plt.title(title)
    plt.xlabel(x_name)
    plt.ylabel(y_name)
    plt.show()

def ThreeDPlot(reports, x_rows = [None, ], x_name = None, y_rows = [None, ], y_name = None,
    z_rows = [None, ], z_name = None, z2_rows = [None, ], z2_name = None,
    z_max = None, z_min = None, output_path = None, title = "", font = ["",""], color = "", colormap = "", size = 12):
    """Represent the 3D scatter plot of some PELE metrics

    PARAMETERS
    ----------
    reports : string
              list of report files to look for data
    x_rows : list of integers
             integers which specify the report columns to represent in the X
             axis
    x_name : string
             label of the X axis
    y_rows : list of integers
             integers which specify the report columns to represent in the Y
             axis
    y_name : string
             label of the Y axis
    z_rows : list of integers
             integers which specify the report columns to represent in the Z
             axis
    z_name : string
             label of the Z axis
    z2_rows : list of integers
             integers which specify the report columns to represent in the colorbar
    z2_name : string
             label of the colorbar
    output_path : string
                  output directory where the resulting plot will be saved
    z_max : float
            it sets the maximum range value of the colorbar
    z_min : float
            it sets the minimum range value of the colorbar
    title: string
       it sets the title name of the plot
    color : string
              color to be used for the figure (when only 3 variables are used)
    size: integer
          The font size of all the labels in the plot
    """

    # The different variables are created and the size of the labels is set
    x_values, y_values, z_values, z2_values, labels, annotations = [], [], [], [], [], []
    plt.rcParams.update({'font.size': size})

    # Set the rows and their labels to perform the scatter plot if they are not specified
    with open(reports[0], 'r') as report_file:
        line = report_file.readline()
        if None in x_rows:
            x_rows = [7, ]
            x_name = "RMSD ($\AA$)"
        if None in y_rows:
            y_rows = [5, ]
            y_name = "Energy ($kcal/mol$)"
        if None in z_rows:
            z_rows = [6, ]
            z_name = "SASA of the ligand"
        if x_name is None:
            x_name = str(line.split("    ")[x_rows[0] - 1])
        if y_name is None:
            y_name = str(line.split("    ")[y_rows[0] - 1])
        if z_name is None:
            z_name = str(line.split("    ")[z_rows[0] - 1])
        if (None not in z2_rows) and (z2_name is None):
            z2_name = str(line.split("    ")[z2_rows[0] - 1])
            z2_name = addUnits(z2_name)

    # Get the report files and save the directory where the report is saved and their number
    for report in reports:
        report_directory = os.path.dirname(report)
        report_number = os.path.basename(report).split('_')[-1].split('.')[0]

        # Open the report file and save the valeus that will be represented in the 3D scatter plot
        with open(report, 'r') as report_file:
            next(report_file)
            for i, line in enumerate(report_file):
                x_total = 0.
                y_total = 0.
                z_total = 0.
                z2_total = 0.

                for x_row in x_rows:
                    x_total += float(line.split()[x_row - 1])

                for y_row in y_rows:
                    y_total += float(line.split()[y_row - 1])

                for z_row in z_rows:
                    z_total += float(line.split()[z_row - 1])

                if None not in z2_rows:
                    for z2_row in z2_rows:
                        z2_total += float(line.split()[z2_row - 1])

                if isnan(x_total) or isnan(y_total) or isnan(z_total) or isnan(z2_total):
                    continue

                x_values.append(x_total)
                y_values.append(y_total)
                z_values.append(z_total)
                z2_values.append(z2_total)

                epoch = report_directory.split('/')[-1]
                if not epoch.isdigit():
                    epoch = '0'

                annotations.append("Epoch: " + epoch + "\n" +
                                   "Trajectory: " + report_number + "\n" +
                                   "Model: " + str(i + 1))

                labels.append(0)

    if z_max is None:
        z_max = max(z2_values)

    if z_min is None:
        z_min = min(z2_values)

    if z_min == z_max:
        cmap = plt.cm.autumn
    else:
        if colormap!="":
            cmap = plt.cm.get_cmap("{}".format(colormap))
        else:
            cmap = plt.cm.plasma

    norm = plt.Normalize(z_min, z_max)

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    if color!="" and (None in z2_rows):
        sc = ax.scatter(x_values, y_values, z_values, c=color)
    else:
        sc = ax.scatter(x_values, y_values, z_values, c=z2_values, cmap=cmap,
                            norm=norm)

    ax.margins(0.05)
    ax.set_facecolor('white')
    if font[0]!="":
        ax.set_ylabel(y_name,Dict_of_fonts[font[0]])
        ax.set_xlabel(x_name,Dict_of_fonts[font[0]])
        ax.set_zlabel(z_name,Dict_of_fonts[font[0]])
    else:
        ax.set_ylabel(y_name)
        ax.set_xlabel(x_name)
        ax.set_zlabel(z_name)
    if font[1]!="":
        ax.set_title(title,Dict_of_fonts[font[1]])
    else:
        ax.set_title(title)

    annot = ax.annotate("", xy=(0, 0), xytext=(20, 20),
                        textcoords="offset points",
                        bbox=dict(boxstyle="round", fc="w"),
                        arrowprops=dict(arrowstyle="->"))
    annot.set_visible(False)

    # Activate the colorbar only if the Z2 axis contains data to plot
    if None not in z2_rows:
        cbar = plt.colorbar(sc, drawedges=False)
        if font[0]!="":
            cbar.set_label(z2_name,Dict_of_fonts[font[0]])
        else:
            cbar.set_label(z2_name)

    def update_annot(ind):
        """Update the information box of the selected point"""
        pos = sc.get_offsets()[ind["ind"][0]]
        annot.xy = pos
        annot.set_text(annotations[int(ind["ind"][0])])
        if color!="" and (None in z2_rows):
            annot.get_bbox_patch().set_facecolor(color)
        else:
            annot.get_bbox_patch().set_facecolor(cmap(norm(
                z2_values[ind["ind"][0]])))

    def hover(event):
        """Action to perform when hovering the mouse on a point"""
        vis = annot.get_visible()
        if event.inaxes == ax:
            cont, ind = sc.contains(event)
            if cont:
                update_annot(ind)
                annot.set_visible(True)
                fig.canvas.draw_idle()
            else:
                if vis:
                    annot.set_visible(False)
                    fig.canvas.draw_idle()

    # Respond to mouse motion
    fig.canvas.mpl_connect("motion_notify_event", hover)

    # Save or display the plot depending on whether an output path was set or not
    if output_path is not None:
        plt.savefig(output_path)
    else:
        plt.show()

def main():
    """Main function

    It is called when this script is the main program called by the interpreter
    """

    # Parse command-line arguments
    reports, x_data, y_data, z_data, z2_data, z_min, z_max, output_path, title, font, color, colormap, size, SP, DDP, DP, PP, TP, BQ = parseArgs()

    # Parse axis data to label it properly
    x_rows, x_name = parseAxisData(x_data)
    y_rows, y_name = parseAxisData(y_data)
    z_rows, z_name = parseAxisData(z_data)
    z2_rows, z2_name = parseAxisData(z2_data)

    # Generate the plot
    if SP:
        scatterPlot(reports,
                    x_rows=x_rows, y_rows=y_rows, z_rows=z_rows,
                    x_name=x_name, y_name=y_name, z_name=z_name,
                    z_min=z_min, z_max=z_max,
                    output_path=output_path,title=title,font=font,color=color,colormap=colormap,size=size)
    if DDP:
        TwoDensityPlot(reports,
                    x_rows=x_rows, y_rows=y_rows,
                    x_name=x_name, y_name=y_name,
                    output_path=output_path,title=title,font=font,color=color,size=size,bestquantile=BQ)
    if DP:
        densityPlot(reports,
                    x_rows=x_rows,x_name=x_name,title=title,color=color,size=size,bestquantile=BQ)
    if PP:
        pointPlot(reports,
                  x_rows=x_rows, y_rows=y_rows,x_name=x_name,y_name=y_name,title=title,color=color,size=size,bestquantile=BQ)
    if TP:
        ThreeDPlot(reports,
                    x_rows=x_rows, y_rows=y_rows, z_rows=z_rows, z2_rows=z2_rows,
                    x_name=x_name, y_name=y_name, z_name=z_name, z2_name=z2_name,
                    z_min=z_min, z_max=z_max,
                    output_path=output_path,title=title,font=font,color=color,colormap=colormap,size=size)


if __name__ == "__main__":
    """Call the main function"""
    main()

# -*- coding: utf-8 -*- # Global imports # Local imports # Script information # Possible font of the axis and the title # Functions Parse arguments from command-line      RETURNS     -------     reports : string               list of report files to look for data     x_data : string              data to parse and assign to the X axis     y_data : string              data to parse and assign to the Y axis     z_data : string              data to parse and assign to the colorbar or the 3rd axis     z_data : string              data to parse and assign to the colorbar of the 3D scatter plot     z_max : float             it sets the maximum range value of the colorbar     z_min : float             it sets the minimum range value of the colorbar     output_path : string                   output directory where the resulting plot will be saved     title : string                   title of the generated figure     font: list of strings 		  key representing the font style in the font dictionary     color : string                   optional color to use for the figure     size: integer           The font size of all the labels in the plot     scatterplot: boolean           Perform the 2D scatter plot of two PELE metrics and a 3rd metric as colorbar (hover function)     twodensityplot: boolean           Perform the 2D density plot of two PELE metrics (with the points represented)     densityplot: boolean           Perform the density plot of one PELE metric     pointplot: boolean           Perform the 2D scatter plot of two PELE metrics and a 3rd metric as colorbar     threeDscatterplot: boolean           Perform the 3D scatter plot of three PELE metrics and a 4th metric as colorbar (hover function)     bestquantile : float           it sets the best quantile out of the data to be plotted Add units according to the input metric      PARAMETERS     ----------     metric_name : string                   name of the metric to plot      RETURNS     -------     label : string             name of the metric to plot with the units that were added to it It sets the columns and label of the data that wants to be plotted      PARAMETERS     ----------     axis_data : string                 axis data to parse      RETURNS     -------     parsed_data : tuple of a list and a string                   the list specifies the report columns that want to be plotted                   in the axis and the string sets the name of the axis Represent the scatter plot      PARAMETERS     ----------     reports : string               list of report files to look for data     x_rows : list of integers              integers which specify the report columns to represent in the X              axis     y_rows : list of integers              integers which specify the report columns to represent in the Y              axis     z_rows : list of integers              integers which specify the report columns to represent in the              colorbar     x_name : string              label of the X axis     y_name : string              label of the Y axis     z_name : string              label of the colorbar     output_path : string                   output directory where the resulting plot will be saved     z_max : float             it sets the maximum range value of the colorbar     z_min : float             it sets the minimum range value of the colorbar     title: string 	   it sets the title name of the plot     color : string               color to be used for the figure (when only 2 variables are used)     size: integer           The font size of all the labels in the plot # The different variables are created and the size of the labels is set # Set the rows and their labels to perform the scatter plot if they are not specified # Get the report files and save the directory where the report is saved and their number # Open the report file and save the valeus that will be represented in the 2D scatter plot # Activate the colorbar only if the Z axis contains data to plot Update the information box of the selected point Action to perform when hovering the mouse on a point # Respond to mouse motion # Save or display the plot depending on whether an output path was set or not Represent the density plot between two variables      PARAMETERS     ----------     reports : string               list of report files to look for data     x_rows : list of integers              integers which specify the report columns to represent in the X              axis     y_rows : list of integers              integers which specify the report columns to represent in the Y              axis     z_rows : list of integers              integers which specify the report columns to represent in the              colorbar     x_name : string              label of the X axis     y_name : string              label of the Y axis     z_name : string              label of the colorbar     output_path : string                   output directory where the resulting plot will be saved     z_max : float             it sets the maximum range value of the colorbar     z_min : float             it sets the minimum range value of the colorbar     title: string        it sets the title name of the plot     color : string         color to be used for the figure     size: integer           The font size of all the labels in the plot # The different variables are created and the size of the labels is # Set the rows and their labels to perform the scatter plot if they # Get the report files and save the directory where the report is saved and their number # Open the report file and save the valeus that will be represented in the 2D scatter plot # Get the best quantile out of the data according to the X axis to be plotted # Save or display the plot depending on whether an output path was set or not Represent the density plot      PARAMETERS     ----------     reports : string               list of report files to look for data     x_rows : list of integers              integers which specify the report columns to represent in the X              axis     x_name : string              label of the X axis     title: string        it sets the title name of the plot     size: integer           The font size of all the labels in the plot     color : string               color to be used for the figure # Get the best quantile out of the data according to the X axis to be plotted Represent the scatter point plot      PARAMETERS     ----------     reports : string               list of report files to look for data     x_rows : list of integers              integers which specify the report columns to represent in the X              axis     x_name : string              label of the X axis     y_rows : list of integers              integers which specify the report columns to represent in the Y              axis     y_name : string              label of the Y axis     title: string        it sets the title name of the plot     color : string               color to be used for the figure     size: integer           The font size of all the labels in the plot # Get the best quantile out of the data according to the X axis to be plotted Represent the 3D scatter plot of some PELE metrics      PARAMETERS     ----------     reports : string               list of report files to look for data     x_rows : list of integers              integers which specify the report columns to represent in the X              axis     x_name : string              label of the X axis     y_rows : list of integers              integers which specify the report columns to represent in the Y              axis     y_name : string              label of the Y axis     z_rows : list of integers              integers which specify the report columns to represent in the Z              axis     z_name : string              label of the Z axis     z2_rows : list of integers              integers which specify the report columns to represent in the colorbar     z2_name : string              label of the colorbar     output_path : string                   output directory where the resulting plot will be saved     z_max : float             it sets the maximum range value of the colorbar     z_min : float             it sets the minimum range value of the colorbar     title: string        it sets the title name of the plot     color : string               color to be used for the figure (when only 3 variables are used)     size: integer           The font size of all the labels in the plot # The different variables are created and the size of the labels is set # Set the rows and their labels to perform the scatter plot if they are not specified # Get the report files and save the directory where the report is saved and their number # Open the report file and save the valeus that will be represented in the 3D scatter plot # Activate the colorbar only if the Z2 axis contains data to plot Update the information box of the selected point Action to perform when hovering the mouse on a point # Respond to mouse motion # Save or display the plot depending on whether an output path was set or not Main function      It is called when this script is the main program called by the interpreter # Parse command-line arguments # Parse axis data to label it properly # Generate the plot Call the main function