# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/transforms/transforms.splitting.ipynb (unless otherwise specified).

__all__ = ['random_split', 'stratified_split', 'split_by_ratio', 'last_session_out_split', 'random_split_v2',
           'chrono_split', 'stratified_split_v2']

# Cell
import numpy as np
import math
import pandas as pd
from sklearn.model_selection import train_test_split

from collections import defaultdict
import functools
import operator
from typing import Any, Iterable, Optional, Tuple, Union

from scipy.sparse import coo_matrix

# Internal Cell
def _subset_interactions(mat, idxs):
    idxs = np.array(idxs)

    coo_mat = coo_matrix(
        (mat.data[idxs], (mat.row[idxs], mat.col[idxs])),
        shape=(mat.shape[0], mat.shape[1])
    )

    return coo_mat


def _validate_val_p_and_test_p(val_p: float, test_p: float) -> None:
    validate_and_test_p = val_p + test_p

    if val_p >= 1 or val_p < 0:
        raise ValueError('``val_p`` must be in the range [0, 1).')
    if test_p >= 1 or test_p < 0:
        raise ValueError('``test_p`` must be in the range [0, 1).')
    if validate_and_test_p >= 1 or validate_and_test_p <= 0:
        raise ValueError('The sum of ``val_p`` and ``test_p`` must be in the range (0, 1).')

# Cell
def random_split(mat, val_p = 0.0, test_p = 0.2, seed = 42):
    """Randomly split interactions into training, validation, and testing sets."""

    _validate_val_p_and_test_p(val_p=val_p, test_p=test_p)

    np.random.seed(seed)

    num_interactions = mat.nnz

    shuffle_indices = np.arange(num_interactions)
    np.random.shuffle(shuffle_indices)

    interactions = _subset_interactions(mat=mat, idxs=shuffle_indices)

    validate_and_test_p = val_p + test_p
    validate_cutoff = int((1.0 - validate_and_test_p) * num_interactions)
    test_cutoff = int((1.0 - test_p) * num_interactions)

    train_idxs = np.arange(validate_cutoff)
    validate_idxs = np.arange(validate_cutoff, test_cutoff)
    test_idxs = np.arange(test_cutoff, num_interactions)

    train_interactions = _subset_interactions(mat=mat, idxs=train_idxs)
    test_interactions = _subset_interactions(mat=mat, idxs=test_idxs)

    if val_p > 0:
        validate_interactions = _subset_interactions(mat=mat, idxs=validate_idxs)

        return train_interactions, validate_interactions, test_interactions
    else:
        return train_interactions, test_interactions

# Internal Cell
def _stratified_split(mat, test_p, seed):
    users = mat.row
    unique_users = set(users)

    # while we should be able to run ``np.where(users == user)[0]`` to find all items each user
    # interacted with, by building up a dictionary to get these values instead, we can achieve the
    # same result in O(N) complexity rather than O(M * N), a nice timesave to have when working with
    # larger datasets
    all_idxs_for_users_dict = defaultdict(list)
    for idx, user in enumerate(users):
        all_idxs_for_users_dict[user].append(idx)

    test_idxs = [
        _stratified_split_parallel_worker(idxs_to_split=all_idxs_for_users_dict[user],
                                            test_p=test_p,
                                            seed=(seed + user))
        for user in unique_users
    ]

    # reduce the list of lists down to a 1-d list
    test_idxs = functools.reduce(operator.iconcat, test_idxs, [])
    # find all indices not in test set - they are now train
    train_idxs = list(set(range(len(users))) - set(test_idxs))

    train_interactions = _subset_interactions(mat=mat,
                                              idxs=train_idxs)
    test_interactions = _subset_interactions(mat=mat,
                                             idxs=test_idxs)

    return train_interactions, test_interactions


def _stratified_split_parallel_worker(idxs_to_split: Iterable[Any],
                                      test_p: float, seed: int) -> np.array:
    _, test_idxs = train_test_split(idxs_to_split,
                                    test_size=test_p,
                                    random_state=seed,
                                    shuffle=True,
                                    stratify=np.ones_like(idxs_to_split))

    return test_idxs

