"""Meta Transformers module

This module has meta-transformers that is build using the pre-existing
transformers as building blocks.
"""
from .base import BaseTransformer
from ..utils.validation import check_ts_array
from ..utils.transformations import tabularize
from sklearn.utils.validation import check_is_fitted
from sklearn.compose import ColumnTransformer
from scipy import sparse
import numpy as np
import pandas as pd


__all__ = ['TSColumnTransformer', 'RowwiseTransformer', 'Tabularizer', 'Tabulariser']


class TSColumnTransformer(ColumnTransformer):
    """Applies transformers to columns of an array or pandas DataFrame. Simply takes the column transformer from sklearn
    and adds capability to handle pandas dataframe.

    This estimator allows different columns or column subsets of the input
    to be transformed separately and the features generated by each transformer
    will be concatenated to form a single feature space.
    This is useful for heterogeneous or columnar data, to combine several
    feature extraction mechanisms or transformations into a single transformer.

    Parameters
    ----------
    transformers : list of tuples
        List of (name, transformer, column(s)) tuples specifying the
        transformer objects to be applied to subsets of the data.
        name : string
            Like in Pipeline and FeatureUnion, this allows the transformer and
            its parameters to be set using ``set_params`` and searched in grid
            search.
        transformer : estimator or {'passthrough', 'drop'}
            Estimator must support `fit` and `transform`. Special-cased
            strings 'drop' and 'passthrough' are accepted as well, to
            indicate to drop the columns or to pass them through untransformed,
            respectively.
        column(s) : string or int, array-like of string or int, slice, \
boolean mask array or callable
            Indexes the data on its second axis. Integers are interpreted as
            positional columns, while strings can reference DataFrame columns
            by name.  A scalar string or int should be used where
            ``transformer`` expects X to be a 1d array-like (vector),
            otherwise a 2d array will be passed to the transformer.
            A callable is passed the input data `X` and can return any of the
            above.
    remainder : {'drop', 'passthrough'} or estimator, default 'drop'
        By default, only the specified columns in `transformers` are
        transformed and combined in the output, and the non-specified
        columns are dropped. (default of ``'drop'``).
        By specifying ``remainder='passthrough'``, all remaining columns that
        were not specified in `transformers` will be automatically passed
        through. This subset of columns is concatenated with the output of
        the transformers.
        By setting ``remainder`` to be an estimator, the remaining
        non-specified columns will use the ``remainder`` estimator. The
        estimator must support `fit` and `transform`.
    sparse_threshold : float, default = 0.3
        If the output of the different transfromers contains sparse matrices,
        these will be stacked as a sparse matrix if the overall density is
        lower than this value. Use ``sparse_threshold=0`` to always return
        dense.  When the transformed output consists of all dense data, the
        stacked result will be dense, and this keyword will be ignored.
    n_jobs : int or None, optional (default=None)
        Number of jobs to run in parallel.
        ``None`` means 1 unless in a :obj:`joblib.parallel_backend` context.
        ``-1`` means using all processors. See :term:`Glossary <n_jobs>`
        for more details.
    transformer_weights : dict, optional
        Multiplicative weights for features per transformer. The output of the
        transformer is multiplied by these weights. Keys are transformer names,
        values the weights.
    preserve_dataframe : boolean
        If True, pandas dataframe is returned.
        If False, numpy array is returned.


    Attributes
    ----------
    transformers_ : list
        The collection of fitted transformers as tuples of
        (name, fitted_transformer, column). `fitted_transformer` can be an
        estimator, 'drop', or 'passthrough'. In case there were no columns
        selected, this will be the unfitted transformer.
        If there are remaining columns, the final element is a tuple of the
        form:
        ('remainder', transformer, remaining_columns) corresponding to the
        ``remainder`` parameter. If there are remaining columns, then
        ``len(transformers_)==len(transformers)+1``, otherwise
        ``len(transformers_)==len(transformers)``.
    named_transformers_ : Bunch object, a dictionary with attribute access
        Read-only attribute to access any transformer by given name.
        Keys are transformer names and values are the fitted transformer
        objects.
    sparse_output_ : boolean
        Boolean flag indicating wether the output of ``transform`` is a
        sparse matrix or a dense numpy array, which depends on the output
        of the individual transformers and the `sparse_threshold` keyword.
    """
    def __init__(
        self,
        transformers,
        remainder="drop",
        sparse_threshold=0.3,
        n_jobs=1,
        transformer_weights=None,
        preserve_dataframe=True,
    ):

        self.preserve_dataframe = preserve_dataframe
        super(TSColumnTransformer, self).__init__(
            transformers=transformers,
            remainder=remainder,
            sparse_threshold=sparse_threshold,
            n_jobs=n_jobs,
            transformer_weights=transformer_weights,
        )

    def _hstack(self, Xs):
        """Stacks X horizontally.

        Supports input types (X): list of numpy arrays, sparse arrays and DataFrames
        """
        types = set(type(X) for X in Xs)

        if self.sparse_output_:
            return sparse.hstack(Xs).tocsr()
        if self.preserve_dataframe and (pd.Series in types or pd.DataFrame in types):
            return pd.concat(Xs, axis="columns")
        return np.hstack(Xs)

    def _validate_output(self, result):
        """
        Ensure that the output of each transformer is 2D. Otherwise
        hstack can raise an error or produce incorrect results.

        Output can also be a pd.Series which is actually a 1D
        """
        names = [name for name, _, _, _ in self._iter(fitted=True,
                                                      replace_strings=True)]
        for Xs, name in zip(result, names):
            if not (getattr(Xs, 'ndim', 0) == 2 or isinstance(Xs, pd.Series)):
                raise ValueError(
                    "The output of the '{0}' transformer should be 2D (scipy " "matrix, array, or pandas DataFrame).".format(name))


