"""
This module gathers tree-based methods, including decision, regression and
randomized trees. Single and multi-output problems are both handled.
"""

# Authors: Gilles Louppe <g.louppe@gmail.com>
#          Peter Prettenhofer <peter.prettenhofer@gmail.com>
#          Brian Holt <bdholt1@gmail.com>
#          Noel Dawe <noel@dawe.me>
#          Satrajit Gosh <satrajit.ghosh@gmail.com>
#          Joly Arnaud <arnaud.v.joly@gmail.com>
#          Fares Hedayati <fares.hedayati@gmail.com>
#
# Licence: BSD 3 clause

from __future__ import division

import numbers
from abc import ABCMeta, abstractmethod

import numpy as np
from scipy.sparse import issparse

from ..base import BaseEstimator, ClassifierMixin, RegressorMixin
from ..externals import six
from ..feature_selection.from_model import _LearntSelectorMixin
from ..utils import check_array, check_random_state, compute_sample_weight
from ..utils.validation import NotFittedError

from ._tree import Criterion
from ._tree import Splitter
from ._tree import DepthFirstTreeBuilder, BestFirstTreeBuilder
from ._tree import Tree
from . import _tree

__all__ = ["DecisionTreeClassifier",
           "DecisionTreeRegressor",
           "ExtraTreeClassifier",
           "ExtraTreeRegressor"]

# =============================================================================
# Types and constants
# =============================================================================

DTYPE = _tree.DTYPE
DOUBLE = _tree.DOUBLE

CRITERIA_CLF = {"gini": _tree.Gini, "entropy": _tree.Entropy}
CRITERIA_REG = {"mse": _tree.MSE, "friedman_mse": _tree.FriedmanMSE}

DENSE_SPLITTERS = {"best": _tree.BestSplitter,
                   "presort-best": _tree.PresortBestSplitter,
                   "random": _tree.RandomSplitter}

SPARSE_SPLITTERS = {"best": _tree.BestSparseSplitter,
                    "random": _tree.RandomSparseSplitter}

# =============================================================================
# Base decision tree
# =============================================================================

class BaseDecisionTree(six.with_metaclass(ABCMeta, BaseEstimator,
                                          _LearntSelectorMixin)):
    """Base class for decision trees.

    Warning: This class should not be used directly.
    Use derived classes instead.
    """

    @abstractmethod
    def __init__(self,
                 criterion,
                 splitter,
                 max_depth,
                 min_samples_split,
                 min_samples_leaf,
                 min_weight_fraction_leaf,
                 max_features,
                 max_leaf_nodes,
                 random_state,
                 class_weight=None):
        self.criterion = criterion
        self.splitter = splitter
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.min_samples_leaf = min_samples_leaf
        self.min_weight_fraction_leaf = min_weight_fraction_leaf
        self.max_features = max_features
        self.random_state = random_state
        self.max_leaf_nodes = max_leaf_nodes
        self.class_weight = class_weight

        self.n_features_ = None
        self.n_outputs_ = None
        self.classes_ = None
        self.n_classes_ = None

        self.tree_ = None
        self.max_features_ = None

    def fit(self, X, y, sample_weight=None, check_input=True):
        """Build a decision tree from the training set (X, y).

        Parameters
        ----------
        X : array-like or sparse matrix, shape = [n_samples, n_features]
            The training input samples. Internally, it will be converted to
            ``dtype=np.float32`` and if a sparse matrix is provided
            to a sparse ``csc_matrix``.

        y : array-like, shape = [n_samples] or [n_samples, n_outputs]
            The target values (class labels in classification, real numbers in
            regression). In the regression case, use ``dtype=np.float64`` and
            ``order='C'`` for maximum efficiency.

        sample_weight : array-like, shape = [n_samples] or None
            Sample weights. If None, then samples are equally weighted. Splits
            that would create child nodes with net zero or negative weight are
            ignored while searching for a split in each node. In the case of
            classification, splits are also ignored if they would result in any
            single class carrying a negative weight in either child node.

        check_input : boolean, (default=True)
            Allow to bypass several input checking.
            Don't use this parameter unless you know what you do.

        Returns
        -------
        self : object
            Returns self.
        """
        random_state = check_random_state(self.random_state)
        if check_input:
            X = check_array(X, dtype=DTYPE, accept_sparse="csc")
            if issparse(X):
                X.sort_indices()

                if X.indices.dtype != np.intc or X.indptr.dtype != np.intc:
                    raise ValueError("No support for np.int64 index based "
                                     "sparse matrices")

        # Determine output settings
        n_samples, self.n_features_ = X.shape
        is_classification = isinstance(self, ClassifierMixin)

        y = np.atleast_1d(y)
        expanded_class_weight = None

        if y.ndim == 1:
            # reshape is necessary to preserve the data contiguity against vs
            # [:, np.newaxis] that does not.
            y = np.reshape(y, (-1, 1))

        self.n_outputs_ = y.shape[1]

        if is_classification:
            y = np.copy(y)

            self.classes_ = []
            self.n_classes_ = []

            if self.class_weight is not None:
                y_original = np.copy(y)

            y_store_unique_indices = np.zeros(y.shape, dtype=np.int)
            for k in range(self.n_outputs_):
                classes_k, y_store_unique_indices[:, k] = np.unique(y[:, k], return_inverse=True)
                self.classes_.append(classes_k)
                self.n_classes_.append(classes_k.shape[0])
            y = y_store_unique_indices

            if self.class_weight is not None:
                expanded_class_weight = compute_sample_weight(
                    self.class_weight, y_original)

        else:
            self.classes_ = [None] * self.n_outputs_
            self.n_classes_ = [1] * self.n_outputs_

        self.n_classes_ = np.array(self.n_classes_, dtype=np.intp)

        if getattr(y, "dtype", None) != DOUBLE or not y.flags.contiguous:
            y = np.ascontiguousarray(y, dtype=DOUBLE)

        # Check parameters
        max_depth = ((2 ** 31) - 1 if self.max_depth is None
                     else self.max_depth)
        max_leaf_nodes = (-1 if self.max_leaf_nodes is None
                          else self.max_leaf_nodes)

        if isinstance(self.max_features, six.string_types):
            if self.max_features == "auto":
                if is_classification:
                    max_features = max(1, int(np.sqrt(self.n_features_)))
                else:
                    max_features = self.n_features_
            elif self.max_features == "sqrt":
                max_features = max(1, int(np.sqrt(self.n_features_)))
            elif self.max_features == "log2":
                max_features = max(1, int(np.log2(self.n_features_)))
            else:
                raise ValueError(
                    'Invalid value for max_features. Allowed string '
                    'values are "auto", "sqrt" or "log2".')
        elif self.max_features is None:
            max_features = self.n_features_
        elif isinstance(self.max_features, (numbers.Integral, np.integer)):
            max_features = self.max_features
        else:  # float
            if self.max_features > 0.0:
                max_features = max(1, int(self.max_features * self.n_features_))
            else:
                max_features = 0

        self.max_features_ = max_features

        if len(y) != n_samples:
            raise ValueError("Number of labels=%d does not match "
                             "number of samples=%d" % (len(y), n_samples))
        if self.min_samples_split <= 0:
            raise ValueError("min_samples_split must be greater than zero.")
        if self.min_samples_leaf <= 0:
            raise ValueError("min_samples_leaf must be greater than zero.")
        if not 0 <= self.min_weight_fraction_leaf <= 0.5:
            raise ValueError("min_weight_fraction_leaf must in [0, 0.5]")
        if max_depth <= 0:
            raise ValueError("max_depth must be greater than zero. ")
        if not (0 < max_features <= self.n_features_):
            raise ValueError("max_features must be in (0, n_features]")
        if not isinstance(max_leaf_nodes, (numbers.Integral, np.integer)):
            raise ValueError("max_leaf_nodes must be integral number but was "
                             "%r" % max_leaf_nodes)
        if -1 < max_leaf_nodes < 2:
            raise ValueError(("max_leaf_nodes {0} must be either smaller than "
                              "0 or larger than 1").format(max_leaf_nodes))

        if sample_weight is not None:
            if (getattr(sample_weight, "dtype", None) != DOUBLE or
                    not sample_weight.flags.contiguous):
                sample_weight = np.ascontiguousarray(
                    sample_weight, dtype=DOUBLE)
            if len(sample_weight.shape) > 1:
                raise ValueError("Sample weights array has more "
                                 "than one dimension: %d" %
                                 len(sample_weight.shape))
            if len(sample_weight) != n_samples:
                raise ValueError("Number of weights=%d does not match "
                                 "number of samples=%d" %
                                 (len(sample_weight), n_samples))

        if expanded_class_weight is not None:
            if sample_weight is not None:
                sample_weight = sample_weight * expanded_class_weight
            else:
                sample_weight = expanded_class_weight

        # Set min_weight_leaf from min_weight_fraction_leaf
        if self.min_weight_fraction_leaf != 0. and sample_weight is not None:
            min_weight_leaf = (self.min_weight_fraction_leaf *
                               np.sum(sample_weight))
        else:
            min_weight_leaf = 0.

        # Set min_samples_split sensibly
        min_samples_split = max(self.min_samples_split,
                                2 * self.min_samples_leaf)

        # Build tree
        criterion = self.criterion
        if not isinstance(criterion, Criterion):
            if is_classification:
                criterion = CRITERIA_CLF[self.criterion](self.n_outputs_,
                                                         self.n_classes_)
            else:
                criterion = CRITERIA_REG[self.criterion](self.n_outputs_)

        SPLITTERS = SPARSE_SPLITTERS if issparse(X) else DENSE_SPLITTERS

        splitter = self.splitter
        if not isinstance(self.splitter, Splitter):
            splitter = SPLITTERS[self.splitter](criterion,
                                                self.max_features_,
                                                self.min_samples_leaf,
                                                min_weight_leaf,
                                                random_state)

        self.tree_ = Tree(self.n_features_, self.n_classes_, self.n_outputs_)

        # Use BestFirst if max_leaf_nodes given; use DepthFirst otherwise
        if max_leaf_nodes < 0:
            builder = DepthFirstTreeBuilder(splitter, min_samples_split,
                                            self.min_samples_leaf,
                                            min_weight_leaf,
                                            max_depth)
        else:
            builder = BestFirstTreeBuilder(splitter, min_samples_split,
                                           self.min_samples_leaf,
                                           min_weight_leaf,
                                           max_depth,
                                           max_leaf_nodes)

        builder.build(self.tree_, X, y, sample_weight)

        if self.n_outputs_ == 1:
            self.n_classes_ = self.n_classes_[0]
            self.classes_ = self.classes_[0]

        return self

    def _validate_X_predict(self, X, check_input):
        """Validate X whenever one tries to predict, apply, predict_proba"""
        if self.tree_ is None:
            raise NotFittedError("Estimator not fitted, "
                                 "call `fit` before exploiting the model.")

        if check_input:
            X = check_array(X, dtype=DTYPE, accept_sparse="csr")
            if issparse(X) and (X.indices.dtype != np.intc or
                                X.indptr.dtype != np.intc):
                raise ValueError("No support for np.int64 index based "
                                 "sparse matrices")

        n_features = X.shape[1]
        if self.n_features_ != n_features:
            raise ValueError("Number of features of the model must "
                             " match the input. Model n_features is %s and "
                             " input n_features is %s "
                             % (self.n_features_, n_features))

        return X

    def predict(self, X, check_input=True):
        """Predict class or regression value for X.

        For a classification model, the predicted class for each sample in X is
        returned. For a regression model, the predicted value based on X is
        returned.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The input samples. Internally, it will be converted to
            ``dtype=np.float32`` and if a sparse matrix is provided
            to a sparse ``csr_matrix``.

        check_input : boolean, (default=True)
            Allow to bypass several input checking.
            Don't use this parameter unless you know what you do.

        Returns
        -------
        y : array of shape = [n_samples] or [n_samples, n_outputs]
            The predicted classes, or the predict values.
        """
        X = self._validate_X_predict(X, check_input)
        proba = self.tree_.predict(X)
        n_samples = X.shape[0]

        # Classification
        if isinstance(self, ClassifierMixin):
            if self.n_outputs_ == 1:
                return self.classes_.take(np.argmax(proba, axis=1), axis=0)

            else:
                predictions = np.zeros((n_samples, self.n_outputs_))

                for k in range(self.n_outputs_):
                    predictions[:, k] = self.classes_[k].take(
                        np.argmax(proba[:, k], axis=1),
                        axis=0)

                return predictions

        # Regression
        else:
            if self.n_outputs_ == 1:
                return proba[:, 0]

            else:
                return proba[:, :, 0]

    def apply(self, X, check_input=True):
        """
        Returns the index of the leaf that each sample is predicted as.

        Parameters
        ----------
        X : array_like or sparse matrix, shape = [n_samples, n_features]
            The input samples. Internally, it will be converted to
            ``dtype=np.float32`` and if a sparse matrix is provided
            to a sparse ``csr_matrix``.

        check_input : boolean, (default=True)
            Allow to bypass several input checking.
            Don't use this parameter unless you know what you do.

        Returns
        -------
        X_leaves : array_like, shape = [n_samples,]
            For each datapoint x in X, return the index of the leaf x
            ends up in. Leaves are numbered within
            ``[0; self.tree_.node_count)``, possibly with gaps in the
            numbering.
        """
        X = self._validate_X_predict(X, check_input)
        return self.tree_.apply(X)

    @property
    def feature_importances_(self):
        """Return the feature importances.

        The importance of a feature is computed as the (normalized) total
        reduction of the criterion brought by that feature.
        It is also known as the Gini importance.

        Returns
        -------
        feature_importances_ : array, shape = [n_features]
        """
        if self.tree_ is None:
            raise NotFittedError("Estimator not fitted, call `fit` before"
                                 " `feature_importances_`.")

        return self.tree_.compute_feature_importances()