# Cell
def stratified_split(mat, val_p = 0.0, test_p = 0.2, seed = 42):
    """Split into train, validate, and test datasets in a stratified
    manner such that each user appears at least once in each of the datasets.
    This split guarantees that every user will be represented in the training, validation, and
    testing datasets given they appear at least three times. If ``val_p ==
    0``, they will appear in the training and testing datasets given they appear at least two times.
    If a user appears fewer than this number of times, a ``ValueError`` will
    be raised."""

    _validate_val_p_and_test_p(val_p=val_p, test_p=test_p)

    np.random.seed(seed)

    train, test = _stratified_split(mat=mat,
                                    test_p=test_p,
                                    seed=seed)

    if val_p > 0:
        train, validate = _stratified_split(mat=train,
                                            test_p=val_p / (1 - test_p),
                                            seed=seed)

        return train, validate, test
    else:
        return train, test

# Internal Cell
def groupby_user(user_indices):
    users, user_position, user_counts = np.unique(user_indices,
                                                  return_inverse=True,
                                                  return_counts=True)
    user_split_indices = np.split(np.argsort(user_position, kind="mergesort"),
                                  np.cumsum(user_counts)[:-1])
    return user_split_indices


def _pad_unknown_item(data_list):
    train_data, test_data = data_list
    n_items = train_data.item.nunique()
    unique_items = set(train_data.item.tolist())
    test_data.loc[~test_data.item.isin(unique_items), "item"] = n_items
    return train_data, test_data


def _filter_unknown_user_item(data_list):
    train_data, test_data = data_list
    unique_values = dict(user=set(train_data.user.tolist()),
                         item=set(train_data.item.tolist()))

    print(f"test data size before filtering: {len(test_data)}")
    out_of_bounds_row_indices = set()
    for col in ["user", "item"]:
        for j, val in enumerate(test_data[col]):
            if val not in unique_values[col]:
                out_of_bounds_row_indices.add(j)

    mask = np.arange(len(test_data))
    test_data_clean = test_data[~np.isin(mask, list(out_of_bounds_row_indices))]
    print(f"test data size after filtering: {len(test_data_clean)}")
    return train_data, test_data_clean

# Cell
def split_by_ratio(data, shuffle=False, test_size=None, pad_unknown=True,
                   filter_unknown=False, seed=42):
    np.random.seed(seed)
    data = data.copy()
    n_users = data.user.nunique()
    user_indices = data.user.to_numpy()
    user_split_indices = groupby_user(user_indices)

    split_indices_all = [[], []]
    for u in range(n_users):
        u_data = user_split_indices[u]
        u_data_len = len(u_data)
        if u_data_len <= 3:   # keep items of rare users in trainset
            split_indices_all[0].extend(u_data)
        else:
            train_threshold = round((1 - test_size) * u_data_len)
            split_indices_all[0].extend(list(u_data[:train_threshold]))
            split_indices_all[1].extend(list(u_data[train_threshold:]))

    if shuffle:
        split_data_all = tuple(
            np.random.permutation(data[idx]) for idx in split_indices_all
        )
    else:
        split_data_all = list(data.iloc[idx] for idx in split_indices_all)

    if pad_unknown:
        split_data_all = _pad_unknown_item(split_data_all)
    elif filter_unknown:
        split_data_all = _filter_unknown_user_item(split_data_all)
    return split_data_all

# Internal Cell
def clean_split(train, test):
    """
    Remove new items from the test set.
    :param train: The training set.
    :param test: The test set.
    :return: The cleaned training and test sets.
    """
    train_items = set()
    train['sequence'].apply(lambda seq: train_items.update(set(seq)))
    test['sequence'] = test['sequence'].apply(lambda seq: [it for it in seq if it in train_items])
    return train, test

# Cell
def last_session_out_split(data,
                           user_key='user_id',
                           session_key='session_id',
                           time_key='ts'):
    """
    Assign the last session of every user to the test set and the remaining ones to the training set
    """
    sessions = data.sort_values(by=[user_key, time_key]).groupby(user_key)[session_key]
    last_session = sessions.last()
    train = data[~data.session_id.isin(last_session.values)].copy()
    test = data[data.session_id.isin(last_session.values)].copy()
    train, test = clean_split(train, test)
    return train, test

