# coding=utf-8
"""
@file decision_tree.py_in

@brief Decision Tree: Driver functions

@namespace decision_tree
"""

from __future__ import division
import plpy
from math import sqrt
from operator import itemgetter
from itertools import groupby
from collections import Iterable

from validation.cross_validation import cross_validation_grouping_w_params
from utilities.control import MinWarning
from utilities.control import EnableOptimizer
from utilities.control import EnableHashagg
from utilities.validate_args import get_cols
from utilities.validate_args import get_cols_and_types
from utilities.validate_args import _get_table_schema_names
from utilities.validate_args import get_expr_type
from utilities.validate_args import table_exists
from utilities.validate_args import table_is_empty
from utilities.validate_args import columns_exist_in_table
from utilities.validate_args import is_var_valid
from utilities.validate_args import unquote_ident
from utilities.utilities import _assert
from utilities.utilities import extract_keyvalue_params
from utilities.utilities import unique_string
from utilities.utilities import add_postfix
from utilities.utilities import split_quoted_delimited_str
from utilities.utilities import _array_to_string
from utilities.utilities import py_list_to_sql_string
# ------------------------------------------------------------


def _tree_validate_args(
        split_criterion, training_table_name, output_table_name,
        id_col_name, list_of_features, dependent_variable,
        list_of_features_to_exclude, grouping_cols, weights, max_depth,
        min_split, min_bucket, n_bins, cp, n_folds, **kwargs):
    """ Validate the arguments
    """
    if not split_criterion:
        split_criterion = 'gini'
    if (split_criterion.lower().strip() not in
            ['mse', 'gini', 'cross-entropy', 'entropy',
             'misclass', 'misclassification']):
        plpy.error("Decision tree error: Invalid split_criterion.")

    _assert(training_table_name and
            training_table_name.strip().lower() not in ('null', ''),
            "Decision tree error: Invalid data table.")
    _assert(table_exists(training_table_name),
            "Decision tree error: Data table is missing.")

    _assert(not table_exists(output_table_name, only_first_schema=True),
            "Decision tree error: Output table already exists.")
    _assert(not table_exists(add_postfix(output_table_name, "_summary"), only_first_schema=True),
            "Decision tree error: Output summary table already exists.")

    _assert(columns_exist_in_table(training_table_name, [id_col_name]),
            "Decision tree error: ID column does not exist.")

    _assert(not (dependent_variable is None or dependent_variable.strip().lower() == ''),
            "Decision tree error: Dependent variable is empty.")

    _assert(is_var_valid(training_table_name, dependent_variable),
            "Decision tree error: Invalid dependent variable ({0}).".
            format(dependent_variable))

    _assert(list_of_features and list_of_features.strip(),
            "Decision tree error: Features to include is empty.")

    if list_of_features.strip() != '*':
        _assert(is_var_valid(training_table_name, list_of_features),
                "Decision tree error: Invalid feature list ({0})".
                format(list_of_features))

    if grouping_cols:
        _assert(is_var_valid(training_table_name, grouping_cols),
                "Decision tree error: Invalid grouping column argument.")

    if weights is not None and weights.strip() != '':
        _assert(is_var_valid(training_table_name, weights),
                "Decision tree error: Invalid weights argument.")
    _assert(max_depth >= 0 and max_depth < 100,
            "Decision tree error: maximum tree depth must be positive and less than 100.")

    _assert(cp >= 0, "Decision tree error: cp must be non-negative.")
    _assert(min_split > 0, "Decision tree error: min_split must be positive.")
    _assert(min_bucket > 0, "Decision tree error: min_bucket must be positive.")
    _assert(n_bins > 1, "Decision tree error: number of bins must be at least 2.")
    _assert(n_folds >= 0, "Decision tree error: number of cross-validation "
                          "folds must be non-negative.")
# ------------------------------------------------------------


def _get_features_to_use(schema_madlib, training_table_name,
                         list_of_features, list_of_features_to_exclude,
                         id_col_name, weights, dependent_variable,
                         grouping_cols=None, **kwargs):
    """ Expand '*' syntax and exclude some features

    Ignore 'list_of_features_to_exclude' if 'list_of_features' is not '*'

    We also exclude from the features all grouping_cols, id_col_name, weights and dependent_variable
    """
    # for some of the sets below we include the quoted name and the unquoted name
    #  to allow user to provide either form. Both forms are added in the exclusion
    #  list.
    if grouping_cols:
        group_set = set(split_quoted_delimited_str(grouping_cols))
        group_set |= set(unquote_ident(i)
                         for i in split_quoted_delimited_str(grouping_cols))
    else:
        group_set = set()
    other_col_set = set([id_col_name, weights, dependent_variable])
    other_col_set |= set(unquote_ident(i)
                         for i in [id_col_name, weights, dependent_variable])

    if list_of_features.strip() == '*':
        all_col_set = set(get_cols(training_table_name, schema_madlib))
        exclude_set = set(split_quoted_delimited_str(list_of_features_to_exclude))
        feature_set = all_col_set - exclude_set
        return list(feature_set - group_set - other_col_set)
    else:
        feature_list = split_quoted_delimited_str(list_of_features)
        feature_exclude = split_quoted_delimited_str(list_of_features_to_exclude)
        return_set = set(feature_list) - set(feature_exclude) - group_set - other_col_set
        # instead of returning list(return_set) we create a list that has
        # elements in same order as original feature_list
        return [feat for feat in feature_list if feat in return_set]
# ------------------------------------------------------------


def _dict_get_quoted(input_dict, col_name):
    """Return value from dict where key could be quoted or unquoted name"""
    return input_dict.get(
        col_name, input_dict.get(unquote_ident(col_name)))
# -------------------------------------------------------------------------


def _classify_features(feature_to_type, features):
    """ Returns
    1) an array of categorical features (all casted to string)
    2) an array of continuous features
    3) an array of boolean categorical variables
    """
    # any column belonging to the following types are categorical
    int_types = ['integer', 'smallint', 'bigint']
    text_types = ['text', 'varchar', 'character varying', 'char', 'character']
    boolean_types = ['boolean']
    cat_types = int_types + text_types + boolean_types
    ordered_cat_types = int_types

    cat_features = [c for c in features
                    if _dict_get_quoted(feature_to_type, c) in cat_types]
    ordered_cat_features = [c for c in features if _dict_get_quoted(
                            feature_to_type, c) in ordered_cat_types]

    cat_features_set = set(cat_features)
    # continuous types - 'real' is cast to 'double precision' for uniformity
    con_types = ['real', 'float8', 'double precision']
    con_features = [c for c in features
                    if (c not in cat_features_set and
                        _dict_get_quoted(feature_to_type, c) in con_types)]

    # In order to be able to form an array, all categorical variables
    # will be cast into TEXT type, but GPDB cannot cast a boolean
    # directly into a text. Thus, boolean categorical variables
    # need special treatment: cast them into integers before casting
    # into text.
    boolean_cats = [c for c in features
                    if _dict_get_quoted(feature_to_type, c) in boolean_types]
    return cat_features, ordered_cat_features, con_features, boolean_cats
# ------------------------------------------------------------


def tree_train_help_message(schema_madlib, message, **kwargs):
    """ Help message for Decision Tree
    """
    if not message:
        help_string = """
------------------------------------------------------------
                        SUMMARY
------------------------------------------------------------
Functionality: Decision Tree

Decision trees use a tree-based predictive model to
predict the value of a target variable based on several input variables.

For more details on the function usage:
    SELECT {schema_madlib}.tree_train('usage');
For an example on using this function:
    SELECT {schema_madlib}.tree_train('example');
        """
    elif message.lower().strip() in ['usage', 'help', '?']:
        help_string = """
------------------------------------------------------------
                        USAGE
------------------------------------------------------------
SELECT {schema_madlib}.tree_train(
    'training_table',       -- Data table name
    'output_table',         -- Table name to store the tree model
    'id_col_name',          -- Row ID, used in tree_predict
    'dependent_variable',   -- The column to fit
    'list_of_features',     -- Comma separated column names to be
                                used as the predictors, can be '*'
                                to include all columns except the
                                dependent_variable
    'features_to_exclude',  -- Comma separated column names to be
                                excluded if list_of_features is '*'
    'split_criterion',      -- How to split a node, options are
                                'gini', 'misclassification' and
                                'entropy' for classification, and
                                'mse' for regression.
    'grouping_cols',        -- Comma separated column names used to
                                group the data. A decision tree model
                                will be created for each group. Default
                                is NULL
    'weights',              -- A Column name containing weights for
                                each observation. Default is NULL
    max_depth,              -- Maximum depth of any node, default is 7
    min_split,              -- Minimum number of observations that must
                                exist in a node for a split to be
                                attemped, default is 20
    min_bucket,             -- Minimum number of observations in any
                                terminal node, default is min_split/3
    n_bins,                 -- Number of bins to find possible node
                                split threshold values for continuous
                                variables, default is 20 (Must be greater than 1)
    pruning_params,         -- A comma-separated text containing
                                key=value pairs of parameters for pruning.
                                Parameters accepted:
                                    'cp' - complexity parameter with default=0.01,
                                    'n_folds' - number of cross-validation folds
                                        with default value of 0 (= no cross-validation)
    surrogate_params,       -- A comma-separated text containing
                                key=value pairs of parameters for computing surrogate splits.
                                Parameters accepted:
                                    'max_surrogates' - Maximum number of surrogates to
                                        compute for each split
    verbose                 -- Boolean, whether to print more info, default is False
);

------------------------------------------------------------
                        OUTPUT
------------------------------------------------------------
The output table ('output_table' above) has the following columns (quoted items
are of type TEXT):
    <grouping columns>      -- Grouping columns, only present when
                                'grouping_cols' is not NULL or ''
    tree                    -- The decision tree model as a binary string
    cat_levels_in_text      -- Distinct levels (casted to text) of all
                                categorical variables combined in a single array
    cat_n_levels            -- Number of distinct levels of all categorical variables
    tree_depth              -- Number of levels in the tree (root has level 0)
    pruning_cp              -- The cost-complexity parameter used for pruning
                                the trained tree(s). This would be different
                                from the input cp value if cross-validation is used.

The output summary table ('output_table_summary') has the following columns:
    'method'                -- Method name: 'tree_train'
    'source_table'          -- Data table name
    'model_table'           -- Tree model table name
    'id_col_name'           -- Name of the 'id' column
    is_classification       -- Boolean value indicating if tree is classification or regression
    'dependent_varname'     -- Response variable column name
    'independent_varnames'  -- Comma-separated feature column names
    'cat_features'          -- Comma-separated column names of categorical variables
    'con_features'          -- Comma-separated column names of continuous variables
    'grouping_cols'         -- Grouping column names
    num_all_groups          -- Number of groups
    num_failed_groups       -- Number of groups for which training failed
    total_rows_processed    -- Number of rows used in the model training
    total_rows_skipped      -- Number of rows skipped because NULL values
    dependent_var_levels    -- For classification, the distinct levels of
                                the dependent variable
    dependent_var_type      -- The type of dependent variable
    input_cp                -- The complexity parameter (cp) used for pruning the
                                 trained tree(s) (before cross-validation is run)
    independent_var_types   -- The types of independent variables, comma-separated

        """
    elif message.lower().strip() in ['example', 'examples']:
        help_string = """
------------------------------------------------------------
                        EXAMPLE
------------------------------------------------------------
DROP TABLE IF EXISTS dummy_dt_con_src CASCADE;
CREATE TABLE dummy_dt_con_src (
    id  INTEGER,
    cat INTEGER[],
    con FLOAT8[],
    y   FLOAT8
);

INSERT INTO dummy_dt_src VALUES
(1, '{{0}}'::INTEGER[], ARRAY[0], 0.5),
(2, '{{0}}'::INTEGER[], ARRAY[1], 0.5),
(3, '{{0}}'::INTEGER[], ARRAY[4], 0.5),
(4, '{{0}}'::INTEGER[], ARRAY[4], 0.5),
(5, '{{0}}'::INTEGER[], ARRAY[4], 0.5),
(6, '{{0}}'::INTEGER[], ARRAY[5], 0.1),
(7, '{{0}}'::INTEGER[], ARRAY[6], 0.1),
(8, '{{1}}'::INTEGER[], ARRAY[9], 0.1);
(9, '{{1}}'::INTEGER[], ARRAY[9], 0.1);
(10, '{{1}}'::INTEGER[], ARRAY[9], 0.1);
(11, '{{1}}'::INTEGER[], ARRAY[9], 0.1);

DROP TABLE IF EXISTS tree_out, tree_out_summary;
SELECT madlib.tree_train(
    'dummy_dt_src',
    'tree_out',
    'id',
    'y',
    'cat, con',
    '',
    'mse',
    NULL::Text,
    NULL::Text,
    3,
    2,
    1,
    5);

SELECT madlib.tree_display('tree_out');
        """
    else:
        help_string = "No such option. Use {schema_madlib}.tree_train('usage')"
    return help_string.format(schema_madlib=schema_madlib)
