import plpy
from elastic_net_generate_result import _elastic_net_generate_result
from elastic_net_utils import _normalize_data
from elastic_net_utils import _tbl_dimension_rownum
from elastic_net_utils import _elastic_net_validate_args
from elastic_net_utils import _compute_average_sq
from elastic_net_utils import _generate_warmup_lambda_sequence
from elastic_net_utils import _process_warmup_lambdas
from utilities.control import MinWarning
from utilities.in_mem_group_control import GroupIterationController
from utilities.utilities import unique_string
from utilities.utilities import extract_keyvalue_params


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


def _fista_params_parser(optimizer_params, lambda_value, tolerance, schema_madlib):
    """
    Parse fista parameters.
    """
    # default values
    defaults_and_types = {
        "max_stepsize": (2., float),
        "eta": (1.2, float),
        "warmup": (False, bool),
        "warmup_lambdas": (None, list),
        "warmup_lambda_no": (15, int),
        "use_active_set": (False, bool),
        "random_stepsize": (False, bool),
        "activeset_tolerance": (tolerance, float),
        "warmup_tolerance": (tolerance, float)
    }
    param_defaults = dict([(k, v[0]) for k, v in defaults_and_types.items()])
    param_types = dict([(k, v[1]) for k, v in defaults_and_types.items()])

    if not optimizer_params:
        return param_defaults

    usage_str = ("\n Run:\n"
                 "   SELECT {0}.elastic_net_train('fista');\n"
                 "   to see the parameters for FISTA algorithm.".
                 format(schema_madlib))
    name_value = extract_keyvalue_params(optimizer_params, param_types,
                                         param_defaults, usage_str=usage_str,
                                         ignore_invalid=True)

    if name_value["warmup"] and name_value['warmup_lambdas'] is not None:
        # errors are handled in _process_warmup_lambdas
        name_value['warmup_lambdas'] = _process_warmup_lambdas(name_value['warmup_lambdas'], lambda_value)

    # validate the parameters
    if name_value["max_stepsize"] <= 0:
        plpy.error("Elastic Net error: backtracking parameter max_stepsize "
                   "must be positive!")

    if name_value["eta"] < 1:
        plpy.error("Elastic Net error: backtracking parameter eta must be "
                   "larger than 1!")

    if (name_value["warmup"] and name_value["warmup_lambdas"] is None and
            name_value["warmup_lambda_no"] < 1):
        plpy.error("Elastic Net error: Number of warm-up lambdas must be a "
                   "positive integer!")

    if name_value["activeset_tolerance"] <= 0:
        plpy.error("Elastic Net error: activeset_tolerance must be positive!")

    if name_value["warmup_tolerance"] <= 0:
        plpy.error("Elastic Net error: warmup_tolerance must be positive!")

    return name_value
# -------------------------------------------------------------------------


def _fista_construct_dict(
        schema_madlib, family, tbl_source, col_ind_var, col_dep_var,
        tbl_result, dimension, row_num, lambda_value, alpha,
        normalization, max_iter, tolerance, outstr_array, optimizer_params_dict):
    """
    Construct the dict used by a series of SQL queries in FISTA optimizer.
    """
    args = dict(schema_madlib=schema_madlib,
                family=family,
                tbl_source=tbl_source,
                tbl_data=tbl_source,  # argument name used in normalization
                col_ind_var=col_ind_var, col_dep_var=col_dep_var,
                col_ind_var_norm_new=unique_string(desp='temp_features_norm'),  # for normalization usage
                col_ind_var_tmp=unique_string(desp='temp_features'),
                col_dep_var_norm_new=unique_string(desp='temp_target_norm'),  # for normalization usage
                col_dep_var_tmp=unique_string(desp='temp_target'),
                tbl_result=tbl_result,
                lambda_value=lambda_value, alpha=alpha,
                dimension=dimension, row_num=row_num,
                max_iter=max_iter, tolerance=tolerance,
                outstr_array=outstr_array,
                normalization=normalization)

    # Add the optimizer parameters
    args.update(optimizer_params_dict)

    # Table names useful when normalizing the original data
    # Note: in order to be consistent with the calling convention
    # of the normalization functions, multiple elements of the dict
    # actually have the same value. This is a price one has to pay
    # to save typing argument names by using **args as the
    # function argument.
    tbl_ind_scales = unique_string(desp='temp_ind_scales')
    tbl_dep_scale = unique_string(desp='temp_dep_scales')
    tbl_data_scaled = unique_string(desp='temp_data_scales')
    args.update(tbl_scale=tbl_dep_scale, tbl_dep_scale=tbl_dep_scale,
                tbl_scales=tbl_ind_scales, tbl_ind_scales=tbl_ind_scales,
                tbl_data_scaled=tbl_data_scaled)

    # Table names used in IGD iterations
    args.update(tbl_fista_state=unique_string(), tbl_fista_args=unique_string())

    # more, for args table
    for name in ('dimension_name', 'lambda_name', 'alpha_name',
                 'total_rows_name', 'max_iter_name', 'tolerance_name',
                 'max_stepsize_name', 'eta_name', 'activeset_name'):
        args[name] = unique_string()
    return args