class RowwiseTransformer(BaseTransformer):
    """A convenience wrapper for row-wise transformers to apply transformation to all rows.

    This estimator allows to create a transformer that works on all rows from a passed transformer that works on a
    single row. This is useful for applying transformers to the time-series in the rows.

    Parameters
    ----------
    transformer : estimator
        An estimator that can work on a row (i.e. a univariate time-series in form of a numpy array or pandas Series.
        must support `fit` and `transform`
    """
    def __init__(self, transformer):
        self.transformer = transformer

    def fit(self, X, y=None):
        """Empty fit function that does nothing.

        Parameters
        ----------
        X : 1D array-like, pandas Series, shape (n_samples, 1)
            The training input samples. Shoould not be a DataFrame.
        y : None, as it is transformer on X

        Returns
        -------
        self : object
            Returns self.
        """

        # check the validity of input
        X = check_ts_array(X)

        # fitting - this transformer needs no fitting
        pass

        # let the model know that it is fitted
        self.is_fitted_ = True
        # `fit` should always return `self`
        return self

    def transform(self, X):
        """Calls the fit_transfor() of the per-row transformer repeatedly
        on each row.

        Parameters
        ----------
        X : 1D array-like, pandas Series, shape (n_samples, 1)
            The training input samples. Shoould not be a DataFrame.

        Returns
        -------
        T : 1D array-like, pandas Series, shape (n_samples, ...)
            The transformed data
        """
        # check the validity of input
        X = check_ts_array(X)
        check_is_fitted(self, 'is_fitted_')

        # Works on single column, but on multiple columns only if columns have equal-length series.
        try:
            Xt = X.apply(self.transformer.fit_transform)

        # Otherwise call apply on each column separately.
        except ValueError as e:
            if str(e) == 'arrays must all be same length':
                Xt = pd.concat([pd.Series(col.apply(self.transformer.fit_transform)) for _, col in X.items()], axis=1)
            else:
                raise

        return Xt


class Tabularizer(BaseTransformer):
    """A transformer that turns time-series/panel data into tabular data.

    This estimator converts nested pandas dataframe containing time-series/panel data with numpy arrays or pandas Series in
    dataframe cells into a tabular pandas dataframe with only primitives in cells. This is useful for transforming
    time-series/panel data into a format that is accepted by standard supervised learning algorithms (as in sklearn).

    Parameters
    ----------
    param check_input: bool, optional (default=True)
        When set to ``True``, inputs will be validated, otherwise inputs are assumed to be valid
        and no checks are performed. Use with caution.
    """
    def __init__(self, check_input=True):
        self.check_input = check_input

    def fit(self, X, y=None):
        """Empty fit function that does nothing. Kept here for consistency.

        Parameters
        ----------
        X : 1D array-like, pandas Series, shape (n_samples, 1)
            The training input samples. Shoould not be a DataFrame.
        y : None, as it is transformer on X

        Returns
        -------
        self : object
            Returns self.
        """

        # check the validity of input
        # TODO check if for each column, all rows have equal-index series
        if self.check_input:
            X = check_ts_array(X)

        # let the model know that it is fitted
        self.is_fitted_ = True
        # `fit` should always return `self`
        return self

    def transform(self, X):
        """
        Transform nested pandas dataframe into tabular pandas dataframe.
        param X : pandas dataframe
            Nested dataframe with pandas series or numpy arrays in cells.

        return : pandas dataframe
            Tabular dataframe with only primitives in cells.
        """

        # check the validity of input
        check_is_fitted(self, 'is_fitted_')

        # TODO check if for each column, all rows have equal-index series
        if self.check_input:
            X = check_ts_array(X)

        Xt = tabularize(X)
        return Xt


Tabulariser = Tabularizer
