# coding=utf-8
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

# m4_changequote(`<!', `!>')

import math
from collections import defaultdict

if __name__ != "__main__":
    import plpy
    from utilities.control import MinWarning
    from utilities.utilities import _assert
    from utilities.utilities import extract_keyvalue_params
    from utilities.utilities import unique_string
    from utilities.utilities import collate_plpy_result
    from utilities.utilities import get_grouping_col_str
    from utilities.validate_args import columns_exist_in_table
    from utilities.validate_args import explicit_bool_to_text
    from utilities.validate_args import get_cols
    from utilities.validate_args import table_exists
    from utilities.validate_args import table_is_empty
else:
    # Used only for Unit Testing
    # FIXME: repeating a function from utilities that is needed by the unit test.
    # This should be removed once a unittest framework in used for testing.
    import random
    import time

    def unique_string(desp='', **kwargs):
        """
        Generate random remporary names for temp table and other names.
        It has a SQL interface so both SQL and Python functions can call it.
        """
        r1 = random.randint(1, 100000000)
        r2 = int(time.time())
        r3 = int(time.time()) % random.randint(1, 100000000)
        u_string = "__madlib_temp_" + desp + str(r1) + "_" + str(r2) + "_" + str(r3) + "__"
        return u_string
# ------------------------------------------------------------------------------

UNIFORM = 'uniform'
UNDERSAMPLE = 'undersample'
OVERSAMPLE = 'oversample'
NOSAMPLE = 'nosample'

NEW_ID_COLUMN = '__madlib_id__'
NULL_IDENTIFIER = '__madlib_null_id__'


def _get_level_frequency_distribution(source_table, class_col,
                                      grp_by_cols=None):
    """ Count the number of rows for each class, partitioned by the grp_by_cols

        Returns a dict containing the number of rows associated with each class
        level. Each class level count is converted to a string using ::text.
        None is a valid key in this dict, capturing NULL value in the database.
    """
    if grp_by_cols and grp_by_cols.lower() != 'null':
        is_grouping = True
        grp_by_cols_comma = grp_by_cols + ', '
        array_grp_by_cols_comma = "ARRAY[{0}]".format(grp_by_cols) + " AS group_values, "
    else:
        is_grouping = False
        grp_by_cols_comma = array_grp_by_cols_comma = ""

    # In below query, the inner query groups the data using grp_by_cols + classes
    # and obtains the count for each combination. The outer query then groups
    # again by the grp_by_cols to collect the classes and counts in an array.
    query_result = plpy.execute("""
        SELECT
            -- For each group get the classes and their rows counts
            {grp_identifier} as group_values,
            array_agg(classes) as classes,
            array_agg(class_count) as class_count
        FROM(
            -- for each group and class combination present in source table
            -- get the count of rows for that combination
            SELECT
                {array_grp_by_cols_comma}
                ({class_col})::TEXT AS classes,
                count(*) AS class_count
            FROM {source_table}
            GROUP BY {grp_by_cols_comma} ({class_col})
        ) q
        {meta_grp_by}
     """.format(grp_identifier="group_values" if is_grouping else "NULL",
                meta_grp_by="GROUP BY group_values" if is_grouping else "",
                **locals()))
    if (len(query_result) > 1) != is_grouping:
        # if is_grouping then query_result should have more than 1 row
        # if not is_grouping then query_result should have only 1 row
        raise RuntimeError("Balance sample: Error during frequency level distribution")

    actual_grp_level_counts = {}
    for each_row in query_result:
        # group_values is a list for each row; convert it to a tuple to use as
        # key in a dictionary
        grp = tuple(each_row['group_values']) if is_grouping else None
        grp_levels, grp_counts = each_row['classes'], each_row['class_count']
        actual_grp_level_counts[grp] = dict(zip(grp_levels, grp_counts))
    return actual_grp_level_counts
# ------------------------------------------------------------------------------