# ------------------------------------------------------------


def _extract_pruning_params(pruning_params_str):
    """
    Args:
        @param pruning_param: str, Parameters used for pruning the tree
                                    cp = Cost-complexity for pruning
                                    n_folds = Number of folds for cross-validation
    Returns:
        dict. A dictionary containing the pruning parameters
    """
    default_dict = dict(cp=0, n_folds=0)
    params_types = dict(cp=float, n_folds=int)
    pruning_params = extract_keyvalue_params(pruning_params_str,
                                             params_types,
                                             default_dict)
    if pruning_params['n_folds'] < 0:
        plpy.error("Decision Tree error: Number of cross-validation folds ({0}) "
                   "must be non-negative".format(pruning_params['n_folds']))
    return pruning_params
# ------------------------------------------------------------


def _get_tree_states(schema_madlib, is_classification, split_criterion,
                     training_table_name, output_table_name, id_col_name,
                     dependent_variable, dep_is_bool,
                     grouping_cols, cat_features, ordered_cat_features,
                     con_features, n_bins, boolean_cats,
                     min_split, min_bucket, weights,
                     max_depth, grp_key_to_cp, compute_cp_list=False,
                     max_n_surr=0, **kwargs):
    """
        Args:
            grp_key_to_cp : Dictionary, mapping from group key to the cp value
                            for each group. For the no grouping case, the
                            key is ''
    """
    filter_dep = _get_filter_str(dependent_variable, grouping_cols)
    # 3)
    if is_classification:
        if split_criterion.lower().strip() == "mse":
            plpy.error("Decision tree error: MSE is not a valid "
                       "split criterion for classification.")
        # For classifications, we also need to map dependent_variable to integers
        n_rows, dep_list = _get_n_and_deplist(training_table_name,
                                              dependent_variable,
                                              filter_dep)
        dep_list.sort()
        if dep_is_bool:
            dep_col_str = ("CASE WHEN {0} THEN 'True' ELSE 'False' END".
                           format(dependent_variable))
        else:
            dep_col_str = dependent_variable
        dep_var_str = ("(CASE " +
                       "\n\t\t".join(["WHEN ({dep_col})::text = $${c}$$ THEN {i}".
                                      format(dep_col=dep_col_str, c=c, i=i)
                                      for i, c in enumerate(dep_list)]) +
                       "\nEND)")
    else:
        if split_criterion.lower().strip() != "mse":
            plpy.warning("Decision tree: Using MSE as split criterion as it "
                         "is the only one supported for regression trees.")
        n_rows = long(plpy.execute("SELECT count(*)::bigint "
                                   "FROM {src} "
                                   "WHERE {filter}".
                                   format(src=training_table_name,
                                          filter=filter_dep))[0]['count'])
        dep_var_str = dependent_variable
        dep_list = []

    dep_n_levels = len(dep_list) if dep_list else 1

    if not grouping_cols:   # non-grouping case
        # 3)  Find the splitting bins, one dict containing two arrays:
        #       categorical bins and continuous bins
        bins = _get_bins(schema_madlib, training_table_name, cat_features,
                         ordered_cat_features, con_features, n_bins,
                         dep_var_str, boolean_cats, n_rows, is_classification,
                         dep_n_levels, filter_dep)
        # some features may be dropped if they have only one value
        cat_features = bins['cat_features']

        # 4) Run tree train till the training is finished
        #  finished: 0 = running, 1 = finished training, 2 = terminated prematurely
        tree = _tree_train_using_bins(**locals())
        tree['grp_key'] = ''
        tree['cp'] = grp_key_to_cp[tree['grp_key']]
        tree_states = [tree]
    else:
        grouping_array_str = get_grouping_array_str(training_table_name, grouping_cols)

        with EnableOptimizer(False):
            # we disable optimizer (ORCA) for platforms that use it
            #  since ORCA doesn't provide an easy way to disable hashagg
            with EnableHashagg(False):
                # we disable hashagg since large number of groups could
                # result in excessive memory usage.
                plpy.notice("Analyzing data to compute split boundaries for variables")
                bins = _get_bins_grps(schema_madlib, training_table_name,
                                      cat_features, ordered_cat_features,
                                      con_features, n_bins,
                                      dep_var_str,
                                      boolean_cats, grouping_cols,
                                      grouping_array_str, n_rows,
                                      is_classification, dep_n_levels,
                                      filter_dep)
                cat_features = bins['cat_features']

                # 3b) Load each group's tree state in memory and set to the initial tree
                tree_states = _tree_train_grps_using_bins(**locals())
                for tree in tree_states:
                    if len(grp_key_to_cp.values()) == 1:
                        # for train w/out CV, the cp value remains the same for
                        # all groups. This is passed as a single-element list.
                        tree['cp'] = grp_key_to_cp.values()[0]
                    else:
                        tree['cp'] = grp_key_to_cp[tree['grp_key']]

    # 5) prune the tree using provided 'cp' value and produce a list of
    #   cp values if cross-validation is required (cp_list = [] if not)
    for tree in tree_states:
        if 'cp' in tree:
            pruned_tree = _prune_and_cplist(schema_madlib, tree,
                                            tree['cp'],
                                            compute_cp_list=compute_cp_list)
            tree['tree_state'] = pruned_tree['tree_state']
            if 'pruned_depth' in pruned_tree:
                tree['pruned_depth'] = pruned_tree['pruned_depth']
            else:
                tree['pruned_depth'] = pruned_tree['tree_depth']
            if 'cp_list' in pruned_tree:
                tree['cp_list'] = pruned_tree['cp_list']

    return tree_states, bins, dep_list, n_rows
# -------------------------------------------------------------------------


def get_grouping_array_str(table_name, grouping_cols, qualifier=None):
    """
    Args:
        @param grouping_cols: list, List of columns used as grouping columns
    """
    if qualifier:
        qualifier_str = qualifier + "."
    else:
        qualifier_str = ''

    all_cols_types = dict(get_cols_and_types(table_name))
    grouping_cols_list = [col.strip() for col in grouping_cols.split(',')]
    grouping_cols_and_types = [(col, _dict_get_quoted(all_cols_types, col))
                               for col in grouping_cols_list]
    grouping_array_str = 'array_to_string(array[' + \
        ','.join("(case when " + col + " then 'True' else 'False' end)::text"
                 if col_type == 'boolean' else '(' + qualifier_str + col + ')::text'
                 for col, col_type in grouping_cols_and_types) + "]::text[], ',')"
    return grouping_array_str
# ------------------------------------------------------------------------------


# XXX This function is used only in decision tree cross-validation.
# So it has some specific code for cv.
def _build_tree(schema_madlib, is_classification, split_criterion,
                training_table_name, output_table_name, id_col_name,
                dependent_variable, dep_is_bool,
                cat_features, ordered_cat_features,
                boolean_cats, con_features, grouping_cols,
                weights, max_depth, min_split, min_bucket, n_bins,
                cp_table, max_n_surr=0, msg_level="warning", k=0, **kwargs):

    compute_cp_list = False
    if grouping_cols:
        grouping_array_str = get_grouping_array_str(cp_table, grouping_cols)
    else:
        grouping_array_str = "''::TEXT"
    sql = """SELECT
                    {0} AS grp_key,
                    explore_value as cp_val
              FROM {1}
          """.format(grouping_array_str, cp_table)
    grp_cp_values = plpy.execute(sql)
    grp_key_to_cp = dict((row['grp_key'], row['cp_val']) for row in grp_cp_values)

    with MinWarning(msg_level):
        plpy.notice("Building tree for cross validation")
        tree_states, bins, dep_list, n_rows = _get_tree_states(**locals())
        all_cols_types = dict(get_cols_and_types(training_table_name))
        n_all_rows = plpy.execute("select count(*) from " + training_table_name
                                  )[0]['count']
        cp = grp_key_to_cp.values()[0]

        # _create_output_tables(...) is a general function
        # we need to let it know that right now it is called from CV.
        use_existing_tables = table_exists(output_table_name)  # create tables if it does not exist
        running_cv = True  # flag to indicate that cv fold ID needs to be included
        _create_output_tables(**locals())
# ------------------------------------------------------------------------------


