def ts_train_test_split(df_input, n, time_colname, ts_id_colnames=None):
    """
    Group data frame by time series ID and split on last n rows for each group.

    :param df_input: input data frame
    :param n: number of observations in the test set
    :param time_colname: time column
    :param ts_id_colnames: (optional) list of grain column names
    :return train and test data frames
    """
    if ts_id_colnames is None:
        ts_id_colnames = []
    ts_id_colnames_original = ts_id_colnames.copy()
    if len(ts_id_colnames) == 0:
        ts_id_colnames = ['Grain']
        df_input[ts_id_colnames[0]] = 'dummy'
    # Sort by ascending time
    df_grouped = (df_input.sort_values(time_colname).groupby(ts_id_colnames, group_keys=False))
    df_head = df_grouped.apply(lambda dfg: dfg.iloc[:-n])
    df_tail = df_grouped.apply(lambda dfg: dfg.iloc[-n:])
    # drop group column name if it was not originally provided
    if len(ts_id_colnames_original) == 0:
        df_head.drop(ts_id_colnames, axis=1, inplace=True)
        df_tail.drop(ts_id_colnames, axis=1, inplace=True)
    return df_head, df_tail