"""
sklearn columntransformer
"""

import pandas as pd
import numpy as np
from shapash.utils.category_encoder_backend import inv_transform_ordinal
from shapash.utils.category_encoder_backend import inv_transform_ce
from shapash.utils.category_encoder_backend import supported_category_encoder
from shapash.utils.category_encoder_backend import dummies_category_encoder
from shapash.utils.category_encoder_backend import category_encoder_binary
from shapash.utils.category_encoder_backend import transform_ordinal
from shapash.utils.model_synoptic import simple_tree_model_sklearn,catboost_model,\
    linear_model,svm_model, xgboost_model, lightgbm_model, dict_model_feature
from shapash.utils.model import extract_features_model

columntransformer = "<class 'sklearn.compose._column_transformer.ColumnTransformer'>"

sklearn_onehot = "<class 'sklearn.preprocessing._encoders.OneHotEncoder'>"
sklearn_ordinal = "<class 'sklearn.preprocessing._encoders.OrdinalEncoder'>"
sklearn_standardscaler = "<class 'sklearn.preprocessing._data.StandardScaler'>"
sklearn_quantiletransformer = "<class 'sklearn.preprocessing._data.QuantileTransformer'>"
sklearn_powertransformer = "<class 'sklearn.preprocessing._data.PowerTransformer'>"

sklearn_model = linear_model + svm_model + simple_tree_model_sklearn

other_model = xgboost_model + catboost_model + lightgbm_model

dummies_sklearn = (sklearn_onehot)

no_dummies_sklearn = (sklearn_ordinal,
                      sklearn_standardscaler,
                      sklearn_quantiletransformer,
                      sklearn_powertransformer)

supported_sklearn = (sklearn_onehot,
                     sklearn_ordinal,
                     sklearn_standardscaler,
                     sklearn_quantiletransformer,
                     sklearn_powertransformer)


def inv_transform_ct(x_in, encoding):
    """
    Inverse transform when using a ColumnsTransformer.

    As ColumnsTransformer output hstack the result of transformers, if the TOP-preprocessed data are re-ordered
    after the ColumnTransformer the inverse transform must return false result.

    We successively inverse the transformers with columns position. That's why inverse colnames
    are prefixed by the transformers names.

    Parameters
    ----------
    x_in : pandas.DataFrame
        Prediction set.
    encoding : list
        The list must contain a single ColumnsTransformer and an optional list of dict.

    Returns
    -------
    pandas.Dataframe
        The reversed transformation for the given list of encoding.
    """
    if str(type(encoding)) == columntransformer:
        # We use inverse tranform from the encoding method base on columns position
        init = 0
        rst = pd.DataFrame()

        for enc in encoding.transformers_:
            name_encoding = enc[0]
            ct_encoding = enc[1]
            col_encoding = enc[2]
            # For Scikit encoding we use the associated inverse transform method
            if str(type(ct_encoding)) in supported_sklearn:
                frame, init = inv_transform_sklearn_in_ct(x_in,
                                                          init,
                                                          name_encoding,
                                                          col_encoding,
                                                          ct_encoding)

            # For category encoding we use the mapping
            elif str(type(ct_encoding)) in supported_category_encoder:
                frame, init = inv_transform_ce_in_ct(x_in,
                                                     init,
                                                     name_encoding,
                                                     col_encoding,
                                                     ct_encoding)

            # columns not encode
            elif name_encoding == 'remainder':
                if ct_encoding == 'passthrough':
                    nb_col = len(col_encoding)
                    frame = x_in.iloc[:, init:init + nb_col]
                else:
                    frame = pd.DataFrame()

            else:
                raise Exception(f'{ct_encoding} is not supported yet.')

            rst = pd.concat([rst, frame], axis=1)

    elif str(type(encoding)) == "<class 'list'>":
        rst = inv_transform_ordinal(x_in, encoding)

    else:
        raise Exception(f"{encoding.__class__.__name__} not supported, no inverse done.")

    return rst


