#AUTOGENERATED! DO NOT EDIT! File to edit: dev/03_ml.ipynb (unless otherwise specified).

__all__ = ['Debugger', 'build_pipeline', 'parse_gs_results']

#Cell
import os
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.compose import make_column_transformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import OneHotEncoder

from htools import save

#Cell
class Debugger(BaseEstimator, TransformerMixin):
    """Transformer to be used to check the output of sklearn pipelines."""

    def __init__(self, verbose=True, path=None):
        """
        Parameters
        -----------
        verbose: bool
            If True, print the shape and null count of the transformed feature
            array.
        path: str (Optional)
            If specified, save the processed data to a pickle file. If left as
            None, output will not be saved. Must use pkl file extension.
        """
        self.verbose = verbose
        self.path = path
        if path: os.makedirs(os.path.dirname(path), exist_ok=True)

    def transform(self, X, y=None):
        if self.verbose:
            print(f'X shape: {X.shape}')
            print(f'# of nulls: {np.isnan(X).sum()}')
            print(pd.DataFrame(X[:5]))
        if self.path: save(X, self.path)
        return X

    def fit(self, X, y=None):
        return self

#Cell
def build_pipeline(x, bin_cols=None, num_cols=None, cat_cols=None,
                    scaler=None, num_impute='median'):
    """Wrapper to build scikit-learn pipelines that impute missing values,
    scale numeric columns, and one hot encode categorical variables. A column
    transformer is returned. No classifier is passed in, so fit and transform
    are available methods but predict is not.

    Parameters
    -----------
    x: pd.DataFrame
        (Optional) A dataframe of features. Lists of numeric and categorical
        features will be generated automatically based on dtypes.
    scaler: sklearn.preprocessing.data
        A scikit-learn scaler object (already instantiated). If none is
        provided, mean normalization will be used.
    num_impute: str
        Strategy to use when imputing missing values in numeric (continuous)
        features. One of ('mean', 'median', 'most_frequent'), default 'median'.
        Note that categorical features will be automatically imputed with
        'most_frequent'.

    Returns
    --------
    sklearn.compose._column_transformer.ColumnTransformer
    """
    scaler = scaler or StandardScaler()

    # Dummy cols need no additional processing.
    bin_cols = bin_cols or [col for col in x.columns
                            if col.startswith('contains_')
                            or col.startswith('is_')]
    num_cols = num_cols or [col for col in x.select_dtypes(np.number).columns
                            if col not in bin_cols]
    cat_cols = cat_cols or x.select_dtypes('category').columns.tolist()

    # Process continuous and categorical variables separately.
    bin_pipe = make_pipeline(SimpleImputer(strategy='most_frequent'))
    num_pipe = make_pipeline(SimpleImputer(strategy=num_impute),
                             scaler)
    cat_pipe = make_pipeline(SimpleImputer(strategy='most_frequent'),
                             OneHotEncoder(drop='first'))

    # Combine into single transformer and add classifier to end of pipeline.
    return make_column_transformer((bin_pipe, bin_cols),
                                   (num_pipe, num_cols),
                                   (cat_pipe, cat_cols),
                                   remainder='passthrough')

#Cell
def parse_gs_results(results, metric, sort_param=None):
    """Parse dictionary returned by grid search cv_results_ so we can examine
    how metrics varied depending on hyperparameter choices.

    Parameters
    -----------
    results: dict
        cv_results_ dict.
    sort_param: str
        Hyperparameter to sort by.
    metric: str

    Returns
    --------
    pd.DataFrame
    """
    sort_by = [metric]
    if sort_param:
        sort_by.insert(0, sort_param)

    return pd.concat([pd.DataFrame(data['params']),
                      pd.DataFrame(data[f'mean_test_{metric}'],
                                   columns=[metric])],
                     axis=1).sort_values(sort_by, ascending=False)