# =============================================================================
# Public estimators
# =============================================================================

class DecisionTreeClassifier(BaseDecisionTree, ClassifierMixin):
    """A decision tree classifier.

    Read more in the :ref:`User Guide <tree>`.

    Parameters
    ----------
    criterion : string, optional (default="gini")
        The function to measure the quality of a split. Supported criteria are
        "gini" for the Gini impurity and "entropy" for the information gain.

    splitter : string, optional (default="best")
        The strategy used to choose the split at each node. Supported
        strategies are "best" to choose the best split and "random" to choose
        the best random split.

    max_features : int, float, string or None, optional (default=None)
        The number of features to consider when looking for the best split:
          - If int, then consider `max_features` features at each split.
          - If float, then `max_features` is a percentage and
            `int(max_features * n_features)` features are considered at each
            split.
          - If "auto", then `max_features=sqrt(n_features)`.
          - If "sqrt", then `max_features=sqrt(n_features)`.
          - If "log2", then `max_features=log2(n_features)`.
          - If None, then `max_features=n_features`.

        Note: the search for a split does not stop until at least one
        valid partition of the node samples is found, even if it requires to
        effectively inspect more than ``max_features`` features.

    max_depth : int or None, optional (default=None)
        The maximum depth of the tree. If None, then nodes are expanded until
        all leaves are pure or until all leaves contain less than
        min_samples_split samples.
        Ignored if ``max_leaf_nodes`` is not None.

    min_samples_split : int, optional (default=2)
        The minimum number of samples required to split an internal node.

    min_samples_leaf : int, optional (default=1)
        The minimum number of samples required to be at a leaf node.

    min_weight_fraction_leaf : float, optional (default=0.)
        The minimum weighted fraction of the input samples required to be at a
        leaf node.

    max_leaf_nodes : int or None, optional (default=None)
        Grow a tree with ``max_leaf_nodes`` in best-first fashion.
        Best nodes are defined as relative reduction in impurity.
        If None then unlimited number of leaf nodes.
        If not None then ``max_depth`` will be ignored.

    class_weight : dict, list of dicts, "balanced" or None, optional
                   (default=None)
        Weights associated with classes in the form ``{class_label: weight}``.
        If not given, all classes are supposed to have weight one. For
        multi-output problems, a list of dicts can be provided in the same
        order as the columns of y.

        The "balanced" mode uses the values of y to automatically adjust
        weights inversely proportional to class frequencies in the input data
        as ``n_samples / (n_classes * np.bincount(y))``

        For multi-output, the weights of each column of y will be multiplied.

        Note that these weights will be multiplied with sample_weight (passed
        through the fit method) if sample_weight is specified.

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    Attributes
    ----------
    classes_ : array of shape = [n_classes] or a list of such arrays
        The classes labels (single output problem),
        or a list of arrays of class labels (multi-output problem).

    feature_importances_ : array of shape = [n_features]
        The feature importances. The higher, the more important the
        feature. The importance of a feature is computed as the (normalized)
        total reduction of the criterion brought by that feature.  It is also
        known as the Gini importance [4]_.

    max_features_ : int,
        The inferred value of max_features.

    n_classes_ : int or list
        The number of classes (for single output problems),
        or a list containing the number of classes for each
        output (for multi-output problems).

    n_features_ : int
        The number of features when ``fit`` is performed.

    n_outputs_ : int
        The number of outputs when ``fit`` is performed.

    tree_ : Tree object
        The underlying Tree object.

    See also
    --------
    DecisionTreeRegressor

    References
    ----------

    .. [1] http://en.wikipedia.org/wiki/Decision_tree_learning

    .. [2] L. Breiman, J. Friedman, R. Olshen, and C. Stone, "Classification
           and Regression Trees", Wadsworth, Belmont, CA, 1984.

    .. [3] T. Hastie, R. Tibshirani and J. Friedman. "Elements of Statistical
           Learning", Springer, 2009.

    .. [4] L. Breiman, and A. Cutler, "Random Forests",
           http://www.stat.berkeley.edu/~breiman/RandomForests/cc_home.htm

    Examples
    --------
    >>> from sklearn.datasets import load_iris
    >>> from sklearn.cross_validation import cross_val_score
    >>> from sklearn.tree import DecisionTreeClassifier
    >>> clf = DecisionTreeClassifier(random_state=0)
    >>> iris = load_iris()
    >>> cross_val_score(clf, iris.data, iris.target, cv=10)
    ...                             # doctest: +SKIP
    ...
    array([ 1.     ,  0.93...,  0.86...,  0.93...,  0.93...,
            0.93...,  0.93...,  1.     ,  0.93...,  1.      ])
    """
    def __init__(self,
                 criterion="gini",
                 splitter="best",
                 max_depth=None,
                 min_samples_split=2,
                 min_samples_leaf=1,
                 min_weight_fraction_leaf=0.,
                 max_features=None,
                 random_state=None,
                 max_leaf_nodes=None,
                 class_weight=None):
        super(DecisionTreeClassifier, self).__init__(
            criterion=criterion,
            splitter=splitter,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            min_weight_fraction_leaf=min_weight_fraction_leaf,
            max_features=max_features,
            max_leaf_nodes=max_leaf_nodes,
            class_weight=class_weight,
            random_state=random_state)

    def predict_proba(self, X, check_input=True):
        """Predict class probabilities of the input samples X.

        The predicted class probability is the fraction of samples of the same
        class in a leaf.

        check_input : boolean, (default=True)
            Allow to bypass several input checking.
            Don't use this parameter unless you know what you do.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The input samples. Internally, it will be converted to
            ``dtype=np.float32`` and if a sparse matrix is provided
            to a sparse ``csr_matrix``.

        Returns
        -------
        p : array of shape = [n_samples, n_classes], or a list of n_outputs
            such arrays if n_outputs > 1.
            The class probabilities of the input samples. The order of the
            classes corresponds to that in the attribute `classes_`.
        """
        X = self._validate_X_predict(X, check_input)
        proba = self.tree_.predict(X)

        if self.n_outputs_ == 1:
            proba = proba[:, :self.n_classes_]
            normalizer = proba.sum(axis=1)[:, np.newaxis]
            normalizer[normalizer == 0.0] = 1.0
            proba /= normalizer

            return proba

        else:
            all_proba = []

            for k in range(self.n_outputs_):
                proba_k = proba[:, k, :self.n_classes_[k]]
                normalizer = proba_k.sum(axis=1)[:, np.newaxis]
                normalizer[normalizer == 0.0] = 1.0
                proba_k /= normalizer
                all_proba.append(proba_k)

            return all_proba

    def predict_log_proba(self, X):
        """Predict class log-probabilities of the input samples X.

        Parameters
        ----------
        X : array-like or sparse matrix of shape = [n_samples, n_features]
            The input samples. Internally, it will be converted to
            ``dtype=np.float32`` and if a sparse matrix is provided
            to a sparse ``csr_matrix``.

        Returns
        -------
        p : array of shape = [n_samples, n_classes], or a list of n_outputs
            such arrays if n_outputs > 1.
            The class log-probabilities of the input samples. The order of the
            classes corresponds to that in the attribute `classes_`.
        """
        proba = self.predict_proba(X)

        if self.n_outputs_ == 1:
            return np.log(proba)

        else:
            for k in range(self.n_outputs_):
                proba[k] = np.log(proba[k])

            return proba