# ------------------------------------------------------------------------


def _fista_cleanup_temp_tbls(**kwargs):
    """
    Drop all temporary tables used by FISTA optimizer,
    including tables used in the possible normalization
    and FISTA iterations.
    """
    plpy.execute("""
                drop table if exists {tbl_ind_scales};
                drop table if exists {tbl_dep_scale};
                drop table if exists {tbl_data_scaled};
                drop table if exists {tbl_fista_args};
                drop table if exists pg_temp.{tbl_fista_state};
                drop table if exists {x_mean_table};
                drop table if exists {y_mean_table};
                """.format(**kwargs))

    return None
# ------------------------------------------------------------------------


def _elastic_net_fista_train(schema_madlib, func_step_aggregate,
                             func_state_diff, family,
                             tbl_source, col_ind_var,
                             col_dep_var, tbl_result, tbl_summary, lambda_value, alpha,
                             normalization, optimizer_params, max_iter,
                             tolerance, outstr_array, grouping_str,
                             grouping_col, **kwargs):
    """
    func_step_aggregate is string, and it is the name of the step function
    """
    _elastic_net_validate_args(tbl_source, col_ind_var, col_dep_var,
                               tbl_result, tbl_summary, lambda_value, alpha,
                               normalization, max_iter, tolerance)
    return _elastic_net_fista_train_compute(schema_madlib,
                                            func_step_aggregate,
                                            func_state_diff,
                                            family,
                                            tbl_source,
                                            col_ind_var, col_dep_var,
                                            tbl_result, tbl_summary,
                                            lambda_value, alpha,
                                            normalization,
                                            optimizer_params, max_iter,
                                            tolerance, outstr_array,
                                            grouping_str, grouping_col,
                                            **kwargs)
# ------------------------------------------------------------------------