def _validate_and_get_sampling_strategy(sampling_strategy_str,
                                        output_table_size,
                                        default=UNIFORM):
    """ Returns the sampling strategy based on the class_sizes input param.
        @param sampling_strategy_str The sampling strategy specified by the
                                         user (class_sizes param)
        @returns:
            Str. One of [UNIFORM, UNDERSAMPLE, OVERSAMPLE]. Default is UNIFORM.
    """
    if not sampling_strategy_str:
        sampling_strategy_str = default
    else:
        if len(sampling_strategy_str) < 3:
            # Require at least 3 characters since UNIFORM and UNDERSAMPLE have
            # common prefix substring
            plpy.error("Sample: Invalid class_sizes parameter")

        supported_strategies = (UNIFORM, UNDERSAMPLE, OVERSAMPLE)
        try:
            # allow user to specify a prefix substring of
            # supported strategies.
            sampling_strategy_str = next(x for x in supported_strategies
                                         if x.startswith(sampling_strategy_str.lower()))
        except StopIteration:
            # next() returns a StopIteration if no element found
            plpy.error("Sample: Invalid class_sizes parameter: "
                       "{0}. Supported class_size parameters are ({1})".
                       format(sampling_strategy_str,
                              ','.join(sorted(supported_strategies))))

    _assert(sampling_strategy_str.lower() in (UNIFORM, UNDERSAMPLE, OVERSAMPLE) or
            (sampling_strategy_str.find('=') > 0),
            "Sample: Invalid class_sizes parameter: "
            "{0}. Supported class_size parameters are ({1})".
            format(sampling_strategy_str,
                   ','.join(sorted(supported_strategies))))

    _assert(not(sampling_strategy_str.lower() == 'oversample' and output_table_size),
            "Sample: Cannot set output_table_size with oversampling.")

    _assert(not(sampling_strategy_str.lower() == 'undersample' and output_table_size),
            "Sample: Cannot set output_table_size with undersampling.")

    return sampling_strategy_str
# ------------------------------------------------------------------------------


def _choose_strategy(actual_count, desired_count):
    """ Choose sampling strategy by comparing actual and desired sample counts

    @param actual_count: Actual number of samples for some level
    @param desired_count: Desired number of sample for the level
    @returns:
        Str. Sampling strategy string (either UNDERSAMPlE or OVERSAMPLE)
    """
    # OVERSAMPLE when the actual count is less than the desired count
    # UNDERSAMPLE when the actual count is more than the desired count

    # If the actual count for a class level is the same as desired count, then
    # we could potentially return the input rows as is.  This, however,
    # precludes the case of bootstrapping (i.e. returning same  number of rows
    # but after sampling with replacement).  Hence, we treat the actual=desired
    # as UNDERSAMPLE.  It's specifically set to UNDERSAMPLE since it provides
    # both 'with' and 'without' replacement  (OVERSAMPLE is always with
    # replacement and NOSAMPLE is always without replacement)
    if actual_count < desired_count:
        return OVERSAMPLE
    else:
        return UNDERSAMPLE
# -------------------------------------------------------------------------


def _get_desired_target_level_counts(desired_level_counts,
                                     actual_grp_level_counts,
                                     output_table_size):
    """ Return the target counts for each group and each class level in the group

        This function is specifically used when the user has provided the
        targets  for all (or a subset) of the levels. The strategy of either
        under or oversampling for each class level is chosen based on the
        desired number of counts for a level, and the actual number of counts
        already present in the input table. This calculation is performed for
        each group (or once if no grouping is used).

        @returns: dict: The key of the dictionary is the group value and the
                        value is another dictionary. The inner dictionary gives
                        the target count for each level present in the group.
    """
    target_grp_level_counts = defaultdict(dict)
    n_grp_values = len(actual_grp_level_counts)
    for each_grp, actual_level_counts in actual_grp_level_counts.items():
        for each_level, desired_count in desired_level_counts.items():
            # actual_grp_level_counts are group specific, whereas
            # desired_level_counts are evenly split among all groups
            try:
                per_grp_desired_count = math.ceil(float(desired_count) / n_grp_values)
                sample_strategy = _choose_strategy(actual_level_counts[each_level],
                                                   per_grp_desired_count)
                target_grp_level_counts[each_grp][each_level] = (
                    per_grp_desired_count, sample_strategy)
            except KeyError:
                plpy.error("Balance sample: Desired class level ({0}) not present in the "
                           "data for each group.".format(each_level))

        # desired levels could contain just a subset of all levels. For the remaining
        # levels in actual_level_counts, compute the desired counts
        remaining_levels = (set(actual_level_counts.keys()) -
                            set(desired_level_counts.keys()))

        #   if 'output_table_size' = NULL, remaining level counts remain as is
        #   if 'output_table_size' = <Integer>, divide remaining count
        #                             uniformly among reamining levels
        if output_table_size:
            # output_table_size is for the whole table and should be split evenly
            # between the groups
            remaining_rows = math.ceil(float(output_table_size -
                                             sum(desired_level_counts.values())) /
                                       n_grp_values)
            if remaining_rows > 0:
                # Uniformly distribute the remaining class levels
                rows_per_level = math.ceil(float(remaining_rows) /
                                           len(remaining_levels))
                for each_level in remaining_levels:
                    sample_strategy = _choose_strategy(
                        actual_level_counts[each_level], rows_per_level)
                    target_grp_level_counts[each_grp][each_level] = (
                        rows_per_level, sample_strategy)
        else:
            # When output_table_size is unspecified, rows from the input table
            # are sampled as is for remaining class levels. This is called as the
            # NOSAMPLE strategy.
            for each_level in remaining_levels:
                target_grp_level_counts[each_grp][each_level] = (
                    actual_level_counts[each_level], NOSAMPLE)
    return target_grp_level_counts
