"""Base and mixin classes for nearest neighbors"""
# Authors: Jake Vanderplas <vanderplas@astro.washington.edu>
#          Fabian Pedregosa <fabian.pedregosa@inria.fr>
#          Alexandre Gramfort <alexandre.gramfort@inria.fr>
#          Sparseness support by Lars Buitinck <L.J.Buitinck@uva.nl>
#          Multi-output support by Arnaud Joly <a.joly@ulg.ac.be>
#
# License: BSD 3 clause (C) INRIA, University of Amsterdam
import warnings
from abc import ABCMeta, abstractmethod

import numpy as np
from scipy.sparse import csr_matrix, issparse

from .ball_tree import BallTree
from .kd_tree import KDTree
from ..base import BaseEstimator
from ..metrics import pairwise_distances
from ..metrics.pairwise import PAIRWISE_DISTANCE_FUNCTIONS
from ..utils import check_X_y, check_array
from ..utils.fixes import argpartition
from ..utils.validation import DataConversionWarning
from ..utils.validation import NotFittedError
from ..externals import six

VALID_METRICS = dict(ball_tree=BallTree.valid_metrics,
                     kd_tree=KDTree.valid_metrics,
                     # The following list comes from the
                     # sklearn.metrics.pairwise doc string
                     brute=(list(PAIRWISE_DISTANCE_FUNCTIONS.keys()) +
                            ['braycurtis', 'canberra', 'chebyshev',
                             'correlation', 'cosine', 'dice', 'hamming',
                             'jaccard', 'kulsinski', 'mahalanobis',
                             'matching', 'minkowski', 'rogerstanimoto',
                             'russellrao', 'seuclidean', 'sokalmichener',
                             'sokalsneath', 'sqeuclidean',
                             'yule', 'wminkowski']))

VALID_METRICS_SPARSE = dict(ball_tree=[],
                            kd_tree=[],
                            brute=PAIRWISE_DISTANCE_FUNCTIONS.keys())

class NeighborsWarning(UserWarning):
    pass

# Make sure that NeighborsWarning are displayed more than once
warnings.simplefilter("always", NeighborsWarning)

def _check_weights(weights):
    """Check to make sure weights are valid"""
    if weights in (None, 'uniform', 'distance'):
        return weights
    elif callable(weights):
        return weights
    else:
        raise ValueError("weights not recognized: should be 'uniform', "
                         "'distance', or a callable function")

def _get_weights(dist, weights):
    """Get the weights from an array of distances and a parameter ``weights``

    Parameters
    ===========
    dist: ndarray
        The input distances
    weights: {'uniform', 'distance' or a callable}
        The kind of weighting used

    Returns
    ========
    weights_arr: array of the same shape as ``dist``
        if ``weights == 'uniform'``, then returns None
    """
    if weights in (None, 'uniform'):
        return None
    elif weights == 'distance':
        # if user attempts to classify a point that was zero distance from one
        # or more training points, those training points are weighted as 1.0
        # and the other points as 0.0
        if dist.dtype is np.dtype(object):
            for point_dist_i, point_dist in enumerate(dist):
                # check if point_dist is iterable
                # (ex: RadiusNeighborClassifier.predict may set an element of
                # dist to 1e-6 to represent an 'outlier')
                if hasattr(point_dist, '__contains__') and 0. in point_dist:
                    dist[point_dist_i] = point_dist == 0.
                else:
                    dist[point_dist_i] = 1. / point_dist
        else:
            with np.errstate(divide='ignore'):
                dist = 1. / dist
            inf_mask = np.isinf(dist)
            inf_row = np.any(inf_mask, axis=1)
            dist[inf_row] = inf_mask[inf_row]
        return dist
    elif callable(weights):
        return weights(dist)
    else:
        raise ValueError("weights not recognized: should be 'uniform', "
                         "'distance', or a callable function")