def inv_transform_ce_in_ct(x_in, init, name_encoding, col_encoding, ct_encoding):
    """
    Inverse transform when using category_encoder in ColumnsTransformer preprocessing.

    Parameters
    ----------
    x_in : pandas.DataFrame
        Data processed.
    init : np.int
        Columns index that give the first column to look at.
    name_encoding : String
        Name of the encoding give by the user.
    col_encoding : list
        Processed features name.
    ct_encoding : category_encoder
        Type of encoding.

    Returns
    -------
    frame : pandas.Dataframe
        The reversed transformation for the given list of encoding.
    init : np.int
        Index of the last column use to make the transformation.
    """
    colname_output = [name_encoding + '_' + val for val in col_encoding]
    colname_input = ct_encoding.get_feature_names()
    nb_col = len(colname_input)
    x_to_inverse = x_in.iloc[:, init:init + nb_col].copy()
    x_to_inverse.columns = colname_input
    frame = inv_transform_ce(x_to_inverse, ct_encoding)
    frame.columns = colname_output
    init += nb_col
    return frame, init


def inv_transform_sklearn_in_ct(x_in, init, name_encoding, col_encoding, ct_encoding):
    """
    Inverse transform when using sklearn in ColumnsTransformer preprocessing.

    Parameters
    ----------
    x_in : pandas.DataFrame
        Data processed.
    init : np.int
        Columns index that give the first column to look at.
    name_encoding : String
        Name of the encoding give by the user.
    col_encoding : list
        Processed features name.
    ct_encoding : sklearn, category_encoder
        Type of encoding.

    Returns
    -------
    frame : pandas.Dataframe
        The reversed transformation for the given list of encoding.
    init : np.int
        Index of the last column use to make the transformation.
    """
    colname_output = [name_encoding + '_' + val for val in col_encoding]
    if str(type(ct_encoding)) in dummies_sklearn:
        colname_input = ct_encoding.get_feature_names(col_encoding)
        nb_col = len(colname_input)
    else:
        nb_col = len(colname_output)
    x_inverse = ct_encoding.inverse_transform(x_in.iloc[:, init:init + nb_col])
    frame = pd.DataFrame(x_inverse, columns=colname_output, index=x_in.index)
    init += nb_col
    return frame, init

def calc_inv_contrib_ct(x_contrib, encoding):
    """
    Reversed contribution when ColumnTransformer is used.

    As columns transformers output hstack the result of transformers, if the TOP-preprocessed data are re-ordered
    after the ColumnTransformer the inverse transform must return false result.

    Parameters
    ----------
    x_contrib : pandas.DataFrame
        Contributions set.
    encoding : ColumnTransformer, list, dict
        The processing apply to the original data.

    Returns
    -------
    pandas.Dataframe
        The aggregate contributions depending on which processing is apply.
    """

    if str(type(encoding)) == columntransformer:
        # We use inverse tranform from the encoding method base on columns position
        init = 0
        rst = pd.DataFrame()

        for enc in encoding.transformers_:
            name_encoding = enc[0]
            ct_encoding = enc[1]
            col_encoding = enc[2]

            if str(type(ct_encoding)) in supported_category_encoder+supported_sklearn:
                # We create new columns names depending on the name of the transformers and the name of the column.
                colname_output = [name_encoding + '_' + val for val in col_encoding]

                # If the processing create multiple columns we find the number of original categories and aggregate
                # the contribution.
                if str(type(ct_encoding)) in dummies_sklearn or str(type(ct_encoding)) in dummies_category_encoder:
                    for i_enc in range(len(colname_output)):
                        if str(type(ct_encoding)) == sklearn_onehot:
                            col_origin = ct_encoding.categories_[i_enc]
                        elif str(type(ct_encoding)) == category_encoder_binary:
                            col_origin = ct_encoding.base_n_encoder.mapping[i_enc].get('mapping').columns.tolist()
                        else:
                            col_origin = ct_encoding.mapping[i_enc].get('mapping').columns.tolist()
                        nb_col = len(col_origin)
                        contrib_inverse = x_contrib.iloc[:, init:init + nb_col].sum(axis=1)
                        frame = pd.DataFrame(contrib_inverse,
                                             columns=[colname_output[i_enc]],
                                             index=contrib_inverse.index)
                        rst = pd.concat([rst, frame], axis=1)
                        init += nb_col
                else:
                    nb_col = len(colname_output)
                    frame = x_contrib.iloc[:, init:init + nb_col]
                    frame.columns = colname_output
                    rst = pd.concat([rst, frame], axis=1)
                    init += nb_col

            elif name_encoding == 'remainder':
                if ct_encoding == 'passthrough':
                    nb_col = len(col_encoding)
                    frame = x_contrib.iloc[:, init:init + nb_col]
                    rst = pd.concat([rst, frame], axis=1)
                    init += nb_col
            else:
                raise Exception(f"{encoding.__class__.__name__} not supported, no inverse done.")
        return rst
    else:
        return x_contrib