class DecisionTreeRegressor(BaseDecisionTree, RegressorMixin):
    """A decision tree regressor.

    Read more in the :ref:`User Guide <tree>`.

    Parameters
    ----------
    criterion : string, optional (default="mse")
        The function to measure the quality of a split. The only supported
        criterion is "mse" for the mean squared error, which is equal to
        variance reduction as feature selection criterion.

    splitter : string, optional (default="best")
        The strategy used to choose the split at each node. Supported
        strategies are "best" to choose the best split and "random" to choose
        the best random split.

    max_features : int, float, string or None, optional (default=None)
        The number of features to consider when looking for the best split:
          - If int, then consider `max_features` features at each split.
          - If float, then `max_features` is a percentage and
            `int(max_features * n_features)` features are considered at each
            split.
          - If "auto", then `max_features=n_features`.
          - If "sqrt", then `max_features=sqrt(n_features)`.
          - If "log2", then `max_features=log2(n_features)`.
          - If None, then `max_features=n_features`.

        Note: the search for a split does not stop until at least one
        valid partition of the node samples is found, even if it requires to
        effectively inspect more than ``max_features`` features.

    max_depth : int or None, optional (default=None)
        The maximum depth of the tree. If None, then nodes are expanded until
        all leaves are pure or until all leaves contain less than
        min_samples_split samples.
        Ignored if ``max_leaf_nodes`` is not None.

    min_samples_split : int, optional (default=2)
        The minimum number of samples required to split an internal node.

    min_samples_leaf : int, optional (default=1)
        The minimum number of samples required to be at a leaf node.

    min_weight_fraction_leaf : float, optional (default=0.)
        The minimum weighted fraction of the input samples required to be at a
        leaf node.

    max_leaf_nodes : int or None, optional (default=None)
        Grow a tree with ``max_leaf_nodes`` in best-first fashion.
        Best nodes are defined as relative reduction in impurity.
        If None then unlimited number of leaf nodes.
        If not None then ``max_depth`` will be ignored.

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    Attributes
    ----------
    feature_importances_ : array of shape = [n_features]
        The feature importances.
        The higher, the more important the feature.
        The importance of a feature is computed as the
        (normalized) total reduction of the criterion brought
        by that feature. It is also known as the Gini importance [4]_.

    max_features_ : int,
        The inferred value of max_features.

    n_features_ : int
        The number of features when ``fit`` is performed.

    n_outputs_ : int
        The number of outputs when ``fit`` is performed.

    tree_ : Tree object
        The underlying Tree object.

    See also
    --------
    DecisionTreeClassifier

    References
    ----------

    .. [1] http://en.wikipedia.org/wiki/Decision_tree_learning

    .. [2] L. Breiman, J. Friedman, R. Olshen, and C. Stone, "Classification
           and Regression Trees", Wadsworth, Belmont, CA, 1984.

    .. [3] T. Hastie, R. Tibshirani and J. Friedman. "Elements of Statistical
           Learning", Springer, 2009.

    .. [4] L. Breiman, and A. Cutler, "Random Forests",
           http://www.stat.berkeley.edu/~breiman/RandomForests/cc_home.htm

    Examples
    --------
    >>> from sklearn.datasets import load_boston
    >>> from sklearn.cross_validation import cross_val_score
    >>> from sklearn.tree import DecisionTreeRegressor
    >>> boston = load_boston()
    >>> regressor = DecisionTreeRegressor(random_state=0)
    >>> cross_val_score(regressor, boston.data, boston.target, cv=10)
    ...                    # doctest: +SKIP
    ...
    array([ 0.61..., 0.57..., -0.34..., 0.41..., 0.75...,
            0.07..., 0.29..., 0.33..., -1.42..., -1.77...])
    """
    def __init__(self,
                 criterion="mse",
                 splitter="best",
                 max_depth=None,
                 min_samples_split=2,
                 min_samples_leaf=1,
                 min_weight_fraction_leaf=0.,
                 max_features=None,
                 random_state=None,
                 max_leaf_nodes=None):
        super(DecisionTreeRegressor, self).__init__(
            criterion=criterion,
            splitter=splitter,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            min_weight_fraction_leaf=min_weight_fraction_leaf,
            max_features=max_features,
            max_leaf_nodes=max_leaf_nodes,
            random_state=random_state)

class ExtraTreeClassifier(DecisionTreeClassifier):
    """An extremely randomized tree classifier.

    Extra-trees differ from classic decision trees in the way they are built.
    When looking for the best split to separate the samples of a node into two
    groups, random splits are drawn for each of the `max_features` randomly
    selected features and the best split among those is chosen. When
    `max_features` is set 1, this amounts to building a totally random
    decision tree.

    Warning: Extra-trees should only be used within ensemble methods.

    Read more in the :ref:`User Guide <tree>`.

    See also
    --------
    ExtraTreeRegressor, ExtraTreesClassifier, ExtraTreesRegressor

    References
    ----------

    .. [1] P. Geurts, D. Ernst., and L. Wehenkel, "Extremely randomized trees",
           Machine Learning, 63(1), 3-42, 2006.
    """
    def __init__(self,
                 criterion="gini",
                 splitter="random",
                 max_depth=None,
                 min_samples_split=2,
                 min_samples_leaf=1,
                 min_weight_fraction_leaf=0.,
                 max_features="auto",
                 random_state=None,
                 max_leaf_nodes=None,
                 class_weight=None):
        super(ExtraTreeClassifier, self).__init__(
            criterion=criterion,
            splitter=splitter,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            min_weight_fraction_leaf=min_weight_fraction_leaf,
            max_features=max_features,
            max_leaf_nodes=max_leaf_nodes,
            class_weight=class_weight,
            random_state=random_state)

class ExtraTreeRegressor(DecisionTreeRegressor):
    """An extremely randomized tree regressor.

    Extra-trees differ from classic decision trees in the way they are built.
    When looking for the best split to separate the samples of a node into two
    groups, random splits are drawn for each of the `max_features` randomly
    selected features and the best split among those is chosen. When
    `max_features` is set 1, this amounts to building a totally random
    decision tree.

    Warning: Extra-trees should only be used within ensemble methods.

    Read more in the :ref:`User Guide <tree>`.

    See also
    --------
    ExtraTreeClassifier, ExtraTreesClassifier, ExtraTreesRegressor

    References
    ----------

    .. [1] P. Geurts, D. Ernst., and L. Wehenkel, "Extremely randomized trees",
           Machine Learning, 63(1), 3-42, 2006.
    """
    def __init__(self,
                 criterion="mse",
                 splitter="random",
                 max_depth=None,
                 min_samples_split=2,
                 min_samples_leaf=1,
                 min_weight_fraction_leaf=0.,
                 max_features="auto",
                 random_state=None,
                 max_leaf_nodes=None):
        super(ExtraTreeRegressor, self).__init__(
            criterion=criterion,
            splitter=splitter,
            max_depth=max_depth,
            min_samples_split=min_samples_split,
            min_samples_leaf=min_samples_leaf,
            min_weight_fraction_leaf=min_weight_fraction_leaf,
            max_features=max_features,
            max_leaf_nodes=max_leaf_nodes,
            random_state=random_state)


"""
MapPlot
-------

Make plots of the cube and interactively connect them to spectrum plotting.
This is really an interactive component of the package; nothing in here is
meant for publication-quality plots, but more for user interactive analysis.

That said, the plotter makes use of `APLpy <https://github.com/aplpy/aplpy>`_,
so it is possible to make publication-quality plots.

:author: Adam Ginsburg
:date: 03/17/2011

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

"""
from __future__ import print_function
import matplotlib
import matplotlib.figure
import numpy as np
import copy
import itertools
import six
try:
    import astropy.wcs as pywcs
    import astropy.io.fits as pyfits
    pywcsOK = True
except ImportError:
    try:
        import pyfits
        import pywcs
        pywcsOK = True
    except ImportError:
        pywcsOK = False
try:
    import aplpy
    icanhasaplpy = True
except: # aplpy fails with generic exceptions instead of ImportError
    icanhasaplpy = False

from . import cubes