def _elastic_net_fista_train_compute(schema_madlib, func_step_aggregate,
                                     func_state_diff, family,
                                     tbl_source, col_ind_var,
                                     col_dep_var, tbl_result, tbl_summary,
                                     lambda_value, alpha,
                                     normalization, optimizer_params, max_iter,
                                     tolerance, outstr_array, grouping_str,
                                     grouping_col, **kwargs):
    """
    Fit linear model with elastic net regularization using FISTA optimization.

    @param tbl_source        Name of data source table
    @param col_ind_var       Name of independent variable column,
                             independent variable is an array
    @param col_dep_var       Name of dependent variable column
    @param tbl_result        Name of the table to store the results,
                             will return fitting coefficients and
                             likelihood
    @param lambda_value      The regularization parameter
    @param alpha             The elastic net parameter, [0, 1]
    @param normalization     Whether to normalize the variables
    @param optimizer_params  Parameters of the above optimizer, the format
                             is '{arg = value, ...}'::varchar[]
    """
    with MinWarning('error'):
        (dimension, row_num) = _tbl_dimension_rownum(schema_madlib, tbl_source, col_ind_var)

        # generate a full dict to ease the following string format
        # including several temporary table names
        args = _fista_construct_dict(schema_madlib, family, tbl_source, col_ind_var,
                                     col_dep_var, tbl_result,
                                     dimension, row_num, lambda_value,
                                     alpha, normalization,
                                     max_iter, tolerance, outstr_array,
                                     _fista_params_parser(optimizer_params,
                                                          lambda_value,
                                                          tolerance,
                                                          schema_madlib))
        args.update({'x_mean_table':unique_string(desp='x_mean_table')})
        args.update({'y_mean_table':unique_string(desp='y_mean_table')})
        args.update({'grouping_col': grouping_col})
        # use normalized data or not
        if normalization:
            _normalize_data(args)
            tbl_used = args["tbl_data_scaled"]
            args["col_ind_var_new"] = args["col_ind_var_norm_new"]
            args["col_dep_var_new"] = args["col_dep_var_norm_new"]
        else:
            tbl_used = tbl_source
            args["col_ind_var_new"] = col_ind_var
            args["col_dep_var_new"] = col_dep_var

        args["tbl_used"] = tbl_used

        if args["warmup_lambdas"] is not None:
            args["warm_no"] = len(args["warmup_lambdas"])

        if args["warmup"] and args["warmup_lambdas"] is None:
            # average squares of each feature
            # used to estimate the largest lambda value
            args["sq"] = _compute_average_sq(**args)
            args["warmup_lambdas"] = \
                _generate_warmup_lambda_sequence(lambda_value,
                args["warmup_lambda_no"])
            args["warm_no"] = len(args["warmup_lambdas"])
        elif args["warmup"] is False:
            args["warm_no"] = 1
            args["warmup_lambdas"] = [lambda_value]  # only one value

        # This update is needed in _elastic_net_generate_result() after
        # _compute_fista is run. Some of these variables are accessed there.
        args.update({
            'rel_state': args["tbl_fista_state"],
            'col_grp_iteration': unique_string(desp='col_grp_iteration'),
            'col_grp_state': unique_string(desp='col_grp_state'),
            'col_grp_key': unique_string(desp='col_grp_key'),
            'col_n_tuples': unique_string(desp='col_n_tuples'),
            'lambda_count': 1,
            'is_active': 0,
            'state_type': "double precision[]",
            'rel_source': tbl_used,
            'grouping_str': grouping_str,
            'tbl_source': tbl_source,
            'tbl_summary': tbl_summary
        })

        # perform the actual calculation
        iteration_run = _compute_fista(
            schema_madlib, func_step_aggregate,
            func_state_diff,
            args["tbl_fista_args"],
            args["tbl_fista_state"],
            tbl_used,
            args["col_ind_var_new"],
            args["col_dep_var_new"],
            grouping_str,
            grouping_col,
            tolerance,
            start_iter=0,
            lambda_name=args["warmup_lambdas"],
            warmup_lambda_value=args.get('warmup_lambdas')[args["lambda_count"] - 1],
            activeset_tolerance=args["activeset_tolerance"],
            warmup_tolerance=args["warmup_tolerance"],
            max_iter=args["max_iter"],
            warm_no=args["warm_no"],
            random_stepsize=args["random_stepsize"],
            use_active_set=args["use_active_set"],
            alpha=args["alpha"],
            row_num=args["row_num"],
            dimension=args["dimension"],
            max_stepsize=args["max_stepsize"],
            eta=args["eta"],
            rel_state=args["tbl_fista_state"],
            col_grp_iteration=args["col_grp_iteration"],
            col_grp_state=args["col_grp_state"],
            col_grp_key=args["col_grp_key"],
            col_n_tuples=args["col_n_tuples"],
            lambda_count=args["lambda_count"],
            is_active=args["is_active"],
            state_type=args["state_type"],
            rel_source=args["rel_source"])

        _elastic_net_generate_result("fista", iteration_run, **args)

        # cleanup
        _fista_cleanup_temp_tbls(**args)
    return None
# ------------------------------------------------------------------------