# -------------------------------------------------------------------------


def _get_supported_target_level_counts(sampling_strategy_str,
                                       actual_grp_level_counts,
                                       output_table_size):
    """ Returns the target level counts for all levels when the class_size param
        is one of [uniform, undersample, oversample]. The strategy of
        (under)oversampling for a specific class level is chosen based on the
        computed number of counts for a level, and the actual number of counts
        already present in the input table.

    """
    target_grp_level_counts = defaultdict(dict)
    n_grp_values = len(actual_grp_level_counts)
    for each_grp, actual_level_counts in actual_grp_level_counts.items():
        if sampling_strategy_str == UNIFORM:
            # UNIFORM: Ensure all level counts are same
            if output_table_size:
                # Ignore actual counts for computing target sizes
                # if output_table_size is specified
                total_ = float(output_table_size) / n_grp_values
            else:
                total_ = sum(actual_level_counts.values())
            target_size_per_level = math.ceil(float(total_) /
                                              len(actual_level_counts))
        else:
            # UNDERSAMPLE: Ensure all level counts are same as the minimum count
            # OVERSAMPLE: Ensure all level counts are same as the maximum count
            if sampling_strategy_str == UNDERSAMPLE:
                target_size_per_level = min(actual_level_counts.values())
            elif sampling_strategy_str == OVERSAMPLE:
                target_size_per_level = max(actual_level_counts.values())
            else:
                raise RuntimeError("Balance sample: Invalid "
                                   "sampling_strategy_str encountered")

        for each_level, actual_count in actual_level_counts.items():
            sample_strategy = _choose_strategy(actual_count, target_size_per_level)
            target_grp_level_counts[each_grp][each_level] = (
                target_size_per_level, sample_strategy)
    return target_grp_level_counts
# -------------------------------------------------------------------------


def _get_target_level_counts(sampling_strategy_str, desired_level_counts,
                             actual_grp_level_counts, output_table_size):
    """
    @param sampling_strategy_str: one of [UNIFORM, UNDERSAMPLE, OVERSAMPLE, None].
                               This is 'None' only if this is user-defined, i.e.,
                               a comma separated list of class levels and number
                               of rows desired pairs.
    @param desired_level_counts: Dict that is defined only when the previous arg
                                 sampling_strategy_str is None. This dict would
                                 then contain the class levels and the
                                 corresponding number of rows specified by
                                 the user.
    @param actual_grp_level_counts: Dictionary that provides for each group the
                                  the count of number of rows for each class
                                  present in the group
    @param output_table_size: Size of the desired output table (NULL or Integer)

    @returns:
        Dict. Number of samples to be drawn, and the sampling strategy to be
              used for each class level.
    """

    if not sampling_strategy_str:
        # This case implies user has provided a desired count for one or more
        # levels. Counts for rest of the levels depend on 'output_table_size'.
        target_level_counts = _get_desired_target_level_counts(
            desired_level_counts, actual_grp_level_counts, output_table_size)
    else:
        # This case imples the user has chosen one of
        # [uniform, undersample, oversample] for the class_size parameter.
        target_level_counts = _get_supported_target_level_counts(
            sampling_strategy_str, actual_grp_level_counts, output_table_size)
    return target_level_counts
# -------------------------------------------------------------------------