class MapPlotter(object):
    """
    Class to plot a spectrum

    See `mapplot` for use documentation; this docstring is only for
    initialization.
    """

    def __init__(self, Cube=None, figure=None, doplot=False, **kwargs):
        """
        Create a map figure for future plotting
        """
        import matplotlib.pyplot
        self._pyplot = matplotlib.pyplot

        # figure out where to put the plot
        if isinstance(figure,matplotlib.figure.Figure):
            self.figure = figure
        elif type(figure) is int:
            self.figure = self._pyplot.figure(figure)
        else:
            self.figure = None
        self.axis = None
        self.FITSFigure = None
        self._click_marks = []
        self._circles = []
        self._clickX = None
        self._clickY = None

        self.overplot_colorcycle = itertools.cycle(['b', 'g', 'r', 'c', 'm', 'y'])
        self.overplot_linestyle = '-'

        self.Cube = Cube
        if self.Cube is not None:
            self.header = cubes.flatten_header(self.Cube.header, delete=True)
        if pywcsOK:
            self.wcs = pywcs.WCS(self.header)

        if doplot: self.mapplot(**kwargs)

    def __call__(self, **kwargs):
        """ see mapplot """
        return self.mapplot(**kwargs)

    def mapplot(self, convention='calabretta', colorbar=True, useaplpy=True,
                vmin=None, vmax=None, cmap=None, plotkwargs={}, **kwargs):
        """
        Plot up a map based on an input data cube.

        The map to be plotted is selected using `makeplane`.
        The `estimator` keyword argument is passed to that function.

        The plotted map, once shown, is interactive.  You can click on it with any
        of the three mouse buttons.

        Button 1 or keyboard '1':
            Plot the selected pixel's spectrum in another window.  Mark the
            clicked pixel with an 'x'
        Button 2 or keyboard 'o':
            Overplot a second (or third, fourth, fifth...) spectrum in the
            external plot window
        Button 3:
            Disconnect the interactive viewer

        You can also click-and-drag with button 1 to average over a circular
        region.  This same effect can be achieved by using the 'c' key to
        set the /c/enter of a circle and the 'r' key to set its /r/adius (i.e.,
        hover over the center and press 'c', then hover some distance away and
        press 'r').

        Parameters
        ----------
        convention : 'calabretta' or 'griesen'
            The default projection to assume for Galactic data when plotting
            with aplpy.
        colorbar : bool
            Whether to show a colorbar
        plotkwargs : dict, optional
            A dictionary of keyword arguments to pass to aplpy.show_colorscale
            or matplotlib.pyplot.imshow
        useaplpy : bool
            Use aplpy if a FITS header is available
        vmin, vmax: float or None
            Override values for the vmin/vmax values.  Will be automatically
            determined if left as None

        .. todo:
            Allow mapplot in subfigure
        """
        if (self.figure is None):
            self.figure = self._pyplot.figure()
        elif (not self._pyplot.fignum_exists(self.figure.number)):
            self.figure = self._pyplot.figure()
        else:
            self._disconnect()
            self.figure.clf()

        # this is where the map is created; everything below this is just plotting
        self.makeplane(**kwargs)

        # have tot pop out estimator so that kwargs can be passed to imshow
        if 'estimator' in kwargs:
            kwargs.pop('estimator')

        # Below here is all plotting stuff

        if vmin is None: vmin = self.plane[self.plane==self.plane].min()
        if vmax is None: vmax = self.plane[self.plane==self.plane].max()

        if icanhasaplpy and useaplpy:
            self.fitsfile = pyfits.PrimaryHDU(data=self.plane,header=self.header)
            self.FITSFigure = aplpy.FITSFigure(self.fitsfile,figure=self.figure,convention=convention)
            self.FITSFigure.show_colorscale(vmin=vmin, vmax=vmax, cmap=cmap, **plotkwargs)
            if hasattr(self.FITSFigure, '_ax1'):
                self.axis = self.FITSFigure._ax1
            else:
                self.axis = self.FITSFigure.ax
            if colorbar:
                try:
                    self.FITSFigure.add_colorbar()
                except Exception as ex:
                    print("ERROR: Could not create colorbar!  Error was %s" % str(ex))
            self._origin = 0 # FITS convention
            # TODO: set _origin to 1 if using PIXEL units, not real wcs
        else:
            self.axis = self.figure.add_subplot(111)
            if hasattr(self,'colorbar') and self.colorbar is not None:
                if self.colorbar.ax in self.axis.figure.axes:
                    self.axis.figure.delaxes(self.colorbar.ax)
            self.axis.imshow(self.plane, vmin=vmin, vmax=vmax, cmap=cmap, **plotkwargs)
            if colorbar:
                try:
                    self.colorbar = self._pyplot.colorbar(self.axis.images[0])
                except Exception as ex:
                    print("ERROR: Could not create colorbar!  Error was %s" % str(ex))
            self._origin = 0 # normal convention

        self.canvas = self.axis.figure.canvas

        self._connect()

    def _connect(self):
        """ Connect click, click up (release click), and key press to events """
        self.clickid = self.canvas.callbacks.connect('button_press_event',self.click)
        self.clickupid = self.canvas.callbacks.connect('button_release_event',self.plot_spectrum)
        self.keyid = self.canvas.callbacks.connect('key_press_event',self.plot_spectrum)

    def _disconnect(self):
        """ Disconnect click, click up (release click), and key press from events """
        if hasattr(self,'canvas'):
            self.canvas.mpl_disconnect(self.clickid)
            self.canvas.mpl_disconnect(self.clickupid)
            self.canvas.mpl_disconnect(self.keyid)

    def makeplane(self, estimator=np.nanmean):
        """
        Create a "plane" view of the cube, either by slicing or projecting it
        or by showing a slice from the best-fit model parameter cube.

        Parameters
        ----------

        estimator : [ function | 'max' | 'int' | FITS filename | integer | slice ]
            A non-pythonic, non-duck-typed variable.  If it's a function, apply that function
            along the cube's spectral axis to obtain an estimate (e.g., mean, min, max, etc.).
            'max' will do the same thing as passing np.max
            'int' will attempt to integrate the image (which is why I didn't duck-type)
            (integrate means sum and multiply by dx)
            a .fits filename will be read using pyfits (so you can make your own cover figure)
            an integer will get the n'th slice in the parcube if it exists
            If it's a slice, slice the input data cube along the Z-axis with this slice

        """
        # THIS IS A HACK!!!  isinstance(a function, function) must be a thing...
        FUNCTION = type(np.max)

        # estimator is NOT duck-typed
        if type(estimator) is FUNCTION:
            self.plane = estimator(self.Cube.cube,axis=0)
        elif isinstance(estimator, six.string_types):
            if estimator == 'max':
                self.plane = self.Cube.cube.max(axis=0)
            elif estimator == 'int':
                dx = np.abs(self.Cube.xarr[1:] - self.Cube.xarr[:-1])
                dx = np.concatenate([dx,[dx[-1]]])
                self.plane = (self.Cube.cube * dx[:,np.newaxis,np.newaxis]).sum(axis=0)
            elif estimator[-5:] == ".fits":
                self.plane = pyfits.getdata(estimator)
        elif type(estimator) is slice:
            self.plane = self.Cube.cube[estimator,:,:]
        elif type(estimator) is int:
            if hasattr(self.Cube,'parcube'):
                self.plane = self.Cube.parcube[estimator,:,:]

        if self.plane is None:
            raise ValueError("Invalid estimator %s" % (str(estimator)))

        if np.sum(np.isfinite(self.plane)) == 0:
            raise ValueError("Map is all NaNs or infs.  Check your estimator or your input cube.")

    def click(self,event):
        """
        Record location of downclick
        """
        if event.inaxes:
            self._clickX = np.round(event.xdata) - self._origin
            self._clickY = np.round(event.ydata) - self._origin

    def plot_spectrum(self, event, plot_fit=True):
        """
        Connects map cube to Spectrum...
        """
        self.event = event
        if event.inaxes:
            clickX = np.round(event.xdata) - self._origin
            clickY = np.round(event.ydata) - self._origin

            # grab toolbar info so that we don't do anything if a tool is selected
            tb = self.canvas.toolbar
            if tb.mode != '':
                return
            elif event.key is not None:
                if event.key == 'c':
                    self._center = (clickX-1,clickY-1)
                    self._remove_circle()
                    self._add_click_mark(clickX,clickY,clear=True)
                elif event.key == 'r':
                    x,y = self._center
                    self._add_circle(x,y,clickX,clickY)
                    self.circle(x,y,clickX-1,clickY-1)
                elif event.key == 'o':
                    clickX,clickY = round(clickX),round(clickY)
                    print("OverPlotting spectrum from point %i,%i" % (clickX-1,clickY-1))
                    color = next(self.overplot_colorcycle)
                    self._add_click_mark(clickX,clickY,clear=False, color=color)
                    self.Cube.plot_spectrum(clickX-1,clickY-1,clear=False, color=color, linestyle=self.overplot_linestyle)
                elif event.key in ('1','2'):
                    event.button = int(event.key)
                    event.key = None
                    self.plot_spectrum(event)
            elif (hasattr(event,'button') and event.button in (1,2)
                    and not (self._clickX == clickX and self._clickY == clickY)):
                if event.button == 1:
                    self._remove_circle()
                    clear=True
                    color = 'k'
                    linestyle = 'steps-mid'
                else:
                    color = next(self.overplot_colorcycle)
                    linestyle = self.overplot_linestyle
                    clear=False
                rad = ( (self._clickX-clickX)**2 + (self._clickY-clickY)**2 )**0.5
                print("Plotting circle from point %i,%i to %i,%i (r=%f)" % (self._clickX,self._clickY,clickX,clickY,rad))
                self._add_circle(self._clickX,self._clickY,clickX,clickY)
                self.circle(self._clickX,self._clickY,clickX,clickY,clear=clear,linestyle=linestyle,color=color)
            elif hasattr(event,'button') and event.button is not None:
                if event.button==1:
                    clickX,clickY = round(clickX),round(clickY)
                    print("Plotting spectrum from point %i,%i" % (clickX,clickY))
                    self._remove_circle()
                    self._add_click_mark(clickX,clickY,clear=True)
                    self.Cube.plot_spectrum(clickX,clickY,clear=True)
                    if plot_fit: self.Cube.plot_fit(clickX, clickY, silent=True)
                elif event.button==2:
                    clickX,clickY = round(clickX),round(clickY)
                    print("OverPlotting spectrum from point %i,%i" % (clickX,clickY))
                    color = next(self.overplot_colorcycle)
                    self._add_click_mark(clickX,clickY,clear=False, color=color)
                    self.Cube.plot_spectrum(clickX,clickY,clear=False, color=color, linestyle=self.overplot_linestyle)
                elif event.button==3:
                    print("Disconnecting GAIA-like tool")
                    self._disconnect()
            else:
                print("Call failed for some reason: ")
                print("event: ",event)
        else:
            pass
            # never really needed... warn("Click outside of axes")

    def _add_click_mark(self,x,y,clear=False,color='k'):
        """
        Add an X at some position
        """
        if clear:
            self._clear_click_marks()
        if self.FITSFigure is not None:
            label = 'xmark%i' % (len(self._click_marks)+1)
            x,y = self.FITSFigure.pixel2world(x,y)
            self.FITSFigure.show_markers(x,y,marker='x',c=color,layer=label)
            self._click_marks.append( label )
        else:
            self._click_marks.append( self.axis.plot(x,y,'kx') )
        self.refresh()

    def _clear_click_marks(self):
        """
        Remove all marks added by previous clicks
        """
        if self.FITSFigure is not None:
            for mark in self._click_marks:
                if mark in self.FITSFigure._layers:
                    self.FITSFigure.remove_layer(mark)
        else:
            for mark in self._click_marks:
                self._click_marks.remove(mark)
                if mark in self.axis.lines:
                    self.axis.lines.remove(mark)
            self.refresh()

    def _add_circle(self,x,y,x2,y2,**kwargs):
        """
        """
        if self.FITSFigure is not None:
            x,y = self.FITSFigure.pixel2world(x,y)
            x2,y2 = self.FITSFigure.pixel2world(x2,y2)
            r = (np.linalg.norm(np.array([x,y])-np.array([x2,y2])))
            #self.FITSFigure.show_markers(x,y,s=r,marker='o',facecolor='none',edgecolor='black',layer='circle')
            layername = "circle%02i" % len(self._circles)
            self.FITSFigure.show_circles(x,y,r,edgecolor='black',facecolor='none',layer=layername,**kwargs)
            self._circles.append(layername)
        else:
            r = np.linalg.norm(np.array([x,y])-np.array([x2,y2]))
            circle = matplotlib.patches.Circle([x,y],radius=r,**kwargs)
            self._circles.append( circle )
            self.axis.patches.append(circle)
            self.refresh()

    def _remove_circle(self):
        """
        """
        if self.FITSFigure is not None:
            for layername in self._circles:
                if layername in self.FITSFigure._layers:
                    self.FITSFigure.remove_layer(layername)
        else:
            for circle in self._circles:
                if circle in self.axis.patches:
                    self.axis.patches.remove(circle)
                self._circles.remove(circle)
            self.refresh()

    def refresh(self):
        if self.axis is not None:
            self.axis.figure.canvas.draw()

    def circle(self,x1,y1,x2,y2,**kwargs):
        """
        Plot the spectrum of a circular aperture
        """

        r = (np.linalg.norm(np.array([x1,y1])-np.array([x2,y2])))
        self.Cube.plot_apspec([x1,y1,r],**kwargs)
        #self.Cube.data = cubes.extract_aperture( self.Cube.cube, [x1,y1,r] , coordsys=None )
        #self.Cube.plotter()

    def copy(self, parent=None):
        """
        Create a copy of the map plotter with blank (uninitialized) axis & figure

        [ parent ]
            A spectroscopic axis instance that is the parent of the specfit
            instance.  This needs to be specified at some point, but defaults
            to None to prevent overwriting a previous plot.
        """

        newmapplot = copy.copy(self)
        newmapplot.Cube = parent
        newmapplot.axis = None
        newmapplot.figure = None

        return newmapplot


# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementations of different data feeders to provide data for TF trainer."""

