""" Data preprocessing helper functions for tabular neural network models """

from collections import OrderedDict
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler, QuantileTransformer, FunctionTransformer  # can also consider: PowerTransformer

from .categorical_encoders import OneHotMergeRaresHandleUnknownEncoder, OrdinalMergeRaresHandleUnknownEncoder

def create_preprocessor(impute_strategy, max_category_levels, unique_category_str, continuous_features, skewed_features, onehot_features, embed_features, bool_features):
    """ Creates sklearn ColumnTransformer that can be fit to training data to preprocess it for tabular neural network. """
    transformers = []  # order of various column transformers in this list is important!
    if continuous_features:
        continuous_transformer = Pipeline(steps=[
            ('imputer', SimpleImputer(strategy=impute_strategy)),
            ('scaler', StandardScaler())])
        transformers.append( ('continuous', continuous_transformer, continuous_features) )
    if skewed_features:
        power_transformer = Pipeline(steps=[
            ('imputer', SimpleImputer(strategy=impute_strategy)),
            ('quantile', QuantileTransformer(output_distribution='normal')) ])  # Or output_distribution = 'uniform'
        transformers.append( ('skewed', power_transformer, skewed_features) )
    if onehot_features:
        onehot_transformer = Pipeline(steps=[
            # TODO: Consider avoiding converting to string for improved memory efficiency
            ('to_str', FunctionTransformer(convert_df_dtype_to_str)),
            ('imputer', SimpleImputer(strategy='constant', fill_value=unique_category_str)),
            ('onehot', OneHotMergeRaresHandleUnknownEncoder(max_levels=max_category_levels, sparse=False))])  # test-time unknown values will be encoded as all zeros vector
        transformers.append( ('onehot', onehot_transformer, onehot_features) )
    if embed_features:  # Ordinal transformer applied to convert to-be-embedded categorical features to integer levels
        ordinal_transformer = Pipeline(steps=[
            ('to_str', FunctionTransformer(convert_df_dtype_to_str)),
            ('imputer', SimpleImputer(strategy='constant', fill_value=unique_category_str)),
            ('ordinal', OrdinalMergeRaresHandleUnknownEncoder(max_levels=max_category_levels))])  # returns 0-n when max_category_levels = n-1. category n is reserved for unknown test-time categories.
        transformers.append( ('ordinal', ordinal_transformer, embed_features) )
    return ColumnTransformer(transformers=transformers, remainder='passthrough')  # numeric features are processed in the same order as in numeric_features vector, so feature-names remain the same.

def convert_df_dtype_to_str(df):
    return df.astype(str)

def get_feature_arraycol_map(processor, max_category_levels):
    """ Returns OrderedDict of feature-name -> list of column-indices in processed data array corresponding to this feature """
    feature_preserving_transforms = set(['continuous','skewed', 'ordinal','bool','remainder'])  # these transforms do not alter dimensionality of feature
    feature_arraycol_map = {}  # unordered version
    current_colindex = 0
    for transformer in processor.transformers_:
        transformer_name = transformer[0]
        transformed_features = transformer[2]
        if transformer_name in feature_preserving_transforms:
            for feature in transformed_features:
                if feature in feature_arraycol_map:
                    raise ValueError("same feature is processed by two different column transformers: %s" % feature)
                feature_arraycol_map[feature] = [current_colindex]
                current_colindex += 1
        elif transformer_name == 'onehot':
            oh_encoder = [step for (name, step) in transformer[1].steps if name == 'onehot'][0]
            for i in range(len(transformed_features)):
                feature = transformed_features[i]
                if feature in feature_arraycol_map:
                    raise ValueError("same feature is processed by two different column transformers: %s" % feature)
                oh_dimensionality = min(len(oh_encoder.categories_[i]), max_category_levels+1)
                feature_arraycol_map[feature] = list(range(current_colindex, current_colindex+oh_dimensionality))
                current_colindex += oh_dimensionality
        else:
            raise ValueError("unknown transformer encountered: %s" % transformer_name)
    return OrderedDict([(key, feature_arraycol_map[key]) for key in feature_arraycol_map])

def get_feature_type_map(feature_arraycol_map, types_of_features):
    """ Returns OrderedDict of feature-name -> feature_type string (options: 'vector', 'embed'). """
    if feature_arraycol_map is None:
        raise ValueError("Must first call get_feature_arraycol_map() to set feature_arraycol_map before calling get_feature_type_map()")
    vector_features = types_of_features['continuous'] + types_of_features['skewed'] + types_of_features['onehot'] + types_of_features['bool']
    feature_type_map = OrderedDict()
    for feature_name in feature_arraycol_map:
        if feature_name in vector_features:
            feature_type_map[feature_name] = 'vector'
        elif feature_name in types_of_features['embed']:
            feature_type_map[feature_name] = 'embed'
        else:
            feature_type_map[feature_name] = 'vector'
    return feature_type_map
