"""
Various preprocessing classes implemented as scikit-learn
transformers compatible with scikit-learn pipelines.
"""

import warnings
import pandas as pd
import numpy as np

from scipy import sparse
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import (
    MultiLabelBinarizer,
    FunctionTransformer,
    LabelEncoder,
    normalize,
)
from sklearn.feature_extraction.text import HashingVectorizer
from sklearn import feature_selection

__all__ = [
    "SelectField",
    "FeatureCast",
    "ImputeNull",
    "DenseTransformer",
    "SparseTransformer",
    "LabelEncoderPipe",
    "SelectorMem",
    "HashingVectorizerChunked",
    "MultihotEncoder",
]

_SELECTOR_LOOKUP = {
    "fpr": feature_selection.SelectFpr,
    "fdr": feature_selection.SelectFdr,
    "kbest": feature_selection.SelectKBest,
    "percentile": feature_selection.SelectPercentile,
    "fwe": feature_selection.SelectFwe,
}


class _CustomFunctionTransformer(TransformerMixin, BaseEstimator):
    """ Base class for function transformers """

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

    def transform(self, X, y=None):
        return self.transformer_.transform(X)


def select_field(X, cols=None, single_dimension=False):
    """
    Select columns from a pandas DataFrame

    Args:
        X (pandas DataFrame): input data
        cols (array-like): list of columns to select
        single_dimension (bool): reduce data to
            one dimension if only one column
            is requested
    Returns:
        X (numpy array)
    """
    if cols is None:
        return X.values
    if len(cols) > 1:
        return X[cols].values
    if len(cols) == 1:
        if single_dimension:
            return X[cols[0]].values
        else:
            return X[cols].values


class SelectField(_CustomFunctionTransformer):
    """
    Applies `select_field` as FunctionTransformer

    Args:
        cols (array-like): list of columns to select
        single_dimension (bool): reduce data to
            one dimension if only one column
            is requested
    """

    def __init__(self, cols=None, single_dimension=False):
        self.cols = cols
        self.single_dimension = single_dimension
        kw_args = {"cols": cols, "single_dimension": single_dimension}
        self.transformer_ = FunctionTransformer(
            select_field, kw_args=kw_args, validate=False
        )


def to_dense(X):
    """ Densify sparse matrix if issparse """
    if sparse.issparse(X):
        return X.todense()
    else:
        return X


class DenseTransformer(_CustomFunctionTransformer):
    """ Applies `to_dense` as a FunctionTransformer """

    def __init__(self):
        self.transformer_ = FunctionTransformer(to_dense, validate=False)


def to_sparse(X):
    """ Sparsify dense matrix if not issparse """
    if sparse.issparse(X):
        return X
    else:
        return sparse.csr_matrix(X)


class SparseTransformer(_CustomFunctionTransformer):
    """ Applies `to_sparse` as a FunctionTransformer """

    def __init__(self):
        self.transformer_ = FunctionTransformer(to_sparse, validate=False)


def feature_cast(X, cast_type=None):
    """
    Casts feature data to requeted type

    Args:
        X (numpy array): input data
        cast_type (type): type to cast data
    Returns:
        X (numpy array)
    """
    if cast_type is None:
        return X
    else:
        return X.astype(cast_type)


class FeatureCast(_CustomFunctionTransformer):
    """
    Applies `feature_cast` as a FunctionTransformer

    Args:
        cast_type (type): type to cast data
    """

    def __init__(self, cast_type=None):
        self.transformer_ = FunctionTransformer(
            feature_cast, kw_args={"cast_type": cast_type}, validate=False
        )


def impute_null(X, impute_val=None):
    """
    Impute null values. Null values are
    determined using `pd.isnull`.

    Args:
        X (numpy array): input data
        impute_val (object): value to impute with
    Returns:
        X (numpy array)
    """
    if impute_val is None:
        return X
    else:
        X[pd.isnull(X)] = impute_val
        return X


class ImputeNull(_CustomFunctionTransformer):
    """
    Applies `impute_null` as a FunctionTransformer

    Args:
        impute_val (object): value to impute with
    """

    def __init__(self, impute_val=None):
        self.transformer_ = FunctionTransformer(
            impute_null, kw_args={"impute_val": impute_val}, validate=False
        )