def _get_sampling_strategy_counts(target_class_sizes):
    """ Return three dicts, one each for undersampling, oversampling, and
        nosampling. The dict contains the number of samples to be drawn for
        each class level.
    """
    undersample_level_counts = {}
    oversample_level_counts = {}
    nosample_level_counts = {}
    for level, (count, strategy) in target_class_sizes.items():
        if strategy == UNDERSAMPLE:
            undersample_level_counts[level] = count
        elif strategy == OVERSAMPLE:
            oversample_level_counts[level] = count
        else:
            nosample_level_counts[level] = count
    return (undersample_level_counts, oversample_level_counts, nosample_level_counts)
# ------------------------------------------------------------------------------


def _get_nosample_subquery(source_table, class_col, nosample_levels,
                           grp_dict=None):
    """ Return the subquery for fetching all rows as is from the input table
        for specific class levels.
    """
    if not nosample_levels:
        return ''
    nosample_level_str = ','.join(["'{0}'".format(level)
                                   for level in nosample_levels if level])
    if grp_dict:
        grp_filter = ' AND ' + ' AND '.join("{0} = '{1}'".format(k, v)
                                            for k, v in grp_dict.items())
    else:
        grp_filter = ''
    return """
        SELECT *
        FROM {source_table}
        WHERE ({class_col} in ({nosample_level_str}) OR
               {class_col} IS NULL)
              {grp_filter}
        """.format(**locals())
# ------------------------------------------------------------------------------


def _get_without_replacement_subquery(schema_madlib, source_table, class_col,
                                      actual_level_counts, desired_level_counts,
                                      grp_dict=None):
    """ Return the subquery for sampling without replacement for specific
        class levels.
    """
    if not desired_level_counts:
        return ''
    class_col_tmp = unique_string(desp='class_col')
    row_number_col = unique_string(desp='row_number')
    desired_count_col = unique_string(desp='desired_count')
    source_table_columns = ','.join(get_cols(source_table))
    null_value_string = "'{0}'".format(NULL_IDENTIFIER)
    desired_level_count_pairs = (','.join("({0}, {1})".
                                 format("'{0}'::text".format(k) if k else null_value_string, v)
                                 for k, v in desired_level_counts.items()))
    desired_level_counts_str = "VALUES " + desired_level_count_pairs
    # Subquery q2 is used to figure out the number of rows to select for each
    # class level. That is used with row_number to order the input rows randomly,
    # and trim the results to the desired number of rows for each class level.
    # q1:
    #    The FROM clause contains information that can be used to figure out the
    #    number of rows to generate for a class level. The tuple basically
    #    contains: (class_level, the_desired_number_of_rows)
    if grp_dict:
        grp_filter = ' AND ' + ' AND '.join("{0} = '{1}'".format(k, v)
                                            for k, v in grp_dict.items())
    else:
        grp_filter = ''
    subquery = """
            SELECT {source_table_columns}
            FROM (
                    SELECT {source_table_columns},
                           row_number() OVER (PARTITION BY {class_col} ORDER BY random()) AS {row_number_col},
                           {desired_count_col}
                    FROM (
                        SELECT {source_table_columns},
                               {desired_count_col}
                        FROM
                            {source_table} s,
                            ({desired_level_counts_str})
                                q1({class_col_tmp}, {desired_count_col})
                        WHERE {class_col_tmp} = coalesce({class_col}::text, '{null_identifier}')
                              {grp_filter}
                    ) q2
                ) q3
            WHERE {row_number_col} <= {desired_count_col}
        """.format(null_identifier=NULL_IDENTIFIER, **locals())
    return subquery
# ------------------------------------------------------------------------------