def tree_train(schema_madlib, training_table_name, output_table_name,
               id_col_name, dependent_variable, list_of_features,
               list_of_features_to_exclude, split_criterion,
               grouping_cols, weights, max_depth,
               min_split, min_bucket, n_bins, pruning_params,
               surrogate_params, verbose_mode, **kwargs):
    """ Decision tree main training function

    Args:
        @param schema_madlib: str, MADlib schema name
        @param training_table_name: str, source table name
        @param output_table_name: str, model table name
        @param id_col_name: str, id column name to uniquely identify each row
        @param dependent_variable: str, dependent variable column name
        @param list_of_features: str, Comma-separated list of feature column names,
                                        can also be '*' implying all columns
                                        except dependent_variable
        @param list_of_features_to_exclude: str, features to exclude if '*' is used
        @param split_criterion: str, Impurity function to use for splitting:
                                  Classification: {'gini', 'entropy', 'misclass'}
                                  Regression: 'mse'
        @param grouping_cols: str, List of grouping columns to group the data
        @param weights: str, Column name for weight for each tuple
        @param max_depth: int, Maximum depth of the tree
        @param min_split: int, Minimum tuples in a node before splitting it
        @param min_bucket: int, Minimum tuples in each child before splitting a node
        @param n_bins: int, Number of bins for quantizing a continuous variables
        @param pruning_params: str, Parameters used for pruning the tree
                                    cp = Cost-complexity for pruning
                                    n_folds = Number of folds for cross-validation
        @param verbose_mode: str, Verbosity of output messages
    """
    msg_level = "notice" if verbose_mode else "warning"

    # Set default values for all arguments
    split_criterion = 'gini' if not split_criterion else split_criterion
    max_depth = 7 if max_depth is None else max_depth
    if min_split is None and min_bucket is None:
        min_split = 20
        min_bucket = 6
    else:
        min_bucket = min_split // 3 if min_bucket is None else min_bucket
        min_split = min_bucket * 3 if min_split is None else min_split
    n_bins = 20 if n_bins is None else n_bins

    # defaults for cp and n_folds set within _extract_pruning_params
    pruning_param_dict = _extract_pruning_params(pruning_params)
    cp = pruning_param_dict['cp']
    n_folds = pruning_param_dict['n_folds']
    surrogate_param_dict = extract_keyvalue_params(surrogate_params,
                                                   dict(max_surrogates=int),
                                                   dict(max_surrogates=0))
    max_n_surr = surrogate_param_dict['max_surrogates']
    _assert(max_n_surr >= 0,
            "Maximum number of surrogates ({0}) should be non-negative".
            format(max_n_surr))

    with MinWarning(msg_level):
        # 1)
        if not grouping_cols or not grouping_cols.strip():
            grouping_cols = ''
        _tree_validate_args(**locals())
        # weights column has be to validated to be present,
        # hence default value allocation happens after the validation
        weights = '1' if not weights or not weights.strip() else weights.strip()

        # expand "*" syntax and exclude some features
        features = _get_features_to_use(**locals())
        _assert(bool(features),
                "Decision tree error: No feature is selected for the model.")

        # 2)
        all_cols_types = dict(get_cols_and_types(training_table_name))
        cat_features, ordered_cat_features, con_features, boolean_cats = \
            _classify_features(all_cols_types, features)
        # get all rows
        n_all_rows = plpy.execute("SELECT count(*) FROM {source_table}".
                                  format(source_table=training_table_name)
                                  )[0]['count']

        is_classification, dep_is_bool = _is_dep_categorical(
            training_table_name, dependent_variable)

        # 4) Build the tree with provided cp value
        compute_cp_list = (n_folds > 1)
        grp_key_to_cp = {'': cp}
        # main training function to get trained decision trees
        plpy.notice("Getting initial tree")
        tree_states, bins, dep_list, n_rows = _get_tree_states(**locals())

        # 5) Perform cross-validation to compute the lowest cp
        dep_n_levels = len(dep_list) if dep_list else 1
        if n_folds > 1:
            plpy.notice("Running cross validation")
            _xvalidate(**locals())

        plpy.notice("Creating output tables")
        _create_output_tables(**locals())
    return None
# ------------------------------------------------------------


def _create_output_tables(schema_madlib, training_table_name, output_table_name,
                          tree_states, bins,
                          split_criterion, cat_features, con_features,
                          id_col_name, dependent_variable,
                          is_classification, n_all_rows, n_rows, dep_list, cp,
                          all_cols_types, grouping_cols=None,
                          use_existing_tables=False, running_cv=False, k=0, **kwargs):
    if not grouping_cols:
        _create_result_table(schema_madlib, tree_states[0],
                             bins['cat_origin'], bins['cat_n'], cat_features,
                             con_features, output_table_name,
                             use_existing_tables, running_cv, k)
    else:
        _create_grp_result_table(
            schema_madlib, tree_states, bins, cat_features,
            con_features, output_table_name, grouping_cols, training_table_name,
            use_existing_tables, running_cv, k)

    failed_groups = sum(row['finished'] != 1 for row in tree_states)
    _create_summary_table(
        schema_madlib, split_criterion, training_table_name,
        output_table_name, id_col_name, cat_features, con_features,
        dependent_variable, failed_groups, is_classification, n_all_rows,
        n_rows, dep_list, all_cols_types, cp, grouping_cols, 1, use_existing_tables,
        running_cv, k)
# -------------------------------------------------------------------------


def _get_n_and_deplist(training_table_name, dependent_variable, filter_null):
    """
    @brief Query the database for the total number of rows and
    levels of dependent variable if the dependent variable is
    categorical.
    """
    sql = """
        SELECT
            sum(n_rows) as n_rows,
            array_agg(dep) as dep
        FROM (
            SELECT
                count(*) as n_rows,
                {dep_var} as dep
            FROM
                {source}
            WHERE {filter_null}
            GROUP BY ({dep_var})
        ) s
        """.format(dep_var=dependent_variable,
                   source=training_table_name,
                   filter_null=filter_null)
    r = plpy.execute(sql)[0]
    r['n_rows'] = 0 if r['n_rows'] is None else r['n_rows']
    return long(r['n_rows']), r['dep']
# ------------------------------------------------------------


def _is_dep_categorical(training_table_name, dependent_variable):
    """
    @brief Sample the dependent variable to check whether it is
    a categorical variable.
    """
    sample_dep = plpy.execute("SELECT " + dependent_variable +
                              " AS dep FROM " +
                              training_table_name + " LIMIT 1")[0]['dep']
    return (not isinstance(sample_dep, float), isinstance(sample_dep, bool))
# ------------------------------------------------------------


def _get_bins(schema_madlib, training_table_name,
              cat_features, ordered_cat_features,
              con_features, n_bins, dependent_variable, boolean_cats, n_rows,
              is_classification, dep_n_levels, filter_null):
    """ Compute the bins of all features

    @param training_table_name Data source table
    @param cat_features A list of strings, categorical column names
    @param con_features A list of strings, continuous column names
    @param n_bins Number of splits equals n_bins - 1
    @param dependent_variable Will be needed when sorting the levels of
    categorical variables
    @param boolean_cats The categorical variables that are of boolean type

    return one dictionary containing two arrays: categorical and continuous
    """
    if len(con_features) > 0:
        if n_bins > n_rows:
            plpy.error("Decision tree error: Number of bins is larger than "
                       "the number of data records.")
        sample_size = n_bins * n_bins  # We use Spark's value here
        # FIXME Hard coded number
        if sample_size < 10000:
            sample_size = 10000

        # Compute the percentage of the sample.
        # We sample a few more values to make sure we can get enough
        # samples, otherwise the number of samples might be smaller
        # than sample_size.
        # Use design doc Eq. (2.1)
        actual_sample_size = sample_size + 14 + sqrt(196 + 28 * sample_size)
        if actual_sample_size > n_rows:
            actual_sample_size = n_rows
        percentage = actual_sample_size / n_rows

        # For continuous variables, use one function to compute
        # the splits for all of them. Similar to the existing
        # _compute_splits function in CoxPH module, but deal with
        # multiple columns together.
        con_features_str = py_list_to_sql_string(con_features, "double precision")

        sample_table_name = unique_string()
        plpy.execute("""
                CREATE TEMP TABLE {sample_table_name} AS
                SELECT *
                FROM (
                    SELECT *, random() AS rand
                    FROM {training_table_name}
                    WHERE {filter_null}
                    AND not {schema_madlib}.array_contains_null({con_features_str})
                ) subq
                WHERE rand <= {percentage}
                """.format(**locals()))

        # The splits for continuous variables
        con_split_str = ("""{schema_madlib}._dst_compute_con_splits(
                                {con_features_str},
                                {actual_sample_size}::integer,
                                {n_bins}::smallint)""".
                         format(**locals()))
        con_splits = plpy.execute("""
                SELECT {con_split_str} as con_splits
                FROM {sample_table_name}
                """.format(**locals()))[0]

        plpy.execute("DROP TABLE {sample_table_name}".format(**locals()))
    else:
        con_splits = {'con_splits': ''}   # no continuous features present

    # For categorical variables, scan the whole table to extract all the
    # levels of the categorical variables, and at the same time
    # sort the levels according to the entropy of the dependent
    # variable.
    # So this aggregate returns a composite type with two columns:
    # col 1 is the array of ordered levels; col 2 is the number of
    # levels in col 1.

    # TODO When n_bins is larger than 2^k - 1, where k is the number
    # of levels of a given categrical feature, we can actually compute
    # all combinations of levels and obtain a complete set of splits
    # instead of using sorting to get an approximate set of splits.
    #
    # We will use integer to represent levels of categorical variables.
    # So before everything, we need to create a mapping from categorical
    # variable levels to integers, and keep this mapping in the memory.
    if len(cat_features) > 0:
        if is_classification:
            # For classifications
            order_fun = ("{madlib}._dst_compute_entropy({dep}, {n})".
                         format(madlib=schema_madlib,
                                dep=dependent_variable,
                                n=dep_n_levels))
        else:
            # For regressions
            order_fun = "AVG({0})".format(dependent_variable)

        # Note that 'sql_cat_levels' goes through two levels of formatting
        # Try to obtain all the levels in one scan of the table.
        # () are needed when casting the categorical variables because
        # they can be expressions.
        sql_cat_levels = """
            SELECT
                '{{col_name}}'::text AS colname,
                levels
            FROM (
                SELECT
                    '{{col_name}}'::text AS colname,
                    array_agg(levels ORDER BY dep_avg) AS levels
                FROM (
                    SELECT
                        ({{col}})::text AS levels,
                        {{order_fun}} AS dep_avg
                    FROM {training_table_name}
                    WHERE {filter_null}
                        AND {{col}} is not NULL
                    GROUP BY {{col}}
                ) s
            ) s1
            WHERE array_upper(levels, 1) > 1
            """.format(training_table_name=training_table_name,
                       filter_null=filter_null)

        sql_all_cats = ' UNION '.join(
            sql_cat_levels.format(
                col="(CASE WHEN " + col + " THEN 'True' ELSE 'False' END)"
                    if col in boolean_cats else col,
                col_name=col,
                order_fun=col if col in ordered_cat_features else order_fun)
            for col in cat_features)
        all_levels = plpy.execute(sql_all_cats)

        if len(all_levels) != len(cat_features):
            plpy.warning("Decision tree warning: Categorical columns with only "
                         "one value are dropped from the tree model.")
            use_cat_features = [row['colname'] for row in all_levels]
            cat_features = [feature for feature in cat_features
                            if feature in use_cat_features]

        col_to_row = dict((row['colname'], i) for i, row in enumerate(all_levels))

        return dict(
            con=con_splits['con_splits'],
            cat_origin=[level for col in cat_features
                        for level in all_levels[col_to_row[col]]['levels']],
            cat_n=[len(all_levels[col_to_row[col]]['levels'])
                   for col in cat_features],
            cat_features=cat_features)
    else:
        # categorical part is empty
        return dict(
            con=con_splits['con_splits'],
            cat_origin=[],
            cat_n=[],
            cat_features=[])