class LabelEncoderPipe(TransformerMixin, BaseEstimator):
    """
    LabelEncoder wrapper with TransformerMixin. Allows LabelEncoder
    to work with a Pipeline.
    """

    def fit(self, X, y=None):
        """ Fit the LabelEncoder """
        self.le = LabelEncoder()
        self.le.fit(X)
        return self

    def transform(self, X, y=None):
        """ Transform the label encoder """
        return np.array([[x] for x in self.le.transform(X)])


class SelectorMem(TransformerMixin, BaseEstimator):
    """
    Memory efficient feature selector. Identical functionality
    to classes in sklearn.feature_selection.univariate_selection but
    with more memory efficient attribute storage.

    Args:
        selector (object): any sklearn.feature_selection.univariate_selection._BaseFilter
            inheritor
        score_func (function): selector scoring function
        threshold (int/float): scoring threshold to use with scoring function and selector
    """

    def __init__(
        self, selector="fpr", score_func=feature_selection.f_classif, threshold=0.05
    ):
        self.selector = selector
        self.score_func = score_func
        self.threshold = threshold

    def fit(self, X, y=None):
        """
        Run score function on (X, y) and get the appropriate features

        Args:
            X (array-like, shape = [n_samples, n_features])
                The training input samples
            y (array-like, shape = [n_samples])
                The target values (class labels in classification, real numbers in
                regression)
        """
        selector = _SELECTOR_LOOKUP[self.selector.lower()](
            self.score_func, self.threshold
        )
        selector.fit(X, y)
        mask_indices = selector.get_support(indices=True)
        mask_bool = selector.get_support(indices=False)
        if np.array(mask_bool).nbytes > np.array(mask_indices).nbytes:
            self.mask = mask_indices
        else:
            self.mask = mask_bool
        return self

    def transform(self, X, y=None):
        """
        Reduce X to the selected features

        Args:
            X (array of shape [n_samples, n_features])
                The input samples

        Returns:
            X_r (array of shape [n_samples, n_selected_features])
                The input samples with only the selected features
        """
        return X[:, self.mask]


class HashingVectorizerChunked(HashingVectorizer):
    """
    Equivalent to HashingVectorizer but with chunked prediction

    Args:
        chunksize (int): size of transform chunks
    """

    def __init__(self, chunksize=100000, **kwargs):
        self.chunksize = chunksize
        HashingVectorizer.__init__(self, **kwargs)

    def transform(self, X):
        """Transform a sequence of documents to a document-term matrix

        Args:
            X (iterable over raw text documents, length = n_samples)
                Samples. Each sample must be a text document (either bytes or
                unicode strings, file name or file object depending on the
                constructor argument) which will be tokenized and hashed.

        Returns:
            X (scipy.sparse matrix, shape = (n_samples, self.n_features))
                Document-term matrix.
        """
        if (len(X) < self.chunksize) or (self.chunksize is None):
            return self._transform(X)
        else:
            return sparse.vstack([self._transform(x) for x in self._chunks(X)])

    def _transform(self, X):
        if isinstance(X, str):
            raise ValueError(
                "Iterable over raw text documents expected, string object received."
            )

        analyzer = self.build_analyzer()
        X = self._get_hasher().transform(analyzer(doc) for doc in X)
        if self.binary:
            X.data.fill(1)
        if self.norm is not None:
            X = normalize(X, norm=self.norm, copy=False)
        return X

    def _chunks(self, l):
        for i in range(0, len(l), self.chunksize):
            yield l[i : i + self.chunksize]


class MultihotEncoder(BaseEstimator, TransformerMixin):
    """
    Wraps `MultiLabelBinarizer` in a pipeline safe transformer

    Args:
        sparse_output (bool): convert output to sparse matrix
    """

    def __init__(self, sparse_output=False):
        self.transformer = MultiLabelBinarizer()
        self.sparse_output = sparse_output

    def fit(self, X, y=None):
        """ Fit MultiLabelBinarizer """
        self.transformer.fit(X)
        return self

    def transform(self, X, y=None):
        """ Transform MultiLabelBinarizer """
        # ignore unseen label warning
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            X_t = self.transformer.transform(X)
        if self.sparse_output:
            return sparse.csr_matrix(X_t)
        else:
            return X_t