# TODO(ipolosukhin): Replace this module with feed-dict queue runners & queues.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools
import math

import numpy as np
import six
from six.moves import xrange  # pylint: disable=redefined-builtin

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.platform import tf_logging as logging

# pylint: disable=g-multiple-import,g-bad-import-order
from .pandas_io import HAS_PANDAS, extract_pandas_data, extract_pandas_matrix, extract_pandas_labels
from .dask_io import HAS_DASK, extract_dask_data, extract_dask_labels

# pylint: enable=g-multiple-import,g-bad-import-order

def _get_in_out_shape(x_shape, y_shape, n_classes, batch_size=None):
  """Returns shape for input and output of the data feeder."""
  x_is_dict, y_is_dict = isinstance(
      x_shape, dict), y_shape is not None and isinstance(y_shape, dict)
  if y_is_dict and n_classes is not None:
    assert (isinstance(n_classes, dict))

  if batch_size is None:
    batch_size = list(x_shape.values())[0][0] if x_is_dict else x_shape[0]
  elif batch_size <= 0:
    raise ValueError('Invalid batch_size %d.' % batch_size)

  if x_is_dict:
    input_shape = {}
    for k, v in list(x_shape.items()):
      input_shape[k] = [batch_size] + (list(v[1:]) if len(v) > 1 else [1])
  else:
    x_shape = list(x_shape[1:]) if len(x_shape) > 1 else [1]
    input_shape = [batch_size] + x_shape

  if y_shape is None:
    return input_shape, None, batch_size

  def out_el_shape(out_shape, num_classes):
    out_shape = list(out_shape[1:]) if len(out_shape) > 1 else []
    # Skip first dimension if it is 1.
    if out_shape and out_shape[0] == 1:
      out_shape = out_shape[1:]
    if num_classes is not None and num_classes > 1:
      return [batch_size] + out_shape + [num_classes]
    else:
      return [batch_size] + out_shape

  if not y_is_dict:
    output_shape = out_el_shape(y_shape, n_classes)
  else:
    output_shape = dict([
        (k, out_el_shape(v, n_classes[k]
                         if n_classes is not None and k in n_classes else None))
        for k, v in list(y_shape.items())
    ])

  return input_shape, output_shape, batch_size

def _data_type_filter(x, y):
  """Filter data types into acceptable format."""
  if HAS_DASK:
    x = extract_dask_data(x)
    if y is not None:
      y = extract_dask_labels(y)
  if HAS_PANDAS:
    x = extract_pandas_data(x)
    if y is not None:
      y = extract_pandas_labels(y)
  return x, y

def _is_iterable(x):
  return hasattr(x, 'next') or hasattr(x, '__next__')

def setup_train_data_feeder(x,
                            y,
                            n_classes,
                            batch_size=None,
                            shuffle=True,
                            epochs=None):
  """Create data feeder, to sample inputs from dataset.

  If `x` and `y` are iterators, use `StreamingDataFeeder`.

  Args:
    x: numpy, pandas or Dask matrix or dictionary of aforementioned. Also
      supports iterables.
    y: numpy, pandas or Dask array or dictionary of aforementioned. Also
      supports
      iterables.
    n_classes: number of classes. Must be None or same type as y. In case, `y`
      is `dict`
      (or iterable which returns dict) such that `n_classes[key] = n_classes for
        y[key]`
    batch_size: size to split data into parts. Must be >= 1.
    shuffle: Whether to shuffle the inputs.
    epochs: Number of epochs to run.

  Returns:
    DataFeeder object that returns training data.

  Raises:
    ValueError: if one of `x` and `y` is iterable and the other is not.
  """
  x, y = _data_type_filter(x, y)
  if HAS_DASK:
    # pylint: disable=g-import-not-at-top
    import dask.dataframe as dd
    if (isinstance(x, (dd.Series, dd.DataFrame)) and
        (y is None or isinstance(y, (dd.Series, dd.DataFrame)))):
      data_feeder_cls = DaskDataFeeder
    else:
      data_feeder_cls = DataFeeder
  else:
    data_feeder_cls = DataFeeder

  if _is_iterable(x):
    if y is not None and not _is_iterable(y):
      raise ValueError('Both x and y should be iterators for '
                       'streaming learning to work.')
    return StreamingDataFeeder(x, y, n_classes, batch_size)
  return data_feeder_cls(
      x, y, n_classes, batch_size, shuffle=shuffle, epochs=epochs)

def _batch_data(x, batch_size=None):
  if (batch_size is not None) and (batch_size <= 0):
    raise ValueError('Invalid batch_size %d.' % batch_size)

  x_first_el = six.next(x)
  x = itertools.chain([x_first_el], x)

  chunk = dict([(k, []) for k in list(x_first_el.keys())]) if isinstance(
      x_first_el, dict) else []
  chunk_filled = False
  for data in x:
    if isinstance(data, dict):
      for k, v in list(data.items()):
        chunk[k].append(v)
        if (batch_size is not None) and (len(chunk[k]) >= batch_size):
          chunk[k] = np.matrix(chunk[k])
          chunk_filled = True
      if chunk_filled:
        yield chunk
        chunk = dict([(k, []) for k in list(x_first_el.keys())]) if isinstance(
            x_first_el, dict) else []
        chunk_filled = False
    else:
      chunk.append(data)
      if (batch_size is not None) and (len(chunk) >= batch_size):
        yield np.matrix(chunk)
        chunk = []

  if isinstance(x_first_el, dict):
    for k, v in list(data.items()):
      chunk[k] = np.matrix(chunk[k])
    yield chunk
  else:
    yield np.matrix(chunk)

def setup_predict_data_feeder(x, batch_size=None):
  """Returns an iterable for feeding into predict step.

  Args:
    x: numpy, pandas, Dask array or dictionary of aforementioned. Also supports
      iterable.
    batch_size: Size of batches to split data into. If `None`, returns one
      batch of full size.

  Returns:
    List or iterator (or dictionary thereof) of parts of data to predict on.

  Raises:
    ValueError: if `batch_size` <= 0.
  """
  if HAS_DASK:
    x = extract_dask_data(x)
  if HAS_PANDAS:
    x = extract_pandas_data(x)
  if _is_iterable(x):
    return _batch_data(x, batch_size)
  if len(x.shape) == 1:
    x = np.reshape(x, (-1, 1))
  if batch_size is not None:
    if batch_size <= 0:
      raise ValueError('Invalid batch_size %d.' % batch_size)
    n_batches = int(math.ceil(float(len(x)) / batch_size))
    return [x[i * batch_size:(i + 1) * batch_size] for i in xrange(n_batches)]
  return [x]

def setup_processor_data_feeder(x):
  """Sets up processor iterable.

  Args:
    x: numpy, pandas or iterable.

  Returns:
    Iterable of data to process.
  """
  if HAS_PANDAS:
    x = extract_pandas_matrix(x)
  return x

def check_array(array, dtype):
  """Checks array on dtype and converts it if different.

  Args:
    array: Input array.
    dtype: Expected dtype.

  Returns:
    Original array or converted.
  """
  # skip check if array is instance of other classes, e.g. h5py.Dataset
  # to avoid copying array and loading whole data into memory
  if isinstance(array, (np.ndarray, list)):
    array = np.array(array, dtype=dtype, order=None, copy=False)
  return array

def _access(data, iloc):
  """Accesses an element from collection, using integer location based indexing.

  Args:
    data: array-like. The collection to access
    iloc: `int` or `list` of `int`s. Location(s) to access in `collection`

  Returns:
    The element of `a` found at location(s) `iloc`.
  """
  if HAS_PANDAS:
    import pandas as pd  # pylint: disable=g-import-not-at-top
    if isinstance(data, pd.Series) or isinstance(data, pd.DataFrame):
      return data.iloc[iloc]
  return data[iloc]

def _check_dtype(dtype):
  if dtypes.as_dtype(dtype) == dtypes.float64:
    logging.warn(
        'float64 is not supported by many models, consider casting to float32.')
  return dtype