def _get_with_replacement_subquery(schema_madlib, source_table, class_col,
                                   actual_level_counts, desired_level_counts,
                                   grp_dict=None):
    """ Return the query for sampling with replacement for specific class
        levels. Always used for oversampling since oversampling will always need
        to use replacement. Used for under sampling only if with_replacement
        flag is set to TRUE.
    """
    if not desired_level_counts:
        return ''

    source_table_columns = ','.join(get_cols(source_table))
    class_col_tmp = unique_string(desp='class_col_with_rep')
    desired_count_col = unique_string(desp='desired_count_with_rep')
    actual_count_col = unique_string(desp='actual_count')
    q1_row_no = unique_string(desp='q1_row')
    q2_row_no = unique_string(desp='q2_row')

    null_value_string = "'{0}'".format(NULL_IDENTIFIER)
    desired_and_actual_count = (','.join("({0}, {1}, {2})".
                                format("'{0}'::text".format(k) if k else null_value_string,
                                       v, actual_level_counts[k])
                                for k, v in desired_level_counts.items()))
    desired_and_actual_level_count_str = "VALUES " + desired_and_actual_count

    # q1 and q2 are two sub queries we create to generate the required number of
    # rows per class level.
    # q1:
    #    The FROM clause contains information that can be used to figure out the
    #    number of rows to generate for a class level. The tuple basically
    #    contains:
    #      (class_level, the_desired_number_of_rows, actual_rows_for_level_in_input_table)
    #    The SELECT clause uses generate series to duplicate a row {q1_row_no}
    #    of times, which is a value between 1 and actual_rows_for_level_in_input_table
    # q2:
    #   Replicates the source_table with row IDs starting from 1 through
    #   actual_rows_for_level_in_input_table.
    #
    # The WHERE clause is used to join the two subqueries to obtain the result.
    if grp_dict:
        grp_filter = 'WHERE ' + ' AND '.join("{0} = '{1}'".format(k, v)
                                             for k, v in grp_dict.items())
    else:
        grp_filter = ''
    subquery = """
            SELECT {source_table_columns}
            FROM
                (
                    SELECT
                         {class_col_tmp},
                         generate_series(1, {desired_count_col}::int) AS _i,
                         ((random()*({actual_count_col}-1)+1)::int) AS {q1_row_no}
                    FROM
                        ({desired_and_actual_level_count_str})
                            q({class_col_tmp}, {desired_count_col}, {actual_count_col})
                ) q1,
                (
                    SELECT
                        *,
                        row_number() OVER(PARTITION BY {class_col}) AS {q2_row_no}
                    FROM
                         {source_table}
                    {grp_filter}
                ) q2
            WHERE {class_col_tmp} = coalesce({class_col}::text, '{null_level_val}') AND
                  q1.{q1_row_no} = q2.{q2_row_no}
        """.format(null_level_val=NULL_IDENTIFIER, **locals())
    return subquery
# ------------------------------------------------------------------------------