# Internal Cell
def process_split_ratio(ratio):
    """Generate split ratio lists.
    Args:
        ratio (float or list): a float number that indicates split ratio or a list of float
        numbers that indicate split ratios (if it is a multi-split).
    Returns:
        tuple:
        - bool: A boolean variable multi that indicates if the splitting is multi or single.
        - list: A list of normalized split ratios.
    """
    if isinstance(ratio, float):
        if ratio <= 0 or ratio >= 1:
            raise ValueError("Split ratio has to be between 0 and 1")

        multi = False
    elif isinstance(ratio, list):
        if any([x <= 0 for x in ratio]):
            raise ValueError(
                "All split ratios in the ratio list should be larger than 0."
            )

        # normalize split ratios if they are not summed to 1
        if math.fsum(ratio) != 1.0:
            ratio = [x / math.fsum(ratio) for x in ratio]

        multi = True
    else:
        raise TypeError("Split ratio should be either float or a list of floats.")

    return multi, ratio


def _get_column_name(name, col_user, col_item):
    if name == "user":
        return col_user
    elif name == "item":
        return col_item
    else:
        raise ValueError("name should be either 'user' or 'item'.")


def split_pandas_data_with_ratios(data, ratios, seed=42, shuffle=False):
    """Helper function to split pandas DataFrame with given ratios
    .. note::
        Implementation referenced from `this source <https://stackoverflow.com/questions/38250710/how-to-split-data-into-3-sets-train-validation-and-test>`_.
    Args:
        data (pandas.DataFrame): Pandas data frame to be split.
        ratios (list of floats): list of ratios for split. The ratios have to sum to 1.
        seed (int): random seed.
        shuffle (bool): whether data will be shuffled when being split.
    Returns:
        list: List of pd.DataFrame split by the given specifications.
    """
    if math.fsum(ratios) != 1.0:
        raise ValueError("The ratios have to sum to 1")

    split_index = np.cumsum(ratios).tolist()[:-1]

    if shuffle:
        data = data.sample(frac=1, random_state=seed)

    splits = np.split(data, [round(x * len(data)) for x in split_index])

    # Add split index (this makes splitting by group more efficient).
    for i in range(len(ratios)):
        splits[i]["split_index"] = i

    return splits


def min_rating_filter_pandas(
    data,
    min_rating=1,
    filter_by="user",
    col_user='USERID',
    col_item='ITEMID',
):
    """Filter rating DataFrame for each user with minimum rating.
    Filter rating data frame with minimum number of ratings for user/item is usually useful to
    generate a new data frame with warm user/item. The warmth is defined by min_rating argument. For
    example, a user is called warm if he has rated at least 4 items.
    Args:
        data (pandas.DataFrame): DataFrame of user-item tuples. Columns of user and item
            should be present in the DataFrame while other columns like rating,
            timestamp, etc. can be optional.
        min_rating (int): minimum number of ratings for user or item.
        filter_by (str): either "user" or "item", depending on which of the two is to
            filter with min_rating.
        col_user (str): column name of user ID.
        col_item (str): column name of item ID.
    Returns:
        pandas.DataFrame: DataFrame with at least columns of user and item that has been filtered by the given specifications.
    """
    split_by_column = _get_column_name(
        filter_by, col_user, col_item
    )

    if min_rating < 1:
        raise ValueError("min_rating should be integer and larger than or equal to 1.")

    return (
        data
        .groupby(split_by_column)
        .filter(lambda x: len(x) >= min_rating)
    )


def _do_stratification(
    data,
    ratio=0.75,
    min_rating=1,
    filter_by="user",
    is_random=True,
    seed=42,
    col_user='USERID',
    col_item='ITEMID',
    col_timestamp='TIMESTAMP',
):
    # A few preliminary checks.
    if not (filter_by == "user" or filter_by == "item"):
        raise ValueError("filter_by should be either 'user' or 'item'.")

    if min_rating < 1:
        raise ValueError("min_rating should be integer and larger than or equal to 1.")

    if col_user not in data.columns:
        raise ValueError("Schema of data not valid. Missing User Col")

    if col_item not in data.columns:
        raise ValueError("Schema of data not valid. Missing Item Col")

    if not is_random:
        if col_timestamp not in data.columns:
            raise ValueError("Schema of data not valid. Missing Timestamp Col")

    multi_split, ratio = process_split_ratio(ratio)

    split_by_column = col_user if filter_by == "user" else col_item

    ratio = ratio if multi_split else [ratio, 1 - ratio]

    if min_rating > 1:
        data = min_rating_filter_pandas(
            data,
            min_rating=min_rating,
            filter_by=filter_by,
            col_user=col_user,
            col_item=col_item,
        )

    # Split by each group and aggregate splits together.
    splits = []

    # If it is for chronological splitting, the split will be performed in a random way.
    df_grouped = (
        data.sort_values(col_timestamp).groupby(split_by_column)
        if is_random is False
        else data.groupby(split_by_column)
    )

    for _, group in df_grouped:
        group_splits = split_pandas_data_with_ratios(
            group, ratio, shuffle=is_random, seed=seed
        )

        # Concatenate the list of split dataframes.
        concat_group_splits = pd.concat(group_splits)

        splits.append(concat_group_splits)

    # Concatenate splits for all the groups together.
    splits_all = pd.concat(splits)

    # Take split by split_index
    splits_list = [
        splits_all[splits_all["split_index"] == x].drop("split_index", axis=1)
        for x in range(len(ratio))
    ]

    return splits_list