# ------------------------------------------------------------


def _create_result_table(schema_madlib, tree_state,
                         cat_origin, cat_n, cat_features,
                         con_features, output_table_name,
                         use_existing_tables=False, running_cv=False,
                         k=0):
    """ Create the output table and the summary table

    In the result table, we need the tree_state and also the categorical
    sorted levels, which will be used in the printing and the prediction
    functions.
    """
    fold = ", " + str(k) + " as k" if running_cv else ""
    if use_existing_tables:
        # plpy.execute("truncate " + output_table_name)
        header = "insert into " + output_table_name + " "
    else:
        header = "create table " + output_table_name + " as "
    depth = (tree_state['pruned_depth'] if 'pruned_depth' in tree_state
             else tree_state['tree_depth'])
    if len(cat_features) > 0:
        sql = header + """
                SELECT
                    {cp} as pruning_cp,
                    $1 as tree,
                    $2 as cat_levels_in_text,
                    $3 as cat_n_levels,
                    {depth} as tree_depth
                    {fold}
            """.format(depth=depth,
                       cp=tree_state['cp'],
                       fold=fold)
        sql_plan = plpy.prepare(sql, ['{0}.bytea8'.format(schema_madlib),
                                      'text[]', 'integer[]'])
        plpy.execute(sql_plan, [tree_state['tree_state'], cat_origin, cat_n])
    else:
        sql = header + """
                SELECT
                    {cp} as pruning_cp,
                    $1 as tree,
                    NULL::text[] as cat_levels_in_text,
                    NULL::integer[] as cat_n_levels,
                    {depth} as tree_depth
                    {fold}
            """.format(depth=depth,
                       cp=tree_state['cp'],
                       fold=fold)
        sql_plan = plpy.prepare(sql, ['{0}.bytea8'.format(schema_madlib)])
        plpy.execute(sql_plan, [tree_state['tree_state']])

# ------------------------------------------------------------


def _get_bins_grps(
        schema_madlib, training_table_name, cat_features, ordered_cat_features,
        con_features, n_bins, dependent_variable, boolean_cats,
        grouping_cols, grouping_array_str, n_rows, is_classification,
        dep_n_levels, filter_null):
    """ Compute the bins for all features in each group

    @brief Similar to _get_bins except that this is for multiple groups.
    So please refer to the comments inside _get_bins for more related
    information.

    returns what _one_step_for_grps needs:
    """
    if len(con_features) > 0:
        if n_bins > n_rows:
            plpy.error("Decision tree error: Number of bins is larger than "
                       "the number of data records.")
        sample_size = n_bins * n_bins  # We use Spark's value here
        # FIXME Hard coded number
        if sample_size < 10000:
            sample_size = 10000

        # Compute the percentage of the sample.
        # We sample a few more values to make sure we can get enough
        # samples, otherwise the number of samples might be smaller
        # than sample_size.
        # Use design doc Eq. (2.1)
        # FIXME Should also use this for the CoxPH module
        sample_size = sample_size + 14 + sqrt(196 + 28 * sample_size)

        grp_key_str = unique_string()
        n_per_seg_str = unique_string()
        grp_size_str = unique_string()
        random_str = unique_string()

        con_features_str = py_list_to_sql_string(con_features, "double precision")
        sample_table_name = unique_string()
        sql = """
                CREATE TEMP TABLE {sample_table_name} AS
                SELECT *
                FROM
                (
                    SELECT
                        *,
                        random() AS rand,
                        CASE WHEN {grp_size_str} < {sample_size}
                            THEN {grp_size_str}
                            ELSE {sample_size}
                        END AS {n_per_seg_str}
                    FROM
                    (
                        SELECT *
                        FROM {training_table_name}
                        WHERE {filter_null}
                        AND not {schema_madlib}.array_contains_null({con_features_str})
                    ) src
                    JOIN
                    (
                        SELECT {grouping_cols}, count(*) AS {grp_size_str}
                        FROM {training_table_name}
                        WHERE {filter_null}
                            AND not {schema_madlib}.array_contains_null({con_features_str})
                        GROUP BY {grouping_cols}
                    ) grp_info
                    USING ({grouping_cols})
                ) subq
                WHERE rand * {grp_size_str} <= {sample_size}
                """.format(**locals())
        plpy.execute(sql)

        # splits is a list, each of whose elements is a dictionary.
        # The dictionary contains 2 items:
        # 1) grp_key - It is an array of text
        # 2) con_splits - continuous split array
        con_split_str = """{schema_madlib}._dst_compute_con_splits(
                {con_features_str},
                {n_per_seg}::integer,
                {n_bins}::smallint)""".format(con_features_str=con_features_str,
                                              schema_madlib=schema_madlib,
                                              n_per_seg=n_per_seg_str,
                                              n_bins=n_bins)
        con_splits_all = plpy.execute(
            """ SELECT
                    {con_split_str} AS con_splits,
                    {grouping_array_str} AS grp_key
                FROM {sample_table_name}
                GROUP BY {grouping_cols}
                """.format(**locals())   # multiple rows
        )

        plpy.execute("DROP TABLE {sample_table_name}".format(**locals()))

    if cat_features:
        if is_classification:
            # For classifications
            order_fun = ("{schema_madlib}._dst_compute_entropy("
                         "{dependent_variable}, {n})".
                         format(schema_madlib=schema_madlib,
                                dependent_variable=dependent_variable,
                                n=dep_n_levels))
        else:
            order_fun = "avg({0})".format(dependent_variable)

        sql_cat_levels = """
                SELECT
                    colname::text,
                    levels::text[],
                    grp_key::text
                from (
                    SELECT
                        grp_key,
                        '{{col_name}}' as colname,
                        array_agg(levels order by dep_avg) as levels
                    from (
                        SELECT
                            {grouping_array_str} as grp_key,
                            ({{col}})::text as levels,
                            {{order_fun}} as dep_avg
                        FROM {training_table_name}
                        WHERE {filter_null}
                            AND {{col}} is not NULL
                        GROUP BY {{col}}, {grouping_cols}
                    ) s
                    GROUP BY grp_key
                ) s1
                where array_upper(levels, 1) > 1
                """.format(**locals())

        sql_all_cats = ' UNION ALL '.join(
            sql_cat_levels.format(
                col=("(CASE WHEN " + col + " THEN 'True' ELSE 'False' END)"
                     if col in boolean_cats else col),
                col_name=col,
                order_fun=col if col in ordered_cat_features else order_fun)
            for col in cat_features)

        all_levels = list(plpy.execute(sql_all_cats))
        all_levels.sort(key=itemgetter('grp_key'))

        use_cat_features = set([row['colname'] for row in all_levels])
        if len(use_cat_features) != len(cat_features):
            plpy.warning("Decision tree warning: Categorical columns with only "
                         "one value are dropped from the tree model.")
            cat_features = [feature for feature in cat_features if feature in use_cat_features]

        grp_to_col_to_row = dict((grp_key, dict(
            (row['colname'], row['levels']) for row in items))
            for grp_key, items in groupby(all_levels, key=itemgetter('grp_key')))

    if cat_features:
        cat_items_list = [rows[col] for col in cat_features
                          for grp_key, rows in grp_to_col_to_row.items() if col in rows]
        cat_n = [len(i) for i in cat_items_list]
        cat_origin = [item for subl in cat_items_list for item in subl]
        grp_key_cat=[grp_key for grp_key in grp_to_col_to_row]
    else:
        cat_n = []
        cat_origin = []
        grp_key_cat=[con_splits['grp_key'] for con_splits in con_splits_all]

    if con_features:
        con = [con_splits['con_splits'] for con_splits in con_splits_all]
        grp_key_con=[con_splits['grp_key'] for con_splits in con_splits_all]
    else:
        grp_key_con=[grp_key for grp_key in grp_to_col_to_row]
        con = [''] * len(grp_key_con)

    return dict(con=con,
                grp_key_con=grp_key_con,
                cat_origin=cat_origin,
                cat_n=cat_n,
                cat_features=cat_features,
                grp_key_cat=grp_key_cat,
                grouping_array_str=grouping_array_str)
# ------------------------------------------------------------


def get_feature_str(schema_madlib, boolean_cats,
                    cat_features, con_features,
                    levels_str, n_levels_str):
    if len(cat_features) > 0:
        unique_val = unique_string()
        cat_features_cast = []
        for col in cat_features:
            if col in boolean_cats:
                cat_features_cast.append(
                    "(CASE WHEN " + col + " THEN 'True' ELSE 'False' END)::text")
            else:
                cat_features_cast.append(
                    "(coalesce(" + col + "::text,'{0}')".format(unique_val) +
                    ")::text")

        cat_features_str = ("{0}._map_catlevel_to_int(array[" +
                            ", ".join(cat_features_cast) + "], {1}, {2})"
                            ).format(schema_madlib, levels_str, n_levels_str)
    else:
        cat_features_str = "NULL"

    if len(con_features) > 0:
        con_features_list = ["coalesce(" + str(con) + ", 'nan'::double precision)"
                             for con in con_features]
        con_features_str = "ARRAY[" + ", ".join(con_features_list) + "]"
    else:
        con_features_str = "NULL"
    return cat_features_str, con_features_str
# -------------------------------------------------------------------------


def _one_step(schema_madlib, training_table_name, cat_features,
              con_features, boolean_cats, bins, n_bins, tree_state, weights,
              dep_var, min_split, min_bucket, max_depth, filter_null,
              dep_n_levels, subsample, n_random_features, max_n_surr=0):
    """ One step of tree training

    @param tree_state A big double precision array that conatins
    (1) internal node: column and split value
    (2) leaf node: the statistics described as a series of numbers
    """
    # The function _map_catlevel_to_int maps a categorical variable value to its
    # integer representation. It returns an integer array.
    # XXX cat_feature_str contains $5 and $2, and a SQL function
    bytea8 = schema_madlib + '.bytea8'
    cat_features_str, con_features_str = get_feature_str(schema_madlib,
                                                         boolean_cats,
                                                         cat_features,
                                                         con_features,
                                                         "$3", "$2")

    # The arguments of the aggregate (in the same order):
    # 1. current tree state, madlib.bytea8
    # 2. categorical features (integer format) in a single array
    # 3. continuous features in a single array
    # 4. weight value
    # 5. categorical sorted levels (integer format) in a combined array
    # 6. continuous splits
    # 7. number of dependent levels
    train_sql = """
        SELECT (result).* from (
            SELECT
                {schema_madlib}._dt_apply($1,
                    {schema_madlib}._compute_leaf_stats(
                        $1,
                        {cat_features_str},
                        {con_features_str},
                        {dep_var},
                        {weights},
                        $2,
                        $4,
                        {dep_n_levels}::smallint,
                        {subsample}::boolean
                    ),
                    $4,
                    {min_split}::smallint,
                    {min_bucket}::smallint,
                    {max_depth}::smallint,
                    {subsample}::boolean,
                    {n_random_features}::integer
                ) as result
            FROM {training_table_name}
            WHERE {filter_null}
        ) s
    """.format(**locals())
    train_sql_plan = plpy.prepare(train_sql, [bytea8, 'integer[]', 'text[]', bytea8])
    # return a new tree state
    updated_tree = plpy.execute(train_sql_plan, [tree_state['tree_state'],
                                                 bins['cat_n'],
                                                 bins['cat_origin'],
                                                 bins['con']
                                                 ])[0]
    # Compute surrogates:
    #   tree_depth outside the scope of dt_apply starts from 0 i.e. 0 depth
    #   implies a single leaf node for the tree. We compute surrogates only for
    #   internal nodes. Hence we at least need the root be an internal node i.e.
    #   we need the tree_depth to be 1 or more.
    if max_n_surr > 0 and updated_tree['tree_depth'] > 0:
        dup_count_expr = weights + "::integer" if subsample else '1'
        surr_sql = """
                SELECT
                    {schema_madlib}._dt_surr_apply(
                        $1,
                        {schema_madlib}._compute_surr_stats(
                            $1,
                            {cat_features_str},
                            {con_features_str},
                            $2,
                            $4,
                            {dup_count_expr}),
                        $4) as result
                FROM {training_table_name}
                WHERE {filter_null}
        """.format(**locals())
        surr_sql_plan = plpy.prepare(surr_sql, [bytea8, 'integer[]', 'text[]', bytea8])
        # return a new tree state
        updated_tree['tree_state'] = plpy.execute(surr_sql_plan,
                                                  [updated_tree['tree_state'],
                                                   bins['cat_n'],
                                                   bins['cat_origin'],
                                                   bins['con']
                                                   ])[0]["result"]
    return updated_tree