def _compute_fista(schema_madlib, func_step_aggregate, func_state_diff,
                   tbl_args, tbl_state, tbl_source, col_ind_var,
                   col_dep_var, grouping_str, grouping_col, tolerance,
                   start_iter, **kwargs):
    """
    Driver function for elastic net using FISTA

    @param schema_madlib Name of the MADlib schema, properly escaped/quoted
    @param tbl_args Name of the (temporary) table containing all non-template
        arguments
    @param tbl_state Name of the (temporary) table containing the inter-iteration
        states
    @param tbl_source Name of the relation containing input points
    @param col_ind_var Name of the independent variables column
    @param col_dep_var Name of the dependent variable column
    @param kwargs We allow the caller to specify additional arguments (all of
        which will be ignored though). The purpose of this is to allow the
        caller to unpack a dictionary whose element set is a superset of
        the required arguments by this function.

    @return The iteration number (i.e., the key) with which to look up the
        result in \c tbl_state
    """
    args = locals()

    for k, v in kwargs.iteritems():
        if k not in args:
            args.update({k: v})
    iterationCtrl = GroupIterationController(args)
    with iterationCtrl as it:
        it.iteration = start_iter
        while True:
            # manually add the intercept term
            if (it.kwargs["lambda_count"] > len(args.get('lambda_name'))):
                break
            it.kwargs["warmup_lambda_value"] = args.get('lambda_name')[it.kwargs["lambda_count"] - 1]
            # Fix for JIRA MADLIB-1092
            # 'col_n_tuples' is supposed to refer to the number of rows in the
            # table, or the number of rows in a group. col_n_tuples gets
            # the right value in in_mem_group_control, so using this instead
            # of row_num (which was used hitherto).
            it.update("""
                    {schema_madlib}.{func_step_aggregate}(
                        ({col_ind_var})::double precision[],
                        ({col_dep_var}),
                        {rel_state}.{col_grp_state},
                        ({warmup_lambda_value})::double precision,
                        ({alpha})::double precision,
                        ({dimension})::integer,
                        ({col_n_tuples})::integer,
                        ({max_stepsize})::double precision,
                        ({eta})::double precision,
                        ({use_active_set})::integer,
                        {is_active}::integer,
                        {random_stepsize}::integer
                    )
                    """)

            if it.kwargs["is_active"] == 1:
                it.kwargs["use_tolerance"] = it.kwargs["activeset_tolerance"]
            elif it.kwargs["lambda_count"] < it.kwargs["warm_no"]:
                it.kwargs["use_tolerance"] = it.kwargs["warmup_tolerance"]
            else:
                it.kwargs["use_tolerance"] = args["tolerance"]

            if it.kwargs["use_active_set"]:
                is_backtracking = it.are_last_state_value_zero()
                if it.test("""
                        {iteration} >= {max_iter}
                        or
                        {schema_madlib}.{func_state_diff}(
                            _state_previous, _state_current) < {use_tolerance}
                        """):
                    if it.iteration < it.kwargs["max_iter"]:
                        if it.kwargs["is_active"] == 0:
                            if (it.kwargs["lambda_count"] < it.kwargs["warm_no"]):
                                it.kwargs["lambda_count"] += 1
                                if (len(args.get('lambda_name')) >=
                                        it.kwargs["lambda_count"]):
                                    break
                            else:
                                break
                        else:
                            it.kwargs["is_active"] = 0
                    else:
                        break
                else:
                    # change active state only outside of backtracking
                    if is_backtracking and it.kwargs["is_active"] == 0:
                        it.kwargs["is_active"] = 1
            else:
                if it.test("""
                        {iteration} >= {max_iter} or
                        {schema_madlib}.{func_state_diff}(
                            _state_previous, _state_current) < {tolerance}
                        """):
                    if (it.iteration < it.kwargs["max_iter"] and
                            it.kwargs["lambda_count"] < it.kwargs["warm_no"]):
                        it.kwargs["lambda_count"] += 1
                    else:
                        break
        it.final()
        if it.kwargs["lambda_count"] < it.kwargs["warm_no"]:
            plpy.error("""
                       Elastic Net error: The final target lambda value is not
                       reached in warm-up iterations. You need more iterations!
                       """)

    return iterationCtrl.iteration