class DataFeeder(object):
  """Data feeder is an example class to sample data for TF trainer."""

  def __init__(self,
               x,
               y,
               n_classes,
               batch_size=None,
               shuffle=True,
               random_state=None,
               epochs=None):
    """Initializes a DataFeeder instance.

    Args:
      x: One feature sample which can either Nd numpy matrix of shape
        `[n_samples, n_features, ...]` or dictionary of Nd numpy matrix.
      y: label vector, either floats for regression or class id for
        classification. If matrix, will consider as a sequence of labels.
        Can be `None` for unsupervised setting. Also supports dictionary of
        labels.
      n_classes: Number of classes, 0 and 1 are considered regression, `None`
        will pass through the input labels without one-hot conversion. Also, if
        `y` is `dict`, then `n_classes` must be `dict` such that
        `n_classes[key] = n_classes for label y[key]`, `None` otherwise.
      batch_size: Mini-batch size to accumulate samples in one mini batch.
      shuffle: Whether to shuffle `x`.
      random_state: Numpy `RandomState` object to reproduce sampling.
      epochs: Number of times to iterate over input data before raising
        `StopIteration` exception.

    Attributes:
      x: Input features (ndarray or dictionary of ndarrays).
      y: Input label (ndarray or dictionary of ndarrays).
      n_classes: Number of classes (if `None`, pass through indices without
        one-hot conversion).
      batch_size: Mini-batch size to accumulate.
      input_shape: Shape of the input (or dictionary of shapes).
      output_shape: Shape of the output (or dictionary of shapes).
      input_dtype: DType of input (or dictionary of shapes).
      output_dtype: DType of output (or dictionary of shapes.
    """
    x_is_dict, y_is_dict = isinstance(x, dict), y is not None and isinstance(
        y, dict)
    if isinstance(y, list):
      y = np.array(y)

    self._x = dict([(k, check_array(v, v.dtype)) for k, v in list(x.items())
                   ]) if x_is_dict else check_array(x, x.dtype)
    self._y = None if y is None else \
      dict([(k, check_array(v, v.dtype)) for k, v in list(y.items())]) if x_is_dict else check_array(y, y.dtype)

    # self.n_classes is not None means we're converting raw target indices to one-hot.
    if n_classes is not None:
      if not y_is_dict:
        y_dtype = (np.int64
                   if n_classes is not None and n_classes > 1 else np.float32)
        self._y = (None if y is None else check_array(y, dtype=y_dtype))

    self.n_classes = n_classes
    self.max_epochs = epochs

    x_shape = dict([(k, v.shape) for k, v in list(self._x.items())
                   ]) if x_is_dict else self._x.shape
    y_shape = dict([(k, v.shape) for k, v in list(self._y.items())
                   ]) if y_is_dict else None if y is None else self._y.shape

    self.input_shape, self.output_shape, self._batch_size = _get_in_out_shape(
        x_shape, y_shape, n_classes, batch_size)

    # Input dtype matches dtype of x.
    self._input_dtype = dict([(k, _check_dtype(v.dtype)) for k, v in list(self._x.items())]) if x_is_dict \
      else _check_dtype(self._x.dtype)

    # note: self._output_dtype = np.float32 when y is None
    self._output_dtype = dict([(k, _check_dtype(v.dtype)) for k, v in list(self._y.items())]) if y_is_dict \
      else _check_dtype(self._y.dtype) if y is not None else np.float32

    # self.n_classes is None means we're passing in raw target indices
    if n_classes is not None and y_is_dict:
      for key in list(n_classes.keys()):
        if key in self._output_dtype:
          self._output_dtype[key] = np.float32

    self._shuffle = shuffle
    self.random_state = np.random.RandomState(
        42) if random_state is None else random_state

    num_samples = list(self._x.values())[0].shape[
        0] if x_is_dict else self._x.shape[0]
    if self._shuffle:
      self.indices = self.random_state.permutation(num_samples)
    else:
      self.indices = np.array(range(num_samples))
    self.offset = 0
    self.epoch = 0
    self._epoch_placeholder = None

  @property
  def x(self):
    return self._x

  @property
  def y(self):
    return self._y

  @property
  def shuffle(self):
    return self._shuffle

  @property
  def input_dtype(self):
    return self._input_dtype

  @property
  def output_dtype(self):
    return self._output_dtype

  @property
  def batch_size(self):
    return self._batch_size

  def make_epoch_variable(self):
    """Adds a placeholder variable for the epoch to the graph.

    Returns:
      The epoch placeholder.
    """
    self._epoch_placeholder = array_ops.placeholder(
        dtypes.int32, [1], name='epoch')
    return self._epoch_placeholder

  def input_builder(self):
    """Builds inputs in the graph.

    Returns:
      Two placeholders for inputs and outputs.
    """

    def get_placeholder(shape, dtype, name_prepend):
      if shape is None:
        return None
      if isinstance(shape, dict):
        placeholder = {}
        for key in list(shape.keys()):
          placeholder[key] = array_ops.placeholder(
              dtypes.as_dtype(dtype[key]), [None] + shape[key][1:],
              name=name_prepend + '_' + key)
      else:
        placeholder = array_ops.placeholder(
            dtypes.as_dtype(dtype), [None] + shape[1:], name=name_prepend)
      return placeholder

    self._input_placeholder = get_placeholder(self.input_shape,
                                              self._input_dtype, 'input')
    self._output_placeholder = get_placeholder(self.output_shape,
                                               self._output_dtype, 'output')
    return self._input_placeholder, self._output_placeholder

  def set_placeholders(self, input_placeholder, output_placeholder):
    """Sets placeholders for this data feeder.

    Args:
      input_placeholder: Placeholder for `x` variable. Should match shape
        of the examples in the x dataset.
      output_placeholder: Placeholder for `y` variable. Should match
        shape of the examples in the y dataset. Can be `None`.
    """
    self._input_placeholder = input_placeholder
    self._output_placeholder = output_placeholder

  def get_feed_params(self):
    """Function returns a `dict` with data feed params while training.

    Returns:
      A `dict` with data feed params while training.
    """
    return {
        'epoch': self.epoch,
        'offset': self.offset,
        'batch_size': self._batch_size
    }

  def get_feed_dict_fn(self):
    """Returns a function that samples data into given placeholders.

    Returns:
      A function that when called samples a random subset of batch size
      from `x` and `y`.
    """
    x_is_dict, y_is_dict = isinstance(
        self._x, dict), self._y is not None and isinstance(self._y, dict)

    # Assign input features from random indices.
    def extract(data, indices):
      return (np.array(_access(data, indices)).reshape((indices.shape[0], 1)) if
              len(data.shape) == 1 else _access(data, indices))

    # assign labels from random indices
    def assign_label(data, shape, dtype, n_classes, indices):
      shape[0] = indices.shape[0]
      out = np.zeros(shape, dtype=dtype)
      for i in xrange(out.shape[0]):
        sample = indices[i]
        # self.n_classes is None means we're passing in raw target indices
        if n_classes is None:
          out[i] = _access(data, sample)
        else:
          if n_classes > 1:
            if len(shape) == 2:
              out.itemset((i, int(_access(data, sample))), 1.0)
            else:
              for idx, value in enumerate(_access(data, sample)):
                out.itemset(tuple([i, idx, value]), 1.0)
          else:
            out[i] = _access(data, sample)
      return out

    def _feed_dict_fn():
      """Function that samples data into given placeholders."""
      if self.max_epochs is not None and self.epoch + 1 > self.max_epochs:
        raise StopIteration
      assert self._input_placeholder is not None
      feed_dict = {}
      if self._epoch_placeholder is not None:
        feed_dict[self._epoch_placeholder.name] = [self.epoch]

      # Take next batch of indices.
      x_len = list(self._x.values())[0].shape[
          0] if x_is_dict else self._x.shape[0]
      end = min(x_len, self.offset + self._batch_size)
      batch_indices = self.indices[self.offset:end]

      # adding input placeholder
      feed_dict.update(
          dict([(self._input_placeholder[k].name, extract(v, batch_indices))
                for k, v in list(self._x.items())]) if x_is_dict else
          {self._input_placeholder.name: extract(self._x, batch_indices)})

      # move offset and reset it if necessary
      self.offset += self._batch_size
      if self.offset >= x_len:
        self.indices = self.random_state.permutation(
            x_len) if self._shuffle else np.array(range(x_len))
        self.offset = 0
        self.epoch += 1

      # return early if there are no labels
      if self._output_placeholder is None:
        return feed_dict

      # adding output placeholders
      if y_is_dict:
        for k, v in list(self._y.items()):
          n_classes = (self.n_classes[k] if k in self.n_classes else
                       None) if self.n_classes is not None else None
          shape, dtype = self.output_shape[k], self._output_dtype[k]
          feed_dict.update({
              self._output_placeholder[k].name:
                  assign_label(v, shape, dtype, n_classes, batch_indices)
          })
      else:
        shape, dtype, n_classes = self.output_shape, self._output_dtype, self.n_classes
        feed_dict.update({
            self._output_placeholder.name:
                assign_label(self._y, shape, dtype, n_classes, batch_indices)
        })

      return feed_dict

    return _feed_dict_fn

class StreamingDataFeeder(DataFeeder):
  """Data feeder for TF trainer that reads data from iterator.

  Streaming data feeder allows to read data as it comes it from disk or
  somewhere else. It's custom to have this iterators rotate infinetly over
  the dataset, to allow control of how much to learn on the trainer side.
  """

  def __init__(self, x, y, n_classes, batch_size):
    """Initializes a StreamingDataFeeder instance.

    Args:
      x: iterator each element of which returns one feature sample. Sample can
        be a Nd numpy matrix or dictionary of Nd numpy matrices.
      y: iterator each element of which returns one label sample. Sample can be
        a Nd numpy matrix or dictionary of Nd numpy matrices with 1 or many
        classes regression values.
      n_classes: indicator of how many classes the corresponding label sample
        has for the purposes of one-hot conversion of label. In case where `y`
        is a dictionary, `n_classes` must be dictionary (with same keys as `y`)
        of how many classes there are in each label in `y`. If key is
        present in `y` and missing in `n_classes`, the value is assumed `None`
        and no one-hot conversion will be applied to the label with that key.
      batch_size: Mini batch size to accumulate samples in one batch. If set
        `None`, then assumes that iterator to return already batched element.

    Attributes:
      x: input features (or dictionary of input features).
      y: input label (or dictionary of output features).
      n_classes: number of classes.
      batch_size: mini batch size to accumulate.
      input_shape: shape of the input (can be dictionary depending on `x`).
      output_shape: shape of the output (can be dictionary depending on `y`).
      input_dtype: dtype of input (can be dictionary depending on `x`).
      output_dtype: dtype of output (can be dictionary depending on `y`).
    """
    # pylint: disable=invalid-name,super-init-not-called
    x_first_el = six.next(x)
    self._x = itertools.chain([x_first_el], x)
    if y is not None:
      y_first_el = six.next(y)
      self._y = itertools.chain([y_first_el], y)
    else:
      y_first_el = None
      self._y = None
    self.n_classes = n_classes

    x_is_dict = isinstance(x_first_el, dict)
    y_is_dict = y is not None and isinstance(y_first_el, dict)
    if y_is_dict and n_classes is not None:
      assert isinstance(n_classes, dict)

    # extract shapes for first_elements
    if x_is_dict:
      x_first_el_shape = dict(
          [(k, [1] + list(v.shape)) for k, v in list(x_first_el.items())])
    else:
      x_first_el_shape = [1] + list(x_first_el.shape)

    if y_is_dict:
      y_first_el_shape = dict(
          [(k, [1] + list(v.shape)) for k, v in list(y_first_el.items())])
    elif y is None:
      y_first_el_shape = None
    else:
      y_first_el_shape = ([1] + list(y_first_el[0].shape if isinstance(
          y_first_el, list) else y_first_el.shape))

    self.input_shape, self.output_shape, self._batch_size = _get_in_out_shape(
        x_first_el_shape, y_first_el_shape, n_classes, batch_size)

    # Input dtype of x_first_el.
    if x_is_dict:
      self._input_dtype = dict(
          [(k, _check_dtype(v.dtype)) for k, v in list(x_first_el.items())])
    else:
      self._input_dtype = _check_dtype(x_first_el.dtype)

    # Output dtype of y_first_el.
    def check_y_dtype(el):
      if isinstance(el, np.ndarray):
        return el.dtype
      elif isinstance(el, list):
        return check_y_dtype(el[0])
      else:
        return _check_dtype(np.dtype(type(el)))

    # Output types are floats, due to both softmaxes and regression req.
    if n_classes is not None and (y is None or not y_is_dict) and n_classes > 0:
      self._output_dtype = np.float32
    elif y_is_dict:
      self._output_dtype = dict(
          [(k, check_y_dtype(v)) for k, v in list(y_first_el.items())])
    elif y is None:
      self._output_dtype = None
    else:
      self._output_dtype = check_y_dtype(y_first_el)

  def get_feed_params(self):
    """Function returns a `dict` with data feed params while training.

    Returns:
      A `dict` with data feed params while training.
    """
    return {'batch_size': self._batch_size}

  def get_feed_dict_fn(self):
    """Returns a function, that will sample data and provide it to placeholders.

    Returns:
      A function that when called samples a random subset of batch size
      from x and y.
    """
    self.stopped = False

    def _feed_dict_fn():
      """Samples data and provides it to placeholders.

      Returns:
        `dict` of input and output tensors.
      """

      def init_array(shape, dtype):
        """Initialize array of given shape or dict of shapes and dtype."""
        if shape is None:
          return None
        elif isinstance(shape, dict):
          return dict([(k, np.zeros(shape[k], dtype[k]))
                       for k in list(shape.keys())])
        else:
          return np.zeros(shape, dtype=dtype)

      def put_data_array(dest, index, source=None, n_classes=None):
        """Puts data array into container."""
        if source is None:
          dest = dest[:index]
        elif n_classes is not None and n_classes > 1:
          if len(self.output_shape) == 2:
            dest.itemset((index, source), 1.0)
          else:
            for idx, value in enumerate(source):
              dest.itemset(tuple([index, idx, value]), 1.0)
        else:
          if len(dest.shape) > 1:
            dest[index, :] = source
          else:
            dest[index] = source[0] if isinstance(source, list) else source
        return dest

      def put_data_array_or_dict(holder, index, data=None, n_classes=None):
        """Puts data array or data dictionary into container."""
        if holder is None:
          return None
        if isinstance(holder, dict):
          if data is None:
            data = {k: None for k in holder.keys()}
          assert isinstance(data, dict)
          for k in holder.keys():
            num_classes = n_classes[k] if (n_classes is not None and
                                           k in n_classes) else None
            holder[k] = put_data_array(holder[k], index, data[k], num_classes)
        else:
          holder = put_data_array(holder, index, data, n_classes)
        return holder

      if self.stopped:
        raise StopIteration

      inp = init_array(self.input_shape, self._input_dtype)
      out = init_array(self.output_shape, self._output_dtype)

      for i in xrange(self._batch_size):
        # Add handling when queue ends.
        try:
          next_inp = six.next(self._x)
          inp = put_data_array_or_dict(inp, i, next_inp, None)
        except StopIteration:
          self.stopped = True
          if i == 0:
            raise
          inp = put_data_array_or_dict(inp, i, None, None)
          out = put_data_array_or_dict(out, i, None, None)
          break

        if self._y is not None:
          next_out = six.next(self._y)
          out = put_data_array_or_dict(out, i, next_out, self.n_classes)

      # creating feed_dict
      if isinstance(inp, dict):
        feed_dict = dict([(self._input_placeholder[k].name, inp[k])
                          for k in list(self._input_placeholder.keys())])
      else:
        feed_dict = {self._input_placeholder.name: inp}
      if self._y is not None:
        if isinstance(out, dict):
          feed_dict.update(
              dict([(self._output_placeholder[k].name, out[k])
                    for k in list(self._output_placeholder.keys())]))
        else:
          feed_dict.update({self._output_placeholder.name: out})

      return feed_dict

    return _feed_dict_fn

