from elastic_net_optimizer_fista import _elastic_net_fista_train
from elastic_net_optimizer_igd import _elastic_net_igd_train
import plpy

# ========================================================================


def _elastic_net_gaussian_fista_train(schema_madlib, 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):
    """
    Use FISTA to solve linear models
    """
    return _elastic_net_fista_train(schema_madlib,
                                    "__gaussian_fista_step",
                                    "__gaussian_fista_state_diff",
                                    "gaussian",
                                    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_gaussian_igd_train(schema_madlib, 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):
    """
    Use IGD to solve linear models
    """
    return _elastic_net_igd_train(schema_madlib,
                                  "__gaussian_igd_step",
                                  "__gaussian_igd_state_diff",
                                  "gaussian",
                                  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_binomial_fista_train(schema_madlib, 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):
    """
    Use FISTA to solve linear models
    """
    return _elastic_net_fista_train(schema_madlib,
                                    "__binomial_fista_step",
                                    "__binomial_fista_state_diff",
                                    "binomial",
                                    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_binomial_igd_train(schema_madlib, 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):
    """
    Use IGD to solve linear models
    """
    return _elastic_net_igd_train(schema_madlib,
                                  "__binomial_igd_step",
                                  "__binomial_igd_state_diff",
                                  "binomial",
                                  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)