# coding=utf-8
# Copyright 2022 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Utility functions for probing experiments."""
import collections
import random
from typing import Dict, Generator, Iterator, List, Set, Text, Tuple

from absl import logging
import numpy as np
import pandas as pd
from scipy import stats
from sklearn import feature_extraction
from sklearn import neighbors



def split_by_length_threshold(
    df, test_set_size):
  """Splits off a test set based on text length.

  Finds a length threshold such that all texts longer than this threshold make
  up a test set that is at max `test_set_size` in size. Note that the exact
  size cannot be guaranteed.

  Args:
    df: DataFrame with 'text_len' column containing the length of every text.
    test_set_size: Number of elements the test set must contain.

  Returns:
    1. Length threshold.
    2. Lengths that exist in the dataset and that are larger than the threshold.
    3. Binary mask in the size of the original data, where True indicates that
      the example must be part of the test set.

  Raises:
    ValueError: In case no threshold can be found such that `test_set_size`
      elements end up in the test set.
  """
  current_count = 0

  # Start from the longest texts.
  for i in range(max(df['text_len']), 0, -1):
    current_count += len(df.loc[df['text_len'] == i, :])

    if current_count > test_set_size:
      return (i,
              set(df['text_len']) & set(range(i + 1,
                                              max(df['text_len']) + 1)),
              df.loc[:, 'text_len'] > i)
  raise ValueError(
      'No length threshold found to create a test set with size {}.'.format(
          test_set_size))


def split_by_random_length(
    df,
    test_set_size,
    size_tolerance = 0.01,
    max_attempts = 100):
  """Creates new train/dev/test split for SentEval data based on text length.

  Finds a random subset of text lengths such that the number of examples having
  these lengths is within an accepted window around `test_set_size`. The
  accepted window is:
    test_set_size * 1 - size_tolerance < window_size < test_set_size * 1 +
    size_tolerance$.

  Since this is a random process, the script can generate multiple such samples.
  Note that the exact size cannot be guaranteed.

  Args:
    df: DataFrame with 'text_len' column containing the length of every text.
    test_set_size: Number of elements the test set must contain. This is the
      middle of the window in which the number of test set examples must fall.
    size_tolerance: Indicates the width of the window according to the above
      formula. This should likely be in [0, 1).
    max_attempts: Number of attempts to try to sample a test set that's size is
      accepted.

  Returns:
    1. Lengths that are reserved for the test set.
    2. Binary mask in the size of the original data, where True indicates that
      the example must be part of the test set.

  Raises:
    RuntimeError: In case no proper length subset can be found such that the
      number of test set examples falls into the window.
  """
  all_lengths = set(df['text_len'])

  for _ in range(max_attempts):
    test_mask = pd.Series([False] * len(df), dtype=bool)
    remaining_lengths = all_lengths.copy()

    while remaining_lengths:
      selected_length = random.choice(tuple(remaining_lengths))
      remaining_lengths.remove(selected_length)
      test_mask.loc[df['text_len'] == selected_length] = True

      # Keep adding lengths until we have a minimum of test examples.
      if test_mask.sum() < test_set_size * (1 - size_tolerance):
        continue
      # If we are still within the tolerance we take this configuration.
      if test_mask.sum() < test_set_size * (1 + size_tolerance):
        return all_lengths - remaining_lengths, test_mask
      # Otherwise we need to start over again.
      break

  raise RuntimeError(
      'No proper split found. Consider increasing the number of attempts or '
      'the size tolerance.')


def split_with_wasserstein(texts, test_set_size,
                           no_of_trials, min_df,
                           leaf_size):
  """Finds test sets by maximizing Wasserstein distances among the given texts.

  This is separating the given texts into training/dev and test sets based on an
  approximate Wasserstein method. First all texts are indexed in a nearest
  neighbors structure. Then a new test centroid is sampled randomly, from which
  the nearest neighbors in Wasserstein space are extracted. Those constitute
  the new test set.
  Similarity is computed based on document-term counts.

  Args:
    texts: Texts to split into training/dev and test sets.
    test_set_size: Number of elements the new test set should contain.
    no_of_trials: Number of test sets requested.
    min_df: Mainly for speed-up and memory efficiency. All tokens must occur at
      least this many times to be considered in the Wasserstein computation.
    leaf_size: Leaf size parameter of the nearest neighbor search. Set high
      values for slower, but less memory-heavy computation.