class DaskDataFeeder(object):
  """Data feeder for that reads data from dask.Series and dask.DataFrame.

  Numpy arrays can be serialized to disk and it's possible to do random seeks
  into them. DaskDataFeeder will remove requirement to have full dataset in the
  memory and still do random seeks for sampling of batches.
  """

  def __init__(self,
               x,
               y,
               n_classes,
               batch_size,
               shuffle=True,
               random_state=None,
               epochs=None):
    """Initializes a DaskDataFeeder instance.

    Args:
      x: iterator that returns for each element, returns features.
      y: iterator that returns for each element, returns 1 or many classes /
        regression values.
      n_classes: indicator of how many classes the label has.
      batch_size: Mini batch size to accumulate.
      shuffle: Whether to shuffle the inputs.
      random_state: random state for RNG. Note that it will mutate so use a
        int value for this if you want consistent sized batches.
      epochs: Number of epochs to run.

    Attributes:
      x: input features.
      y: input label.
      n_classes: number of classes.
      batch_size: mini batch size to accumulate.
      input_shape: shape of the input.
      output_shape: shape of the output.
      input_dtype: dtype of input.
      output_dtype: dtype of output.

    Raises:
      ValueError: if `x` or `y` are `dict`, as they are not supported currently.
    """

    if isinstance(x, dict) or isinstance(y, dict):
      raise ValueError(
          'DaskDataFeeder does not support dictionaries at the moment.')

    # pylint: disable=invalid-name,super-init-not-called
    import dask.dataframe as dd  # pylint: disable=g-import-not-at-top
    # TODO(terrytangyuan): check x and y dtypes in dask_io like pandas
    self._x = x
    self._y = y
    # save column names
    self._x_columns = list(x.columns)
    if isinstance(y.columns[0], str):
      self._y_columns = list(y.columns)
    else:
      # deal with cases where two DFs have overlapped default numeric colnames
      self._y_columns = len(self._x_columns) + 1
      self._y = self._y.rename(columns={y.columns[0]: self._y_columns})

    # TODO(terrytangyuan): deal with unsupervised cases
    # combine into a data frame
    self.df = dd.multi.concat([self._x, self._y], axis=1)
    self.n_classes = n_classes

    x_count = x.count().compute()[0]
    x_shape = (x_count, len(self._x.columns))
    y_shape = (x_count, len(self._y.columns))
    # TODO(terrytangyuan): Add support for shuffle and epochs.
    self._shuffle = shuffle
    self.epochs = epochs
    self.input_shape, self.output_shape, self._batch_size = _get_in_out_shape(
        x_shape, y_shape, n_classes, batch_size)
    self.sample_fraction = self._batch_size / float(x_count)
    self._input_dtype = _check_dtype(self._x.dtypes[0])
    self._output_dtype = _check_dtype(self._y.dtypes[self._y_columns])
    if random_state is None:
      self.random_state = 66
    else:
      self.random_state = random_state

  def get_feed_params(self):
    """Function returns a `dict` with data feed params while training.

    Returns:
      A `dict` with data feed params while training.
    """
    return {'batch_size': self._batch_size}

  def get_feed_dict_fn(self, input_placeholder, output_placeholder):
    """Returns a function, that will sample data and provide it to placeholders.

    Args:
      input_placeholder: tf.Placeholder for input features mini batch.
      output_placeholder: tf.Placeholder for output labels.

    Returns:
      A function that when called samples a random subset of batch size
      from x and y.
    """

    def _feed_dict_fn():
      """Samples data and provides it to placeholders."""
      # TODO(ipolosukhin): option for with/without replacement (dev version of
      # dask)
      sample = self.df.random_split(
          [self.sample_fraction, 1 - self.sample_fraction],
          random_state=self.random_state)
      inp = extract_pandas_matrix(sample[0][self._x_columns].compute()).tolist()
      out = extract_pandas_matrix(sample[0][self._y_columns].compute())
      # convert to correct dtype
      inp = np.array(inp, dtype=self._input_dtype)
      # one-hot encode out for each class for cross entropy loss
      if HAS_PANDAS:
        import pandas as pd  # pylint: disable=g-import-not-at-top
        if not isinstance(out, pd.Series):
          out = out.flatten()
      out_max = self._y.max().compute().values[0]
      encoded_out = np.zeros((out.size, out_max + 1), dtype=self._output_dtype)
      encoded_out[np.arange(out.size), out] = 1
      return {input_placeholder.name: inp, output_placeholder.name: encoded_out}

    return _feed_dict_fn


"""
Read a directory of expression counts in ht-seq format. Each sample
should be an individual file in the directory. File names and
sample order are specified in the config file (order is determined
by order IN the config.)

This class is intended to return the raw dataframe of samples with
missing sample columns as NaN.
"""

import pandas as pd
from pathos.multiprocessing import ProcessPool
import pathlib

try:
    from functools import reduce # for py3 compatibility
except ImportError:
    pass

class HtSeqParser(object):

    def __init__(self, nodes=2):

        self.nodes = nodes

    def parse_data(self, data_paths, sample_names):
        """
        Read the input files from the config file and load in to a
        pandas dataframe.

        params
            data_paths: list of file paths specified in the config. Returned
            from config parse sample_names: list of sample names specified in
            the config returned from config parse

        """
        output = self.load_data(data_paths, sample_names)
        data, ercc_df = (self.merge_dfs(output)
                         .pipe(self.df_cleanup)
                         .pipe(self.split_on_ercc))

        return data, ercc_df

    def load_data(self, data_paths, sample_names):
        " Multiprocess load of files in to a list of dfs "
        pool = ProcessPool(nodes=self.nodes)
        dfs = pool.map(self.load_func, zip(data_paths, sample_names))
        return dfs

    @staticmethod
    def load_func(data_tuple):
        path, sample_name = data_tuple
        return pd.read_csv(path, sep='\t', names=['gene', sample_name])

    def merge_dfs(self, dfs):
        return reduce(lambda x, y: pd.merge(x, y, on='gene', how='outer'), dfs)

    def df_cleanup(self, df_old):
        " Clean away unwanted columns, reset index, and fillna "
        df = df_old.copy()
        df = df[df['gene'].str.startswith('__') == False]
        df.set_index('gene', inplace=True)
        df.fillna(value='Nan', inplace=True)
        return df

    def split_on_ercc(self, df):
        " Extract the ERCC data "
        ercc_cols = df.index.str.startswith('ERCC-')

        ercc_df = df[ercc_cols]
        data = df[~ercc_cols]

        return data, ercc_df


# -*- coding: utf-8 -*-
import numpy as np
import xgboost as xgb
import testing as tm
import pytest

try:
    import pandas as pd
except ImportError:
    pass

pytestmark = pytest.mark.skipif(**tm.no_pandas())

dpath = 'demo/data/'
rng = np.random.RandomState(1994)