def balance_sample(schema_madlib, source_table, output_table, class_col,
                   class_sizes, output_table_size, grouping_cols,
                   with_replacement, keep_null, **kwargs):
    """
    Balance sampling function
    Args:
        @param schema_madlib      Schema that MADlib is installed on.
        @param source_table       Input table name.
        @param output_table       Output table name.
        @param class_col          Name of the column containing the class to be
                                  balanced.
        @param class_sizes        Parameter to define the size of the different
                                  class values.
        @param output_table_size  Desired size of the output data set.
        @param grouping_cols      The columns that define the grouping.
        @param with_replacement   The sampling method.
        @param keep_null          Flag to include rows with class level values
                                  NULL. Default is False.

    """
    with MinWarning("warning"):

        desired_sample_per_class = unique_string(desp='desired_sample_per_class')
        desired_counts = unique_string(desp='desired_counts')

        # set all default values
        if not class_sizes:
            class_sizes = UNIFORM
        if not with_replacement:
            with_replacement = False
        keep_null = False if not keep_null else True
        if class_sizes:
            class_sizes = class_sizes.strip()

        _validate_strs(source_table, output_table, class_col,
                       output_table_size, grouping_cols)

        # If keep_null=False, create a view of the input table ignoring NULL
        # values for class levels.
        if keep_null:
            new_source_table = source_table
        else:
            new_source_table = unique_string(desp='source_table')
            plpy.execute("""
                CREATE VIEW {new_source_table} AS
                SELECT * FROM {source_table}
                WHERE {class_col} IS NOT NULL
            """.format(**locals()))
        # class_sizes can be of two forms:
        #   1. A string describing sampling strategy (as described in
        #       _validate_and_get_sampling_strategy).
        #       In this case, 'sampling_strategy_str' is set to one of
        #       [UNIFORM, UNDERSAMPLE, OVERSAMPLE]
        #   2. Class sizes for all (or a subset) of the class levels
        #       In this case, sampling_strategy_str = None and parsed_class_sizes
        #       is used for the sampling.
        parsed_class_sizes = extract_keyvalue_params(class_sizes,
                                                     allow_duplicates=False,
                                                     lower_case_names=False)

        # GREENPLUM 4.3.X does not support bool-to-text cast which is relied
        #  upon for class_col in multiple queries. The explicit_bool_to_text
        #  function wraps the class_col with a MADlib function that provides the
        #  cast just for those platforms that don't provide the cast. For
        #  platforms that provide the cast, class_col is unchanged below.
        class_col = explicit_bool_to_text(source_table,
                                          [class_col],
                                          schema_madlib)[0]
        distinct_sql = ("SELECT DISTINCT ({0})::TEXT as levels FROM {1} ".
                        format(class_col,
                               source_table))
        distinct_levels = collate_plpy_result(plpy.execute(distinct_sql))['levels']
        if not parsed_class_sizes:
            sampling_strategy_str = _validate_and_get_sampling_strategy(
                class_sizes, output_table_size)
        else:
            sampling_strategy_str = None
            try:
                for each_level, each_class_size in parsed_class_sizes.items():
                    _assert(each_level in distinct_levels,
                            "Sample: Invalid class value specified ({0})".
                            format(each_level))
                    each_class_size = int(each_class_size)
                    _assert(each_class_size >= 1,
                            "Sample: Class size has to be greater than zero")
                    parsed_class_sizes[each_level] = each_class_size
            except ValueError:
                plpy.error("Sample: Invalid value for class_sizes ({0})".
                           format(class_sizes))

        # Get the number of rows to be sampled for each class level, based on
        # the input table, class_sizes, and output_table_size params. This also
        # includes info about the resulting sampling strategy, i.e., one of
        # UNDERSAMPLE, OVERSAMPLE, or NOSAMPLE for each level.
        grp_col_str, grp_cols = get_grouping_col_str(
            schema_madlib, 'Balance sample', [NEW_ID_COLUMN, class_col],
            source_table, grouping_cols)
        actual_grp_level_counts = _get_level_frequency_distribution(
            new_source_table, class_col, grp_col_str)
        target_grp_class_sizes = _get_target_level_counts(
            sampling_strategy_str, parsed_class_sizes,
            actual_grp_level_counts, output_table_size)

        is_output_created = False
        grp_cols_list = grp_col_str.split(',') if grp_cols else []

        # for each group
        for grp_vals, actual_level_counts in actual_grp_level_counts.items():
            target_class_sizes = target_grp_class_sizes[grp_vals]

            (undersample_level_counts,
             oversample_level_counts,
             nosample_level_counts) = _get_sampling_strategy_counts(target_class_sizes)

            # grp_dict represents each grouping column and its value
            # for current iteration
            grp_dict = dict(zip(grp_cols_list, grp_vals)) if grp_cols else None

            # Get subqueries for each sampling strategy and union together in
            # one big query.

            # NOSAMPLE are levels that are to be retained in output without sampling
            nosample_subquery = _get_nosample_subquery(
                new_source_table, class_col, nosample_level_counts.keys(), grp_dict)
            # OVERSAMPLE are levels that to be sampled more than existing count
            # (always with replacement)
            oversample_subquery = _get_with_replacement_subquery(
                schema_madlib, new_source_table, class_col,
                actual_level_counts, oversample_level_counts, grp_dict)
            # UNDERSAMPLE are levels that are to be sampled to less than
            # existing count. Undersampling supports both with and without
            # replacement.
            if with_replacement:
                undersample_subquery = _get_with_replacement_subquery(
                    schema_madlib, new_source_table, class_col,
                    actual_level_counts, undersample_level_counts, grp_dict)
            else:
                undersample_subquery = _get_without_replacement_subquery(
                    schema_madlib, new_source_table, class_col,
                    actual_level_counts, undersample_level_counts, grp_dict)

            # Merge the three subqueries using a UNION ALL clause.
            sampling_queries = (undersample_subquery, oversample_subquery, nosample_subquery)
            union_all_subquery = ' UNION ALL '.join(
                ['({0})'.format(subquery)
                 for subquery in sampling_queries if subquery])

            # Populate the output table
            if is_output_created:
                table_header = "INSERT INTO {0}".format(output_table)
            else:
                table_header = "CREATE TABLE {0} AS ".format(output_table)
                is_output_created = True
            final_query = """
                {table_header}
                SELECT row_number() OVER() AS {id_col_name}, *
                FROM (
                    {union_all_subquery}
                ) union_query
            """.format(id_col_name=NEW_ID_COLUMN, **locals())
            plpy.execute(final_query)
        # end of grouping loop

        if not keep_null:
            plpy.execute("DROP VIEW IF EXISTS {0}".format(new_source_table))