# ------------------------------------------------------------


def _one_step_for_grps(
        schema_madlib, training_table_name, cat_features,
        con_features, boolean_cats, bins, n_bins, tree_states, weights,
        grouping_cols, grouping_array_str, dep_var, min_split, min_bucket,
        max_depth, filter_null, dep_n_levels, subsample, n_random_features,
        max_n_surr=0):
    """ One step of trees training with grouping support
    """
    # The function _map_catlevel_to_int maps a categorical variable value to its
    # integer representation. It returns an integer array.
    # XXX cat_feature_str contains $5 and $2, and a SQL function

    bytea8arr = schema_madlib + '.bytea8[]'
    # avoid name conflicts
    grp_key = unique_string()
    tree_state = unique_string()
    con_splits = unique_string()
    cat_n_levels = unique_string()
    finished = unique_string()
    cat_levels_in_text = unique_string()

    cat_features_str, con_features_str = get_feature_str(
        schema_madlib, boolean_cats, cat_features, con_features,
        cat_levels_in_text, cat_n_levels)

    train_apply_func = """
        {schema_madlib}._dt_apply(
                {tree_state},
                agg_result,
                {con_splits},
                {min_split}::smallint,
                {min_bucket}::smallint,
                {max_depth}::smallint,
                {subsample}::boolean,
                {n_random_features}::integer)
        """.format(**locals())

    train_aggregate = """
            {schema_madlib}._compute_leaf_stats(
            {tree_state},
            {cat_features_str},
            {con_features_str},
            {dep_var},
            {weights},
            {cat_n_levels},
            {con_splits},
            {dep_n_levels}::smallint,
            {subsample}::boolean)
        """.format(**locals())

    sql = """
            SELECT
                s1.grp_key,
                {apply_func} as result
            FROM (  SELECT
                        {grouping_array_str} AS grp_key,
                        {aggregate} AS agg_result
                    FROM
                        {training_table_name} as src,
                        ( SELECT
                            grp_key            AS {grp_key},
                            finished           AS {finished},
                            tree_state         AS {tree_state},
                            con_splits         AS {con_splits},
                            cat_n_levels       AS {cat_n_levels},
                            cat_levels_in_text AS {cat_levels_in_text}
                          FROM
                            (   SELECT
                                    unnest($1) AS grp_key,
                                    unnest($2) AS finished,
                                    unnest($3) AS tree_state
                            ) AS tree_state_set
                            JOIN (
                                SELECT
                                    unnest($4) AS grp_key,
                                    unnest($9) AS con_splits
                            ) AS con_splits
                            USING (grp_key)
                            JOIN
                                {schema_madlib}._gen_cat_levels_set($5, $6, $7, $8) AS cat_levels
                            USING (grp_key)
                        ) AS needed_data
                    WHERE {grouping_array_str} = {grp_key}
                        AND {filter_null}
                    GROUP BY {grouping_cols}
                 ) s1
                JOIN (  SELECT
                            grp_key,
                            tree_state AS {tree_state},
                            con_splits AS {con_splits}
                        FROM (  SELECT
                                    unnest($1) AS grp_key,
                                    unnest($3) AS tree_state
                             ) AS tree_state_set
                            JOIN
                             (  SELECT
                                    unnest($4) AS grp_key,
                                    unnest($9) AS con_splits
                             ) AS con_splits
                            USING (grp_key)
                    ) s2
            USING (grp_key)
    """
    train_sql = "SELECT grp_key, (result).* from (" + sql + ") sub"
    train_sql = train_sql.format(aggregate=train_aggregate,
                                 apply_func=train_apply_func,
                                 # check_finished="AND " + finished + " = 0",
                                 **locals())
    train_sql_plan = plpy.prepare(train_sql,
                                  ['text[]', 'integer[]', bytea8arr, 'text[]',
                                   'text[]', 'integer[]', 'integer', 'text[]',
                                   bytea8arr])

    unfinished_trees = [t for t in tree_states if t['finished'] == 0]
    finished_trees = [t for t in tree_states if t['finished'] != 0]

    # Need to cast the output of plpy.execute to 'list' since newer versions
    # of Postgres return a PlyResult object instead of a list
    updated_unfinished = list(plpy.execute(train_sql_plan, [
        [t['grp_key'] for t in unfinished_trees],
        [t['finished'] for t in unfinished_trees],
        [t['tree_state'] for t in unfinished_trees],
        bins['grp_key_con'],
        bins['grp_key_cat'],
        bins['cat_n'],
        len(cat_features),
        bins['cat_origin'],
        bins['con']]))

    if max_n_surr > 0:
        surr_apply_func = """
            {schema_madlib}._dt_surr_apply(
                    {tree_state},
                    agg_result,
                    {con_splits})
            """.format(schema_madlib=schema_madlib,
                       tree_state=tree_state,
                       con_splits=con_splits)

        dup_count_expr = weights + "::integer" if subsample else '1'
        surr_aggregate = """
                {schema_madlib}._compute_surr_stats(
                    {tree_state},
                    {cat_features_str},
                    {con_features_str},
                    {cat_n_levels},
                    {con_splits},
                    {dup_count_expr})
            """.format(**locals())

        # update the tree to compute surrogates for the layer above the just
        # computed leaves
        surr_sql = sql.format(aggregate=surr_aggregate,
                              apply_func=surr_apply_func,
                              check_finished="",
                              **locals())
        surr_sql_plan = plpy.prepare(surr_sql,
                                     ['text[]', 'integer[]', bytea8arr,
                                      'text[]', 'text[]', 'integer[]', 'integer',
                                      'text[]', bytea8arr])
        surr_trees = list(plpy.execute(surr_sql_plan, [
            [t['grp_key'] for t in updated_unfinished],
            [t['finished'] for t in updated_unfinished],
            [t['tree_state'] for t in updated_unfinished],
            bins['grp_key_con'],
            bins['grp_key_cat'],
            bins['cat_n'],
            len(cat_features),
            bins['cat_origin'],
            bins['con']]))

        surr_dict = dict()
        for each_tree in surr_trees:
            surr_dict[each_tree['grp_key']] = each_tree['result']
        for each_tree in updated_unfinished:
            if each_tree['grp_key'] in surr_dict:
                each_tree['tree_state'] = surr_dict[each_tree['grp_key']]

    # return all tree states (finished and unfinished)
    updated_unfinished.extend(finished_trees)

    return updated_unfinished
# ------------------------------------------------------------


def _create_grp_result_table(
        schema_madlib, tree_states, bins, cat_features,
        con_features, output_table_name, grouping_cols,
        training_table_name, use_existing_tables=False,
        running_cv=False, k=0):
    """ Create the output table for grouping case.
    """
    grp_key = unique_string()
    tree_state = unique_string()
    tree_depth = unique_string()
    cp_col = unique_string()
    cat_levels_in_text = unique_string()
    cat_n_levels = unique_string()
    cat_levels_val = cat_levels_in_text if cat_features else "NULL::TEXT[]"
    cat_n_levels_val = cat_n_levels if cat_features else "NULL::INTEGER[]"
    grouping_array_str = bins['grouping_array_str']

    fold = ", " + str(k) + " as k" if running_cv else ""
    if use_existing_tables:
        # plpy.execute("truncate " + output_table_name)
        header = "insert into " + output_table_name + " "
    else:
        header = "create table " + output_table_name + " as "
    sql = header + """
            SELECT
                {grouping_cols},
                {tree_state} as tree,
                {cat_levels_val} as cat_levels_in_text,
                {cat_n_levels_val} as cat_n_levels,
                {tree_depth} as tree_depth,
                {cp_col} as pruning_cp
                {fold}
            FROM (
                SELECT
                    {grouping_cols},
                    {grouping_array_str} as {grp_key}
                FROM {training_table_name}
                group by {grouping_cols}
            ) s1
            JOIN (
                SELECT
                    unnest($1) as {grp_key},
                    unnest($2) as {tree_state},
                    unnest($3) as {tree_depth},
                    unnest($4) as {cp_col}
            ) s2
            USING ({grp_key})
            """
    if cat_features:
        sql += """
                JOIN (
                    SELECT
                        grp_key as {grp_key},
                        cat_n_levels as {cat_n_levels},
                        cat_levels_in_text as {cat_levels_in_text}
                    FROM
                        {schema_madlib}._gen_cat_levels_set($5, $6, $7, $8)

                ) s3
                USING ({grp_key})
            """
    sql = sql.format(**locals())
    if cat_features:
        sql_plan = plpy.prepare(
            sql, ['text[]',
                  '{schema_madlib}.bytea8[]'.format(schema_madlib=schema_madlib),
                  'integer[]', 'double precision[]',
                  'text[]', 'integer[]', 'integer', 'text[]'])
        plpy.execute(sql_plan, [
            [t['grp_key'] for t in tree_states],
            [t['tree_state'] for t in tree_states],
            [t['pruned_depth'] if 'pruned_depth' in t else t['tree_depth']
             for t in tree_states],
            [t['cp'] for t in tree_states],
            bins['grp_key_cat'],
            bins['cat_n'],
            len(cat_features),
            bins['cat_origin']])
    else:
        sql_plan = plpy.prepare(sql, [
            'text[]',
            '{schema_madlib}.bytea8[]'.format(schema_madlib=schema_madlib),
            'integer[]',
            'double precision[]'])
        plpy.execute(sql_plan, [
            [t['grp_key'] for t in tree_states],
            [t['tree_state'] for t in tree_states],
            [t['pruned_depth'] if 'pruned_depth' in t else t['tree_depth']
             for t in tree_states],
            [t['cp'] for t in tree_states]
        ])