class TestPandas:
    def test_pandas(self):
        df = pd.DataFrame([[1, 2., True], [2, 3., False]],
                          columns=['a', 'b', 'c'])
        dm = xgb.DMatrix(df, label=pd.Series([1, 2]))
        assert dm.feature_names == ['a', 'b', 'c']
        assert dm.feature_types == ['int', 'float', 'i']
        assert dm.num_row() == 2
        assert dm.num_col() == 3
        np.testing.assert_array_equal(dm.get_label(), np.array([1, 2]))

        # overwrite feature_names and feature_types
        dm = xgb.DMatrix(df, label=pd.Series([1, 2]),
                         feature_names=['x', 'y', 'z'],
                         feature_types=['q', 'q', 'q'])
        assert dm.feature_names == ['x', 'y', 'z']
        assert dm.feature_types == ['q', 'q', 'q']
        assert dm.num_row() == 2
        assert dm.num_col() == 3

        # incorrect dtypes
        df = pd.DataFrame([[1, 2., 'x'], [2, 3., 'y']],
                          columns=['a', 'b', 'c'])
        with pytest.raises(ValueError):
            xgb.DMatrix(df)

        # numeric columns
        df = pd.DataFrame([[1, 2., True], [2, 3., False]])
        dm = xgb.DMatrix(df, label=pd.Series([1, 2]))
        assert dm.feature_names == ['0', '1', '2']
        assert dm.feature_types == ['int', 'float', 'i']
        assert dm.num_row() == 2
        assert dm.num_col() == 3
        np.testing.assert_array_equal(dm.get_label(), np.array([1, 2]))

        df = pd.DataFrame([[1, 2., 1], [2, 3., 1]], columns=[4, 5, 6])
        dm = xgb.DMatrix(df, label=pd.Series([1, 2]))
        assert dm.feature_names == ['4', '5', '6']
        assert dm.feature_types == ['int', 'float', 'int']
        assert dm.num_row() == 2
        assert dm.num_col() == 3

        df = pd.DataFrame({'A': ['X', 'Y', 'Z'], 'B': [1, 2, 3]})
        dummies = pd.get_dummies(df)
        #    B  A_X  A_Y  A_Z
        # 0  1    1    0    0
        # 1  2    0    1    0
        # 2  3    0    0    1
        result, _, _ = xgb.data._transform_pandas_df(dummies,
                                                     enable_categorical=False)
        exp = np.array([[1., 1., 0., 0.],
                        [2., 0., 1., 0.],
                        [3., 0., 0., 1.]])
        np.testing.assert_array_equal(result, exp)
        dm = xgb.DMatrix(dummies)
        assert dm.feature_names == ['B', 'A_X', 'A_Y', 'A_Z']
        assert dm.feature_types == ['int', 'int', 'int', 'int']
        assert dm.num_row() == 3
        assert dm.num_col() == 4

        df = pd.DataFrame({'A=1': [1, 2, 3], 'A=2': [4, 5, 6]})
        dm = xgb.DMatrix(df)
        assert dm.feature_names == ['A=1', 'A=2']
        assert dm.feature_types == ['int', 'int']
        assert dm.num_row() == 3
        assert dm.num_col() == 2

        df_int = pd.DataFrame([[1, 1.1], [2, 2.2]], columns=[9, 10])
        dm_int = xgb.DMatrix(df_int)
        df_range = pd.DataFrame([[1, 1.1], [2, 2.2]], columns=range(9, 11, 1))
        dm_range = xgb.DMatrix(df_range)
        assert dm_int.feature_names == ['9', '10']  # assert not "9 "
        assert dm_int.feature_names == dm_range.feature_names

        # test MultiIndex as columns
        df = pd.DataFrame(
            [
                (1, 2, 3, 4, 5, 6),
                (6, 5, 4, 3, 2, 1)
            ],
            columns=pd.MultiIndex.from_tuples((
                ('a', 1), ('a', 2), ('a', 3),
                ('b', 1), ('b', 2), ('b', 3),
            ))
        )
        dm = xgb.DMatrix(df)
        assert dm.feature_names == ['a 1', 'a 2', 'a 3', 'b 1', 'b 2', 'b 3']
        assert dm.feature_types == ['int', 'int', 'int', 'int', 'int', 'int']
        assert dm.num_row() == 2
        assert dm.num_col() == 6

    def test_slice(self):
        rng = np.random.RandomState(1994)
        rows = 100
        X = rng.randint(3, 7, size=rows)
        X = pd.DataFrame({'f0': X})
        y = rng.randn(rows)
        ridxs = [1, 2, 3, 4, 5, 6]
        m = xgb.DMatrix(X, y)
        sliced = m.slice(ridxs)

        assert m.feature_types == sliced.feature_types

    def test_pandas_categorical(self):
        rng = np.random.RandomState(1994)
        rows = 100
        X = rng.randint(3, 7, size=rows)
        X = pd.Series(X, dtype="category")
        X = pd.DataFrame({'f0': X})
        y = rng.randn(rows)
        m = xgb.DMatrix(X, y, enable_categorical=True)
        assert m.feature_types[0] == 'categorical'

    def test_pandas_sparse(self):
        import pandas as pd
        rows = 100
        X = pd.DataFrame(
            {"A": pd.arrays.SparseArray(np.random.randint(0, 10, size=rows)),
             "B": pd.arrays.SparseArray(np.random.randn(rows)),
             "C": pd.arrays.SparseArray(np.random.permutation(
                 [True, False] * (rows // 2)))}
        )
        y = pd.Series(pd.arrays.SparseArray(np.random.randn(rows)))
        dtrain = xgb.DMatrix(X, y)
        booster = xgb.train({}, dtrain, num_boost_round=4)
        predt_sparse = booster.predict(xgb.DMatrix(X))
        predt_dense = booster.predict(xgb.DMatrix(X.sparse.to_dense()))
        np.testing.assert_allclose(predt_sparse, predt_dense)

    def test_pandas_label(self):
        # label must be a single column
        df = pd.DataFrame({'A': ['X', 'Y', 'Z'], 'B': [1, 2, 3]})
        with pytest.raises(ValueError):
            xgb.data._transform_pandas_df(df, False, None, None, 'label', 'float')

        # label must be supported dtype
        df = pd.DataFrame({'A': np.array(['a', 'b', 'c'], dtype=object)})
        with pytest.raises(ValueError):
            xgb.data._transform_pandas_df(df, False, None, None, 'label', 'float')

        df = pd.DataFrame({'A': np.array([1, 2, 3], dtype=int)})
        result, _, _ = xgb.data._transform_pandas_df(df, False, None, None,
                                                     'label', 'float')
        np.testing.assert_array_equal(result, np.array([[1.], [2.], [3.]],
                                                       dtype=float))
        dm = xgb.DMatrix(np.random.randn(3, 2), label=df)
        assert dm.num_row() == 3
        assert dm.num_col() == 2

    def test_pandas_weight(self):
        kRows = 32
        kCols = 8

        X = np.random.randn(kRows, kCols)
        y = np.random.randn(kRows)
        w = np.random.uniform(size=kRows).astype(np.float32)
        w_pd = pd.DataFrame(w)
        data = xgb.DMatrix(X, y, w_pd)

        assert data.num_row() == kRows
        assert data.num_col() == kCols

        np.testing.assert_array_equal(data.get_weight(), w)

    def test_cv_as_pandas(self):
        dm = xgb.DMatrix(dpath + 'agaricus.txt.train')
        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic', 'eval_metric': 'error'}

        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10)
        assert isinstance(cv, pd.DataFrame)
        exp = pd.Index([u'test-error-mean', u'test-error-std',
                        u'train-error-mean', u'train-error-std'])
        assert len(cv.columns.intersection(exp)) == 4

        # show progress log (result is the same as above)
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    verbose_eval=True)
        assert isinstance(cv, pd.DataFrame)
        exp = pd.Index([u'test-error-mean', u'test-error-std',
                        u'train-error-mean', u'train-error-std'])
        assert len(cv.columns.intersection(exp)) == 4
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    verbose_eval=True, show_stdv=False)
        assert isinstance(cv, pd.DataFrame)
        exp = pd.Index([u'test-error-mean', u'test-error-std',
                        u'train-error-mean', u'train-error-std'])
        assert len(cv.columns.intersection(exp)) == 4

        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic', 'eval_metric': 'auc'}
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, as_pandas=True)
        assert 'eval_metric' in params
        assert 'auc' in cv.columns[0]

        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic', 'eval_metric': ['auc']}
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10, as_pandas=True)
        assert 'eval_metric' in params
        assert 'auc' in cv.columns[0]

        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic', 'eval_metric': ['auc']}
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    as_pandas=True, early_stopping_rounds=1)
        assert 'eval_metric' in params
        assert 'auc' in cv.columns[0]
        assert cv.shape[0] < 10

        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic'}
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    as_pandas=True, metrics='auc')
        assert 'auc' in cv.columns[0]

        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic'}
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    as_pandas=True, metrics=['auc'])
        assert 'auc' in cv.columns[0]

        params = {'max_depth': 2, 'eta': 1, 'verbosity': 0,
                  'objective': 'binary:logistic', 'eval_metric': ['auc']}
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    as_pandas=True, metrics='error')
        assert 'eval_metric' in params
        assert 'auc' not in cv.columns[0]
        assert 'error' in cv.columns[0]

        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    as_pandas=True, metrics=['error'])
        assert 'eval_metric' in params
        assert 'auc' not in cv.columns[0]
        assert 'error' in cv.columns[0]

        params = list(params.items())
        cv = xgb.cv(params, dm, num_boost_round=10, nfold=10,
                    as_pandas=True, metrics=['error'])
        assert isinstance(params, list)
        assert 'auc' not in cv.columns[0]
        assert 'error' in cv.columns[0]


from mnist import MNIST
import  numpy as np
import tensorflow as tf
from src.my.lib.utils import montage
import matplotlib.pyplot as plt
from PIL import Image

src = '../../../../data/mnist/'
output='./content/1/%s.jpg'

mndata = MNIST(src)
data = np.array(mndata.load_testing())
X = data[0]
Y = data[1]

items = 100
imgs = np.array([i for i in np.array(X[:items])]).reshape(items,28,28)
n_features = 784

n_input = n_features

Y = imgs.reshape(items,n_features).astype(float)
current_input = imgs.reshape(items,n_features).astype(float)

Ws = []
Bs = []
dimensions = [512,256,128,64]

for layer_i,n_ouputs in enumerate(dimensions):
    with tf.variable_scope("encoder/variable/%s" % layer_i):
        W = tf.get_variable(name="weight%s" % layer_i, dtype=tf.float64,
                            initializer=tf.contrib.layers.xavier_initializer(),
                            shape=[n_input, n_ouputs])
        #B = tf.get_variable(name='bias%s' % layer_i, dtype=tf.float64,
        #                    initializer=tf.random_normal_initializer(mean=0.0, stddev=1.1),
        #                    shape=[n_ouputs])
        #h = tf.nn.bias_add(value=tf.matmul(current_input, W),
        #                   bias=B)

        h = tf.matmul(current_input, W)
        current_input = h
        current_input = tf.nn.relu(current_input)

        n_input = n_ouputs
        Ws.append(W)
        #Bs.append()

Ws = Ws[::-1]#reverse
Bs = Bs[::-1]#reverse

#dimensions = dimensions[::1][1:].append(n_features)
dimensions = dimensions[::-1][1:] +[n_features]

#Build DECODER

for layer_i,n_ouputs in enumerate(dimensions):
    with tf.variable_scope("encoder/variable/%s" % layer_i):

                                                            ##128x64 -> 64x128
        h = value=tf.matmul(current_input,tf.transpose(Ws[layer_i]))
        if layer_i + 1 < len(Bs):
            h = tf.nn.bias_add(h,bias=Bs[layer_i + 1])
        current_input = h
        current_input = tf.nn.relu(current_input)

        n_input = n_ouputs

loss_func = tf.reduce_mean(tf.squared_difference(current_input, Y), 1)
optimizer = tf.train.AdamOptimizer(learning_rate=0.00001)
train = optimizer.minimize(loss_func)

counter = 0
with tf.Session() as sess:

    sess.run(tf.initialize_all_variables())

    for i in range(50000):
        sess.run(train)
        if i % 15 == 0:

            Image.fromarray(montage(sess.run(current_input).reshape(items,28,28)).astype(np.uint8)) \
                .save(output % ("0"*(5 - len(str(counter))) +  str(counter)))
            print(sess.run(tf.reduce_mean(loss_func)))
            counter += 1