# ------------------------------------------------------------------------------


def _validate_strs(source_table, output_table, class_col,
                   output_table_size, grouping_cols):
    _assert(source_table and table_exists(source_table),
            "Sample: Source table ({source_table}) does not exist.".format(**locals()))
    _assert(not table_is_empty(source_table),
            "Sample: Source table ({source_table}) is empty.".format(**locals()))

    _assert(output_table,
            "Sample: Output table name is missing.".format(**locals()))
    _assert(not table_exists(output_table),
            "Sample: Output table ({output_table}) already exists.".format(**locals()))

    _assert(class_col,
            "Sample: Class column name is missing.".format(**locals()))
    _assert(columns_exist_in_table(source_table, [class_col]),
            ("""Sample: Class column ({class_col}) does not exist in""" +
             """ table ({source_table}).""").format(**locals()))

    _assert(not columns_exist_in_table(source_table, [NEW_ID_COLUMN]),
            ("""Sample: Please ensure the source table ({0})""" +
             """ does not contain a column named {1}""").format(source_table, NEW_ID_COLUMN))

    _assert((not output_table_size) or (output_table_size > 0),
            "Sample: Invalid output table size ({output_table_size}).".format(
            **locals()))
# ------------------------------------------------------------------------------


def balance_sample_help(schema_madlib, message, **kwargs):
    """
    Help function for balance_sample

    Args:
        @param schema_madlib
        @param message: string, Help message string
        @param kwargs

    Returns:
        String. Help/usage information
    """
    if not message:
        help_string = """
-----------------------------------------------------------------------
                            SUMMARY
-----------------------------------------------------------------------
Given a table with varying set of records for each class label,
this function will create an output table with a varying types (by
default: uniform) of sampling distributions of each class label. It is
possible to use with or without replacement sampling methods, specify
different proportions of each class, multiple grouping columns and/or
output table size.

For more details on function usage:
    SELECT {schema_madlib}.balance_sample('usage');
            """
    elif message.lower() in ['usage', 'help', '?']:
        help_string = """

Given a table, stratified sampling returns a proportion of records for
each group (strata). It is possible to use with or without replacement
sampling methods, specify a set of target columns, and assume the
whole table is a single strata.

----------------------------------------------------------------------------
                            USAGE
----------------------------------------------------------------------------

 SELECT {schema_madlib}.balance_sample(
    source_table      TEXT,     -- Input table name.
    output_table      TEXT,     -- Output table name.
    class_col         TEXT,     -- Name of column containing the class to be
                                -- balanced.
    class_size        TEXT,     -- (Default: NULL) Parameter to define the size
                                -- of the different class values.
    output_table_size INTEGER,  -- (Default: NULL) Desired size of the output
                                -- data set.
    grouping_cols     TEXT,     -- (Default: NULL) The columns columns that
                                -- defines the grouping.
    with_replacement  BOOLEAN   -- (Default: FALSE) The sampling method.
    keep_null         BOOLEAN   -- (Default: FALSE) Consider class levels with
                                    NULL values or not.

If class_size is NULL, the source table is uniformly sampled.

If output_table_size is NULL, the resulting output table size will depend on
the settings for the ‘class_size’ parameter. It is ignored if ‘class_size’
parameter is set to either ‘oversample’ or ‘undersample’.

If grouping_cols is NULL, the whole table is treated as a single group and
sampled accordingly.

If with_replacement is TRUE, each sample is independent (the same row may
be selected in the sample set more than once). Else (if with_replacement
is FALSE), a row can be selected at most once.
);

The output_table would contain the required number of samples, along with a
new column named __madlib_id__, that contain unique numbers for all
sampled rows.
"""
    else:
        help_string = "No such option. Use {schema_madlib}.balance_sample()"

    return help_string.format(schema_madlib=schema_madlib)


import unittest