# ------------------------------------------------------------


def _get_dep_type(data_table, dep):
    """
    @brief Obtain the dependent_variable type
    """
    table_schema_str, table_name = _get_table_schema_names(data_table)
    dep_type = plpy.execute("""
            SELECT data_type from information_schema.columns
            where
                table_name = '{table_name}' and
                table_schema in {table_schema} and
                column_name = '{dep}'
            """.format(table_name=table_name,
                       table_schema=table_schema_str,
                       dep=dep))
    if dep_type:
        return dep_type[0]['data_type']
    else:
        return get_expr_type(dep, data_table)
# ------------------------------------------------------------


def _create_summary_table(
        schema_madlib, split_criterion,
        training_table_name, output_table_name, id_col_name,
        cat_features, con_features, dependent_variable,
        num_failed_groups, is_classification, n_all_rows, n_rows,
        dep_list, all_cols_types, cp, grouping_cols=None, n_groups=1,
        use_existing_tables=False, running_cv=False, k=0):
    # dependent variables
    features = ', '.join(cat_features + con_features)
    if dep_list:
        dep_list_str = ("$dep_list$" +
                        ','.join('"{0}"'.format(str(dep)) for dep in dep_list) +
                        "$dep_list$")
    else:
        dep_list_str = "NULL"
    indep_type = ', '.join(_dict_get_quoted(all_cols_types, col)
                           for col in cat_features + con_features)
    dep_type = _get_dep_type(training_table_name, dependent_variable)
    cat_features_str = ','.join(cat_features)
    con_features_str = ','.join(con_features)
    if grouping_cols:
        grouping_cols_str = "'{0}'".format(grouping_cols)
    else:
        grouping_cols_str = "NULL"
    n_rows_skipped = n_all_rows - n_rows

    if isinstance(cp, Iterable):
        cp_str = py_list_to_sql_string(cp, 'double precision')
    else:
        cp_str = str(cp) + "::double precision"

    fold = ", " + str(k) + " as k" if running_cv else ""
    if use_existing_tables:
        # plpy.execute("truncate " + output_table_name + "_summary")
        header = "insert into " + add_postfix(output_table_name, "_summary") + " "
    else:
        header = "create table " + add_postfix(output_table_name, "_summary") + " as "

    sql = header + """
            SELECT
                'tree_train'::text             AS method,
                '{is_classification}'::boolean AS is_classification,
                '{training_table_name}'::text  AS source_table,
                '{output_table_name}'::text    AS model_table,
                '{id_col_name}'::text          AS id_col_name,
                '{dependent_variable}'::text   AS dependent_varname,
                '{features}'::text             AS independent_varnames,
                '{cat_features_str}'::text     AS cat_features,
                '{con_features_str}'::text     AS con_features,
                {grouping_cols_str}::text      AS grouping_cols,
                {n_groups}::integer            AS num_all_groups,
                {num_failed_groups}::integer   AS num_failed_groups,
                {n_rows}::integer              AS total_rows_processed,
                {n_rows_skipped}::integer      AS total_rows_skipped,
                {dep_list_str}::text           AS dependent_var_levels,
                '{dep_type}'::text             AS dependent_var_type,
                {cp_str}                       AS input_cp,
                '{indep_type}'::text           AS independent_var_types
                {fold}
        """.format(**locals())
    plpy.execute(sql)
# ------------------------------------------------------------


def _get_filter_str(dependent_variable, grouping_cols):
    """ Return a 'WHERE' clause string that filters out all rows that contain a
    NULL.
    """
    if grouping_cols:
        group_filter = ' and '.join('({0}) is not NULL'.format(g.strip())
                                    for g in grouping_cols.split(','))
    else:
        group_filter = None
    dep_filter = '(' + dependent_variable + ") is not NULL"
    return ' and '.join(filter(None, [group_filter, dep_filter]))
# -------------------------------------------------------------------------


def _validate_predict(schema_madlib, model, source, output, use_existing_tables):
    # validations for inputs
    _assert(source and source.strip().lower() not in ('null', ''),
            "Decision tree error: Invalid data table name: {0}".format(source))
    _assert(table_exists(source),
            "Decision tree error: Data table ({0}) does not exist".format(source))
    _assert(not table_is_empty(source),
            "Decision tree error: Data table ({0}) is empty".format(source))
    _assert(model and
            model.strip().lower() not in ('null', ''),
            "Decision tree error: Invalid model table name: {0}".format(model))
    _assert(table_exists(model),
            "Decision tree error: Model table ({0}) does not exist".format(model))
    _assert(not table_is_empty(model),
            "Decision tree error: Model table ({0}) is empty".format(model))
    model_summary = add_postfix(model, "_summary")
    _assert(table_exists(model_summary),
            "Decision tree error: Model summary table ({0}) does not exist".format(model_summary))
    _assert(not table_is_empty(model_summary),
            "Decision tree error: Model summary table ({0}) is empty".format(model_summary))
    _assert(output and
            output.strip().lower() not in ('null', ''),
            "Decision tree error: Invalid output table name: {0}".format(output))
    if not use_existing_tables:
        _assert(not table_exists(output),
                "Decision tree error: Output table ({0}) already exists".format(output))
    _assert(
        columns_exist_in_table(
            model,
            ["tree", "cat_levels_in_text", "cat_n_levels"],
            schema_madlib),
        "Decision tree error: Invalid model table ({0})".format(model))
    _assert(
        columns_exist_in_table(
            model_summary,
            ["grouping_cols", "id_col_name", "dependent_varname",
             "cat_features", "con_features", "is_classification"],
            schema_madlib),
        "Decision tree error: Invalid model summary table ({0})".format(model_summary))
# -------------------------------------------------------------------------


def tree_predict(schema_madlib, model, source, output, pred_type='response',
                 use_existing_tables=False, k=0, **kwargs):
    """
    Args:
        @param schema_madlib: str, Name of MADlib schema
        @param model: str, Name of table containing the tree model
        @param source: str, Name of table containing prediction data
        @param output: str, Name of table to output the results
        @param pred_type: str, The type of output required:
                            'response' gives the actual response values,
                            'prob' gives the probability of the classes in a
                             classification tree.
                          For regression tree, only type='response' is defined.

    Returns:
        None

    Side effect:
        Creates an output table containing the prediction for given source table

    Throws:
        None
    """
    _validate_predict(schema_madlib, model, source, output, use_existing_tables)
    model_summary = add_postfix(model, "_summary")
    # obtain the cat_features and con_features from model table
    summary_elements = plpy.execute("SELECT * FROM {0}".format(model_summary))[0]

    cat_features = split_quoted_delimited_str(summary_elements["cat_features"])
    con_features = split_quoted_delimited_str(summary_elements["con_features"])
    id_col_name = summary_elements["id_col_name"]
    grouping_cols_str = summary_elements["grouping_cols"]
    dep_varname = summary_elements["dependent_varname"]
    dep_levels = summary_elements["dependent_var_levels"]
    is_classification = summary_elements["is_classification"]
    dep_type = summary_elements['dependent_var_type']

    # find which columns are of type boolean
    boolean_cats = set([key for key, value in get_cols_and_types(source)
                        if value == 'boolean'])

    cat_features_str, con_features_str = get_feature_str(
        schema_madlib, boolean_cats, cat_features, con_features,
        "m.cat_levels_in_text", "m.cat_n_levels")

    if use_existing_tables and table_exists(output):
        plpy.execute("truncate " + output)
        header = "INSERT INTO " + output + " "
        use_fold = 'WHERE k = ' + str(k)
    else:
        header = "CREATE TABLE " + output + " AS "
        use_fold = ''

    if not grouping_cols_str:
        using_str = ""
        join_str = ","
    else:
        using_str = "USING ( " + grouping_cols_str + ")"
        join_str = "LEFT OUTER JOIN"

    pred_name = ('"prob_{0}"' if pred_type == "prob" else
                 '"estimated_{0}"').format(dep_varname.replace('"', '').strip())

    if not is_classification:
        sql = header + """
            SELECT {id_col_name},
                   {schema_madlib}._predict_dt_response(
                        tree,
                        {cat_features_str}::INTEGER[],
                        {con_features_str}::DOUBLE PRECISION[]) as {pred_name}
            FROM {source} as s {join_str} {model} as m {using_str}
            {use_fold}
            """
    else:
        if dep_type.lower() == "boolean":
            # some platforms don't have text to boolean cast. We manually check the string.
            dep_cast_str = ("(case {pred_name} when 'True' then "
                            "True else False end)::BOOLEAN as {pred_name}")
        else:
            dep_cast_str = "{pred_name}::{dep_type}"
        if pred_type == "response":
            sql = header + """
                    SELECT
                        {id_col_name}
                       , %s
                    FROM (
                         SELECT
                            {id_col_name}
                           ,
                           ($sql${{ {dep_levels} }}$sql$::varchar[])[
                               {schema_madlib}._predict_dt_response (
                                   tree,
                                   {cat_features_str}::INTEGER[],
                                   {con_features_str}::DOUBLE PRECISION[]) + 1]::TEXT
                           as {pred_name}
                        FROM {source} as s {join_str} {model} as m {using_str}
                        {use_fold}
                    ) q
                """ % (dep_cast_str)
        else:
            intermediate_col = unique_string()
            score_format = ', \n'.join([
                '{interim}[{j}] as "estimated_prob_{c}"'.
                format(j=i+1, c=c.strip(' "'), interim=intermediate_col)
                for i, c in enumerate(split_quoted_delimited_str(dep_levels))])
            sql = header + """
                SELECT
                    {id_col_name},
                    {score_format}
                    FROM (
                        SELECT {id_col_name},
                               {schema_madlib}._predict_dt_prob(tree,
                                    {cat_features_str}::INTEGER[],
                                    {con_features_str}::DOUBLE PRECISION[])
                                        AS {intermediate_col}
                        FROM {source} as s {join_str} {model} as m {using_str}
                        {use_fold}
                    ) q
                """
    sql = sql.format(**locals())
    with MinWarning('warning'):
        with EnableOptimizer(False):
            # we disable optimizer (ORCA) for platforms that use it
            #  since ORCA doesn't provide an easy way to disable hashagg
            with EnableHashagg(False):
                # we disable hashagg since large number of groups could
                # result in excessive memory usage.
                plpy.execute(sql)
    return None
# -------------------------------------------------------------------------


def _get_display_header(table_name, dep_levels, is_regression, dot_format=True):
    if dot_format:
        tree_type = ("Classification", "Regression")[is_regression]
        return ("digraph {0} {{".format('"' + tree_type + ' tree for ' +
                                        str(table_name) + '"'))
    else:
        return_str = """-------------------------------------
    - Each node represented by 'id' inside ().
    - Each internal nodes has the split condition at the end, while each
        leaf node has a * at the end.
    - For each internal node (i), its child nodes are indented by 1 level
        with ids (2i+1) for True node and (2i+2) for False node.
    """
        if is_regression:
            return_str += ("- Number of rows and average response value inside []. "
                           "For a leaf node, this is the prediction.\n")
        else:
            return_str += """- Number of (weighted) rows for each response variable inside [].'
        The response label order is given as {0}.
        For each leaf, the prediction is given after the '-->'
        """.format(str(dep_levels))
        return_str += "\n-------------------------------------"
    return return_str