# Cell
def random_split_v2(data, ratio=0.8, seed=42):
    """Pandas random splitter.
    The splitter randomly splits the input data.
    Args:
        data (pandas.DataFrame): Pandas DataFrame to be split.
        ratio (float or list): Ratio for splitting data. If it is a single float number
            it splits data into two halves and the ratio argument indicates the ratio
            of training data set; if it is a list of float numbers, the splitter splits
            data into several portions corresponding to the split ratios. If a list is
            provided and the ratios are not summed to 1, they will be normalized.
        seed (int): Seed.
    Returns:
        list: Splits of the input data as pandas.DataFrame.
    """
    multi_split, ratio = process_split_ratio(ratio)

    if multi_split:
        splits = split_pandas_data_with_ratios(data, ratio, shuffle=True, seed=seed)
        splits_new = [x.drop("split_index", axis=1) for x in splits]

        return splits_new
    else:
        return train_test_split(data, test_size=None, train_size=ratio, random_state=seed)

# Cell
def chrono_split(
    data,
    ratio=0.75,
    min_rating=1,
    filter_by="user",
    col_user='USERID',
    col_item='ITEMID',
    col_timestamp='TIMESTAMP',
):
    """Pandas chronological splitter.
    This function splits data in a chronological manner. That is, for each user / item, the
    split function takes proportions of ratings which is specified by the split ratio(s).
    The split is stratified.
    Args:
        data (pandas.DataFrame): Pandas DataFrame to be split.
        ratio (float or list): Ratio for splitting data. If it is a single float number
            it splits data into two halves and the ratio argument indicates the ratio of
            training data set; if it is a list of float numbers, the splitter splits
            data into several portions corresponding to the split ratios. If a list is
            provided and the ratios are not summed to 1, they will be normalized.
        seed (int): Seed.
        min_rating (int): minimum number of ratings for user or item.
        filter_by (str): either "user" or "item", depending on which of the two is to
            filter with min_rating.
        col_user (str): column name of user IDs.
        col_item (str): column name of item IDs.
        col_timestamp (str): column name of timestamps.
    Returns:
        list: Splits of the input data as pandas.DataFrame.
    """
    return _do_stratification(
        data,
        ratio=ratio,
        min_rating=min_rating,
        filter_by=filter_by,
        col_user=col_user,
        col_item=col_item,
        col_timestamp=col_timestamp,
        is_random=False,
    )

# Cell
def stratified_split_v2(
    data,
    ratio=0.75,
    min_rating=1,
    filter_by="user",
    col_user='USERID',
    col_item='ITEMID',
    seed=42,
):
    """Pandas stratified splitter.
    For each user / item, the split function takes proportions of ratings which is
    specified by the split ratio(s). The split is stratified.
    Args:
        data (pandas.DataFrame): Pandas DataFrame to be split.
        ratio (float or list): Ratio for splitting data. If it is a single float number
            it splits data into two halves and the ratio argument indicates the ratio of
            training data set; if it is a list of float numbers, the splitter splits
            data into several portions corresponding to the split ratios. If a list is
            provided and the ratios are not summed to 1, they will be normalized.
        seed (int): Seed.
        min_rating (int): minimum number of ratings for user or item.
        filter_by (str): either "user" or "item", depending on which of the two is to
            filter with min_rating.
        col_user (str): column name of user IDs.
        col_item (str): column name of item IDs.
    Returns:
        list: Splits of the input data as pandas.DataFrame.
    """
    return _do_stratification(
        data,
        ratio=ratio,
        min_rating=min_rating,
        filter_by=filter_by,
        col_user=col_user,
        col_item=col_item,
        is_random=True,
        seed=seed,
    )