class NeighborsBase(six.with_metaclass(ABCMeta, BaseEstimator)):
    """Base class for nearest neighbors estimators."""

    @abstractmethod
    def __init__(self):
        pass

    def _init_params(self, n_neighbors=None, radius=None,
                     algorithm='auto', leaf_size=30, metric='minkowski',
                     p=2, metric_params=None, **kwargs):
        if kwargs:
            warnings.warn("Passing additional arguments to the metric "
                          "function as **kwargs is deprecated "
                          "and will no longer be supported in 0.18. "
                          "Use metric_params instead.",
                          DeprecationWarning, stacklevel=3)
            if metric_params is None:
                metric_params = {}
            metric_params.update(kwargs)

        self.n_neighbors = n_neighbors
        self.radius = radius
        self.algorithm = algorithm
        self.leaf_size = leaf_size
        self.metric = metric
        self.metric_params = metric_params
        self.p = p

        if algorithm not in ['auto', 'brute',
                             'kd_tree', 'ball_tree']:
            raise ValueError("unrecognized algorithm: '%s'" % algorithm)

        if algorithm == 'auto':
            alg_check = 'ball_tree'
        else:
            alg_check = algorithm

        if callable(metric):
            if algorithm == 'kd_tree':
                # callable metric is only valid for brute force and ball_tree
                raise ValueError(
                    "kd_tree algorithm does not support callable metric '%s'"
                    % metric)
        elif metric not in VALID_METRICS[alg_check]:
            raise ValueError("Metric '%s' not valid for algorithm '%s'"
                             % (metric, algorithm))

        if self.metric_params is not None and 'p' in self.metric_params:
            warnings.warn("Parameter p is found in metric_params. "
                          "The corresponding parameter from __init__ "
                          "is ignored.", SyntaxWarning, stacklevel=3)
            effective_p = metric_params['p']
        else:
            effective_p = self.p

        if self.metric in ['wminkowski', 'minkowski'] and effective_p < 1:
            raise ValueError("p must be greater than one for minkowski metric")

        self._fit_X = None
        self._tree = None
        self._fit_method = None

    def _fit(self, X):
        if self.metric_params is None:
            self.effective_metric_params_ = {}
        else:
            self.effective_metric_params_ = self.metric_params.copy()

        effective_p = self.effective_metric_params_.get('p', self.p)
        if self.metric in ['wminkowski', 'minkowski']:
            self.effective_metric_params_['p'] = effective_p

        self.effective_metric_ = self.metric
        # For minkowski distance, use more efficient methods where available
        if self.metric == 'minkowski':
            p = self.effective_metric_params_.pop('p', 2)
            if p < 1:
                raise ValueError("p must be greater than one "
                                 "for minkowski metric")
            elif p == 1:
                self.effective_metric_ = 'manhattan'
            elif p == 2:
                self.effective_metric_ = 'euclidean'
            elif p == np.inf:
                self.effective_metric_ = 'chebyshev'
            else:
                self.effective_metric_params_['p'] = p

        if isinstance(X, NeighborsBase):
            self._fit_X = X._fit_X
            self._tree = X._tree
            self._fit_method = X._fit_method
            return self

        elif isinstance(X, BallTree):
            self._fit_X = X.data
            self._tree = X
            self._fit_method = 'ball_tree'
            return self

        elif isinstance(X, KDTree):
            self._fit_X = X.data
            self._tree = X
            self._fit_method = 'kd_tree'
            return self

        X = check_array(X, accept_sparse='csr')

        n_samples = X.shape[0]
        if n_samples == 0:
            raise ValueError("n_samples must be greater than 0")

        if issparse(X):
            if self.algorithm not in ('auto', 'brute'):
                warnings.warn("cannot use tree with sparse input: "
                              "using brute force")
            if self.effective_metric_ not in VALID_METRICS_SPARSE['brute']:
                raise ValueError("metric '%s' not valid for sparse input"
                                 % self.effective_metric_)
            self._fit_X = X.copy()
            self._tree = None
            self._fit_method = 'brute'
            return self

        self._fit_method = self.algorithm
        self._fit_X = X

        if self._fit_method == 'auto':
            # A tree approach is better for small number of neighbors,
            # and KDTree is generally faster when available
            if (self.n_neighbors is None
                    or self.n_neighbors < self._fit_X.shape[0] // 2):
                if self.effective_metric_ in VALID_METRICS['kd_tree']:
                    self._fit_method = 'kd_tree'
                else:
                    self._fit_method = 'ball_tree'
            else:
                self._fit_method = 'brute'

        if self._fit_method == 'ball_tree':
            self._tree = BallTree(X, self.leaf_size,
                                  metric=self.effective_metric_,
                                  **self.effective_metric_params_)
        elif self._fit_method == 'kd_tree':
            self._tree = KDTree(X, self.leaf_size,
                                metric=self.effective_metric_,
                                **self.effective_metric_params_)
        elif self._fit_method == 'brute':
            self._tree = None
        else:
            raise ValueError("algorithm = '%s' not recognized"
                             % self.algorithm)
        return self

class KNeighborsMixin(object):
    """Mixin for k-neighbors searches"""

    def kneighbors(self, X=None, n_neighbors=None, return_distance=True):
        """Finds the K-neighbors of a point.

        Returns distance

        Parameters
        ----------
        X : array-like, last dimension same as that of fit data, optional
            The query point or points.
            If not provided, neighbors of each indexed point are returned.
            In this case, the query point is not considered its own neighbor.

        n_neighbors : int
            Number of neighbors to get (default is the value
            passed to the constructor).

        return_distance : boolean, optional. Defaults to True.
            If False, distances will not be returned

        Returns
        -------
        dist : array
            Array representing the lengths to points, only present if
            return_distance=True

        ind : array
            Indices of the nearest points in the population matrix.

        Examples
        --------
        In the following example, we construct a NeighborsClassifier
        class from an array representing our data set and ask who's
        the closest point to [1,1,1]

        >>> samples = [[0., 0., 0.], [0., .5, 0.], [1., 1., .5]]
        >>> from sklearn.neighbors import NearestNeighbors
        >>> neigh = NearestNeighbors(n_neighbors=1)
        >>> neigh.fit(samples) # doctest: +ELLIPSIS
        NearestNeighbors(algorithm='auto', leaf_size=30, ...)
        >>> print(neigh.kneighbors([1., 1., 1.])) # doctest: +ELLIPSIS
        (array([[ 0.5]]), array([[2]]...))

        As you can see, it returns [[0.5]], and [[2]], which means that the
        element is at distance 0.5 and is the third element of samples
        (indexes start at 0). You can also query for multiple points:

        >>> X = [[0., 1., 0.], [1., 0., 1.]]
        >>> neigh.kneighbors(X, return_distance=False) # doctest: +ELLIPSIS
        array([[1],
               [2]]...)

        """
        if self._fit_method is None:
            raise NotFittedError("Must fit neighbors before querying.")

        if n_neighbors is None:
            n_neighbors = self.n_neighbors

        if X is not None:
            query_is_train = False
            X = check_array(X, accept_sparse='csr')
        else:
            query_is_train = True
            X = self._fit_X
            # Include an extra neighbor to account for the sample itself being
            # returned, which is removed later
            n_neighbors += 1

        train_size = self._fit_X.shape[0]
        if n_neighbors > train_size:
            raise ValueError(
                "Expected n_neighbors <= n_samples, "
                " but n_samples = %d, n_neighbors = %d" %
                (train_size, n_neighbors)
            )
        n_samples, _ = X.shape
        sample_range = np.arange(n_samples)[:, None]

        if self._fit_method == 'brute':
            # for efficiency, use squared euclidean distances
            if self.effective_metric_ == 'euclidean':
                dist = pairwise_distances(X, self._fit_X, 'euclidean',
                                          squared=True)
            else:
                dist = pairwise_distances(X, self._fit_X,
                                          self.effective_metric_,
                                          **self.effective_metric_params_)

            neigh_ind = argpartition(dist, n_neighbors - 1, axis=1)
            neigh_ind = neigh_ind[:, :n_neighbors]
            # argpartition doesn't guarantee sorted order, so we sort again
            neigh_ind = neigh_ind[
                sample_range, np.argsort(dist[sample_range, neigh_ind])]

            if return_distance:
                if self.effective_metric_ == 'euclidean':
                    result = np.sqrt(dist[sample_range, neigh_ind]), neigh_ind
                else:
                    result = dist[sample_range, neigh_ind], neigh_ind
            else:
                result = neigh_ind

        elif self._fit_method in ['ball_tree', 'kd_tree']:
            if issparse(X):
                raise ValueError(
                    "%s does not work with sparse matrices. Densify the data, "
                    "or set algorithm='brute'" % self._fit_method)
            result = self._tree.query(X, n_neighbors,
                                      return_distance=return_distance)
        else:
            raise ValueError("internal: _fit_method not recognized")

        if not query_is_train:
            return result

        else:
            # If the query data is the same as the indexed data, we would like
            # to ignore the first nearest neighbor of every sample, i.e
            # the sample itself.
            if return_distance:
                dist, neigh_ind = result
            else:
                neigh_ind = result

            sample_mask = neigh_ind != sample_range

            # Corner case: When the number of duplicates are more
            # than the number of neighbors, the first NN will not
            # be the sample, but a duplicate.
            # In that case mask the first duplicate.
            dup_gr_nbrs = np.all(sample_mask, axis=1)
            sample_mask[:, 0][dup_gr_nbrs] = False

            neigh_ind = np.reshape(
                neigh_ind[sample_mask], (n_samples, n_neighbors - 1))

            if return_distance:
                dist = np.reshape(
                    dist[sample_mask], (n_samples, n_neighbors - 1))
                return dist, neigh_ind
            return neigh_ind

    def kneighbors_graph(self, X=None, n_neighbors=None,
                         mode='connectivity'):
        """Computes the (weighted) graph of k-Neighbors for points in X

        Parameters
        ----------
        X : array-like, last dimension same as that of fit data, optional
            The query point or points.
            If not provided, neighbors of each indexed point are returned.
            In this case, the query point is not considered its own neighbor.

        n_neighbors : int
            Number of neighbors for each sample.
            (default is value passed to the constructor).

        mode : {'connectivity', 'distance'}, optional
            Type of returned matrix: 'connectivity' will return the
            connectivity matrix with ones and zeros, in 'distance' the
            edges are Euclidean distance between points.

        Returns
        -------
        A : sparse matrix in CSR format, shape = [n_samples, n_samples_fit]
            n_samples_fit is the number of samples in the fitted data
            A[i, j] is assigned the weight of edge that connects i to j.

        Examples
        --------
        >>> X = [[0], [3], [1]]
        >>> from sklearn.neighbors import NearestNeighbors
        >>> neigh = NearestNeighbors(n_neighbors=2)
        >>> neigh.fit(X) # doctest: +ELLIPSIS
        NearestNeighbors(algorithm='auto', leaf_size=30, ...)
        >>> A = neigh.kneighbors_graph(X)
        >>> A.toarray()
        array([[ 1.,  0.,  1.],
               [ 0.,  1.,  1.],
               [ 1.,  0.,  1.]])

        See also
        --------
        NearestNeighbors.radius_neighbors_graph
        """
        if n_neighbors is None:
            n_neighbors = self.n_neighbors

        # kneighbors does the None handling.
        if X is not None:
            X = check_array(X, accept_sparse='csr')
            n_samples1 = X.shape[0]
        else:
            n_samples1 = self._fit_X.shape[0]

        n_samples2 = self._fit_X.shape[0]
        n_nonzero = n_samples1 * n_neighbors
        A_indptr = np.arange(0, n_nonzero + 1, n_neighbors)

        # construct CSR matrix representation of the k-NN graph
        if mode == 'connectivity':
            A_data = np.ones(n_samples1 * n_neighbors)
            A_ind = self.kneighbors(X, n_neighbors, return_distance=False)

        elif mode == 'distance':
            A_data, A_ind = self.kneighbors(
                X, n_neighbors, return_distance=True)
            A_data = np.ravel(A_data)

        else:
            raise ValueError(
                'Unsupported mode, must be one of "connectivity" '
                'or "distance" but got "%s" instead' % mode)

        kneighbors_graph = csr_matrix((A_data, A_ind.ravel(), A_indptr),
                                      shape=(n_samples1, n_samples2))

        return kneighbors_graph

class RadiusNeighborsMixin(object):
    """Mixin for radius-based neighbors searches"""

    def radius_neighbors(self, X=None, radius=None, return_distance=True):
        """Finds the neighbors within a given radius of a point or points.

        Return the indices and distances of each point from the dataset
        lying in a ball with size ``radius`` around the points of the query
        array. Points lying on the boundary are included in the results.

        The result points are *not* necessarily sorted by distance to their
        query point.

        Parameters
        ----------
        X : array-like, (n_samples, n_features), optional
            The query point or points.
            If not provided, neighbors of each indexed point are returned.
            In this case, the query point is not considered its own neighbor.

        radius : float
            Limiting distance of neighbors to return.
            (default is the value passed to the constructor).

        return_distance : boolean, optional. Defaults to True.
            If False, distances will not be returned

        Returns
        -------
        dist : array, shape (n_samples,) of arrays
            Array representing the distances to each point, only present if
            return_distance=True. The distance values are computed according
            to the ``metric`` constructor parameter.

        ind : array, shape (n_samples,) of arrays
            An array of arrays of indices of the approximate nearest points
            from the population matrix that lie within a ball of size
            ``radius`` around the query points.

        Examples
        --------
        In the following example, we construct a NeighborsClassifier
        class from an array representing our data set and ask who's
        the closest point to [1, 1, 1]:

        >>> import numpy as np
        >>> samples = [[0., 0., 0.], [0., .5, 0.], [1., 1., .5]]
        >>> from sklearn.neighbors import NearestNeighbors
        >>> neigh = NearestNeighbors(radius=1.6)
        >>> neigh.fit(samples) # doctest: +ELLIPSIS
        NearestNeighbors(algorithm='auto', leaf_size=30, ...)
        >>> rng = neigh.radius_neighbors([1., 1., 1.])
        >>> print(np.asarray(rng[0][0])) # doctest: +ELLIPSIS
        [ 1.5  0.5]
        >>> print(np.asarray(rng[1][0])) # doctest: +ELLIPSIS
        [1 2]

        The first array returned contains the distances to all points which
        are closer than 1.6, while the second array returned contains their
        indices.  In general, multiple points can be queried at the same time.

        Notes
        -----
        Because the number of neighbors of each point is not necessarily
        equal, the results for multiple query points cannot be fit in a
        standard data array.
        For efficiency, `radius_neighbors` returns arrays of objects, where
        each object is a 1D array of indices or distances.
        """
        if self._fit_method is None:
            raise NotFittedError("Must fit neighbors before querying.")

        if X is not None:
            query_is_train = False
            X = check_array(X, accept_sparse='csr')
        else:
            query_is_train = True
            X = self._fit_X

        if radius is None:
            radius = self.radius

        n_samples = X.shape[0]
        if self._fit_method == 'brute':
            # for efficiency, use squared euclidean distances
            if self.effective_metric_ == 'euclidean':
                dist = pairwise_distances(X, self._fit_X, 'euclidean',
                                          squared=True)
                radius *= radius
            else:
                dist = pairwise_distances(X, self._fit_X,
                                          self.effective_metric_,
                                          **self.effective_metric_params_)

            neigh_ind_list = [np.where(d <= radius)[0] for d in dist]

            # See https://github.com/numpy/numpy/issues/5456
            # if you want to understand why this is initialized this way.
            neigh_ind = np.empty(n_samples, dtype='object')
            neigh_ind[:] = neigh_ind_list

            if return_distance:
                dist_array = np.empty(n_samples, dtype='object')
                if self.effective_metric_ == 'euclidean':
                    dist_list = [np.sqrt(d[neigh_ind[i]])
                                 for i, d in enumerate(dist)]
                else:
                    dist_list = [d[neigh_ind[i]]
                                 for i, d in enumerate(dist)]
                dist_array[:] = dist_list

                results = dist_array, neigh_ind
            else:
                results = neigh_ind

        elif self._fit_method in ['ball_tree', 'kd_tree']:
            if issparse(X):
                raise ValueError(
                    "%s does not work with sparse matrices. Densify the data, "
                    "or set algorithm='brute'" % self._fit_method)
            results = self._tree.query_radius(X, radius,
                                              return_distance=return_distance)
            if return_distance:
                results = results[::-1]
        else:
            raise ValueError("internal: _fit_method not recognized")

        if not query_is_train:
            return results
        else:
            # If the query data is the same as the indexed data, we would like
            # to ignore the first nearest neighbor of every sample, i.e
            # the sample itself.
            if return_distance:
                dist, neigh_ind = results
            else:
                neigh_ind = results

            for ind, ind_neighbor in enumerate(neigh_ind):
                mask = ind_neighbor != ind

                neigh_ind[ind] = ind_neighbor[mask]
                if return_distance:
                    dist[ind] = dist[ind][mask]

            if return_distance:
                return dist, neigh_ind
            return neigh_ind

    def radius_neighbors_graph(self, X=None, radius=None, mode='connectivity'):
        """Computes the (weighted) graph of Neighbors for points in X

        Neighborhoods are restricted the points at a distance lower than
        radius.

        Parameters
        ----------
        X : array-like, shape = [n_samples, n_features], optional
            The query point or points.
            If not provided, neighbors of each indexed point are returned.
            In this case, the query point is not considered its own neighbor.

        radius : float
            Radius of neighborhoods.
            (default is the value passed to the constructor).

        mode : {'connectivity', 'distance'}, optional
            Type of returned matrix: 'connectivity' will return the
            connectivity matrix with ones and zeros, in 'distance' the
            edges are Euclidean distance between points.

        Returns
        -------
        A : sparse matrix in CSR format, shape = [n_samples, n_samples]
            A[i, j] is assigned the weight of edge that connects i to j.

        Examples
        --------
        >>> X = [[0], [3], [1]]
        >>> from sklearn.neighbors import NearestNeighbors
        >>> neigh = NearestNeighbors(radius=1.5)
        >>> neigh.fit(X) # doctest: +ELLIPSIS
        NearestNeighbors(algorithm='auto', leaf_size=30, ...)
        >>> A = neigh.radius_neighbors_graph(X)
        >>> A.toarray()
        array([[ 1.,  0.,  1.],
               [ 0.,  1.,  0.],
               [ 1.,  0.,  1.]])

        See also
        --------
        kneighbors_graph
        """
        if X is not None:
            X = check_array(X, accept_sparse=['csr', 'csc', 'coo'])

        n_samples2 = self._fit_X.shape[0]
        if radius is None:
            radius = self.radius

        # construct CSR matrix representation of the NN graph
        if mode == 'connectivity':
            A_ind = self.radius_neighbors(X, radius,
                                          return_distance=False)
            A_data = None
        elif mode == 'distance':
            dist, A_ind = self.radius_neighbors(X, radius,
                                                return_distance=True)
            A_data = np.concatenate(list(dist))
        else:
            raise ValueError(
                'Unsupported mode, must be one of "connectivity", '
                'or "distance" but got %s instead' % mode)

        n_samples1 = A_ind.shape[0]
        n_neighbors = np.array([len(a) for a in A_ind])
        A_ind = np.concatenate(list(A_ind))
        if A_data is None:
            A_data = np.ones(len(A_ind))
        A_indptr = np.concatenate((np.zeros(1, dtype=int),
                                   np.cumsum(n_neighbors)))

        return csr_matrix((A_data, A_ind, A_indptr),
                          shape=(n_samples1, n_samples2))

class SupervisedFloatMixin(object):
    def fit(self, X, y):
        """Fit the model using X as training data and y as target values

        Parameters
        ----------
        X : {array-like, sparse matrix, BallTree, KDTree}
            Training data. If array or matrix, shape = [n_samples, n_features]

        y : {array-like, sparse matrix}
            Target values, array of float values, shape = [n_samples]
             or [n_samples, n_outputs]
        """
        if not isinstance(X, (KDTree, BallTree)):
            X, y = check_X_y(X, y, "csr", multi_output=True)
        self._y = y
        return self._fit(X)

class SupervisedIntegerMixin(object):
    def fit(self, X, y):
        """Fit the model using X as training data and y as target values

        Parameters
        ----------
        X : {array-like, sparse matrix, BallTree, KDTree}
            Training data. If array or matrix, shape = [n_samples, n_features]

        y : {array-like, sparse matrix}
            Target values of shape = [n_samples] or [n_samples, n_outputs]

        """
        if not isinstance(X, (KDTree, BallTree)):
            X, y = check_X_y(X, y, "csr", multi_output=True)

        if y.ndim == 1 or y.ndim == 2 and y.shape[1] == 1:
            if y.ndim != 1:
                warnings.warn("A column-vector y was passed when a 1d array "
                              "was expected. Please change the shape of y to "
                              "(n_samples, ), for example using ravel().",
                              DataConversionWarning, stacklevel=2)

            self.outputs_2d_ = False
            y = y.reshape((-1, 1))
        else:
            self.outputs_2d_ = True

        self.classes_ = []
        self._y = np.empty(y.shape, dtype=np.int)
        for k in range(self._y.shape[1]):
            classes, self._y[:, k] = np.unique(y[:, k], return_inverse=True)
            self.classes_.append(classes)

        if not self.outputs_2d_:
            self.classes_ = self.classes_[0]
            self._y = self._y.ravel()

        return self._fit(X)

class UnsupervisedMixin(object):
    def fit(self, X, y=None):
        """Fit the model using X as training data

        Parameters
        ----------
        X : {array-like, sparse matrix, BallTree, KDTree}
            Training data. If array or matrix, shape = [n_samples, n_features]
        """
        return self._fit(X)


"""
Pyicoteo is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""
import sys, os
import math
import random

from core import Cluster, Region, InvalidLine, InsufficientData, ConversionNotSupported
from defaults import *
import utils
import bam

from regions import AnnotationGene, AnnotationTranscript, AnnotationExon, RegionWriter, read_gff_file, get_exons, get_introns, gene_slide

import warnings

try:
    from shutil import move
except:
    from os import rename as move

"""
Differential expression and MA plot visualization module.
"""
def _region_from_dual(self, line):
    try:
        self.cluster_aux.clear()
        self.cluster_aux.read_line(line)
        strand = None
        if self.stranded_analysis:
            strand = self.cluster_aux.strand
        ret = Region(self.cluster_aux.name, self.cluster_aux.start, self.cluster_aux.end, name2=self.cluster_aux.name2, strand=strand)
        self.cluster_aux.clear()
        return ret

    except ValueError:
        pass #discarding header

def __calc_reg_write(self, region_file, count, calculated_region):
    if count > self.region_mintags:
        region_file.write(calculated_region.write())

def calculate_region(self):
    """
    Calculate a region file using the reads present in the both main files to analyze.
    """
    self.logger.info('Generating regions...')
    self.sorted_region_path = '%s/calcregion_%s.bed'%(self._output_dir(), os.path.basename(self.current_output_path))
    region_file = open(self.sorted_region_path, 'wb')

    if self.region_magic:
        regwriter = RegionWriter(self.gff_file, region_file, self.region_magic, no_sort=self.no_sort, logger=self.logger, write_as=BED, galaxy_workarounds=self.galaxy_workarounds)
        regwriter.write_regions()

    dual_reader = utils.DualSortedReader(self.current_experiment_path, self.current_control_path, self.experiment_format, self.logger)
    if self.stranded_analysis:
        calculate_region_stranded(self, dual_reader, region_file)
    else:
        calculate_region_notstranded(self, dual_reader, region_file)

    region_file.flush()

def __cr_append(self, regions, region):
    regions.append(region)

def calculate_region_notstranded(self, dual_reader, region_file):
    calculated_region = Region()
    readcount = 1
    for line in dual_reader:
        if not calculated_region: #first region only
            calculated_region = _region_from_dual(self, line)
            calculated_region.end += self.proximity
        else:
            new_region = _region_from_dual(self, line)
            new_region.end += self.proximity
            if calculated_region.overlap(new_region):
                calculated_region.join(new_region)
                readcount += 1
            else:
                calculated_region.end -= self.proximity
                __calc_reg_write(self, region_file, readcount, calculated_region)
                calculated_region = new_region.copy()
                readcount = 1

    if calculated_region:
        calculated_region.end -= self.proximity
        __calc_reg_write(self, region_file, readcount, calculated_region)

def calculate_region_stranded(self, dual_reader, region_file):
    temp_region_file = open(self.sorted_region_path, 'wb')
    region_plus = Region()
    region_minus = Region()
    regions = []
    numreads_plus = 1
    numreads_minus = 1
    dual_reader = utils.DualSortedReader(self.current_experiment_path, self.current_control_path, self.experiment_format, self.logger)
    for line in dual_reader:
        new_region = _region_from_dual(self, line)
        new_region.end += self.proximity
        if not (region_plus and new_region.strand == PLUS_STRAND):
            region_plus = _region_from_dual(self, line)

        elif not (region_plus and new_region.strand == PLUS_STRAND):
            region_minus = _region_from_dual(self, line)

        else:
            if region_plus.overlap(new_region) and region_plus.strand == new_region.strand:
                region_plus.join(new_region)
                numreads_plus += 1
            elif region_minus.overlap(new_region) and region_minus.strand == new_region.strand:
                region_minus.join(new_region)
                numreads_minus += 1
            else:
                if new_region.strand == region_plus.strand:
                    region_plus.end -= self.proximity
                    self.__calc_reg_write(region_file, numreads_plus, region_plus)
                    region_plus = new_region.copy()
                    numreads_plus = 1
                else:
                    region_minus.end -= self.proximity
                    self.__calc_reg_write(region_file, numreads_minus, region_minus)
                    region_minus = new_region.copy()
                    numreads_minus = 1

    if region_plus:
        region_plus.end -= self.proximity
        regions.append(region_plus)

    if region_minus:
        region_minus.end -= self.proximity
        regions.append(region_minus)

    regions.sort(key=lambda x:(x.name, x.start, x.end, x.strand))
    for region in regions:
        region_file.write(region.write())

def get_zscore(x, mean, sd):
    if sd > 0:
        return float(x-mean)/sd
    else:
        return 0 #This points are weird anyway

def read_interesting_regions(self, file_path):
    regs = []
    try:
        regs_file = open(file_path, 'r')
        for line in regs_file:
            regs.append(line.strip())
    except IOError as ioerror:
        self.logger.warning("Interesting regions file not found")
    return regs # memory inefficient if there's a large number of interesting regions

def plot_enrichment(self, file_path):
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        try:
            if self.postscript:
                import matplotlib
                matplotlib.use("PS")

            from matplotlib.pyplot import *
            from matplotlib import rcParams
            rcParams.update({'font.size': 22})
            rcParams['legend.fontsize'] = 14
            #decide labels
            if self.label1:
                label_main = self.label1
            else:
                if self.real_control_path and self.real_experiment_path:
                    label_main = '%s VS %s'%(os.path.basename(self.real_experiment_path), os.path.basename(self.real_control_path))
                else:
                    label_main = "A VS B"

            if self.label2:
                label_control = self.label2
            else:
                if self.replica_path:
                    label_control = '%s(A) VS %s(A)'%(os.path.basename(self.real_experiment_path), os.path.basename(self.replica_path))
                else:
                    label_control = 'Background distribution'

            #self.logger.info("Interesting regions path: %s" % (self.interesting_regions))
            interesting_regs = []
            if self.interesting_regions:
                self.logger.info("Reading interesting regions...")
                interesting_regs = read_interesting_regions(self, self.interesting_regions)
            #self.logger.info("Interesting regions: %s" % (interesting_regs))
            #self.logger.info("Plot path: %s" % (file_path))
            interesting_A = []
            interesting_M = []
            #self.logger.info("disable_significant: %s" % (self.disable_significant_color))

            A = []
            A_prime = []
            M = []
            M_significant = []
            A_significant = []
            M_prime = []
            A_medians = []
            points = []
            minus_points = []
            all_points = []
            figure(figsize=(14,22))
            biggest_A  = -sys.maxint #for drawing
            smallest_A = sys.maxint #for drawing
            biggest_M = 0 #for drawing
            self.logger.info("Loading table...")
            for line in open(file_path):
                sline = line.split()
                try:
                    enrich = dict(zip(enrichment_keys, sline))
                    # WARNING: for slide inter and slide intra: name2 = 'start:end' (no gene_id, FIXME?)
                    name2 = enrich['name2'].split(':')
                    gene_id = name2[0]
                    if len(name2) >= 2:
                        transcript_id = name2[1] # consider transcript_id? (exons)
                    else:
                        transcript_id = None

                    if gene_id in interesting_regs or transcript_id in interesting_regs:
                        interesting_M.append(float(enrich["M"]))
                        interesting_A.append(float(enrich["A"]))

                    biggest_A = max(biggest_A, float(enrich["A"]))
                    smallest_A = min(smallest_A, float(enrich["A"]))
                    biggest_M = max(biggest_M, abs(float(enrich["M"])))
                    biggest_A = max(biggest_A, float(enrich["A_prime"]))
                    smallest_A = min(smallest_A, float(enrich["A_prime"]))
                    biggest_M = max(biggest_M, abs(float(enrich["M_prime"])))
                    positive_point = self.zscore*float(enrich["sd"])+float(enrich["mean"])
                    negative_point = -self.zscore*float(enrich["sd"])+float(enrich["mean"])
                    A_median = float(enrich["A_median"])
                    all_points.append((A_median, positive_point, negative_point))
                    if abs(float(enrich["zscore"])) < self.zscore:
                        M.append(float(enrich["M"]))
                        A.append(float(enrich["A"]))
                    else:
                        M_significant.append(float(enrich["M"]))
                        A_significant.append(float(enrich["A"]))

                    M_prime.append(float(enrich["M_prime"]))
                    A_prime.append(float(enrich["A_prime"]))
                except ValueError:
                    pass #to skip the header
            all_points.sort(key= lambda x:x[0])

            for t in all_points:
                (A_medians.append(t[0]), points.append(t[1]), minus_points.append(t[2]))

            if points:
                margin = 1.1
                A_medians.append(biggest_A*margin)
                points.append(points[-1])
                minus_points.append(minus_points[-1])
                A_medians.insert(0, smallest_A)
                points.insert(0, points[0])
                minus_points.insert(0, minus_points[0])
                self.logger.info("Plotting points...")
                #Background plot
                subplot(211, axisbg="lightyellow")
                xlabel('Average', fontsize=30)
                ylabel('Log2 ratio', fontsize=30)
                axis([smallest_A*margin, biggest_A*margin, -biggest_M*margin, biggest_M*margin])
                plot(A_prime, M_prime, '.', label=label_control, color = '#666666')
                plot(A_medians, points, 'r--', label="Z-score (%s)"%self.zscore)
                plot(A_medians, minus_points,  'r--')
                axhline(0, linestyle='--', color="grey", alpha=0.75)
                leg = legend(fancybox=True, scatterpoints=1, numpoints=1, loc=2, ncol=4, mode="expand")
                leg.get_frame().set_alpha(0.5)
                #Experiment plot
                subplot(212, axisbg="lightyellow")
                axis([smallest_A*margin, biggest_A*margin, -biggest_M*margin, biggest_M*margin])
                plot(A, M, 'k.', label=label_main)
                if self.disable_significant_color:
                    significant_marker = 'ko'
                else:
                    significant_marker = 'ro'

                plot(A_significant, M_significant, significant_marker, label="%s (significant)"%label_main)
                plot(A_medians, points, 'r--', label="Z-score (%s)"%self.zscore)
                plot(A_medians, minus_points, 'r--')
                if self.interesting_regions:
                    interesting_label = label_main + ' (interesting)'
                    plot(interesting_A, interesting_M, 'H', label=interesting_label, color='#00EE00') # plotting "interesting" regions

                axhline(0, linestyle='--', color="grey", alpha=0.75)
                xlabel('Average', fontsize=30)
                ylabel('Log2 ratio', fontsize=30)
                leg2 = legend(fancybox=True, scatterpoints=1, numpoints=1, loc=2, ncol=4)
                leg2.get_frame().set_alpha(0.7)
                self._save_figure("enrichment_MA", width=500, height=2800)
            else:
                self.logger.warning("Nothing to plot.")

        except ImportError:
            if self.debug:
                raise
            __matplotlibwarn(self)

def __matplotlibwarn(self):
    #FIXME move to utils.py or plotting module
    self.logger.warning('Pyicos can not find an installation of matplotlib, so no plot will be drawn. If you want to get a plot with the correlation values, install the matplotlib library.')

def __calc_M(signal_a, signal_b):
    return math.log(float(signal_a)/float(signal_b), 2)

def __calc_A(signal_a, signal_b):
    return (math.log(float(signal_a), 2)+math.log(float(signal_b), 2))/2

def _calculate_MA(self, region_path, read_counts, factor = 1, replica_factor = 1, file_a_reader=None, file_b_reader=None, replica_reader=None):
    tags_a = []
    tags_b = []
    numreads_background_1 = 0
    numreads_background_2 = 0
    total_reads_background_1 = 0
    total_reads_background_2 = 0
    self.logger.debug("Inside _calculate_MA")
    self.regions_analyzed_count = 0
    enrichment_result = [] #This will hold the name, start and end of the region, plus the A, M, 'A and 'M
    if NOWRITE not in self.operations:
        out_file = open(self.current_output_path, 'wb')

    for region_line in open(region_path):
        sline = region_line.split()
        region_of_interest = self._region_from_sline(sline)
        if region_of_interest:
            region_a = None
            replica = None
            replica_tags = None
            signal_a = -1
            signal_b = -1
            signal_background_1 = -1
            signal_background_2 = -1
            swap1 = Region()
            swap2 = Region()
            if read_counts:
                signal_a = float(sline[6])
                signal_b = float(sline[7])*factor
                signal_background_1 = float(sline[8])
                signal_background_2 = float(sline[9])*replica_factor
                if CHECK_REPLICAS in self.operations:
                    self.experiment_values.append(signal_background_1)
                    self.replica_values.append(signal_background_2)

            else:
                self.logger.debug("Reading tags for %s ..."%region_of_interest)
                if self.experiment_format == BAM:

                    tags_a = len(file_a_reader.get_overlaping_clusters(region_of_interest, overlap=self.overlap))
                    tags_b = len(file_b_reader.get_overlaping_clusters(region_of_interest, overlap=self.overlap))
                else:
                    tags_a = file_a_reader.get_overlaping_counts(region_of_interest, overlap=self.overlap)
                    tags_b = file_b_reader.get_overlaping_counts(region_of_interest, overlap=self.overlap)

                if self.use_replica:
                    if self.experiment_format == BAM:
                        replica_tags = len(replica_reader.get_overlaping_clusters(region_of_interest, overlap=self.overlap))
                    else:
                        replica_tags = replica_reader.get_overlaping_counts(region_of_interest, overlap=self.overlap)

                self.logger.debug("... done. tags_a: %s tags_b: %s"%(tags_a, tags_b))

                #if we are using pseudocounts, use the union, use the intersection otherwise
                if (self.pseudocount and (tags_a or tags_b)) or (not self.pseudocount and tags_a and tags_b):
                    signal_a = region_of_interest.normalized_counts(self.len_norm, self.n_norm, self.total_regions, self.pseudocount, factor, self.total_reads_a, tags_a)
                    signal_b = region_of_interest.normalized_counts(self.len_norm, self.n_norm, self.total_regions, self.pseudocount, factor, self.total_reads_b, tags_b)
                    self.already_norm = True

            if not self.counts_file:
                if (self.pseudocount and (tags_a or tags_b)) or (not self.pseudocount and tags_a and tags_b):
                    if self.use_replica:
                        replica = region_of_interest.copy()
                        #replica.add_tags(replica_tags)
                        numreads_background_1 = tags_a
                        numreads_background_2 = replica_tags
                        total_reads_background_1 = self.total_reads_a
                        total_reads_background_2 = self.total_reads_replica
                        signal_background_1 = signal_a
                        signal_background_2 = region_of_interest.normalized_counts(self.len_norm, self.n_norm, self.total_regions, self.pseudocount,
                                                                        replica_factor, self.total_reads_replica, replica_tags)

                    else:
                        numreads_background_1 = 0
                        numreads_background_2 = 0
                        for i in range(0, tags_a+tags_b):
                            if random.uniform(0,2) > 1:
                                numreads_background_1 += 1
                            else:
                                numreads_background_2 += 1

                        total_reads_background_1 = total_reads_background_2 = self.average_total_reads
                        signal_background_1 = region_of_interest.normalized_counts(self.len_norm, self.n_norm, self.total_regions, self.pseudocount,
                                                                                replica_factor, self.average_total_reads, numreads_background_1)
                        signal_background_2 = region_of_interest.normalized_counts(self.len_norm, self.n_norm, self.total_regions, self.pseudocount,
                                                                                replica_factor, self.average_total_reads, numreads_background_2)

            #if there is no data in the replica or in the swap and we are not using pseudocounts, dont write the data
            if signal_a > 0 and signal_b > 0 and signal_background_1 > 0 and signal_background_2 > 0 or self.use_MA:
                if self.use_MA and not self.already_norm:
                    A = float(sline[10])
                    M = float(sline[11])
                    A_prime = float(sline[16])
                    M_prime = float(sline[17])
                else:
                    if not self.already_norm: #TODO refractor
                        if self.len_norm: #read per kilobase in region
                            signal_a = 1e3*(float(signal_a)/len(region_of_interest))
                            signal_b = 1e3*(float(signal_b)/len(region_of_interest))
                            signal_background_1 = 1e3*(float(signal_background_1)/len(region_of_interest))
                            signal_background_2 = 1e3*(float(signal_background_2)/len(region_of_interest))

                        if self.n_norm: #per million reads in the sample
                            signal_a = 1e6*(float(signal_a)/self.total_reads_a)
                            signal_b = 1e6*(float(signal_b)/self.total_reads_b)
                            if self.use_replica:
                                signal_background_1 = signal_a
                                signal_background_2 = 1e6*(float(signal_background_2)/self.total_reads_replica)
                            else:
                                signal_background_1 = 1e6*(float(signal_background_1)/self.average_total_reads)
                                signal_background_2 = 1e6*(float(signal_background_2)/self.average_total_reads)

                    A = __calc_A(signal_a, signal_b)
                    M = __calc_M(signal_a, signal_b)
                    A_prime = __calc_A(signal_background_1, signal_background_2)
                    M_prime = __calc_M(signal_background_1, signal_background_2)
                    if CHECK_REPLICAS in self.operations:
                        self.experiment_values.append(signal_background_1)
                        self.replica_values.append(signal_background_2)

                if NOWRITE not in self.operations:
                    out_file.write("%s\n"%("\t".join([region_of_interest.write().rstrip("\n"), str(signal_a), str(signal_b), str(signal_background_1), str(signal_background_2), str(A), str(M), str(self.total_reads_a), str(self.total_reads_b), str(tags_a), str(tags_b),  str(A_prime), str(M_prime), str(total_reads_background_1), str(total_reads_background_2), str(numreads_background_1), str(numreads_background_2)])))
                self.regions_analyzed_count += 1

    self.logger.debug("LEAVING _calculate_MA")
    if NOWRITE in self.operations:
        return ""
    else:
        out_file.flush()
        out_file.close()
        # Outputting to HTML (if specified)
        if self.html_output is not None:
            self.logger.info("Generating HTML")
            try:
                from jinja2 import Environment, PackageLoader, Markup
            except:
                self.logger.error("Could not find the jinja2 library")
                return out_file.name

            loadr = PackageLoader('pyicoteolib', 'templates')
            env = Environment(loader=loadr)
            template = env.get_template('enrich_html.html')

            def jinja_read_file(filename):
                f = open(filename, 'r')
                #for line in f:
                #    print line
                txt = ''.join(f.readlines())
                f.close()
                return txt
            env.globals['jinja_read_file'] = jinja_read_file

            if self.galaxy_workarounds: # Galaxy changes the working directory when outputting multiple files
                parent_dir = "./"
            else:
                parent_dir = os.sep.join(out_file.name.split(os.sep)[0:-1]) + "/"
            plot_path = parent_dir + "enrichment_MA_" + out_file.name.split(os.sep)[-1] + ".png"
            bed_path = parent_dir + out_file.name.split(os.sep)[-1]

            html_file = open(self.html_output, 'w')
            html_file.write(template.render({'page_title': 'Enrichment results', 'results_output': jinja_read_file(out_file.name), 'plot_path': plot_path, 'bed_path': bed_path}))
            html_file.flush()
            html_file.close()

        return out_file.name

def _calculate_total_lengths(self):
    msg = "Calculating enrichment in regions"
    if self.counts_file:
        self.sorted_region_path = self.counts_file
        if (not self.total_reads_a or not self.total_reads_b or (not self.total_reads_replica and self.use_replica)) and not self.use_MA:
            self.logger.info("... counting from counts file...")
            self.total_reads_a = 0
            self.total_reads_b = 0
            if self.total_reads_replica:
                self.total_reads_replica = 0
            else:
                self.total_reads_replica = 1
            for line in open(self.counts_file):
                try:
                    enrich = dict(zip(enrichment_keys, line.split()))
                    self.total_reads_a += float(enrich["signal_a"])
                    self.total_reads_b += float(enrich["signal_b"])
                    if self.use_replica:
                        self.total_reads_replica += float(enrich["signal_prime_2"])
                except ValueError:
                    self.logger.debug("(Counting) skip header...")
    else:
        self.logger.info("... counting number of lines in files...")
        if not self.total_reads_a:
            if self.experiment_format == BAM:
                self.total_reads_a = bam.size(self.current_experiment_path)
            else:
                self.total_reads_a = sum(1 for line in utils.open_file(self.current_experiment_path, self.experiment_format, logger=self.logger))
        if not self.total_reads_b:
            if self.experiment_format == BAM:
                self.total_reads_b = bam.size(self.current_control_path)
            else:
                self.total_reads_b = sum(1 for line in utils.open_file(self.current_control_path, self.control_format, logger=self.logger))
        if self.use_replica and not self.total_reads_replica:
            if self.experiment_format  == BAM:
                self.total_reads_replica = bam.size(self.replica_path)
            else:
                self.total_reads_replica = sum(1 for line in utils.open_file(self.replica_path, self.experiment_format, logger=self.logger))

        self.logger.debug("Number lines in experiment A: %s Experiment B: %s"%(self.total_reads_a, self.total_reads_b))
        if self.use_replica:
            msg = "%s using replicas..."%msg
        else:
            msg = "%s using swap..."%msg

        self.logger.info(msg)

    self.average_total_reads = (self.total_reads_a+self.total_reads_b)/2

def enrichment(self):
    file_a_reader = file_b_reader = replica_reader = None
    self.use_replica = (bool(self.replica_path) or (bool(self.counts_file) and self.use_replica_flag))
    self.logger.debug("Use replica: %s"%self.use_replica)
    if not USE_MA in self.operations:
        _calculate_total_lengths(self)
    if not self.counts_file:
        file_a_reader = utils.read_fetcher(self.current_experiment_path, self.experiment_format, cached=self.cached, logger=self.logger, use_samtools=self.use_samtools, access_sequential=self.access_sequential, only_counts=True)
        file_b_reader = utils.read_fetcher(self.current_control_path, self.experiment_format, cached=self.cached, logger=self.logger, use_samtools=self.use_samtools, access_sequential=self.access_sequential, only_counts=True)
        if self.use_replica:
            replica_reader = utils.read_fetcher(self.current_replica_path, self.experiment_format, cached=self.cached, logger=self.logger, use_samtools=self.use_samtools, access_sequential=self.access_sequential, only_counts=True)

        if self.sorted_region_path:
            self.logger.info('Using region file %s (%s)'%(self.region_path, self.region_format))
        else:
            calculate_region(self) #create region file semi automatically

        self.total_regions = sum(1 for line in open(self.sorted_region_path))

    self.logger.info("... analyzing regions, calculating normalized counts, A / M and replica or swap...")
    self.already_norm = False
    if self.use_MA:
        ma_path = self.counts_file
    else:
        ma_path = self.sorted_region_path

    out_path = _calculate_MA(self, ma_path, bool(self.counts_file), 1, 1, file_a_reader, file_b_reader, replica_reader)
    self.already_norm = True
    self.logger.debug("Already normalized: %s"%self.already_norm)
    if self.tmm_norm:
        if CHECK_REPLICAS in self.operations:
            self.experiment_values = []
            self.replica_values = []

        self.logger.info("TMM Normalizing...")
        tmm_factor = calc_tmm_factor(self, out_path, self.regions_analyzed_count, False)
        replica_tmm_factor = 1
        if self.use_replica:
            replica_tmm_factor = calc_tmm_factor(self, out_path, self.regions_analyzed_count, True)
        #move output file to old output
        #use as input
        old_output = '%s/notnormalized_%s'%(self._current_directory(), os.path.basename(self.current_output_path))
        move(os.path.abspath(self.current_output_path), old_output)
        out_path = _calculate_MA(self, old_output, True, tmm_factor, replica_tmm_factor, True) #recalculate with the new factor, using the counts again

    if self.quant_norm:
        self.logger.info("Full quantile normalization...")
        signal_a = []
        signal_prime_1 = []
        enrich = []
        for line in open(out_path):
            sline = line.split()
            enrich_line = dict(zip(enrichment_keys, sline))
            enrich.append(enrich_line)
            signal_a.append(float(enrich_line['signal_a']))
            signal_prime_1.append(float(enrich_line['signal_prime_1']))

        #full quantile normalization
        signal_a.sort()
        enrich.sort(key=lambda x:float(x['signal_b']))
        quant_counts = open('%s/quantcounts_%s'%(self._current_directory(), os.path.basename(self.current_output_path)), 'w')
        for i in range(len(enrich)):
            enrich[i]['signal_b'] = signal_a[i]

        self.logger.info("Full quantile normalization replica...")
        #full quantile normalization of the replica
        signal_prime_1.sort()
        enrich.sort(key=lambda x:float(x['signal_prime_2']))
        for i in range(len(enrich)):
            enrich[i]['signal_prime_2'] = signal_prime_1[i]
            quant_counts.write("%s\n"%"\t".join(str(enrich[i][key]) for key in enrichment_keys[:20])) #write the lines

        quant_counts.flush()
        out_path = _calculate_MA(self, quant_counts.name, True, 1, 1, True) #recalculate with the new factor, using the counts again
        self._manage_temp_file(quant_counts.name)

    self.logger.info("%s regions analyzed."%self.regions_analyzed_count)
    if not NOWRITE in self.operations:
        self.logger.info("Enrichment result saved to %s"%self.current_output_path)

    if CHECK_REPLICAS in self.operations:
        check_replica(self)

    return out_path

def _sub_tmm(counts_a, counts_b, reads_a, reads_b):
    return (counts_a-reads_a)/(counts_a*reads_a) + (counts_b-reads_b)/(counts_b*reads_b)

def calc_tmm_factor(self, file_counts, total_regions, replica):
    if replica:
        signal_1 = "signal_prime_1"
        signal_2 = "signal_prime_2"
        M = "M_prime"
        reads_2 = self.total_reads_replica
    else:
        signal_1 = "signal_a"
        signal_2 = "signal_b"
        M = "M"
        reads_2 = self.total_reads_b

    values_list = []
    #read the file inside the values_list
    for line in open(file_counts):
        sline = line.split()
        values_list.append(dict(zip(enrichment_keys, sline)))

    a_trim_number = int(round(total_regions*self.a_trim))
    #discard the bad A
    self.logger.debug("Removing the worst A (%s regions, %s percent)"%(a_trim_number, self.a_trim*100))
    values_list.sort(key=lambda x:float(x["A"])) #sort by A
    for i in range (0, a_trim_number):
        values_list.pop(0)

    values_list.sort(key=lambda x:float(x[M])) #sort by M
    m_trim_number = int(round(total_regions*(self.m_trim/2))) #this number is half the value of the flag, because we will trim half below, and half over
    #remove on the left
    for i in range(0, m_trim_number):
        values_list.pop(0)
    #remove on the right
    for i in range(0, m_trim_number):
        values_list.pop(-1)

    #now calculate the normalization factor
    arriba = 0
    abajo = 0
    for value in values_list:
        w = _sub_tmm(float(value[signal_1]), float(value[signal_2]), self.total_reads_a, reads_2)
        arriba += w*float(value[M])
        abajo += w

    try:
        factor = 2**(arriba/abajo)
    except ZeroDivisionError:
        self.logger.warning("Division by zero, TMM factor could not be calculated.")
        factor = 1

    if replica:
        self.logger.info("Replica TMM Normalization Factor: %s"%factor)
    else:
        self.logger.info("TMM Normalization Factor: %s"%factor)

    return factor

def __load_enrichment_result(values_path):
    ret = []
    for line in open(values_path):
        sline = line.split()
        try:
            float(sline[1])
            ret.append(dict(zip(enrichment_keys, sline)))
        except ValueError:
            pass

    return ret

def calculate_zscore(self, values_path):
    num_regions = sum(1 for line in open(values_path))
    bin_size = int(self.binsize*num_regions)
    if bin_size < 50:
        self.logger.warning("The bin size results in a sliding window smaller than 50, adjusting window to 50 in order to get statistically meaningful results.")
        bin_size = 50

    bin_step = max(1, int(round(self.bin_step*bin_size)))
    self.logger.info("Enrichment window calculation using a sliding window size of %s, sliding with a step of %s"%(bin_size, bin_step))
    self.logger.info("... calculating zscore...")
    enrichment_result = __load_enrichment_result(values_path)
    enrichment_result.sort(key= lambda x:(float(x["A_prime"])))
    self.logger.debug("Number of loaded counts: %s"%len(enrichment_result))
    self.points = []
    #get the standard deviations
    for i in range(0, num_regions-bin_size+bin_step, bin_step):
        #get the slice
        if i+bin_size < num_regions:
            result_chunk = enrichment_result[i:i+bin_size]
        else:
            result_chunk = enrichment_result[i:]  #last chunk

        #retrieve the values
        mean_acum = 0
        a_acum = 0
        Ms_replica = []
        for entry in result_chunk:
            mean_acum += float(entry["M_prime"])
            a_acum += float(entry["A_prime"])
            Ms_replica.append(float(entry["M_prime"]))

        #add them to the points of mean and sd
        mean = mean_acum/len(result_chunk)
        sd = math.sqrt((sum((x - mean)**2 for x in Ms_replica))/len(Ms_replica))
        #the A median
        A_median = a_acum / len(result_chunk)
        self.points.append([A_median, mean, sd]) #The A asigned to the window, the mean and the standard deviation
        #self.logger.debug("Window of %s length, with A median: %s mean: %s sd: %s"%(len(result_chunk), self.points[-1][0], self.points[-1][1], self.points[-1][2], len(self.points)))

    #update z scores
    for entry in enrichment_result:
        entry["A_median"] = 0
        entry["mean"] = 0
        entry["sd"] = 0
        entry["zscore"] = 0
        closest_A = sys.maxint
        sd_position = 0
        for i in range(0, len(self.points)):
            new_A = self.points[i][0]
            if new_A != closest_A: #skip repeated points
                if abs(closest_A - float(entry["A"])) >= abs(new_A - float(entry["A"])):
                    closest_A = new_A
                    sd_position = i
                else:
                    break #already found, no need to go further since the points are ordered

        entry["A_median"] = closest_A
        if self.points: #only calculate if there where windows...
            __sub_zscore(self.sdfold, entry, self.points[sd_position])

    if not self.points: # ... otherwise give a warning
        self.logger.warning("Insufficient number of regions analyzed (%s), z-score values could not be calculated"%num_regions)

    enrichment_result.sort(key=lambda x:(x["name"], int(x["start"]), int(x["end"])))
    old_file_path = '%s/before_zscore_%s'%(self._current_directory(), os.path.basename(values_path)) #create path for the outdated file
    move(os.path.abspath(values_path), old_file_path) #move the file
    new_file = file(values_path, 'w') #open a new file in the now empty file space
    if not self.skip_header:
        new_file.write('\t'.join(enrichment_keys))
        new_file.write('\n')

    for entry in enrichment_result:
        new_file.write("\t".join(str(entry[key]) for key in enrichment_keys)+"\n")

    self._manage_temp_file(old_file_path)
    return values_path

def __sub_zscore(sdfold, entry, point):
    entry["mean"] = str(point[1])
    entry["sd"] = str(point[2])
    entry["zscore"] = str(get_zscore(float(entry["M"]), float(entry["mean"]), sdfold*float(entry["sd"])))

def check_replica(self):
    #discard everything below the flag
    new_experiment = []
    new_replica = []
    min_value = sys.maxint
    max_value = -sys.maxint
    for i in range(len(self.replica_values)):
        if self.experiment_values[i] > self.count_filter and self.replica_values[i] > self.count_filter:
            new_experiment.append(math.log(self.experiment_values[i], 2))
            new_replica.append(math.log(self.replica_values[i], 2))
            min_value = min(min_value, math.log(self.experiment_values[i], 2), math.log(self.replica_values[i], 2))
            max_value = max(max_value, math.log(self.experiment_values[i], 2), math.log(self.replica_values[i], 2))
    #print self.replica_values
    self.experiment_values = new_experiment
    self.replica_values = new_replica
    try:
        if self.postscript:
            import matplotlib
            matplotlib.use("PS")
        from matplotlib.pyplot import plot, show, xlabel, ylabel, axhline, axis, clf, text, title, xlim, ylim
    except:
        __matplotlibwarn(self)
        return 0
    clf()
    r_squared = utils.pearson(self.experiment_values, self.replica_values)**2
    text(min_value+abs(max_value)*0.1, max_value-abs(max_value)*0.2, r'Pearson $R^2$= %s'%round(r_squared, 3), fontsize=18, bbox={'facecolor':'yellow', 'alpha':0.5, 'pad':10})
    xlabel("log2(%s)"%self.experiment_label, fontsize=18)
    ylabel("log2(%s)"%self.replica_label, fontsize=18)
    xlim(min_value, max_value)
    ylim(min_value, max_value)
    title(self.title_label, fontsize=24)
    plot(self.experiment_values, self.replica_values, '.')

    self._save_figure("check_replica")

def check_replica_correlation(self):
    "No usado, de momento"
    min_tags = 20
    experiment_reader = utils.read_fetcher(self.current_experiment_path, self.experiment_format, cached=self.cached, logger=self.logger, use_samtools=self.use_samtools, access_sequential=self.access_sequential)
    replica_reader = utils.read_fetcher(self.current_replica_path, self.experiment_format, cached=self.cached, logger=self.logger, use_samtools=self.use_samtools, access_sequential=self.access_sequential)
    correlations_acum = 0
    num_correlations = 0
    for region_line in open(self.region_path):
        sline = region_line.split()
        region_experiment = self._region_from_sline(sline)
        region_replica = region_experiment.copy()
        tags_experiment = experiment_reader.get_overlaping_clusters(region_experiment, overlap=1)
        tags_replica = replica_reader.get_overlaping_clusters(region_experiment, overlap=1)
        count_experiment = len(tags_experiment)
        count_replica = len(tags_replica)
        correlations = []
        if count_experiment+count_replica > min_tags:
            region_experiment.add_tags(tags_experiment, clusterize=True)
            region_replica.add_tags(tags_replica, clusterize=True)
            num_correlations += 1
            correlation = utils.pearson(region_experiment.get_array(), region_replica.get_array())
            correlations_acum += max(0, correlation)
            correlations.append(correlation)

    print correlations_acum/num_correlations
    try:
        if self.postscript:
            import matplotlib
            matplotlib.use("PS")
        from matplotlib.pyplot import plot, boxplot, show, legend, figure, xlabel, ylabel, subplot, axhline, axis
    except:
        __matplotlibwarn(self)
        return 0

    print correlations
    boxplot(correlations)
    self._save_figure("check_replica")


import logging
import os

import numpy as np
import pandas as pd

from tardis.plasma.properties.base import (
    PreviousIterationProperty,
    ProcessingPlasmaProperty,
)
from tardis.plasma.properties.ion_population import PhiSahaNebular

__all__ = [
    "PreviousElectronDensities",
    "PreviousBetaSobolev",
    "HeliumNLTE",
    "HeliumNumericalNLTE",
]

logger = logging.getLogger(__name__)

class PreviousElectronDensities(PreviousIterationProperty):
    """
    Attributes
    ----------
    previous_electron_densities : The values for the electron densities converged upon in the previous iteration.
    """

    outputs = ("previous_electron_densities",)

    def set_initial_value(self, kwargs):
        initial_value = pd.Series(
            1000000.0,
            index=kwargs["abundance"].columns,
        )
        self._set_initial_value(initial_value)

class PreviousBetaSobolev(PreviousIterationProperty):
    """
    Attributes
    ----------
    previous_beta_sobolev : The beta sobolev values converged upon in the previous iteration.
    """

    outputs = ("previous_beta_sobolev",)

    def set_initial_value(self, kwargs):
        initial_value = pd.DataFrame(
            1.0,
            index=kwargs["atomic_data"].lines.index,
            columns=kwargs["abundance"].columns,
        )
        self._set_initial_value(initial_value)

class HeliumNLTE(ProcessingPlasmaProperty):
    outputs = ("helium_population",)

    @staticmethod
    def calculate(
        level_boltzmann_factor,
        ionization_data,
        beta_rad,
        g,
        g_electron,
        w,
        t_rad,
        t_electrons,
        delta,
        zeta_data,
        number_density,
        partition_function,
    ):
        """
        Updates all of the helium level populations according to the helium NLTE recomb approximation.
        """
        helium_population = level_boltzmann_factor.loc[2].copy()
        # He I excited states
        he_one_population = HeliumNLTE.calculate_helium_one(
            g_electron, beta_rad, ionization_data, level_boltzmann_factor, g, w
        )
        helium_population.loc[0].update(he_one_population)
        # He I ground state
        helium_population.loc[0, 0] = 0.0
        # He II excited states
        he_two_population = level_boltzmann_factor.loc[2, 1].mul(
            (g.loc[2, 1, 0] ** (-1.0))
        )
        helium_population.loc[1].update(he_two_population)
        # He II ground state
        helium_population.loc[1, 0] = 1.0
        # He III states
        helium_population.loc[2, 0] = HeliumNLTE.calculate_helium_three(
            t_rad,
            w,
            zeta_data,
            t_electrons,
            delta,
            g_electron,
            beta_rad,
            ionization_data,
            g,
        )
        #        unnormalised = helium_population.sum()
        #        normalised = helium_population.mul(number_density.ix[2] / unnormalised)
        #        helium_population.update(normalised)
        return helium_population

    @staticmethod
    def calculate_helium_one(
        g_electron, beta_rad, ionization_data, level_boltzmann_factor, g, w
    ):
        """
        Calculates the He I level population values, in equilibrium with the He II ground state.
        """
        return (
            level_boltzmann_factor.loc[2, 0]
            * (1.0 / (2 * g.loc[2, 1, 0]))
            * (1 / g_electron)
            * (1 / (w ** 2.0))
            * np.exp(ionization_data.loc[2, 1] * beta_rad)
        )

    @staticmethod
    def calculate_helium_three(
        t_rad,
        w,
        zeta_data,
        t_electrons,
        delta,
        g_electron,
        beta_rad,
        ionization_data,
        g,
    ):
        """
        Calculates the He III level population values.
        """
        zeta = PhiSahaNebular.get_zeta_values(zeta_data, 2, t_rad)[1]
        he_three_population = (
            2
            * (float(g.loc[2, 2, 0]) / g.loc[2, 1, 0])
            * g_electron
            * np.exp(-ionization_data.loc[2, 2] * beta_rad)
            * w
            * (delta.loc[2, 2] * zeta + w * (1.0 - zeta))
            * (t_electrons / t_rad) ** 0.5
        )
        return he_three_population

class HeliumNumericalNLTE(ProcessingPlasmaProperty):
    """
    IMPORTANT: This particular property requires a specific numerical NLTE
    solver and a specific atomic dataset (neither of which are distributed
    with Tardis) to work.
    """

    outputs = ("helium_population",)

    def __init__(self, plasma_parent, heating_rate_data_file):
        super(HeliumNumericalNLTE, self).__init__(plasma_parent)
        self._g_upper = None
        self._g_lower = None
        self.heating_rate_data = np.loadtxt(heating_rate_data_file, unpack=True)

    def calculate(
        self,
        ion_number_density,
        electron_densities,
        t_electrons,
        w,
        lines,
        j_blues,
        levels,
        level_boltzmann_factor,
        t_rad,
        zeta_data,
        g_electron,
        delta,
        partition_function,
        ionization_data,
        beta_rad,
        g,
        time_explosion,
    ):
        logger.info("Performing numerical NLTE He calculations.")
        if len(j_blues) == 0:
            return None
        # Outputting data required by SH module
        for zone, _ in enumerate(electron_densities):
            with open(
                f"He_NLTE_Files/shellconditions_{zone}.txt", "w"
            ) as output_file:
                output_file.write(ion_number_density.loc[2].sum()[zone])
                output_file.write(electron_densities[zone])
                output_file.write(t_electrons[zone])
                output_file.write(self.heating_rate_data[zone])
                output_file.write(w[zone])
                output_file.write(time_explosion)
                output_file.write(t_rad[zone])
                output_file.write(self.plasma_parent.v_inner[zone])
                output_file.write(self.plasma_parent.v_outer[zone])

        for zone, _ in enumerate(electron_densities):
            with open(
                f"He_NLTE_Files/abundances_{zone}.txt", "w"
            ) as output_file:
                for element in range(1, 31):
                    try:
                        number_density = (
                            ion_number_density[zone].loc[element].sum()
                        )
                    except:
                        number_density = 0.0
                    output_file.write(number_density)

            helium_lines = lines[lines["atomic_number"] == 2]
            helium_lines = helium_lines[helium_lines["ion_number"] == 0]
        for zone, _ in enumerate(electron_densities):
            with open(
                f"He_NLTE_Files/discradfield_{zone}.txt", "w"
            ) as output_file:
                j_blues = pd.DataFrame(j_blues, index=lines.index)
                helium_j_blues = j_blues[zone].loc[helium_lines.index]
                for value in helium_lines.index:
                    if helium_lines.level_number_lower.loc[value] < 35:
                        output_file.write(
                            int(helium_lines.level_number_lower.loc[value] + 1),
                            int(helium_lines.level_number_upper.loc[value] + 1),
                            j_blues[zone].loc[value],
                        )
        # Running numerical simulations
        for zone, _ in enumerate(electron_densities):
            os.rename(
                f"He_NLTE_Files/abundances{zone}.txt",
                "He_NLTE_Files/abundances_current.txt",
            )
            os.rename(
                f"He_NLTE_Files/shellconditions{zone}.txt",
                "He_NLTE_Files/shellconditions_current.txt",
            )
            os.rename(
                f"He_NLTE_Files/discradfield{zone}.txt",
                "He_NLTE_Files/discradfield_current.txt",
            )
            os.system("nlte-solver-module/bin/nlte_solvertest >/dev/null")
            os.rename(
                "He_NLTE_Files/abundances_current.txt",
                f"He_NLTE_Files/abundances{zone}.txt",
            )
            os.rename(
                "He_NLTE_Files/shellconditions_current.txt",
                f"He_NLTE_Files/shellconditions{zone}.txt",
            )
            os.rename(
                "He_NLTE_Files/discradfield_current.txt",
                f"He_NLTE_Files/discradfield{zone}.txt",
            )
            os.rename("debug_occs.dat", f"He_NLTE_Files/occs{zone}.txt")
        # Reading in populations from files
        helium_population = level_boltzmann_factor.loc[2].copy()
        for zone, _ in enumerate(electron_densities):
            with open(
                f"He_NLTE_Files/discradfield{zone}.txt", "r"
            ) as read_file:
                for level in range(0, 35):
                    level_population = read_file.readline()
                    level_population = float(level_population)
                    helium_population[zone].loc[0, level] = level_population
                helium_population[zone].loc[1, 0] = float(read_file.readline())
        # Performing He LTE level populations (upper two energy levels,
        # He II excited states, He III)
        he_one_population = HeliumNLTE.calculate_helium_one(
            g_electron,
            beta_rad,
            partition_function,
            ionization_data,
            level_boltzmann_factor,
            electron_densities,
            g,
            w,
            t_rad,
            t_electrons,
        )
        helium_population.loc[0, 35].update(he_one_population.loc[35])
        helium_population.loc[0, 36].update(he_one_population.loc[36])

        he_two_population = level_boltzmann_factor.loc[2, 1, 1:].mul(
            (g.loc[2, 1, 0] ** (-1)) * helium_population.loc[s1, 0]
        )
        helium_population.loc[1, 1:].update(he_two_population)

        helium_population.loc[2, 0] = HeliumNLTE.calculate_helium_three(
            t_rad,
            w,
            zeta_data,
            t_electrons,
            delta,
            g_electron,
            beta_rad,
            partition_function,
            ionization_data,
            electron_densities,
        )
        unnormalised = helium_population.sum()
        normalised = helium_population.mul(
            ion_number_density.loc[2].sum() / unnormalised
        )
        helium_population.update(normalised)
        return helium_population


#!/usr/bin/env python

"""
Demos of MIMO time encoding and decoding algorithms that use IAF
neurons with delays.
"""

# Copyright (c) 2009-2015, Lev Givon
# All rights reserved.
# Distributed under the terms of the BSD license:
# http://www.opensource.org/licenses/bsd-license

import numpy as np

# Set matplotlib backend so that plots can be generated without a
# display:
import matplotlib
matplotlib.use('AGG')

from bionet.utils.misc import func_timer
import bionet.utils.band_limited as bl
import bionet.utils.plotting as pl
import bionet.ted.iaf as iaf

# For determining output plot file names:
output_name = 'iaf_delay_demo_'
output_count = 0
output_ext = '.png'

# Define input signal generation parameters:
T = 0.05
dur = 2*T
dt = 1e-6
f = 100
bw = 2*np.pi*f

np.random.seed(0)
noise_power = None
comps = 8

if noise_power == None:
    fig_title = 'IAF Input Signal with No Noise'
else:
    fig_title = 'IAF Input Signal with %d dB of Noise' % noise_power

M = 3 # number of input signals
N = 9 # number of neurons

# Starting and ending points of interval that is encoded:
t_start = 0.02
t_end = t_start+T
if t_end > dur:
    raise ValueError('t_start is too large')
k_start = int(np.round(t_start/dt))
k_end = int(np.round(t_end/dt))
t_enc = np.arange(k_start, k_end, dtype=np.float)*dt

u_list = []
for i in xrange(M):
    fig_title_in = fig_title + ' (Signal #' + str(i+1) + ')'
    print fig_title_in
    u = func_timer(bl.gen_band_limited)(dur, dt, f, noise_power, comps)
    u /= max(u)
    u *= 1.5
    pl.plot_signal(t_enc, u[k_start:k_end], fig_title_in,
                   output_name + str(output_count) + output_ext)
    u_list.append(u)
    output_count += 1

t = np.arange(len(u_list[0]), dtype=np.float)*dt

# Define neuron parameters:
def randu(a, b, *d):
    """Create an array of the given shape and propagate it with random
    samples from a uniform distribution over ``[a, b)``."""

    if a >= b:
        raise ValueError('b must exceed a')
    return a+(b-a)*np.random.rand(*d)

b_list = list(randu(2.3, 3.3, N))
d_list = list(randu(0.15, 0.25, N))
k_list = list(0.01*np.ones(N))
a_list = map(list, np.reshape(np.random.exponential(0.003, N*M), (N, M)))
w_list = map(list, np.reshape(randu(0.5, 1.0, N*M), (N, M)))

fig_title = 'Signal Encoded Using Delayed IAF Encoder'
print fig_title
s_list = func_timer(iaf.iaf_encode_delay)(u_list, t_start, dt, b_list, d_list,
                                          k_list, a_list, w_list)

for i in xrange(M):
    for j in xrange(N):
        fig_title_out = fig_title + '\n(Signal #' + str(i+1) + \
                        ', Neuron #' + str(j+1) + ')'
        pl.plot_encoded(t_enc, u_list[i][k_start:k_end],
                        s_list[j][np.cumsum(s_list[j])<T],
                        fig_title_out,
                        output_name + str(output_count) + output_ext)
        output_count += 1

fig_title = 'Signal Decoded Using Delayed IAF Decoder'
print fig_title
u_rec_list = func_timer(iaf.iaf_decode_delay)(s_list, T, dt,
                                              b_list, d_list, k_list,
                                              a_list, w_list)

for i in xrange(M):
    fig_title_out = fig_title + ' (Signal #' + str(i+1) + ')'
    pl.plot_compare(t_enc, u_list[i][k_start:k_end],
                    u_rec_list[i][0:k_end-k_start], fig_title_out,
                    output_name + str(output_count) + output_ext)
    output_count += 1


# Author: Lars Buitinck <L.J.Buitinck@uva.nl>
# License: BSD 3 clause

import numbers

import numpy as np
import scipy.sparse as sp

from . import _hashing
from ..base import BaseEstimator, TransformerMixin

def _iteritems(d):
    """Like d.iteritems, but accepts any collections.Mapping."""
    return d.iteritems() if hasattr(d, "iteritems") else d.items()

class FeatureHasher(BaseEstimator, TransformerMixin):
    """Implements feature hashing, aka the hashing trick.

    This class turns sequences of symbolic feature names (strings) into
    scipy.sparse matrices, using a hash function to compute the matrix column
    corresponding to a name. The hash function employed is the signed 32-bit
    version of Murmurhash3.

    Feature names of type byte string are used as-is. Unicode strings are
    converted to UTF-8 first, but no Unicode normalization is done.
    Feature values must be (finite) numbers.

    This class is a low-memory alternative to DictVectorizer and
    CountVectorizer, intended for large-scale (online) learning and situations
    where memory is tight, e.g. when running prediction code on embedded
    devices.

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

    Parameters
    ----------
    n_features : integer, optional
        The number of features (columns) in the output matrices. Small numbers
        of features are likely to cause hash collisions, but large numbers
        will cause larger coefficient dimensions in linear learners.
    dtype : numpy type, optional
        The type of feature values. Passed to scipy.sparse matrix constructors
        as the dtype argument. Do not set this to bool, np.boolean or any
        unsigned integer type.
    input_type : string, optional
        Either "dict" (the default) to accept dictionaries over
        (feature_name, value); "pair" to accept pairs of (feature_name, value);
        or "string" to accept single strings.
        feature_name should be a string, while value should be a number.
        In the case of "string", a value of 1 is implied.
        The feature_name is hashed to find the appropriate column for the
        feature. The value's sign might be flipped in the output (but see
        non_negative, below).
    non_negative : boolean, optional, default np.float64
        Whether output matrices should contain non-negative values only;
        effectively calls abs on the matrix prior to returning it.
        When True, output values can be interpreted as frequencies.
        When False, output values will have expected value zero.

    Examples
    --------
    >>> from sklearn.feature_extraction import FeatureHasher
    >>> h = FeatureHasher(n_features=10)
    >>> D = [{'dog': 1, 'cat':2, 'elephant':4},{'dog': 2, 'run': 5}]
    >>> f = h.transform(D)
    >>> f.toarray()
    array([[ 0.,  0., -4., -1.,  0.,  0.,  0.,  0.,  0.,  2.],
           [ 0.,  0.,  0., -2., -5.,  0.,  0.,  0.,  0.,  0.]])

    See also
    --------
    DictVectorizer : vectorizes string-valued features using a hash table.
    sklearn.preprocessing.OneHotEncoder : handles nominal/categorical features
      encoded as columns of integers.
    """

    def __init__(self, n_features=(2 ** 20), input_type="dict",
                 dtype=np.float64, non_negative=False):
        self._validate_params(n_features, input_type)

        self.dtype = dtype
        self.input_type = input_type
        self.n_features = n_features
        self.non_negative = non_negative

    @staticmethod
    def _validate_params(n_features, input_type):
        # strangely, np.int16 instances are not instances of Integral,
        # while np.int64 instances are...
        if not isinstance(n_features, (numbers.Integral, np.integer)):
            raise TypeError("n_features must be integral, got %r (%s)."
                            % (n_features, type(n_features)))
        elif n_features < 1 or n_features >= 2 ** 31:
            raise ValueError("Invalid number of features (%d)." % n_features)

        if input_type not in ("dict", "pair", "string"):
            raise ValueError("input_type must be 'dict', 'pair' or 'string',"
                             " got %r." % input_type)

    def fit(self, X=None, y=None):
        """No-op.

        This method doesn't do anything. It exists purely for compatibility
        with the scikit-learn transformer API.

        Returns
        -------
        self : FeatureHasher

        """
        # repeat input validation for grid search (which calls set_params)
        self._validate_params(self.n_features, self.input_type)
        return self

    def transform(self, raw_X, y=None):
        """Transform a sequence of instances to a scipy.sparse matrix.

        Parameters
        ----------
        raw_X : iterable over iterable over raw features, length = n_samples
            Samples. Each sample must be iterable an (e.g., a list or tuple)
            containing/generating feature names (and optionally values, see
            the input_type constructor argument) which will be hashed.
            raw_X need not support the len function, so it can be the result
            of a generator; n_samples is determined on the fly.
        y : (ignored)

        Returns
        -------
        X : scipy.sparse matrix, shape = (n_samples, self.n_features)
            Feature matrix, for use with estimators or further transformers.

        """
        raw_X = iter(raw_X)
        if self.input_type == "dict":
            raw_X = (_iteritems(d) for d in raw_X)
        elif self.input_type == "string":
            raw_X = (((f, 1) for f in x) for x in raw_X)
        indices, indptr, values = \
            _hashing.transform(raw_X, self.n_features, self.dtype)
        n_samples = indptr.shape[0] - 1

        if n_samples == 0:
            raise ValueError("Cannot vectorize empty sequence.")

        X = sp.csr_matrix((values, indices, indptr), dtype=self.dtype,
                          shape=(n_samples, self.n_features))
        X.sum_duplicates()  # also sorts the indices
        if self.non_negative:
            np.abs(X.data, X.data)
        return X


#!/usr/bin/python
import h5py
import matplotlib.pylab as plt
import sys
import numpy as np;

################ Node # 2 Displacement #############################
#######################################
## Analytical Solution
#######################################

finput = h5py.File('Analytical_Solution.feioutput')

plt.figure()

# Read the time and displacement
times = finput["time"][:]
normal_strain = finput["/Model/Elements/Element_Outputs"][6,:]
normal_stress = -finput["/Model/Elements/Element_Outputs"][9,:];

# plt.plot(normal_strain,normal_stress,'-r',Linewidth=4,label='Analytical Solution')
plt.hold(True)

#######################################
## Current Solution
#######################################

# Go over each feioutput and plot each one.
thefile = "Monotonic_Contact_Behaviour_Adding_Normal_Load.h5.feioutput";
finput = h5py.File(thefile)

# Read the time and displacement
times = finput["time"][:]
normal_strain = finput["/Model/Elements/Element_Outputs"][6,:]
normal_stress = -finput["/Model/Elements/Element_Outputs"][9,:];

# Configure the figure filename, according to the input filename.
outfig=thefile.replace("_","-")
outfigname=outfig.replace("h5.feioutput","pdf")

# Plot the figure. Add labels and titles.
plt.plot(normal_strain,normal_stress,'-k',Linewidth=4,label='Numerical Solution')
plt.xlabel(r"Normal Strain $\epsilon$")
plt.ylabel(r"Normal Stress $\sigma$")
plt.legend()
plt.savefig("Contact_Normal_Interface_Behavour.pdf",  bbox_inches='tight')
# plt.show()
# #####################################################################