""" Module to sort a DataFrame of spike times
"""

import numpy as np


def by_condition(df, condition="event_label", index="event_index"):
    """
    Sorts a DataFrame of spike times by a condition column

    Assumes that there is one column tracking the trial index, and
    one sortable column tracking the trial condition, usually
    generated by the `align.to_events` method.

    Parameters
    ----------
    df : pd.DataFrame
        DataFrame containing spike times, as well as trial
        index and condition info
    condition : str
        Name of the column with the trial condition
    index : str
        Name of the column with the trial index

    Returns
    -------
    sorted_df : pd.DataFrame
        DataFrame sorted by condition; the `index` column
        will be overwritten with new trial indices.

    """

    sorted_df = df.sort_values(by=[condition, index])

    new_inds = []
    total_trials = 0

    for condition in sorted_df.event_label.unique():
        temp_df = sorted_df[sorted_df.event_label == condition]
        t, unique_values = np.unique(temp_df.event_index, return_inverse=True)
        new_inds.append(unique_values + total_trials)
        total_trials += len(t)

    sorted_df["event_index"] = np.concatenate(new_inds)

    return sorted_df


sort_by_condition = by_condition
""" Alias for `by_condition` """