class UtilitiesTestCase(unittest.TestCase):
    """
        Comment "import plpy" and replace plpy.error calls with appropriate
        Python Exceptions to successfully run the test cases
    """

    def setUp(self):
        self.input_class_level_counts1 = {None: {'a': 20, 'b': 30, 'c': 25}}
        self.level1a = 'a'
        self.level1a_cnt1 = 15
        self.level1a_cnt2 = 25
        self.level1a_cnt3 = 20

        self.sampling_strategy_str0 = ''
        self.sampling_strategy_str1 = 'uniform'
        self.sampling_strategy_str2 = 'oversample'
        self.sampling_strategy_str3 = 'undersample'
        self.user_specified_class_size0 = ''
        self.user_specified_class_size1 = {'a': 25, 'b': 25}
        self.user_specified_class_size2 = {'b': 25}
        self.user_specified_class_size3 = {'a': 30}
        self.output_table_size1 = None
        self.output_table_size2 = 60
        # self.input_class_level_counts2 = {'a':100, 'b':100, 'c':100}

    def test__choose_strategy(self):
        self.assertEqual(UNDERSAMPLE, _choose_strategy(35, 25))
        self.assertEqual(OVERSAMPLE, _choose_strategy(15, 25))
        self.assertEqual(UNDERSAMPLE, _choose_strategy(25, 25))

    def test__get_target_level_counts(self):
        # Test cases for user defined class level samples, without output table size
        self.assertEqual({None: {'a': (25, OVERSAMPLE), 'b': (25, UNDERSAMPLE), 'c': (25, NOSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str0,
                                                  self.user_specified_class_size1,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size1))
        self.assertEqual({None: {'a': (20, NOSAMPLE), 'b': (25, UNDERSAMPLE), 'c': (25, NOSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str0,
                                                  self.user_specified_class_size2,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size1))
        self.assertEqual({None: {'a': (30, OVERSAMPLE), 'b': (30, NOSAMPLE), 'c': (25, NOSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str0,
                                                  self.user_specified_class_size3,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size1))
        # Test cases for user defined class level samples, with output table size
        self.assertEqual({None: {'a': (25, OVERSAMPLE), 'b': (25, UNDERSAMPLE), 'c': (10, UNDERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str0,
                                                  self.user_specified_class_size1,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size2))
        self.assertEqual({None: {'a': (18, UNDERSAMPLE), 'b': (25, UNDERSAMPLE), 'c': (18, UNDERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str0,
                                                  self.user_specified_class_size2,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size2))
        self.assertEqual({None: {'a': (30, OVERSAMPLE), 'b': (15, UNDERSAMPLE), 'c': (15, UNDERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str0,
                                                  self.user_specified_class_size3,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size2))
        # Test cases for UNIFORM, OVERSAMPLE, and UNDERSAMPLE without any output table size
        self.assertEqual({None: {'a': (25, OVERSAMPLE), 'b': (25, UNDERSAMPLE), 'c': (25, UNDERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str1,
                                                  self.user_specified_class_size0,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size1))
        self.assertEqual({None: {'a': (30, OVERSAMPLE), 'b': (30, UNDERSAMPLE), 'c': (30, OVERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str2,
                                                  self.user_specified_class_size0,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size1))
        self.assertEqual({None: {'a': (20, UNDERSAMPLE), 'b': (20, UNDERSAMPLE), 'c': (20, UNDERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str3,
                                                  self.user_specified_class_size0,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size1))
        # Test cases for UNIFORM with output table size
        self.assertEqual({None: {'a': (20, UNDERSAMPLE), 'b': (20, UNDERSAMPLE), 'c': (20, UNDERSAMPLE)}},
                         _get_target_level_counts(self.sampling_strategy_str1,
                                                  self.user_specified_class_size0,
                                                  self.input_class_level_counts1,
                                                  self.output_table_size2))

    def test__get_sampling_strategy_specific_dict(self):
        # Test cases for getting sampling strategy specific counts
        target_level_counts_1 = {'a': (25, OVERSAMPLE), 'b': (25, UNDERSAMPLE), 'c': (25, NOSAMPLE)}
        target_level_counts_2 = {'a': (25, OVERSAMPLE), 'b': (25, UNDERSAMPLE)}
        target_level_counts_3 = {'a': (25, OVERSAMPLE), 'b': (25, NOSAMPLE), 'c': (25, NOSAMPLE)}
        self.assertEqual(({'b': 25}, {'a': 25}, {'c': 25}),
                         _get_sampling_strategy_counts(target_level_counts_1))
        self.assertEqual(({'b': 25}, {'a': 25}, {}),
                         _get_sampling_strategy_counts(target_level_counts_2))
        self.assertEqual(({}, {'a': 25}, {'c': 25, 'b': 25}),
                         _get_sampling_strategy_counts(target_level_counts_3))


if __name__ == '__main__':
    unittest.main()