dataframes.py

import numpy as np
from warnings import warn

# external libraries
from tabulate import tabulate

# local imports
from .plotting import quick_plot


def print_dataframe(dataframe, max_langth=20):
    """
    Prints out nicely `dataframe`. Truncates column names to length `max_langth`. Not suitable for large dataframes (it makes a copy of the original dataframe).

    Parameters
    ----------
    dataframe : DataFrame
        Data frame to print.
    max_langth : int, optional
        Maximum length of the data frame column names (if longer, it gets truncated).

    Returns
    -------
    None

    """

    df = dataframe.copy()

    # truncate too long column names
    for col_name in df.columns:
        if len(col_name) > max_langth:
            df = df.rename(columns={col_name: (col_name[:max_langth - 2] + '..')})

    # print out the data frame
    print(tabulate(df, headers='keys', tablefmt='psql'))


def get_varied_columns_names(df):
    """
    Returns a list of data frame columns' names whose values are different from each other (i.e. there are at least two distinct values in each of the columns whose names are returned).

    Parameters
    ----------
    df : DataFrame
        Data frame whose columns are searched for different values.

    Returns
    -------
    varied_columns : list
        A list containing names of columns in which different values can be found.

    """

    varied_columns = []
    error_columns = []
    for c in df:
        # check if the value of the first row is not identical to all the other row
        try:
            check = not np.all(df[c] == df[c].iloc[0])
        except ValueError:
            error_columns.append(c)
        else:
            if check:
                varied_columns.append(c)

    if len(error_columns) > 0:
        warn(f'failed to to deduce if columns are varied for the following: {error_columns}')

    return varied_columns


def series_to_string(series, delimiter=' = '):
    """ Returns a string of pairs of `series` indices and values in a form 'index_0 = value_0, index_1 = value_1, ...' (if `delimiter` is ' = '). """

    s = ''
    for name in series.index:
        try:
            value = f'{series[name]:.3g}'
        except ValueError:  # if value is not a number
            value = str(series[name])

        s += f'{name}{delimiter}{value}, '

    s = s[:-2]

    return s


def plot_from_df(df, x_label, y_label, y2_label=None, grouping_label=None, plot_with='matplotlib', **kwargs):
    """ 
    Returns a plot with data from `df` based on selected labels (df's columns).
    To retrieve data it uses get_plot_dict_from_df(). 
    kwargs are passed on to quick_plot(). 
    """

    plot_data = get_plot_dict_from_df(df, x_label, y_label, y2_label=y2_label, grouping_label=grouping_label)

    # replace '_' with white spaces
    x_label = " ".join(x_label.split("_"))
    y_label = " ".join(y_label.split("_"))
    y2_label = " ".join(y2_label.split("_")) if y2_label else None

    if grouping_label:
        plot = quick_plot(plot_data, x_label=x_label, y_label=y_label, y2_label=y2_label, plot_with=plot_with, legend_title=grouping_label, **kwargs)
    else:
        show_legend = True if y2_label else False
        plot = quick_plot(plot_data, x_label=x_label, y_label=y_label, y2_label=y2_label, plot_with=plot_with, show_legend=show_legend, **kwargs)

    return plot


def get_plot_dict_from_df(df, x_label, y_label, y2_label=None, grouping_label=None):
    """ 
    Returns a dictionary that can be passed to quick_plot() based on data from dataframe `df` and chosen labels (dataframe's columns).
    It looks also in `df` for uncertainties as (label + '_error'), ignores if they are not found.
    If `grouping_label` is provided, it returns many data sets corresponding to distinct values found in the dataframe's column `grouping_label`.
    """

    if grouping_label:
        grouping_unique_values = df[grouping_label].unique()
            
        plot_data = {}
        for value in grouping_unique_values:
            label = f'{value}'  # legend labels are the distinct values from `grouping_label` column
            # label = value if not y2_label else f'{value}; {y_label}'

            plot_data[label] = {}

            plot_data[label]['x'] = df.loc[df[grouping_label] == value][x_label]
            # check if the uncertainties are provided
            try: 
                plot_data[label]['x_err'] = df.loc[df[grouping_label] == value][x_label + '_error']
            except KeyError:
                pass

            plot_data[label]['y'] = df.loc[df[grouping_label] == value][y_label]
            # check if the uncertainties are provided
            try: 
                plot_data[label]['y_err'] = df.loc[df[grouping_label] == value][y_label + '_error']
            except KeyError:
                pass

            if y2_label:
                label_2 = f'{value} - {" ".join(y2_label.split("_"))}'  # replace '_' with white spaces

                plot_data[label_2] = {}
                plot_data[label_2]['second_y'] = True

                plot_data[label_2]['x'] = plot_data[label]['x']
                try:
                    plot_data[label_2]['x_err'] = plot_data[label]['x_err']
                except KeyError:
                    pass

                plot_data[label_2]['y'] = df.loc[df[grouping_label] == value][y2_label]
                try: 
                    plot_data[label_2]['y_err'] = df.loc[df[grouping_label] == value][y2_label + '_error']
                except KeyError:
                    pass
                    
    else:
        label = " ".join(y_label.split("_"))  # replace '_' with white spaces

        plot_data = {}
        plot_data[label] = {}

        plot_data[label]['x'] = df[x_label]
        try: 
            plot_data[label]['x_err'] = df[x_label + '_error']
        except KeyError:
            pass

        plot_data[label]['y'] = df[y_label]
        try: 
            plot_data[label]['y_err'] = df[y_label + '_error']
        except KeyError:
            pass

        if y2_label:
            label_2 = " ".join(y2_label.split("_"))  # replace '_' with white spaces

            plot_data[label_2] = {}
            plot_data[label_2]['second_y'] = True

            plot_data[label_2]['x'] = plot_data[label]['x']
            try:
                plot_data[label_2]['x_err'] = plot_data[label]['x_err']
            except KeyError:
                pass

            plot_data[label_2]['y'] = df[y2_label]
            try: 
                plot_data[label_2]['y_err'] = df[y2_label + '_error']
            except KeyError:
                pass

    return plot_data