# ------------------------------------------------------------------------------


def tree_display(schema_madlib, model_table, dot_format=True, verbose=False,
                 disp_surr=False, **kwargs):

    if dot_format:
        disp_surr = False  # surrogates cannot be displayed in dot format
    bytea8 = schema_madlib + '.bytea8'
    summary_table = add_postfix(model_table, "_summary")
    summary = plpy.execute("SELECT * FROM {summary_table}".
                           format(summary_table=summary_table))[0]
    dep_levels = summary["dependent_var_levels"]
    dep_levels = [''] if not dep_levels else split_quoted_delimited_str(dep_levels)
    table_name = summary["source_table"]
    is_regression = not summary["is_classification"]
    cat_features_str = split_quoted_delimited_str(summary['cat_features'])
    con_features_str = split_quoted_delimited_str(summary['con_features'])

    grouping_cols = summary["grouping_cols"]
    grouping_cols = '' if grouping_cols is None else grouping_cols + ','
    grouped_trees = plpy.execute("SELECT {grouping_cols} "
                                 "tree, cat_levels_in_text, cat_n_levels "
                                 "FROM {model_table}".
                                 format(model_table=model_table,
                                        grouping_cols=grouping_cols))

    return_str_list = []
    if not disp_surr:
        return_str_list.append(_get_display_header(table_name, dep_levels,
                                                   is_regression, dot_format))
    else:
        return_str_list.append("""
  -------------------------------------
      Surrogates for internal nodes
  -------------------------------------
        """)

    with MinWarning('warning'):
        for index, dtree in enumerate(grouped_trees):
            grouping_keys = [k + ' = ' + str(dtree[k]) for k in dtree.keys()
                             if k not in ('tree', 'cat_features', 'con_features',
                                          'cat_n_levels', 'cat_levels_in_text')]
            if grouping_keys:
                group_name = "(" + ','.join(grouping_keys) + ")"
            else:
                group_name = ''
            tree = dtree['tree']
            if dtree['cat_levels_in_text']:
                cat_levels_in_text = dtree['cat_levels_in_text']
                cat_n_levels = dtree['cat_n_levels']
            else:
                cat_levels_in_text = []
                cat_n_levels = []

            if disp_surr:
                if group_name:
                    return_str_list.append("--- Surrogates for tree {0} ---".format(group_name))
                sql = """SELECT {0}._display_decision_tree_surrogate(
                                    $1, $2, $3, $4, $5) as display_tree
                      """.format(schema_madlib)
                # execute sql to get display string
                sql_plan = plpy.prepare(sql, [bytea8,
                                              'text[]', 'text[]', 'text[]',
                                              'int[]'])
                tree_display = plpy.execute(
                    sql_plan, [tree, cat_features_str, con_features_str,
                               cat_levels_in_text, cat_n_levels])[0]
            else:
                if dot_format:
                    return_str_list.append('\t subgraph "cluster{0}"{{'.format(index))
                    return_str_list.append('\t label="{0}"'.format(group_name.replace('"', '\\"')))
                    sql = """
                            SELECT {0}._display_decision_tree(
                                        $1, $2, $3, $4, $5, $6, '{1}', {2}
                                    ) as display_tree
                          """.format(schema_madlib, "g" + str(index) + "_", verbose)
                else:
                    if group_name:
                        return_str_list.append("--- Tree for {0} ---".format(group_name))
                    sql = """
                            SELECT {0}._display_text_decision_tree(
                                        $1, $2, $3, $4, $5, $6) as display_tree
                          """.format(schema_madlib)
                sql_plan = plpy.prepare(sql, [bytea8,
                                              'text[]', 'text[]', 'text[]',
                                              'int[]', 'text[]'])
                tree_display = plpy.execute(
                    sql_plan, [tree, cat_features_str, con_features_str,
                               cat_levels_in_text, cat_n_levels, dep_levels])[0]

            return_str_list.append(tree_display["display_tree"])
            if dot_format:
                return_str_list.append("\t } //--- end of subgraph------------")
            else:
                return_str_list.append("-------------------------------------")
    if dot_format:
        return_str_list.append("} //---end of digraph--------- ")
    return ("\n".join(return_str_list))
# -------------------------------------------------------------------------


def tree_predict_help_message(schema_madlib, message, **kwargs):
    """ Help message for Decision Tree predict
    """
    if not message:
        help_string = """
------------------------------------------------------------
                        SUMMARY
------------------------------------------------------------
Functionality: Decision Tree Prediction

Prediction for a decision tree (trained using {schema_madlib}.tree_train) can
be performed on a new data table.

For more details on the function usage:
    SELECT {schema_madlib}.tree_predict('usage');
For an example on using this function:
    SELECT {schema_madlib}.tree_predict('example');
        """
    elif message.lower().strip() in ['usage', 'help', '?']:
        help_string = """
------------------------------------------------------------
                        USAGE
------------------------------------------------------------
SELECT {schema_madlib}.tree_predict(
    'tree_model',           -- Model table name (output of tree_train)
    'new_data_table',       -- Prediction source table
    'output_table',         -- Table name to store the predictions
    'type'                  -- Type of prediction output
);

Note: The 'new_data_table' should have the same 'id_col_name' column as used
in the training function. This is used to corelate the prediction data row with
the actual prediction in the output table.

------------------------------------------------------------
                        OUTPUT
------------------------------------------------------------
The output table ('output_table' above) has the '<id_col_name>' column giving
the 'id' for each prediction and the prediction columns for the response
variable (also called as dependent variable).

If prediction type = 'response', then the table has a single column with the
prediction value of the response. The type of this column depends on the type
of the response variable used during training.

If prediction type = 'prob', then the table has multiple columns, one for each
possible value of the response variable. The columns are labeled as
'estimated_prob_<dep value>', where <dep value> represents for each value
of the response.
        """
    elif message.lower().strip() in ['example', 'examples']:
        help_string = """
------------------------------------------------------------
                        EXAMPLE
------------------------------------------------------------
-- Assuming the example of tree_train() has been run
SELECT {schema_madlib}.tree_predict(
    'tree_out',
    'dummy_dt_src',
    'tree_predict_out',
    'response'
);

SELECT * FROM tree_predict_out;
        """
    else:
        help_string = "No such option. Use {schema_madlib}.tree_predict('usage')"
    return help_string.format(schema_madlib=schema_madlib)
# ------------------------------------------------------------


def _prune_and_cplist(schema_madlib, tree, cp, compute_cp_list=False):
    """ Prune tree with given cost-complexity parameters
        and return a list of cp values at which tree can be pruned

        Args:
            @param schema_madlib: str, MADlib schema name
            @param tree: Tree data to prune
            @param cp: float, cost-complexity parameter, all splits that have a
                                complexity lower than 'cp' will be pruned
            @param compute_cp_list: bool, optionally return a list of cp values that
                                are the various complexity boundaries for the splits
                                in the tree. This list can be used by cross-validation
                                to explore different cp values.

        Returns:
            Dictionary containing following keys:
                tree_state: pruned tree state
                pruned_depth: depth of tree after pruning
                cp_list: list of cp values at which tree can be pruned
                         (returned only if compute_cp_list=True)
    """
    sql = """
        SELECT (pruned_tree).*
        FROM (
            SELECT {madlib}._prune_and_cplist(
                        $1,
                        ({cp})::double precision,
                        ({compute_cp_list})::boolean
                    ) as pruned_tree
        ) q
    """.format(madlib=schema_madlib, cp=cp,
               compute_cp_list=bool(compute_cp_list))

    sql_plan = plpy.prepare(sql, [schema_madlib + '.bytea8'])
    pruned_tree = plpy.execute(sql_plan, [tree['tree_state']])[0]
    return pruned_tree
# -------------------------------------------------------------------------


