text: # Copyright (c) Microsoft Corporation. All rights reserved.
  # Licensed under the MIT License.
  import numpy as np
  import pandas as pd
  from sklearn.model_selection import train_test_split as sk_split
  
  from recommenders.utils.constants import (
      DEFAULT_ITEM_COL,
      DEFAULT_USER_COL,
      DEFAULT_TIMESTAMP_COL,
  )
  from recommenders.datasets.split_utils import (
      process_split_ratio,
      min_rating_filter_pandas,
      split_pandas_data_with_ratios,
  )
  
  
  def python_random_split(data, ratio=0.75, 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 sk_split(data, test_size=None, train_size=ratio, random_state=seed)
  
  
  def _do_stratification(
      data,
      ratio=0.75,
      min_rating=1,
      filter_by="user",
      is_random=True,
      seed=42,
      col_user=DEFAULT_USER_COL,
      col_item=DEFAULT_ITEM_COL,
      col_timestamp=DEFAULT_TIMESTAMP_COL,
  ):
      # 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
  
  
  def python_chrono_split(
      data,
      ratio=0.75,
      min_rating=1,
      filter_by="user",
      col_user=DEFAULT_USER_COL,
      col_item=DEFAULT_ITEM_COL,
      col_timestamp=DEFAULT_TIMESTAMP_COL,
  ):
      """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,
      )
  
  
  def python_stratified_split(
      data,
      ratio=0.75,
      min_rating=1,
      filter_by="user",
      col_user=DEFAULT_USER_COL,
      col_item=DEFAULT_ITEM_COL,
      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,
      )
  
  
  def numpy_stratified_split(X, ratio=0.75, seed=42):
      """Split the user/item affinity matrix (sparse matrix) into train and test set matrices while maintaining
      local (i.e. per user) ratios.
  
      Main points :
  
      1. In a typical recommender problem, different users rate a different number of items,
      and therefore the user/affinity matrix has a sparse structure with variable number
      of zeroes (unrated items) per row (user). Cutting a total amount of ratings will
      result in a non-homogeneous distribution between train and test set, i.e. some test
      users may have many ratings while other very little if none.
  
      2. In an unsupervised learning problem, no explicit answer is given. For this reason
      the split needs to be implemented in a different way then in supervised learningself.
      In the latter, one typically split the dataset by rows (by examples), ending up with
      the same number of features but different number of examples in the train/test setself.
      This scheme does not work in the unsupervised case, as part of the rated items needs to
      be used as a test set for fixed number of users.
  
      Solution:
  
      1. Instead of cutting a total percentage, for each user we cut a relative ratio of the rated
      items. For example, if user1 has rated 4 items and user2 10, cutting 25% will correspond to
      1 and 2.6 ratings in the test set, approximated as 1 and 3 according to the round() function.
      In this way, the 0.75 ratio is satisfied both locally and globally, preserving the original
      distribution of ratings across the train and test set.
  
      2. It is easy (and fast) to satisfy this requirements by creating the test via element subtraction
      from the original dataset X. We first create two copies of X; for each user we select a random
      sample of local size ratio (point 1) and erase the remaining ratings, obtaining in this way the
      train set matrix Xtst. The train set matrix is obtained in the opposite way.
  
      Args:
          X (numpy.ndarray, int): a sparse matrix to be split
          ratio (float): fraction of the entire dataset to constitute the train set
          seed (int): random seed
  
      Returns:
          numpy.ndarray, numpy.ndarray:
          - Xtr: The train set user/item affinity matrix.
          - Xtst: The test set user/item affinity matrix.
      """
  
      np.random.seed(seed)  # set the random seed
      test_cut = int((1 - ratio) * 100)  # percentage of ratings to go in the test set
  
      # initialize train and test set matrices
      Xtr = X.copy()
      Xtst = X.copy()
  
      # find the number of rated movies per user
      rated = np.sum(Xtr != 0, axis=1)
  
      # for each user, cut down a test_size% for the test set
      tst = np.around((rated * test_cut) / 100).astype(int)
  
      for u in range(X.shape[0]):
          # For each user obtain the index of rated movies
          idx = np.asarray(np.where(Xtr[u] != 0))[0].tolist()
  
          # extract a random subset of size n from the set of rated movies without repetition
          idx_tst = np.random.choice(idx, tst[u], replace=False)
          idx_train = list(set(idx).difference(set(idx_tst)))
  
          # change the selected rated movies to unrated in the train set
          Xtr[u, idx_tst] = 0
          # set the movies that appear already in the train set as 0
          Xtst[u, idx_train] = 0
  
      del idx, idx_train, idx_tst
  
      return Xtr, Xtst
  