def transform_ct(x_in, model, encoding):
    """
    Transform when using a ColumnsTransformer.

    As ColumnsTransformer output hstack the result of transformers, if the TOP-preprocessed data are re-ordered
    after the ColumnTransformer the inverse transform must return false result.

    We successively apply the transformers with columns position. That's why colnames
    are prefixed by the transformers names.

    Parameters
    ----------
    x_in : pandas.DataFrame
        Raw dataset to apply preprocessing
    model: model object
        model used to check the different values of target estimate predict_proba
    encoding : list
        The list must contain a single ColumnsTransformer and an optional list of dict.

    Returns
    -------
    pandas.Dataframe
        The data preprocessed for the given list of encoding.
    """
    if str(type(encoding)) == columntransformer:
        # We use inverse tranform from the encoding method base on columns position
        if str(type(model)) in sklearn_model:
            rst = pd.DataFrame(encoding.transform(x_in),
                               index=x_in.index)
            rst.columns = ["col_" + str(feature) for feature in rst.columns]

        elif str(type(model)) in other_model:
            rst = pd.DataFrame(encoding.transform(x_in),
                                columns=extract_features_model(model, dict_model_feature[str(type(model))]),
                                index=x_in.index)
        else:
            raise ValueError("Model specified isn't supported by Shapash.")

    elif str(type(encoding)) == "<class 'list'>":
        rst = transform_ordinal(x_in, encoding)

    else:
        raise Exception(f"{encoding.__class__.__name__} not supported, no preprocessing done.")

    return rst

def get_names(name, trans, column, column_transformer):
    """
    Allow to extract features names from one encoder of the ColumnTransformer.
    If the right names aren't available, It creates a list with customized names.

    Parameters
    ----------
    name: string
        String which indicates the name of the transformer.
    trans: sklearn encoder or category_encoders
        One of the encoder fitted through the ColumnTransformer.
    column: list
        List of features impacted by the specific transformer.
    column_transformer: sklearn ColumnTransformer
        The fitted ColumnTransformer containing the specific encoder.

    Returns
    -------
    list:
        List of returned features when specific transformer is applied.
    """
    if trans == 'drop' or (
            hasattr(column, '__len__') and not len(column)):
        return []
    if trans == 'passthrough':
        if hasattr(column_transformer, '_df_columns'):
            if ((not isinstance(column, slice))
                    and all(isinstance(col, str) for col in column)):
                return column
            else:
                return column_transformer._df_columns[column]
        else:
            indices = np.arange(column_transformer._n_features)
            return ['x%d' % i for i in indices[column]]
    if not hasattr(trans, 'get_feature_names'):
        if column is None:
            return []
        else:
            return [name + "__" + f for f in column]

    return [name + "__" + f for f in trans.get_feature_names()]

def get_feature_names(column_transformer):
    """
    Allow to extract all features names from encoders of the ColumnTransformer once it has been applied.
    If the right names aren't available, It creates a list with customized names.

    Parameters
    ----------
    column_transformer: sklearn ColumnTransformer
        The fitted ColumnTransformer containing the specific encoder.

    Returns
    -------
    feature_names: list
        List of returned features names when ColumnTransformer is applied.
    """
    feature_names = []
    l_transformers = list(column_transformer._iter(fitted=True))

    for name, trans, column, _ in l_transformers:
        feature_names.extend(get_names(name, trans, column, column_transformer))

    return feature_names

def get_list_features_names(list_preprocessing, columns_dict):
    """
    Allow to extract all features names from encoders when a list of preprocessing is uesd once it has been applied.
    If the right names aren't available, It creates a list with customized names.

    Parameters
    ----------
    list_preprocessing: list
        The fitted list_preprocessing containing the specific encoders.

    Returns
    -------
    feature_names: list
        List of returned features names when list_preprocessing is applied.

    """
    feature_expected = [value for value in columns_dict.values()]

    for enc in list_preprocessing:
        if str(type(enc)) in supported_category_encoder:
            feature_expected = enc.feature_names

        elif str(type(enc)) in columntransformer:
            feature_expected = get_feature_names(enc)

        elif str(type(enc)) in ("<class 'list'>"):
            feature_expected = feature_expected

    return feature_expected