def _xvalidate(schema_madlib, tree_states, training_table_name, output_table_name,
               id_col_name, dependent_variable,
               list_of_features, list_of_features_to_exclude,
               cat_features, ordered_cat_features, boolean_cats, con_features,
               split_criterion, grouping_cols, weights, max_depth,
               min_split, min_bucket, n_bins, is_classification,
               dep_is_bool, dep_n_levels, n_folds, n_rows,
               max_n_surr, msg_level='warning', **kwargs):
    """
    Run cross validation for decision tree over multiple cp values

    Returns:
        cp value corresponding to lowest cross validation error

    Side effect:
        Creates a table containing cross-validation error for each cp value
    """
    # 1) create group_to_param_list_table for CV
    group_to_param_list_table = unique_string()
    param_list_name = unique_string()
    if not grouping_cols:
        plan = plpy.prepare("""
                CREATE TEMP TABLE {group_to_param_list_table} AS
                SELECT $1 AS {param_list_name}
                """.format(**locals()), ["float8[]"])
        plpy.execute(plan, [tree_states[0]['cp_list']])
        plpy.notice("Running cross validation for cp values: {0}".format(
                    str(tree_states[0]['cp_list'])))
    else:
        grp_list = []
        cp_list = []
        plpy.notice("Running cross validation for cp values:")
        for tree in tree_states:
            grp_list.extend([tree['grp_key']]*len(tree['cp_list']))
            cp_list.extend(tree['cp_list'])
            plpy.notice("{0} -> {1}".format(tree['grp_key'], str(tree['cp_list'])))

        grouping_array_str = get_grouping_array_str(training_table_name,
                                                    grouping_cols)
        grp_key_str = unique_string()
        plan = plpy.prepare("""
                CREATE TEMP TABLE {group_to_param_list_table} AS
                SELECT
                    {grouping_cols},
                    {param_list_name}
                FROM
                (
                    SELECT {grp_key_str},
                           array_agg(cp_value) as {param_list_name}
                    FROM
                    (
                        SELECT unnest($1) as {grp_key_str},
                               unnest($2) as cp_value
                    ) unnested_grp_cp
                    GROUP BY {grp_key_str}
                ) grp_key_to_cp_list
                JOIN
                (
                    SELECT
                        {grouping_cols},
                        {grouping_array_str} as {grp_key_str}
                    FROM {training_table_name}
                ) grp_key_to_grouping_cols
                USING ({grp_key_str})
                """.format(**locals()), ["text[]", "float8[]"])
        plpy.execute(plan, [grp_list, cp_list])

    # 2) call CV function to actually cross-validate _build_tree
    # expects output table model_cv({grouping_cols), cp, avg, stddev)
    model_cv = output_table_name + "_cv"
    metric_function = "_tree_misclassified" if is_classification else "_tree_rmse"
    pred_name = '"estimated_{0}"'.format(dependent_variable.strip(' "'))
    grouping_str = 'NULL' if not grouping_cols else '"' + grouping_cols + '"'
    cat_feature_str = _array_to_string(cat_features)
    ordered_cat_feature_str = _array_to_string(ordered_cat_features)
    boolean_cat_str = _array_to_string(boolean_cats)
    con_feature_str = _array_to_string(con_features)
    modeling_params = [str(i) for i in
                       (is_classification,
                        split_criterion, "%data%", "%model%", id_col_name,
                        dependent_variable, dep_is_bool,
                        cat_feature_str, ordered_cat_feature_str,
                        boolean_cat_str, con_feature_str,
                        grouping_str, weights, max_depth,
                        min_split, min_bucket, n_bins,
                        "%explore%", max_n_surr, msg_level)]
    modeling_param_types = (["BOOLEAN"] + ["TEXT"] * 5 + ["BOOLEAN"] +
                            ["VARCHAR[]"] * 4 + ["TEXT"] * 2 + ["INTEGER"] * 4 +
                            ["TEXT", "SMALLINT", "TEXT"])

    cross_validation_grouping_w_params(
        schema_madlib,
        schema_madlib + '.__build_tree',
        modeling_params,
        modeling_param_types,
        schema_madlib + '.__tree_predict',
        ["%model%", "%data%", "%prediction%", "response", "True"],
        ["VARCHAR", "VARCHAR", "VARCHAR", "VARCHAR", "BOOLEAN"],
        schema_madlib + "." + metric_function,
        ["%data%", dependent_variable, "%prediction%", pred_name, id_col_name, grouping_cols, "%error%", "True"],
        ["VARCHAR", "VARCHAR", "VARCHAR", "VARCHAR", "VARCHAR", "TEXT", "VARCHAR", "BOOLEAN"],
        group_to_param_list_table, param_list_name, grouping_cols,
        training_table_name, id_col_name, False, model_cv, 'cp', None, n_folds)

    # 3) find the best cp for each group from table {model_cv}
    if not grouping_cols:
        grouping_array_str_comma = "''::TEXT AS grp_key,"
        group_by_str = ''
        qualified_group_by_str = ''
        group_select_str = ''
    else:
        grouping_array_qualified_str = get_grouping_array_str(training_table_name,
                                                              grouping_cols,
                                                              qualifier="min_cv_error")
        grouping_array_str_comma = grouping_array_qualified_str + " AS grp_key,"
        grouping_cols_list = grouping_cols.split(",")
        qualified_group_by_str = "GROUP BY " + ','.join("min_cv_error." + i for i in grouping_cols_list)
        group_by_str = "GROUP BY " + grouping_cols
        group_select_str = grouping_cols + ","

    plpy.notice(str(list(plpy.execute("SELECT * FROM {0}".format(model_cv)))))
    validation_result_query = """
            SELECT
                {grouping_array_str_comma}
                max(cp) AS cp
            FROM
                {model_cv}
            NATURAL JOIN
            ( SELECT
                    {group_select_str}
                    min(cv_error_avg) AS cv_error_avg
                FROM {model_cv}
                {group_by_str}
            ) min_cv_error
            {qualified_group_by_str}
        """.format(**locals())
    plpy.notice("validation_result_query:")
    plpy.notice(validation_result_query)
    validation_result = plpy.execute(validation_result_query)
    plpy.notice("finished validation_result_query, validation_result = " + str(list(validation_result)))

    grp_key_to_best_cp = dict((row['grp_key'], row['cp']) for row in validation_result)

    # 4) update tree_states to have the best cp cross-validated
    for tree in tree_states:
        best_cp = grp_key_to_best_cp[tree['grp_key']]
        if best_cp > tree['cp']:
            tree['cp'] = best_cp
            # we prune each tree further using a higher cp value
            # giving the optimal pruned tree.
            # This time we don't need the cp_list.
            pruned_tree = _prune_and_cplist(schema_madlib,
                                            tree,
                                            tree['cp'],
                                            compute_cp_list=False)
            tree['tree_state'] = pruned_tree['tree_state']
            if 'pruned_depth' in pruned_tree:
                tree['pruned_depth'] = pruned_tree['pruned_depth']
            elif 'tree_depth' in pruned_tree:
                tree['pruned_depth'] = pruned_tree['tree_depth']
            else:
                tree['pruned_depth'] = 0

    plpy.execute("DROP TABLE {group_to_param_list_table}".format(**locals()))
# ------------------------------------------------------------


def _tree_train_using_bins(
        schema_madlib, bins, training_table_name,
        cat_features, con_features, boolean_cats, n_bins, weights,
        dep_var_str, min_split, min_bucket, max_depth, filter_dep,
        dep_n_levels, is_classification, split_criterion,
        subsample=False, n_random_features=1, max_n_surr=0, **kwargs):
    """Trains a tree without grouping columns"""
    # Iterations for training the tree
    tree_state = plpy.execute(
        """
        SELECT {schema_madlib}._initialize_decision_tree(
            {is_regression_tree},
            '{split_criterion}'::text,
            {dep_n_levels}::smallint,
            {max_n_surr}::smallint)
            as tree_state, False as finished
        """.format(schema_madlib=schema_madlib,
                   is_regression_tree=not is_classification,
                   split_criterion=split_criterion,
                   dep_n_levels=dep_n_levels,
                   max_n_surr=max_n_surr))[0]
    plpy.notice("Starting tree building")
    tree_depth = -1
    while tree_state['finished'] == 0:
        #  finished: 0 = running, 1 = finished training, 2 = terminated prematurely
        tree_depth += 1
        tree_state = _one_step(
            schema_madlib, training_table_name,
            cat_features, con_features, boolean_cats, bins,
            n_bins, tree_state, weights, dep_var_str,
            min_split, min_bucket, max_depth, filter_dep,
            dep_n_levels, subsample, n_random_features, max_n_surr)
        plpy.notice("Completed training of level {0}".format(tree_depth))

    return tree_state
# ------------------------------------------------------------------------------


def _tree_train_grps_using_bins(
        schema_madlib, bins, training_table_name, cat_features, con_features,
        boolean_cats, n_bins, weights, grouping_cols, grouping_array_str, dep_var_str,
        min_split, min_bucket, max_depth, filter_dep, dep_n_levels,
        is_classification, split_criterion, subsample=False,
        n_random_features=1, tree_terminated=None, max_n_surr=0, **kwargs):

    """Trains a tree with grouping columns included """
    # Iterations for training the tree
    initialized_tree_state = plpy.execute(
        """
        SELECT
            {schema_madlib}._initialize_decision_tree(
                {is_regression_tree},
                '{split_criterion}'::text,
                {dep_n_levels}::smallint,
                {max_n_surr}::smallint
                )
            AS tree_state,
            0 AS finished
        """.format(schema_madlib=schema_madlib,
                   is_regression_tree=not is_classification,
                   split_criterion=split_criterion,
                   dep_n_levels=dep_n_levels,
                   max_n_surr=max_n_surr))[0]
    tree_states = []
    for key in bins['grp_key_con']:
        group_copy = dict(initialized_tree_state, grp_key=key)
        tree_states.append(group_copy)

    # The following is only used in random forest:
    # If a group already has a terminated tree (not finished
    # properly), then we do not continue to compute more trees
    # for that specific group of data.
    if tree_terminated is not None:
        for item in tree_states:
            if item['grp_key'] in tree_terminated and \
                    tree_terminated[item['grp_key']] == 2:  # terminated
                item['finished'] = 2  # won't continue in _one_step_for_grps

    plpy.notice("Started tree building for all groups")
    level = 0
    while not all(t['finished'] for t in tree_states):
        tree_states = _one_step_for_grps(
            schema_madlib, training_table_name, cat_features,
            con_features, boolean_cats, bins, n_bins,
            tree_states, weights, grouping_cols,
            grouping_array_str, dep_var_str, min_split, min_bucket,
            max_depth, filter_dep, dep_n_levels, subsample,
            n_random_features, max_n_surr)
        level += 1
        plpy.notice("Finished training for level " + str(level))

    return tree_states
# ------------------------------------------------------------


def _tree_rmse(schema_madlib, source_table, dependent_varname, prediction_table,
               pred_dep_name, id_col_name, grouping_cols, output_table,
               use_existing_tables=False, k=0, **kwargs):
    old_messages = plpy.execute("SELECT setting FROM pg_settings "
                                "WHERE name = 'client_min_messages'")[0]['setting']
    plpy.execute('SET client_min_messages TO warning')

    fold = ", " + str(k) + " as k"
    if use_existing_tables and table_exists(output_table):
        # plpy.execute("truncate " + output_table)
        header = "INSERT INTO " + output_table + " "
    else:
        header = "CREATE TABLE " + output_table + " AS "

    grouping_str = '' if not grouping_cols else "GROUP BY " + grouping_cols
    grouping_col_str = "" if not grouping_cols else grouping_cols + ','

    sql = header + """
        SELECT
            {grouping_col_str}
            sqrt(avg(
                    ({prediction_table}.{pred_dep_name} -
                     {source_table}.{dependent_varname})^2
                    )
                ) as cv_error
            {fold}
        FROM {prediction_table}, {source_table}
        WHERE {prediction_table}.{id_col_name} = {source_table}.{id_col_name}
        {grouping_str}
        """.format(output_table=output_table,
                   prediction_table=prediction_table,
                   pred_dep_name=pred_dep_name,
                   source_table=source_table,
                   dependent_varname=dependent_varname,
                   grouping_col_str=grouping_col_str,
                   grouping_str=grouping_str,
                   id_col_name=id_col_name,
                   fold=fold)
    plpy.execute(sql)
    plpy.execute('SET client_min_messages TO ' + old_messages)
# ------------------------------------------------------------


def _tree_misclassified(schema_madlib, source_table, dependent_varname,
                        prediction_table, pred_dep_name, id_col_name,
                        grouping_cols, output_table, use_existing_tables=False,
                        k=0, **kwargs):
    old_messages = plpy.execute(
        "SELECT setting FROM pg_settings WHERE name = 'client_min_messages'")[0]['setting']
    plpy.execute('SET client_min_messages TO warning')

    fold = ", " + str(k) + " as k"
    if use_existing_tables and table_exists(output_table):
        # plpy.execute("truncate " + output_table)
        header = "INSERT INTO " + output_table + " "
    else:
        header = "CREATE TABLE " + output_table + " AS "

    grouping_str = '' if not grouping_cols else "GROUP BY " + grouping_cols
    grouping_col_str = "" if not grouping_cols else grouping_cols + ','

    sql = header + """
        SELECT
            {grouping_col_str}
            1.0 * sum(CASE WHEN ({prediction_table}.{pred_dep_name} =
                           {source_table}.{dependent_varname})
                THEN 0 ELSE 1 END) / count(*) as cv_error
            {fold}
        FROM
            {prediction_table}, {source_table}
        WHERE
          {prediction_table}.{id_col_name} = {source_table}.{id_col_name}
        {grouping_str}
        """.format(output_table=output_table,
                   prediction_table=prediction_table,
                   pred_dep_name=pred_dep_name,
                   source_table=source_table,
                   dependent_varname=dependent_varname,
                   grouping_col_str=grouping_col_str,
                   grouping_str=grouping_str,
                   id_col_name=id_col_name,
                   fold=fold)
    plpy.execute(sql)
    plpy.execute('SET client_min_messages TO ' + old_messages)