from statsmodels.compat.python import range
import numpy as np

def _make_index(prob,size):
    """
    Returns a boolean index for given probabilities.

    Notes
    ---------
    prob = [.75,.25] means that there is a 75% chance of the first column
    being True and a 25% chance of the second column being True. The
    columns are mutually exclusive.
    """
    rv = np.random.uniform(size=(size,1))
    cumprob = np.cumsum(prob)
    return np.logical_and(np.r_[0,cumprob[:-1]] <= rv, rv < cumprob)

def mixture_rvs(prob, size, dist, kwargs=None):
    """
    Sample from a mixture of distributions.

    Parameters
    ----------
    prob : array-like
        Probability of sampling from each distribution in dist
    size : int
        The length of the returned sample.
    dist : array-like
        An iterable of distributions objects from scipy.stats.
    kwargs : tuple of dicts, optional
        A tuple of dicts.  Each dict in kwargs can have keys loc, scale, and
        args to be passed to the respective distribution in dist.  If not
        provided, the distribution defaults are used.

    Examples
    --------
    Say we want 5000 random variables from mixture of normals with two
    distributions norm(-1,.5) and norm(1,.5) and we want to sample from the
    first with probability .75 and the second with probability .25.

    >>> from scipy import stats
    >>> prob = [.75,.25]
    >>> Y = mixture_rvs(prob, 5000, dist=[stats.norm, stats.norm], kwargs =
                (dict(loc=-1,scale=.5),dict(loc=1,scale=.5)))
    """
    if len(prob) != len(dist):
        raise ValueError("You must provide as many probabilities as distributions")
    if not np.allclose(np.sum(prob), 1):
        raise ValueError("prob does not sum to 1")

    if kwargs is None:
        kwargs = ({},)*len(prob)

    idx = _make_index(prob,size)
    sample = np.empty(size)
    for i in range(len(prob)):
        sample_idx = idx[...,i]
        sample_size = sample_idx.sum()
        loc = kwargs[i].get('loc',0)
        scale = kwargs[i].get('scale',1)
        args = kwargs[i].get('args',())
        sample[sample_idx] = dist[i].rvs(*args, **dict(loc=loc,scale=scale,
            size=sample_size))
    return sample

class MixtureDistribution(object):
    '''univariate mixture distribution

    for simple case for now (unbound support)
    does not yet inherit from scipy.stats.distributions

    adding pdf to mixture_rvs, some restrictions on broadcasting
    Currently it does not hold any state, all arguments included in each method.
    '''

    #def __init__(self, prob, size, dist, kwargs=None):

    def rvs(self, prob, size, dist, kwargs=None):
        return mixture_rvs(prob, size, dist, kwargs=kwargs)

    def pdf(self, x, prob, dist, kwargs=None):
        """
        pdf a mixture of distributions.

        Parameters
        ----------
        prob : array-like
            Probability of sampling from each distribution in dist
        dist : array-like
            An iterable of distributions objects from scipy.stats.
        kwargs : tuple of dicts, optional
            A tuple of dicts.  Each dict in kwargs can have keys loc, scale, and
            args to be passed to the respective distribution in dist.  If not
            provided, the distribution defaults are used.

        Examples
        --------
        Say we want 5000 random variables from mixture of normals with two
        distributions norm(-1,.5) and norm(1,.5) and we want to sample from the
        first with probability .75 and the second with probability .25.

        >>> from scipy import stats
        >>> prob = [.75,.25]
        >>> Y = mixture.pdf(x, prob, dist=[stats.norm, stats.norm], kwargs =
                    (dict(loc=-1,scale=.5),dict(loc=1,scale=.5)))
        """
        if len(prob) != len(dist):
            raise ValueError("You must provide as many probabilities as distributions")
        if not np.allclose(np.sum(prob), 1):
            raise ValueError("prob does not sum to 1")

        if kwargs is None:
            kwargs = ({},)*len(prob)

        for i in range(len(prob)):
            loc = kwargs[i].get('loc',0)
            scale = kwargs[i].get('scale',1)
            args = kwargs[i].get('args',())
            if i == 0:  #assume all broadcast the same as the first dist
                pdf_ = prob[i] * dist[i].pdf(x, *args, loc=loc, scale=scale)
            else:
                pdf_ += prob[i] * dist[i].pdf(x, *args, loc=loc, scale=scale)
        return pdf_

    def cdf(self, x, prob, dist, kwargs=None):
        """
        cdf of a mixture of distributions.

        Parameters
        ----------
        prob : array-like
            Probability of sampling from each distribution in dist
        size : int
            The length of the returned sample.
        dist : array-like
            An iterable of distributions objects from scipy.stats.
        kwargs : tuple of dicts, optional
            A tuple of dicts.  Each dict in kwargs can have keys loc, scale, and
            args to be passed to the respective distribution in dist.  If not
            provided, the distribution defaults are used.

        Examples
        --------
        Say we want 5000 random variables from mixture of normals with two
        distributions norm(-1,.5) and norm(1,.5) and we want to sample from the
        first with probability .75 and the second with probability .25.

        >>> from scipy import stats
        >>> prob = [.75,.25]
        >>> Y = mixture.pdf(x, prob, dist=[stats.norm, stats.norm], kwargs =
                    (dict(loc=-1,scale=.5),dict(loc=1,scale=.5)))
        """
        if len(prob) != len(dist):
            raise ValueError("You must provide as many probabilities as distributions")
        if not np.allclose(np.sum(prob), 1):
            raise ValueError("prob does not sum to 1")

        if kwargs is None:
            kwargs = ({},)*len(prob)

        for i in range(len(prob)):
            loc = kwargs[i].get('loc',0)
            scale = kwargs[i].get('scale',1)
            args = kwargs[i].get('args',())
            if i == 0:  #assume all broadcast the same as the first dist
                cdf_ = prob[i] * dist[i].cdf(x, *args, loc=loc, scale=scale)
            else:
                cdf_ += prob[i] * dist[i].cdf(x, *args, loc=loc, scale=scale)
        return cdf_

def mv_mixture_rvs(prob, size, dist, nvars, **kwargs):
    """
    Sample from a mixture of multivariate distributions.

    Parameters
    ----------
    prob : array-like
        Probability of sampling from each distribution in dist
    size : int
        The length of the returned sample.
    dist : array-like
        An iterable of distributions instances with callable method rvs.
    nvargs : int
        dimension of the multivariate distribution, could be inferred instead
    kwargs : tuple of dicts, optional
        ignored

    Examples
    --------
    Say we want 2000 random variables from mixture of normals with two
    multivariate normal distributions, and we want to sample from the
    first with probability .4 and the second with probability .6.

    import statsmodels.sandbox.distributions.mv_normal as mvd

    cov3 = np.array([[ 1.  ,  0.5 ,  0.75],
                       [ 0.5 ,  1.5 ,  0.6 ],
                       [ 0.75,  0.6 ,  2.  ]])

    mu = np.array([-1, 0.0, 2.0])
    mu2 = np.array([4, 2.0, 2.0])
    mvn3 = mvd.MVNormal(mu, cov3)
    mvn32 = mvd.MVNormal(mu2, cov3/2., 4)
    rvs = mix.mv_mixture_rvs([0.4, 0.6], 2000, [mvn3, mvn32], 3)

    """
    if len(prob) != len(dist):
        raise ValueError("You must provide as many probabilities as distributions")
    if not np.allclose(np.sum(prob), 1):
        raise ValueError("prob does not sum to 1")

    if kwargs is None:
        kwargs = ({},)*len(prob)

    idx = _make_index(prob,size)
    sample = np.empty((size, nvars))
    for i in range(len(prob)):
        sample_idx = idx[...,i]
        sample_size = sample_idx.sum()
        #loc = kwargs[i].get('loc',0)
        #scale = kwargs[i].get('scale',1)
        #args = kwargs[i].get('args',())
        # use int to avoid numpy bug with np.random.multivariate_normal
        sample[sample_idx] = dist[i].rvs(size=int(sample_size))
    return sample

if __name__ == '__main__':

    from scipy import stats

    obs_dist = mixture_rvs([.25,.75], size=10000, dist=[stats.norm, stats.beta],
                kwargs=(dict(loc=-1,scale=.5),dict(loc=1,scale=1,args=(1,.5))))

    nobs = 10000
    mix = MixtureDistribution()
##    mrvs = mixture_rvs([1/3.,2/3.], size=nobs, dist=[stats.norm, stats.norm],
##                   kwargs = (dict(loc=-1,scale=.5),dict(loc=1,scale=.75)))

    mix_kwds = (dict(loc=-1,scale=.25),dict(loc=1,scale=.75))
    mrvs = mix.rvs([1/3.,2/3.], size=nobs, dist=[stats.norm, stats.norm],
                   kwargs=mix_kwds)

    grid = np.linspace(-4,4, 100)
    mpdf = mix.pdf(grid, [1/3.,2/3.], dist=[stats.norm, stats.norm],
                   kwargs=mix_kwds)
    mcdf = mix.cdf(grid, [1/3.,2/3.], dist=[stats.norm, stats.norm],
                   kwargs=mix_kwds)

    doplot = 1
    if doplot:
        import matplotlib.pyplot as plt
        plt.figure()
        plt.hist(mrvs, bins=50, normed=True, color='red')
        plt.title('histogram of sample and pdf')
        plt.plot(grid, mpdf, lw=2, color='black')

        plt.figure()
        plt.hist(mrvs, bins=50, normed=True, cumulative=True, color='red')
        plt.title('histogram of sample and pdf')
        plt.plot(grid, mcdf, lw=2, color='black')

        plt.show()


#!/usr/bin/env python
"""
=============================================
Joint feature selection with multi-task Lasso
=============================================

The multi-task lasso allows to fit multiple regression problems
jointly enforcing the selected features to be the same across
tasks. This example simulates sequential measurements, each task
is a time instant, and the relevant features vary in amplitude
over time while being the same. The multi-task lasso imposes that
features that are selected at one time point are select for all time
point. This makes feature selection by the Lasso more stable.

"""
print(__doc__)

# Author: Alexandre Gramfort <alexandre.gramfort@inria.fr>
# License: BSD 3 clause

import matplotlib.pyplot as plt
import numpy as np

from sklearn.linear_model import MultiTaskLasso, Lasso

rng = np.random.RandomState(42)

# Generate some 2D coefficients with sine waves with random frequency and phase
n_samples, n_features, n_tasks = 100, 30, 40
n_relevant_features = 5
coef = np.zeros((n_tasks, n_features))
times = np.linspace(0, 2 * np.pi, n_tasks)
for k in range(n_relevant_features):
    coef[:, k] = np.sin((1. + rng.randn(1)) * times + 3 * rng.randn(1))

X = rng.randn(n_samples, n_features)
Y = np.dot(X, coef.T) + rng.randn(n_samples, n_tasks)

coef_lasso_ = np.array([Lasso(alpha=0.5).fit(X, y).coef_ for y in Y.T])
coef_multi_task_lasso_ = MultiTaskLasso(alpha=1.).fit(X, Y).coef_

###############################################################################
# Plot support and time series
fig = plt.figure(figsize=(8, 5))
plt.subplot(1, 2, 1)
plt.spy(coef_lasso_)
plt.xlabel('Feature')
plt.ylabel('Time (or Task)')
plt.text(10, 5, 'Lasso')
plt.subplot(1, 2, 2)
plt.spy(coef_multi_task_lasso_)
plt.xlabel('Feature')
plt.ylabel('Time (or Task)')
plt.text(10, 5, 'MultiTaskLasso')
fig.suptitle('Coefficient non-zero location')

feature_to_plot = 0
plt.figure()
lw = 2
plt.plot(coef[:, feature_to_plot], color='seagreen', linewidth=lw,
         label='Ground truth')
plt.plot(coef_lasso_[:, feature_to_plot], color='cornflowerblue', linewidth=lw,
         label='Lasso')
plt.plot(coef_multi_task_lasso_[:, feature_to_plot], color='gold', linewidth=lw,
         label='MultiTaskLasso')
plt.legend(loc='upper center')
plt.axis('tight')
plt.ylim([-1.1, 1.1])
plt.show()


"""
The :mod:`sklearn.covariance` module includes methods and algorithms to
robustly estimate the covariance of features given a set of points. The
precision matrix defined as the inverse of the covariance is also estimated.
Covariance estimation is closely related to the theory of Gaussian Graphical
Models.
"""

from .empirical_covariance_ import empirical_covariance, EmpiricalCovariance, \
    log_likelihood
from .shrunk_covariance_ import shrunk_covariance, ShrunkCovariance, \
    ledoit_wolf, ledoit_wolf_shrinkage, \
    LedoitWolf, oas, OAS
from .robust_covariance import fast_mcd, MinCovDet
from .graph_lasso_ import graph_lasso, GraphLasso, GraphLassoCV
from .outlier_detection import EllipticEnvelope

__all__ = ['EllipticEnvelope',
           'EmpiricalCovariance',
           'GraphLasso',
           'GraphLassoCV',
           'LedoitWolf',
           'MinCovDet',
           'OAS',
           'ShrunkCovariance',
           'empirical_covariance',
           'fast_mcd',
           'graph_lasso',
           'ledoit_wolf',
           'ledoit_wolf_shrinkage',
           'log_likelihood',
           'oas',
           'shrunk_covariance']


from .pls_ import _PLS

__all__ = ['CCA']

class CCA(_PLS):
    """CCA Canonical Correlation Analysis.

    CCA inherits from PLS with mode="B" and deflation_mode="canonical".

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

    Parameters
    ----------
    n_components : int, (default 2).
        number of components to keep.

    scale : boolean, (default True)
        whether to scale the data?

    max_iter : an integer, (default 500)
        the maximum number of iterations of the NIPALS inner loop

    tol : non-negative real, default 1e-06.
        the tolerance used in the iterative algorithm

    copy : boolean
        Whether the deflation be done on a copy. Let the default value
        to True unless you don't care about side effects

    Attributes
    ----------
    x_weights_ : array, [p, n_components]
        X block weights vectors.

    y_weights_ : array, [q, n_components]
        Y block weights vectors.

    x_loadings_ : array, [p, n_components]
        X block loadings vectors.

    y_loadings_ : array, [q, n_components]
        Y block loadings vectors.

    x_scores_ : array, [n_samples, n_components]
        X scores.

    y_scores_ : array, [n_samples, n_components]
        Y scores.

    x_rotations_ : array, [p, n_components]
        X block to latents rotations.

    y_rotations_ : array, [q, n_components]
        Y block to latents rotations.

    n_iter_ : array-like
        Number of iterations of the NIPALS inner loop for each
        component.

    Notes
    -----
    For each component k, find the weights u, v that maximizes
    max corr(Xk u, Yk v), such that ``|u| = |v| = 1``

    Note that it maximizes only the correlations between the scores.

    The residual matrix of X (Xk+1) block is obtained by the deflation on the
    current X score: x_score.

    The residual matrix of Y (Yk+1) block is obtained by deflation on the
    current Y score.

    Examples
    --------
    >>> from sklearn.cross_decomposition import CCA
    >>> X = [[0., 0., 1.], [1.,0.,0.], [2.,2.,2.], [3.,5.,4.]]
    >>> Y = [[0.1, -0.2], [0.9, 1.1], [6.2, 5.9], [11.9, 12.3]]
    >>> cca = CCA(n_components=1)
    >>> cca.fit(X, Y)
    ... # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    CCA(copy=True, max_iter=500, n_components=1, scale=True, tol=1e-06)
    >>> X_c, Y_c = cca.transform(X, Y)

    References
    ----------

    Jacob A. Wegelin. A survey of Partial Least Squares (PLS) methods, with
    emphasis on the two-block case. Technical Report 371, Department of
    Statistics, University of Washington, Seattle, 2000.

    In french but still a reference:
    Tenenhaus, M. (1998). La regression PLS: theorie et pratique. Paris:
    Editions Technic.

    See also
    --------
    PLSCanonical
    PLSSVD
    """

    def __init__(self, n_components=2, scale=True,
                 max_iter=500, tol=1e-06, copy=True):
        _PLS.__init__(self, n_components=n_components, scale=scale,
                      deflation_mode="canonical", mode="B",
                      norm_y_weights=True, algorithm="nipals",
                      max_iter=max_iter, tol=tol, copy=copy)


# Author: Alexander Fabisch  -- <afabisch@informatik.uni-bremen.de>
# Author: Christopher Moody <chrisemoody@gmail.com>
# Author: Nick Travers <nickt@squareup.com>
# License: BSD 3 clause (C) 2014

# This is the exact and Barnes-Hut t-SNE implementation. There are other
# modifications of the algorithm:
# * Fast Optimization for t-SNE:
#   http://cseweb.ucsd.edu/~lvdmaaten/workshops/nips2010/papers/vandermaaten.pdf

import numpy as np
from scipy import linalg
import scipy.sparse as sp
from scipy.spatial.distance import pdist
from scipy.spatial.distance import squareform
from ..neighbors import BallTree
from ..base import BaseEstimator
from ..utils import check_array
from ..utils import check_random_state
from ..utils.extmath import _ravel
from ..decomposition import RandomizedPCA
from ..metrics.pairwise import pairwise_distances
from . import _utils
from . import _barnes_hut_tsne
from ..utils.fixes import astype

MACHINE_EPSILON = np.finfo(np.double).eps

def _joint_probabilities(distances, desired_perplexity, verbose):
    """Compute joint probabilities p_ij from distances.

    Parameters
    ----------
    distances : array, shape (n_samples * (n_samples-1) / 2,)
        Distances of samples are stored as condensed matrices, i.e.
        we omit the diagonal and duplicate entries and store everything
        in a one-dimensional array.

    desired_perplexity : float
        Desired perplexity of the joint probability distributions.

    verbose : int
        Verbosity level.

    Returns
    -------
    P : array, shape (n_samples * (n_samples-1) / 2,)
        Condensed joint probability matrix.
    """
    # Compute conditional probabilities such that they approximately match
    # the desired perplexity
    distances = astype(distances, np.float32, copy=False)
    conditional_P = _utils._binary_search_perplexity(
        distances, None, desired_perplexity, verbose)
    P = conditional_P + conditional_P.T
    sum_P = np.maximum(np.sum(P), MACHINE_EPSILON)
    P = np.maximum(squareform(P) / sum_P, MACHINE_EPSILON)
    return P

def _joint_probabilities_nn(distances, neighbors, desired_perplexity, verbose):
    """Compute joint probabilities p_ij from distances using just nearest
    neighbors.

    This method is approximately equal to _joint_probabilities. The latter
    is O(N), but limiting the joint probability to nearest neighbors improves
    this substantially to O(uN).

    Parameters
    ----------
    distances : array, shape (n_samples * (n_samples-1) / 2,)
        Distances of samples are stored as condensed matrices, i.e.
        we omit the diagonal and duplicate entries and store everything
        in a one-dimensional array.

    desired_perplexity : float
        Desired perplexity of the joint probability distributions.

    verbose : int
        Verbosity level.

    Returns
    -------
    P : array, shape (n_samples * (n_samples-1) / 2,)
        Condensed joint probability matrix.
    """
    # Compute conditional probabilities such that they approximately match
    # the desired perplexity
    distances = astype(distances, np.float32, copy=False)
    neighbors = astype(neighbors, np.int64, copy=False)
    conditional_P = _utils._binary_search_perplexity(
        distances, neighbors, desired_perplexity, verbose)
    m = "All probabilities should be finite"
    assert np.all(np.isfinite(conditional_P)), m
    P = conditional_P + conditional_P.T
    sum_P = np.maximum(np.sum(P), MACHINE_EPSILON)
    P = np.maximum(squareform(P) / sum_P, MACHINE_EPSILON)
    assert np.all(np.abs(P) <= 1.0)
    return P

def _kl_divergence(params, P, degrees_of_freedom, n_samples, n_components,
                   skip_num_points=0):
    """t-SNE objective function: gradient of the KL divergence
    of p_ijs and q_ijs and the absolute error.

    Parameters
    ----------
    params : array, shape (n_params,)
        Unraveled embedding.

    P : array, shape (n_samples * (n_samples-1) / 2,)
        Condensed joint probability matrix.

    degrees_of_freedom : float
        Degrees of freedom of the Student's-t distribution.

    n_samples : int
        Number of samples.

    n_components : int
        Dimension of the embedded space.

    skip_num_points : int (optional, default:0)
        This does not compute the gradient for points with indices below
        `skip_num_points`. This is useful when computing transforms of new
        data where you'd like to keep the old data fixed.

    Returns
    -------
    kl_divergence : float
        Kullback-Leibler divergence of p_ij and q_ij.

    grad : array, shape (n_params,)
        Unraveled gradient of the Kullback-Leibler divergence with respect to
        the embedding.
    """
    X_embedded = params.reshape(n_samples, n_components)

    # Q is a heavy-tailed distribution: Student's t-distribution
    n = pdist(X_embedded, "sqeuclidean")
    n += 1.
    n /= degrees_of_freedom
    n **= (degrees_of_freedom + 1.0) / -2.0
    Q = np.maximum(n / (2.0 * np.sum(n)), MACHINE_EPSILON)

    # Optimization trick below: np.dot(x, y) is faster than
    # np.sum(x * y) because it calls BLAS

    # Objective: C (Kullback-Leibler divergence of P and Q)
    kl_divergence = 2.0 * np.dot(P, np.log(P / Q))

    # Gradient: dC/dY
    grad = np.ndarray((n_samples, n_components))
    PQd = squareform((P - Q) * n)
    for i in range(skip_num_points, n_samples):
        np.dot(_ravel(PQd[i]), X_embedded[i] - X_embedded, out=grad[i])
    grad = grad.ravel()
    c = 2.0 * (degrees_of_freedom + 1.0) / degrees_of_freedom
    grad *= c

    return kl_divergence, grad

def _kl_divergence_error(params, P, neighbors, degrees_of_freedom, n_samples,
                         n_components):
    """t-SNE objective function: the absolute error of the
    KL divergence of p_ijs and q_ijs.

    Parameters
    ----------
    params : array, shape (n_params,)
        Unraveled embedding.

    P : array, shape (n_samples * (n_samples-1) / 2,)
        Condensed joint probability matrix.

    neighbors : array (n_samples, K)
        The neighbors is not actually required to calculate the
        divergence, but is here to match the signature of the
        gradient function

    degrees_of_freedom : float
        Degrees of freedom of the Student's-t distribution.

    n_samples : int
        Number of samples.

    n_components : int
        Dimension of the embedded space.

    Returns
    -------
    kl_divergence : float
        Kullback-Leibler divergence of p_ij and q_ij.

    grad : array, shape (n_params,)
        Unraveled gradient of the Kullback-Leibler divergence with respect to
        the embedding.
    """
    X_embedded = params.reshape(n_samples, n_components)

    # Q is a heavy-tailed distribution: Student's t-distribution
    n = pdist(X_embedded, "sqeuclidean")
    n += 1.
    n /= degrees_of_freedom
    n **= (degrees_of_freedom + 1.0) / -2.0
    Q = np.maximum(n / (2.0 * np.sum(n)), MACHINE_EPSILON)

    # Optimization trick below: np.dot(x, y) is faster than
    # np.sum(x * y) because it calls BLAS

    # Objective: C (Kullback-Leibler divergence of P and Q)
    if len(P.shape) == 2:
        P = squareform(P)
    kl_divergence = 2.0 * np.dot(P, np.log(P / Q))

    return kl_divergence

def _kl_divergence_bh(params, P, neighbors, degrees_of_freedom, n_samples,
                      n_components, angle=0.5, skip_num_points=0,
                      verbose=False):
    """t-SNE objective function: KL divergence of p_ijs and q_ijs.

    Uses Barnes-Hut tree methods to calculate the gradient that
    runs in O(NlogN) instead of O(N^2)

    Parameters
    ----------
    params : array, shape (n_params,)
        Unraveled embedding.

    P : array, shape (n_samples * (n_samples-1) / 2,)
        Condensed joint probability matrix.

    neighbors: int64 array, shape (n_samples, K)
        Array with element [i, j] giving the index for the jth
        closest neighbor to point i.

    degrees_of_freedom : float
        Degrees of freedom of the Student's-t distribution.

    n_samples : int
        Number of samples.

    n_components : int
        Dimension of the embedded space.

    angle : float (default: 0.5)
        This is the trade-off between speed and accuracy for Barnes-Hut T-SNE.
        'angle' is the angular size (referred to as theta in [3]) of a distant
        node as measured from a point. If this size is below 'angle' then it is
        used as a summary node of all points contained within it.
        This method is not very sensitive to changes in this parameter
        in the range of 0.2 - 0.8. Angle less than 0.2 has quickly increasing
        computation time and angle greater 0.8 has quickly increasing error.

    skip_num_points : int (optional, default:0)
        This does not compute the gradient for points with indices below
        `skip_num_points`. This is useful when computing transforms of new
        data where you'd like to keep the old data fixed.

    verbose : int
        Verbosity level.

    Returns
    -------
    kl_divergence : float
        Kullback-Leibler divergence of p_ij and q_ij.

    grad : array, shape (n_params,)
        Unraveled gradient of the Kullback-Leibler divergence with respect to
        the embedding.
    """
    params = astype(params, np.float32, copy=False)
    X_embedded = params.reshape(n_samples, n_components)
    neighbors = astype(neighbors, np.int64, copy=False)
    if len(P.shape) == 1:
        sP = squareform(P).astype(np.float32)
    else:
        sP = P.astype(np.float32)

    grad = np.zeros(X_embedded.shape, dtype=np.float32)
    error = _barnes_hut_tsne.gradient(sP, X_embedded, neighbors,
                                      grad, angle, n_components, verbose,
                                      dof=degrees_of_freedom)
    c = 2.0 * (degrees_of_freedom + 1.0) / degrees_of_freedom
    grad = grad.ravel()
    grad *= c

    return error, grad

def _gradient_descent(objective, p0, it, n_iter, objective_error=None,
                      n_iter_check=1, n_iter_without_progress=50,
                      momentum=0.5, learning_rate=1000.0, min_gain=0.01,
                      min_grad_norm=1e-7, min_error_diff=1e-7, verbose=0,
                      args=None, kwargs=None):
    """Batch gradient descent with momentum and individual gains.

    Parameters
    ----------
    objective : function or callable
        Should return a tuple of cost and gradient for a given parameter
        vector. When expensive to compute, the cost can optionally
        be None and can be computed every n_iter_check steps using
        the objective_error function.

    p0 : array-like, shape (n_params,)
        Initial parameter vector.

    it : int
        Current number of iterations (this function will be called more than
        once during the optimization).

    n_iter : int
        Maximum number of gradient descent iterations.

    n_iter_check : int
        Number of iterations before evaluating the global error. If the error
        is sufficiently low, we abort the optimization.

    objective_error : function or callable
        Should return a tuple of cost and gradient for a given parameter
        vector.

    n_iter_without_progress : int, optional (default: 30)
        Maximum number of iterations without progress before we abort the
        optimization.

    momentum : float, within (0.0, 1.0), optional (default: 0.5)
        The momentum generates a weight for previous gradients that decays
        exponentially.

    learning_rate : float, optional (default: 1000.0)
        The learning rate should be extremely high for t-SNE! Values in the
        range [100.0, 1000.0] are common.

    min_gain : float, optional (default: 0.01)
        Minimum individual gain for each parameter.

    min_grad_norm : float, optional (default: 1e-7)
        If the gradient norm is below this threshold, the optimization will
        be aborted.

    min_error_diff : float, optional (default: 1e-7)
        If the absolute difference of two successive cost function values
        is below this threshold, the optimization will be aborted.

    verbose : int, optional (default: 0)
        Verbosity level.

    args : sequence
        Arguments to pass to objective function.

    kwargs : dict
        Keyword arguments to pass to objective function.

    Returns
    -------
    p : array, shape (n_params,)
        Optimum parameters.

    error : float
        Optimum.

    i : int
        Last iteration.
    """
    if args is None:
        args = []
    if kwargs is None:
        kwargs = {}

    p = p0.copy().ravel()
    update = np.zeros_like(p)
    gains = np.ones_like(p)
    error = np.finfo(np.float).max
    best_error = np.finfo(np.float).max
    best_iter = 0

    for i in range(it, n_iter):
        new_error, grad = objective(p, *args, **kwargs)
        grad_norm = linalg.norm(grad)

        inc = update * grad >= 0.0
        dec = np.invert(inc)
        gains[inc] += 0.05
        gains[dec] *= 0.95
        np.clip(gains, min_gain, np.inf)
        grad *= gains
        update = momentum * update - learning_rate * grad
        p += update

        if (i + 1) % n_iter_check == 0:
            if new_error is None:
                new_error = objective_error(p, *args)
            error_diff = np.abs(new_error - error)
            error = new_error

            if verbose >= 2:
                m = "[t-SNE] Iteration %d: error = %.7f, gradient norm = %.7f"
                print(m % (i + 1, error, grad_norm))

            if error < best_error:
                best_error = error
                best_iter = i
            elif i - best_iter > n_iter_without_progress:
                if verbose >= 2:
                    print("[t-SNE] Iteration %d: did not make any progress "
                          "during the last %d episodes. Finished."
                          % (i + 1, n_iter_without_progress))
                break
            if grad_norm <= min_grad_norm:
                if verbose >= 2:
                    print("[t-SNE] Iteration %d: gradient norm %f. Finished."
                          % (i + 1, grad_norm))
                break
            if error_diff <= min_error_diff:
                if verbose >= 2:
                    m = "[t-SNE] Iteration %d: error difference %f. Finished."
                    print(m % (i + 1, error_diff))
                break

        if new_error is not None:
            error = new_error

    return p, error, i

def trustworthiness(X, X_embedded, n_neighbors=5, precomputed=False):
    """Expresses to what extent the local structure is retained.

    The trustworthiness is within [0, 1]. It is defined as

    .. math::

        T(k) = 1 - \frac{2}{nk (2n - 3k - 1)} \sum^n_{i=1}
            \sum_{j \in U^{(k)}_i (r(i, j) - k)}

    where :math:`r(i, j)` is the rank of the embedded datapoint j
    according to the pairwise distances between the embedded datapoints,
    :math:`U^{(k)}_i` is the set of points that are in the k nearest
    neighbors in the embedded space but not in the original space.

    * "Neighborhood Preservation in Nonlinear Projection Methods: An
      Experimental Study"
      J. Venna, S. Kaski
    * "Learning a Parametric Embedding by Preserving Local Structure"
      L.J.P. van der Maaten

    Parameters
    ----------
    X : array, shape (n_samples, n_features) or (n_samples, n_samples)
        If the metric is 'precomputed' X must be a square distance
        matrix. Otherwise it contains a sample per row.

    X_embedded : array, shape (n_samples, n_components)
        Embedding of the training data in low-dimensional space.

    n_neighbors : int, optional (default: 5)
        Number of neighbors k that will be considered.

    precomputed : bool, optional (default: False)
        Set this flag if X is a precomputed square distance matrix.

    Returns
    -------
    trustworthiness : float
        Trustworthiness of the low-dimensional embedding.
    """
    if precomputed:
        dist_X = X
    else:
        dist_X = pairwise_distances(X, squared=True)
    dist_X_embedded = pairwise_distances(X_embedded, squared=True)
    ind_X = np.argsort(dist_X, axis=1)
    ind_X_embedded = np.argsort(dist_X_embedded, axis=1)[:, 1:n_neighbors + 1]

    n_samples = X.shape[0]
    t = 0.0
    ranks = np.zeros(n_neighbors)
    for i in range(n_samples):
        for j in range(n_neighbors):
            ranks[j] = np.where(ind_X[i] == ind_X_embedded[i, j])[0][0]
        ranks -= n_neighbors
        t += np.sum(ranks[ranks > 0])
    t = 1.0 - t * (2.0 / (n_samples * n_neighbors *
                          (2.0 * n_samples - 3.0 * n_neighbors - 1.0)))
    return t

class TSNE(BaseEstimator):
    """t-distributed Stochastic Neighbor Embedding.

    t-SNE [1] is a tool to visualize high-dimensional data. It converts
    similarities between data points to joint probabilities and tries
    to minimize the Kullback-Leibler divergence between the joint
    probabilities of the low-dimensional embedding and the
    high-dimensional data. t-SNE has a cost function that is not convex,
    i.e. with different initializations we can get different results.

    It is highly recommended to use another dimensionality reduction
    method (e.g. PCA for dense data or TruncatedSVD for sparse data)
    to reduce the number of dimensions to a reasonable amount (e.g. 50)
    if the number of features is very high. This will suppress some
    noise and speed up the computation of pairwise distances between
    samples. For more tips see Laurens van der Maaten's FAQ [2].

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

    Parameters
    ----------
    n_components : int, optional (default: 2)
        Dimension of the embedded space.

    perplexity : float, optional (default: 30)
        The perplexity is related to the number of nearest neighbors that
        is used in other manifold learning algorithms. Larger datasets
        usually require a larger perplexity. Consider selecting a value
        between 5 and 50. The choice is not extremely critical since t-SNE
        is quite insensitive to this parameter.

    early_exaggeration : float, optional (default: 4.0)
        Controls how tight natural clusters in the original space are in
        the embedded space and how much space will be between them. For
        larger values, the space between natural clusters will be larger
        in the embedded space. Again, the choice of this parameter is not
        very critical. If the cost function increases during initial
        optimization, the early exaggeration factor or the learning rate
        might be too high.

    learning_rate : float, optional (default: 1000)
        The learning rate can be a critical parameter. It should be
        between 100 and 1000. If the cost function increases during initial
        optimization, the early exaggeration factor or the learning rate
        might be too high. If the cost function gets stuck in a bad local
        minimum increasing the learning rate helps sometimes.

    n_iter : int, optional (default: 1000)
        Maximum number of iterations for the optimization. Should be at
        least 200.

    n_iter_without_progress : int, optional (default: 30)
        Maximum number of iterations without progress before we abort the
        optimization.

        .. versionadded:: 0.17
           parameter *n_iter_without_progress* to control stopping criteria.

    min_grad_norm : float, optional (default: 1E-7)
        If the gradient norm is below this threshold, the optimization will
        be aborted.

    metric : string or callable, optional
        The metric to use when calculating distance between instances in a
        feature array. If metric is a string, it must be one of the options
        allowed by scipy.spatial.distance.pdist for its metric parameter, or
        a metric listed in pairwise.PAIRWISE_DISTANCE_FUNCTIONS.
        If metric is "precomputed", X is assumed to be a distance matrix.
        Alternatively, if metric is a callable function, it is called on each
        pair of instances (rows) and the resulting value recorded. The callable
        should take two arrays from X as input and return a value indicating
        the distance between them. The default is "euclidean" which is
        interpreted as squared euclidean distance.

    init : string, optional (default: "random")
        Initialization of embedding. Possible options are 'random' and 'pca'.
        PCA initialization cannot be used with precomputed distances and is
        usually more globally stable than random initialization.

    verbose : int, optional (default: 0)
        Verbosity level.

    random_state : int or RandomState instance or None (default)
        Pseudo Random Number generator seed control. If None, use the
        numpy.random singleton. Note that different initializations
        might result in different local minima of the cost function.

    method : string (default: 'barnes_hut')
        By default the gradient calculation algorithm uses Barnes-Hut
        approximation running in O(NlogN) time. method='exact'
        will run on the slower, but exact, algorithm in O(N^2) time. The
        exact algorithm should be used when nearest-neighbor errors need
        to be better than 3%. However, the exact method cannot scale to
        millions of examples.

        .. versionadded:: 0.17
           Approximate optimization *method* via the Barnes-Hut.

    angle : float (default: 0.5)
        Only used if method='barnes_hut'
        This is the trade-off between speed and accuracy for Barnes-Hut T-SNE.
        'angle' is the angular size (referred to as theta in [3]) of a distant
        node as measured from a point. If this size is below 'angle' then it is
        used as a summary node of all points contained within it.
        This method is not very sensitive to changes in this parameter
        in the range of 0.2 - 0.8. Angle less than 0.2 has quickly increasing
        computation time and angle greater 0.8 has quickly increasing error.

    Attributes
    ----------
    embedding_ : array-like, shape (n_samples, n_components)
        Stores the embedding vectors.

    kl_divergence_ : float
        Kullback-Leibler divergence after optimization.

    Examples
    --------

    >>> import numpy as np
    >>> from sklearn.manifold import TSNE
    >>> X = np.array([[0, 0, 0], [0, 1, 1], [1, 0, 1], [1, 1, 1]])
    >>> model = TSNE(n_components=2, random_state=0)
    >>> np.set_printoptions(suppress=True)
    >>> model.fit_transform(X) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
    array([[ 0.00017599,  0.00003993],
           [ 0.00009891,  0.00021913],
           [ 0.00018554, -0.00009357],
           [ 0.00009528, -0.00001407]])

    References
    ----------

    [1] van der Maaten, L.J.P.; Hinton, G.E. Visualizing High-Dimensional Data
        Using t-SNE. Journal of Machine Learning Research 9:2579-2605, 2008.

    [2] van der Maaten, L.J.P. t-Distributed Stochastic Neighbor Embedding
        http://homepage.tudelft.nl/19j49/t-SNE.html

    [3] L.J.P. van der Maaten. Accelerating t-SNE using Tree-Based Algorithms.
        Journal of Machine Learning Research 15(Oct):3221-3245, 2014.
        http://lvdmaaten.github.io/publications/papers/JMLR_2014.pdf
    """

    def __init__(self, n_components=2, perplexity=30.0,
                 early_exaggeration=4.0, learning_rate=1000.0, n_iter=1000,
                 n_iter_without_progress=30, min_grad_norm=1e-7,
                 metric="euclidean", init="random", verbose=0,
                 random_state=None, method='barnes_hut', angle=0.5):
        if init not in ["pca", "random"] or isinstance(init, np.ndarray):
            msg = "'init' must be 'pca', 'random' or a NumPy array"
            raise ValueError(msg)
        self.n_components = n_components
        self.perplexity = perplexity
        self.early_exaggeration = early_exaggeration
        self.learning_rate = learning_rate
        self.n_iter = n_iter
        self.n_iter_without_progress = n_iter_without_progress
        self.min_grad_norm = min_grad_norm
        self.metric = metric
        self.init = init
        self.verbose = verbose
        self.random_state = random_state
        self.method = method
        self.angle = angle
        self.embedding_ = None

    def _fit(self, X, skip_num_points=0):
        """Fit the model using X as training data.

        Note that sparse arrays can only be handled by method='exact'.
        It is recommended that you convert your sparse array to dense
        (e.g. `X.toarray()`) if it fits in memory, or otherwise using a
        dimensionality reduction technique (e.g. TruncatedSVD).

        Parameters
        ----------
        X : array, shape (n_samples, n_features) or (n_samples, n_samples)
            If the metric is 'precomputed' X must be a square distance
            matrix. Otherwise it contains a sample per row. Note that this
            when method='barnes_hut', X cannot be a sparse array and if need be
            will be converted to a 32 bit float array. Method='exact' allows
            sparse arrays and 64bit floating point inputs.

        skip_num_points : int (optional, default:0)
            This does not compute the gradient for points with indices below
            `skip_num_points`. This is useful when computing transforms of new
            data where you'd like to keep the old data fixed.
        """
        if self.method not in ['barnes_hut', 'exact']:
            raise ValueError("'method' must be 'barnes_hut' or 'exact'")
        if self.angle < 0.0 or self.angle > 1.0:
            raise ValueError("'angle' must be between 0.0 - 1.0")
        if self.method == 'barnes_hut' and sp.issparse(X):
            raise TypeError('A sparse matrix was passed, but dense '
                            'data is required for method="barnes_hut". Use '
                            'X.toarray() to convert to a dense numpy array if '
                            'the array is small enough for it to fit in '
                            'memory. Otherwise consider dimensionality '
                            'reduction techniques (e.g. TruncatedSVD)')
        else:
            X = check_array(X, accept_sparse=['csr', 'csc', 'coo'], dtype=np.float64)
        random_state = check_random_state(self.random_state)

        if self.early_exaggeration < 1.0:
            raise ValueError("early_exaggeration must be at least 1, but is "
                             "%f" % self.early_exaggeration)

        if self.n_iter < 200:
            raise ValueError("n_iter should be at least 200")

        if self.metric == "precomputed":
            if self.init == 'pca':
                raise ValueError("The parameter init=\"pca\" cannot be used "
                                 "with metric=\"precomputed\".")
            if X.shape[0] != X.shape[1]:
                raise ValueError("X should be a square distance matrix")
            distances = X
        else:
            if self.verbose:
                print("[t-SNE] Computing pairwise distances...")

            if self.metric == "euclidean":
                distances = pairwise_distances(X, metric=self.metric,
                                               squared=True)
            else:
                distances = pairwise_distances(X, metric=self.metric)

        if not np.all(distances >= 0):
            raise ValueError("All distances should be positive, either "
                             "the metric or precomputed distances given "
                             "as X are not correct")

        # Degrees of freedom of the Student's t-distribution. The suggestion
        # degrees_of_freedom = n_components - 1 comes from
        # "Learning a Parametric Embedding by Preserving Local Structure"
        # Laurens van der Maaten, 2009.
        degrees_of_freedom = max(self.n_components - 1.0, 1)
        n_samples = X.shape[0]
        # the number of nearest neighbors to find
        k = min(n_samples - 1, int(3. * self.perplexity + 1))

        neighbors_nn = None
        if self.method == 'barnes_hut':
            if self.verbose:
                print("[t-SNE] Computing %i nearest neighbors..." % k)
            if self.metric == 'precomputed':
                # Use the precomputed distances to find
                # the k nearest neighbors and their distances
                neighbors_nn = np.argsort(distances, axis=1)[:, :k]
            else:
                # Find the nearest neighbors for every point
                bt = BallTree(X)
                # LvdM uses 3 * perplexity as the number of neighbors
                # And we add one to not count the data point itself
                # In the event that we have very small # of points
                # set the neighbors to n - 1
                distances_nn, neighbors_nn = bt.query(X, k=k + 1)
                neighbors_nn = neighbors_nn[:, 1:]
            P = _joint_probabilities_nn(distances, neighbors_nn,
                                        self.perplexity, self.verbose)
        else:
            P = _joint_probabilities(distances, self.perplexity, self.verbose)
        assert np.all(np.isfinite(P)), "All probabilities should be finite"
        assert np.all(P >= 0), "All probabilities should be zero or positive"
        assert np.all(P <= 1), ("All probabilities should be less "
                                "or then equal to one")

        if self.init == 'pca':
            pca = RandomizedPCA(n_components=self.n_components,
                                random_state=random_state)
            X_embedded = pca.fit_transform(X)
        elif isinstance(self.init, np.ndarray):
            X_embedded = self.init
        elif self.init == 'random':
            X_embedded = None
        else:
            raise ValueError("Unsupported initialization scheme: %s"
                             % self.init)

        return self._tsne(P, degrees_of_freedom, n_samples, random_state,
                          X_embedded=X_embedded,
                          neighbors=neighbors_nn,
                          skip_num_points=skip_num_points)

    def _tsne(self, P, degrees_of_freedom, n_samples, random_state,
              X_embedded=None, neighbors=None, skip_num_points=0):
        """Runs t-SNE."""
        # t-SNE minimizes the Kullback-Leiber divergence of the Gaussians P
        # and the Student's t-distributions Q. The optimization algorithm that
        # we use is batch gradient descent with three stages:
        # * early exaggeration with momentum 0.5
        # * early exaggeration with momentum 0.8
        # * final optimization with momentum 0.8
        # The embedding is initialized with iid samples from Gaussians with
        # standard deviation 1e-4.

        if X_embedded is None:
            # Initialize embedding randomly
            X_embedded = 1e-4 * random_state.randn(n_samples,
                                                   self.n_components)
        params = X_embedded.ravel()

        opt_args = {}
        opt_args = {"n_iter": 50, "momentum": 0.5, "it": 0,
                    "learning_rate": self.learning_rate,
                    "verbose": self.verbose, "n_iter_check": 25,
                    "kwargs": dict(skip_num_points=skip_num_points)}
        if self.method == 'barnes_hut':
            m = "Must provide an array of neighbors to use Barnes-Hut"
            assert neighbors is not None, m
            obj_func = _kl_divergence_bh
            objective_error = _kl_divergence_error
            sP = squareform(P).astype(np.float32)
            neighbors = neighbors.astype(np.int64)
            args = [sP, neighbors, degrees_of_freedom, n_samples,
                    self.n_components]
            opt_args['args'] = args
            opt_args['min_grad_norm'] = 1e-3
            opt_args['n_iter_without_progress'] = 30
            # Don't always calculate the cost since that calculation
            # can be nearly as expensive as the gradient
            opt_args['objective_error'] = objective_error
            opt_args['kwargs']['angle'] = self.angle
            opt_args['kwargs']['verbose'] = self.verbose
        else:
            obj_func = _kl_divergence
            opt_args['args'] = [P, degrees_of_freedom, n_samples,
                                self.n_components]
            opt_args['min_error_diff'] = 0.0
            opt_args['min_grad_norm'] = 0.0

        # Early exaggeration
        P *= self.early_exaggeration

        params, kl_divergence, it = _gradient_descent(obj_func, params,
                                                      **opt_args)
        opt_args['n_iter'] = 100
        opt_args['momentum'] = 0.8
        opt_args['it'] = it + 1
        params, kl_divergence, it = _gradient_descent(obj_func, params,
                                                      **opt_args)
        if self.verbose:
            print("[t-SNE] KL divergence after %d iterations with early "
                  "exaggeration: %f" % (it + 1, kl_divergence))
        # Save the final number of iterations
        self.n_iter_final = it

        # Final optimization
        P /= self.early_exaggeration
        opt_args['n_iter'] = self.n_iter
        opt_args['it'] = it + 1
        params, error, it = _gradient_descent(obj_func, params, **opt_args)

        if self.verbose:
            print("[t-SNE] Error after %d iterations: %f"
                  % (it + 1, kl_divergence))

        X_embedded = params.reshape(n_samples, self.n_components)
        self.kl_divergence_ = kl_divergence

        return X_embedded

    def fit_transform(self, X, y=None):
        """Fit X into an embedded space and return that transformed
        output.

        Parameters
        ----------
        X : array, shape (n_samples, n_features) or (n_samples, n_samples)
            If the metric is 'precomputed' X must be a square distance
            matrix. Otherwise it contains a sample per row.

        Returns
        -------
        X_new : array, shape (n_samples, n_components)
            Embedding of the training data in low-dimensional space.
        """
        embedding = self._fit(X)
        self.embedding_ = embedding
        return self.embedding_

    def fit(self, X, y=None):
        """Fit X into an embedded space.

        Parameters
        ----------
        X : array, shape (n_samples, n_features) or (n_samples, n_samples)
            If the metric is 'precomputed' X must be a square distance
            matrix. Otherwise it contains a sample per row. If the method
            is 'exact', X may be a sparse matrix of type 'csr', 'csc'
            or 'coo'.
        """
        self.fit_transform(X)
        return self


from rpy2.robjects import Formula
from rpy2.robjects import pandas2ri
from rpy2.robjects.packages import importr

class LM(object):
  '''Class for fitting (simple) linear models using rpy2.  When extracting
the coefficients for a model (lmerMod or glmerMod) that uses orthogonal
polynomials (poly in R syntax), it is necessary to fit a linear model
that maps from the original data to the fitted polynomial.  The mermod
class uses this class to fir such linear models.

  '''
  def __init__(self, formula=None, response=None, predictors=None):
    self.__stats = importr('stats')
    self.formula = formula
    self.response = response
    self.predictors = predictors
    self._coefs = None

  @property
  def formula(self):
    return self._formula

  @formula.setter
  def formula(self, f):
    self._formula = Formula(f)
    self.env = self._formula.environment

  def fit(self):
    '''Fit the linear model and extract the coefficients.
FIXME: This function assumes the model has a single predictor variable (x), but may appear multiple times with different exponents.  That is, the equation must be of the form

    y ~ x + I(x^2) + I(x^3)
'''
    if self.formula is None or self.response is None or self.predictors is None:
      raise RuntimeError('set formula, response, and predictor variables')
    ## FIXME: This is a quick and dirty hack.
    self.env['y'] = self.response
    self.env['x'] = self.predictors.loc[:, 'x']
    fit = self.__stats.lm(self.formula)
    self._coefs = pandas2ri.ri2py(fit.rx('coefficients')[0])

  @property
  def coefs(self):
    if self._coefs == None:
      self.fit()
    return self._coefs


# force floating point division. Can still use integer with //
from __future__ import division
# This file is used for importing the common utilities classes.
import numpy as np
import matplotlib.pyplot as plt
# import the patrick-specific utilities
import GenUtilities  as pGenUtil
import PlotUtilities as pPlotUtil
import CheckpointUtilities as pCheckUtil
from scipy.signal import savgol_filter
DEF_FILTER_CONST = 0.005 # 0.5%

BASE_GROUP = "/Volumes/group/4Patrick/"
SUBDIR_BINARIES = "PRH_AFM_Databases/BinaryFilesTimeSeparationForce/"

def getDatabaseFolder():
    """
    Returns the location of the database binary folder location

    Args:
        None
    Returns:
        Where the database is, as a string.
    """
    # XXX TODO: right now assumes mac-style mounting...
    return BASE_GROUP + SUBDIR_BINARIES

def getDatabaseFile(fileName,extension=".hdf"):
    """
    Returns the absolute path to a previously-saved file with the given filename
    Path is *not* guaranteed to exist, if the file hasn't been saved already.

    Args:
        fileName: the name of the file (usually according to the "TraceData"
        table, field "FileTimSepFor")

        extension: the recquired extension
    Returns:
        Where the file is located, an absolute path. Doesn't guarantee the file
        *does* exist, just that *if* it does, it would be there.
    """
    fileWithExt = pGenUtil.ensureEnds(fileName,extension)
    return  getDatabaseFolder() + fileWithExt

def DemoDir():
    '''
    :return: the absolute path to the demo directory
    '''
    return BASE_GROUP + "DemoData/IgorDemos/"

# all demo directories should have an input and output directory
def GetDemoInOut(demoName,baseDir=DemoDir(),raiseOnError=True):
    """
    Returns the demo input and output directories, given a path baseDir and
    name demoName. Recquires files to exist at "<baseDir><demoName>". If
    encountering an error (e.g. permissions, something isn't mounted), raises
    an error.

    Args:
        demoName: The name of the demo. Assumed to be the subdir under "basedir"
        we want to use

        baseDir: the base directory. Input and output directories are
        "<baseDir><demoName>Input/" and "<baseDir><demoName>Output/", resp.

        raiseOnError : if true, raises an error on an OS. otherwise, just
        prints a warning that something went wrong.
    Returns:
        tuple of <inputDir>,<outputDir>
    """
    fullBase =  baseDir + demoName
    inputV = pGenUtil.getSanitaryPath(fullBase + "/Input/")
    outputV = pGenUtil.getSanitaryPath(fullBase + "/Output/")
    try:
        pGenUtil.ensureDirExists(inputV)
        pGenUtil.ensureDirExists(outputV)
    except OSError as e:
        if (raiseOnError):
            raise(e)
        print("Warning, couldn't open demo directories based in " + fullBase +
              ". Most likely, not connected to JILA network")
    return inputV,outputV

def DemoJilaOrLocal(demoName,localPath):
    """
    Looks for the demo dir in the default (jila-hosted) space. If nothing is
    found, looks in the paths specified by localpath (where it puts input
    and output directories according to its name)

    Args:
        demoName: see GetDemoInOut

        localPath: equivalent of baseDir in GetDemoInOut. Where we put the input        and Output directories for the unit test if JILA can't be found.

    Returns:
        tuple of <inputDir>,<outputDir>
    """
    inDir,outDir = GetDemoInOut(demoName,raiseOnError=False)
    if (not pGenUtil.dirExists(inDir)):
        print("Warning: Couldn't connect to JILA's Network. Using local data.")
        # get "sanitary paths" which as OS-indepdent (in theory..)
        localPath = pGenUtil.ensureEnds(localPath,"/")
        inDir = pGenUtil.getSanitaryPath(localPath)
        outDir = pGenUtil.getSanitaryPath(localPath + "Output" + demoName +"/")
        pGenUtil.ensureDirExists(outDir)
        if (not pGenUtil.dirExists(inDir)):
            # whoops...
            raise IOError("Demo Directory {:s} not found anywhere.".\
                          format(inDir))
    return inDir,outDir

# read a txt or similarly formatted file
def readIgorWave(mFile,skip_header=3,skip_footer=1,comments="X "):
    data = np.genfromtxt(mFile,comments=comments,skip_header=skip_header,
                         skip_footer=skip_footer)
    return data

def savitskyFilter(inData,nSmooth = None,degree=2):
    if (nSmooth is None):
        nSmooth = int(len(inData)/200)
    # POST: have an nSmooth
    if (nSmooth % 2 == 0):
        # must be odd
        nSmooth += 1
    # get the filtered version of the data
    return savgol_filter(inData,nSmooth,degree)

def SplitIntoApproachAndRetract(sep,force,sepToSplit=None):
    '''
    Given a full force/sep curve, returns the approach/retract
    according to before/after sepToSplit, cutting out the surface (assumed
    at minimm separation )
    :param sep: the separation, units not important. minimum is surface
    :param force: the force, units not important
    :param sepToSplot: the separation where we think the surface is. same units
    as sep
    '''
    # find where sep is closest to sepToSplit before/after minIdx (surface)
    if (sepToSplit is None):
        sepToSplit = np.min(sep)
    surfIdx = np.argmin(sep)
    sepAppr = sep[:surfIdx]
    sepRetr = sep[surfIdx:]
    apprIdx = np.argmin(np.abs(sepAppr-sepToSplit))
    retrIdx = surfIdx + np.argmin(np.abs(sepRetr-sepToSplit))
    forceAppr = force[:apprIdx]
    forceRetr = force[retrIdx:]
    sepAppr = sep[:apprIdx]
    sepRetr = sep[retrIdx:]
    return sepAppr,sepRetr,forceAppr,forceRetr

def NormalizeSepForce(sep,force,surfIdx=None,normalizeSep=True,
                      normalizeFor=True,sensibleUnits=True):
    if (sensibleUnits):
        sepUnits = sep * 1e9
        forceUnits = force * 1e12
    else:
        sepUnits = sep
        forceUnits= force
    if (surfIdx is None):
        surfIdx = np.argmin(sep)
    if (normalizeSep):
        sepUnits -= sepUnits[surfIdx]
    if (normalizeFor):
        # reverse: sort low to high
        sortIdx = np.argsort(sep)[::-1]
        # get the percentage of points we want
        percent = 0.05
        nPoints = int(percent*sortIdx.size)
        idxForMedian = sortIdx[:nPoints]
        # get the median force at these indices
        forceMedUnits = np.median(forceUnits[idxForMedian])
        # correct the force
        forceUnits -= forceMedUnits
        # multiply it by -1 (flip)
        forceUnits *= -1
    return sepUnits,forceUnits

# plot a force extension curve with approach and retract
def PlotFec(sep,force,surfIdx = None,normalizeSep=True,normalizeFor=True,
            filterN=None,sensibleUnits=True):
    """
    Plot a force extension curve

    :param sep: The separation in meters
    :param force: The force in meters
    :param surfIdx: The index between approach and retract. if not present,
    intuits approximate index from minmmum Sep
    :param normalizeSep: If true, then zeros sep to its minimum
    :paran normalizeFor: If true, then zeros force to the median-filtered last
    5% of data, by separation (presummably, already detached)
    :param filterT: Plots the raw data in grey, and filters
    the force to the Number of points given. If none, assumes default % of curve
    :param sensibleUnits: Plots in nm and pN, defaults to true
    """
    if (surfIdx is None):
        surfIdx = np.argmin(sep)
    sepUnits,forceUnits = NormalizeSepForce(sep,force,surfIdx,normalizeSep,
                                            normalizeFor,sensibleUnits)
    if (filterN is None):
        filterN = int(np.ceil(DEF_FILTER_CONST*sepUnits.size))
    # POST: go ahead and normalize/color
    sepAppr = sepUnits[:surfIdx]
    sepRetr = sepUnits[surfIdx:]
    forceAppr = forceUnits[:surfIdx]
    forceRetr = forceUnits[surfIdx:]
    PlotFilteredSepForce(sepAppr,forceAppr,filterN=filterN,color='r',
                         label="Approach")
    PlotFilteredSepForce(sepRetr,forceRetr,filterN=filterN,color='b',
                         label="Retract")
    plt.xlim([min(sepUnits),max(sepUnits)])
    pPlotUtil.lazyLabel("Separation [nm]","Force [pN]","Force Extension Curve")
    return sepUnits,forceUnits

def filterForce(force,filterN=None):
    if (filterN is None):
        filterN = int(np.ceil(DEF_FILTER_CONST*force.size))
    return savitskyFilter(force,filterN)

def PlotFilteredSepForce(sep,force,filterN=None,labelRaw=None,
                         linewidthFilt=2.0,color='r',**kwargs):
    forceFilt =filterForce(force,filterN)
    plt.plot(sep,forceFilt,color=color,lw=linewidthFilt,**kwargs)
    # plot the raw data as grey
    plt.plot(sep,force,color='k',label=labelRaw,alpha=0.3)
    return forceFilt


#! /usr/bin/env python
#
# Copyright (C) 2007-2009 Cournapeau David <cournape@gmail.com>
#               2010 Fabian Pedregosa <fabian.pedregosa@inria.fr>
# License: 3-clause BSD

descr = """A set of python modules for machine learning and data mining"""

import sys
import os
import shutil
from distutils.command.clean import clean as Clean

if sys.version_info[0] < 3:
    import __builtin__ as builtins
else:
    import builtins

# This is a bit (!) hackish: we are setting a global variable so that the main
# sklearn __init__ can detect if it is being loaded by the setup routine, to
# avoid attempting to load components that aren't built yet:
# the numpy distutils extensions that are used by scikit-learn to recursively
# build the compiled extensions in sub-packages is based on the Python import
# machinery.
builtins.__SKLEARN_SETUP__ = True

DISTNAME = 'scikit-learn'
DESCRIPTION = 'A set of python modules for machine learning and data mining'
with open('README.rst') as f:
    LONG_DESCRIPTION = f.read()
MAINTAINER = 'Andreas Mueller'
MAINTAINER_EMAIL = 'amueller@ais.uni-bonn.de'
URL = 'http://scikit-learn.org'
LICENSE = 'new BSD'
DOWNLOAD_URL = 'http://sourceforge.net/projects/scikit-learn/files/'

# We can actually import a restricted version of sklearn that
# does not need the compiled code
import sklearn
VERSION = sklearn.__version__

# Optional setuptools features
# We need to import setuptools early, if we want setuptools features,
# as it monkey-patches the 'setup' function
# For some commands, use setuptools
SETUPTOOLS_COMMANDS = set([
    'develop', 'release', 'bdist_egg', 'bdist_rpm',
    'bdist_wininst', 'install_egg_info', 'build_sphinx',
    'egg_info', 'easy_install', 'upload', 'bdist_wheel',
    '--single-version-externally-managed',
])
if SETUPTOOLS_COMMANDS.intersection(sys.argv):
    import setuptools
    extra_setuptools_args = dict(
        zip_safe=False,  # the package can run out of an .egg file
        include_package_data=True,
    )
else:
    extra_setuptools_args = dict()

# Custom clean command to remove build artifacts

class CleanCommand(Clean):
    description = "Remove build artifacts from the source tree"

    def run(self):
        Clean.run(self)
        if os.path.exists('build'):
            shutil.rmtree('build')
        for dirpath, dirnames, filenames in os.walk('sklearn'):
            for filename in filenames:
                if (filename.endswith('.so') or filename.endswith('.pyd')
                        or filename.endswith('.dll')
                        or filename.endswith('.pyc')):
                    os.unlink(os.path.join(dirpath, filename))
            for dirname in dirnames:
                if dirname == '__pycache__':
                    shutil.rmtree(os.path.join(dirpath, dirname))

cmdclass = {'clean': CleanCommand}

# Optional wheelhouse-uploader features
# To automate release of binary packages for scikit-learn we need a tool
# to download the packages generated by travis and appveyor workers (with
# version number matching the current release) and upload them all at once
# to PyPI at release time.
# The URL of the artifact repositories are configured in the setup.cfg file.

WHEELHOUSE_UPLOADER_COMMANDS = set(['fetch_artifacts', 'upload_all'])
if WHEELHOUSE_UPLOADER_COMMANDS.intersection(sys.argv):
    import wheelhouse_uploader.cmd
    cmdclass.update(vars(wheelhouse_uploader.cmd))

def configuration(parent_package='', top_path=None):
    if os.path.exists('MANIFEST'):
        os.remove('MANIFEST')

    from numpy.distutils.misc_util import Configuration
    config = Configuration(None, parent_package, top_path)

    # Avoid non-useful msg:
    # "Ignoring attempt to set 'name' (from ... "
    config.set_options(ignore_setup_xxx_py=True,
                       assume_default_configuration=True,
                       delegate_options_to_subpackages=True,
                       quiet=True)

    config.add_subpackage('sklearn')

    return config

def is_scipy_installed():
    try:
        import scipy
    except ImportError:
        return False
    return True

def is_numpy_installed():
    try:
        import numpy
    except ImportError:
        return False
    return True

def setup_package():
    metadata = dict(name=DISTNAME,
                    maintainer=MAINTAINER,
                    maintainer_email=MAINTAINER_EMAIL,
                    description=DESCRIPTION,
                    license=LICENSE,
                    url=URL,
                    version=VERSION,
                    download_url=DOWNLOAD_URL,
                    long_description=LONG_DESCRIPTION,
                    classifiers=['Intended Audience :: Science/Research',
                                 'Intended Audience :: Developers',
                                 'License :: OSI Approved',
                                 'Programming Language :: C',
                                 'Programming Language :: Python',
                                 'Topic :: Software Development',
                                 'Topic :: Scientific/Engineering',
                                 'Operating System :: Microsoft :: Windows',
                                 'Operating System :: POSIX',
                                 'Operating System :: Unix',
                                 'Operating System :: MacOS',
                                 'Programming Language :: Python :: 2',
                                 'Programming Language :: Python :: 2.6',
                                 'Programming Language :: Python :: 2.7',
                                 'Programming Language :: Python :: 3',
                                 'Programming Language :: Python :: 3.3',
                                 'Programming Language :: Python :: 3.4',
                                 ],
                    cmdclass=cmdclass,
                    **extra_setuptools_args)

    if (len(sys.argv) >= 2
            and ('--help' in sys.argv[1:] or sys.argv[1]
                 in ('--help-commands', 'egg_info', '--version', 'clean'))):

        # For these actions, NumPy is not required.
        #
        # They are required to succeed without Numpy for example when
        # pip is used to install Scikit-learn when Numpy is not yet present in
        # the system.
        try:
            from setuptools import setup
        except ImportError:
            from distutils.core import setup

        metadata['version'] = VERSION
    else:
        if is_numpy_installed() is False:
            raise ImportError("Numerical Python (NumPy) is not installed.\n"
                             "scikit-learn requires NumPy.\n"
                             "Installation instructions are available on scikit-learn website: "
                             "http://scikit-learn.org/stable/install.html\n")
        if is_scipy_installed() is False:
            raise ImportError("Scientific Python (SciPy) is not installed.\n"
                             "scikit-learn requires SciPy.\n"
                             "Installation instructions are available on scikit-learn website: "
                             "http://scikit-learn.org/stable/install.html\n")
        from numpy.distutils.core import setup

        metadata['configuration'] = configuration

    setup(**metadata)

if __name__ == "__main__":
    setup_package()


import matplotlib.pyplot as plt
import numpy as np
from sklearn.svm import SVC
from sklearn.datasets import make_blobs
from .plot_2d_separator import plot_2d_separator

def make_handcrafted_dataset():
    # a carefully hand-designed dataset lol
    X, y = make_blobs(centers=2, random_state=4, n_samples=30)
    y[np.array([7, 27])] = 0
    mask = np.ones(len(X), dtype=np.bool)
    mask[np.array([0, 1, 5, 26])] = 0
    X, y = X[mask], y[mask]
    return X, y

def plot_rbf_svm_parameters():
    X, y = make_handcrafted_dataset()

    fig, axes = plt.subplots(1, 3, figsize=(12, 4))
    for ax, C in zip(axes, [1e0, 5, 10, 100]):
        ax.scatter(X[:, 0], X[:, 1], s=150, c=np.array(['red', 'blue'])[y])

        svm = SVC(kernel='rbf', C=C).fit(X, y)
        plot_2d_separator(svm, X, ax=ax, eps=.5)
        ax.set_title("C = %f" % C)

    fig, axes = plt.subplots(1, 4, figsize=(15, 3))
    for ax, gamma in zip(axes, [0.1, .5, 1, 10]):
        ax.scatter(X[:, 0], X[:, 1], s=150, c=np.array(['red', 'blue'])[y])
        svm = SVC(gamma=gamma, kernel='rbf', C=1).fit(X, y)
        plot_2d_separator(svm, X, ax=ax, eps=.5)
        ax.set_title("gamma = %f" % gamma)

def plot_svm(log_C, log_gamma):
    X, y = make_handcrafted_dataset()
    C = 10. ** log_C
    gamma = 10. ** log_gamma
    svm = SVC(kernel='rbf', C=C, gamma=gamma).fit(X, y)
    ax = plt.gca()
    plot_2d_separator(svm, X, ax=ax, eps=.5)
    # plot data
    ax.scatter(X[:, 0], X[:, 1], s=150, c=np.array(['red', 'blue'])[y])
    # plot support vectors
    sv = svm.support_vectors_
    ax.scatter(sv[:, 0], sv[:, 1], s=230, facecolors='none', zorder=10, linewidth=3)
    ax.set_title("C = %.4f gamma = %.4f" % (C, gamma))

def plot_svm_interactive():
    from IPython.html.widgets import interactive, FloatSlider
    C_slider = FloatSlider(min=-3, max=3, step=.1, value=0, readout=False)
    gamma_slider = FloatSlider(min=-2, max=2, step=.1, value=0, readout=False)
    return interactive(plot_svm, log_C=C_slider, log_gamma=gamma_slider)


"""
Our experiment environment.
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation

GRAVITY = -9.81

class Environment:

    def __init__(self, dimensions=None, dynamic_objects=None,
                 static_objects=None, robot=None):
        """An environment to run experiments in

        :param dimensions: (Optional) The dimensions of env
        :type dimensions: 1x3 numpy.array or None

        :param dynamic_objects: (Optional) A list of objects that can move
        :type dynamic_objects: list of maddux.objects.DynamicObject or None

        :param static_objects: (Optional) A list of stationary objects
        :type static_objects: list of maddux.objects.StaticObject or None

        :param robot: (Optional) A robot to simulate
        :type robot: maddux.robot.Arm or None

        :rtype: None
        """

        if dimensions is not None:
            self.dimensions = np.array(dimensions)
        else:
            self.dimensions = np.array([10.0, 10.0, 100.0])
        self.dynamic_objects = dynamic_objects if dynamic_objects else []
        self.static_objects = static_objects if static_objects else []
        self.robot = robot

    def run(self, duration):
        """Run for a certain duration

        :param duration: duration to run environment in seconds
        :type duration: integer

        :rtype: None
        """
        duration_ms = int(duration * 1000)

        for _ in xrange(duration_ms):
            map(lambda obj: obj.step(), self.dynamic_objects)
            if self.collision():
                break

    def animate(self, duration=None, save_path=None):
        """Animates the running of the program

        :param duration: (Optional) Duration of animation in seconds
        :type duration: int or None

        :param save_path: (Optional) Path to save mp4 in instead of displaying
        :type save_path: String or None

        :rtype: None
        """
        fps = 15
        dynamic_iter_per_frame = 10 * fps

        if duration is None:
            if self.robot is None:
                # Sensible Default
                frames = fps * 5
            else:
                frames = len(self.robot.qs)
        else:
            frames = int(fps * duration)

        def update(i):
            ax.clear()
            for _ in xrange(dynamic_iter_per_frame):
                map(lambda obj: obj.step(), self.dynamic_objects)
                # Check for collisions
                self.collision()
            if self.robot is not None:
                next_q = self.robot.qs[i]
                self.robot.update_angles(next_q)
            self.plot(ax=ax, show=False)

        fig = plt.figure(figsize=(8, 8))
        ax = Axes3D(fig)
        self.plot(ax=ax, show=False)

        # If we don't assign its return to something, it doesn't run.
        # Seems like really weird behavior..
        ani = animation.FuncAnimation(fig, update, frames=frames, blit=False)
        if save_path is None:
            plt.show()
        else:
            Writer = animation.writers['ffmpeg']
            writer = Writer(
                fps=fps, metadata=dict(
                    artist='Maddux'), bitrate=1800)
            ani.save(save_path, writer=writer)

    def hypothetical_landing_position(self):
        """Find the position that the ball would land (or hit a wall)

        :returns: Position (x, y, z) of hypothetical landing position of a
                  thrown object based on end effector velocity.
        :rtype: numpy.ndarray or None
        """
        pos = self.robot.end_effector_position().copy()
        # Only need linear velocity
        v = self.robot.end_effector_velocity()[0:3]

        for t in np.linspace(0, 15, 5000):
            # Check if it hit a target
            for static in self.static_objects:
                if static.is_hit(pos):
                    return pos.copy()

            # Or a wall
            for i in range(len(pos)):
                in_negative_space = pos[i] <= 0
                past_boundary = pos[i] >= self.dimensions[i]

                if in_negative_space or past_boundary:
                    return pos.copy()

            # Otherwise step forward
            v[2] += t * GRAVITY
            pos += t * v
        # If we never hit anything (which is completely impossible (TM))
        # return None
        return None

    def collision(self):
        """Check if any dynamic objects collide with any static
        objects or walls.

        :return: Whether there was a collision
        :rtype: bool
        """
        for dynamic in self.dynamic_objects:
            if dynamic.attached:
                continue

            for static in self.static_objects:
                if static.is_hit(dynamic.position):
                    dynamic.attach()
                    return True

            for i in range(len(dynamic.position)):
                in_negative_space = dynamic.position[i] <= 0
                past_boundary = (dynamic.position[i] >=
                                 self.dimensions[i])
                if in_negative_space or past_boundary:
                    dynamic.attach()
                    return True

        return False

    def plot(self, ax=None, show=True):
        """Plot throw trajectory and ball

        :param ax: Current axis if a figure already exists
        :type ax: matplotlib.axes

        :param show: (Default: True) Whether to show the figure
        :type show: bool

        :rtype: None
        """
        if ax is None:
            fig = plt.figure(figsize=(12, 12))
            ax = Axes3D(fig)

        # Set the limits to be environment ranges
        ax.set_xlim([0, self.dimensions[0]])
        ax.set_ylim([0, self.dimensions[1]])

        if self.dynamic_objects:
            zmax = max([o.positions[:, 2].max()
                        for o in self.dynamic_objects])
        else:
            zmax = 10
        ax.set_zlim([0, max(10, zmax)])

        # And set our labels
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')

        for dynamic in self.dynamic_objects:
            # Plot Trajectory
            ax.plot(dynamic.positions[:, 0], dynamic.positions[:, 1],
                    dynamic.positions[:, 2], 'r--', label='Trajectory')

        # Plot objects
        map(lambda obj: obj.plot(ax), self.dynamic_objects)
        map(lambda obj: obj.plot(ax), self.static_objects)

        if self.robot:
            self.robot.plot(ax)

        if show:
            plt.show()


#!/usr/bin/env python

# hacktasic viz of the quadratic surface of advantage around the max output
# for a couple of clear block on right / left / center cases

import agents
import argparse
import base_network
import Image
import numpy as np
import models
import sys
import tensorflow as tf
import replay_memory
import util
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm

np.set_printoptions(precision=5, threshold=10000, suppress=True, linewidth=10000)

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--width', type=int, default=160, help="render width")
parser.add_argument('--height', type=int, default=120, help="render height")
agents.add_opts(parser)
models.add_opts(parser)
replay_memory.add_opts(parser)
util.add_opts(parser)
opts = parser.parse_args()
#opts.ckpt_dir = "runs/14/d/ckpts"  # last known good
print >>sys.stderr, "OPTS", opts

# init our rl_agent
agent_cstr = eval("agents.NafAgent")
agent = agent_cstr(opts)
an = agent.network

# prepare three plots; one for each of block on left, in center, or on right
fig = plt.figure(figsize=plt.figaspect(0.3))
plt.title(opts.ckpt_dir)
R = np.arange(-1, 1.25, 0.25)
X, Y = np.meshgrid(R, R)
for plot_idx, (img_file, desc) in enumerate([("runs/14/d/imgs/ep_00007/e0000.png", "on left"),
                                             ("runs/14/d/imgs/ep_00007/e0019.png", "center"),
                                             ("runs/14/d/imgs/ep_00007/e0034.png", "on right")]):
  print "calculating for", desc, "..."

  # slurp in bitmap
  img = Image.open(img_file)
  img = np.array(img)[:,:,:3]

  # collect q-value for all x, y values in one hit
  all_x_y_pairs = np.stack(zip(np.ravel(X), np.ravel(Y)))
  img_repeated = [img] * all_x_y_pairs.shape[0]
  q_values = agent.sess.run(an.q_value,
                            feed_dict={an.input_state: img_repeated,
                                       an.input_action: all_x_y_pairs,
                                       base_network.FLIP_HORIZONTALLY: False})
  Z = q_values.reshape(X.shape)

  # plot as surface
  ax = fig.add_subplot(1,3,plot_idx+1, projection='3d')
  ax.plot_surface(X, Y, Z, rstride=1, cstride=1, color='b', cmap=cm.coolwarm, linewidth=1)
  ax.set_title(desc)
  ax.set_xlabel("turn")
  ax.set_ylabel("move")
  ax.set_zlabel("q")

  # include single vertical line where q was maximised (according to output_action)
  output = agent.sess.run(an.output_action,
                          feed_dict={an.input_state: [img],
                                     base_network.FLIP_HORIZONTALLY: False})
  turn, move = np.squeeze(output)
  q_value = agent.sess.run(an.q_value,
                           feed_dict={an.input_state: [img],
                                      an.input_action: [[turn, move]],
                                      base_network.FLIP_HORIZONTALLY: False})
  print "turn", turn, "move", move, "=> q", np.squeeze(q_value), "Zmin=", np.min(Z), "Zmax=", np.max(Z)
  ax.plot([turn, turn], [move, move], [np.min(Z), np.max(Z)], linewidth=5)

# render
plt.savefig("/tmp/test.png")
plt.show()


import pandas as pd
import folium.folium as folium
import itertools
import numpy as np
import logging
import geojson as gj
import copy
import attrdict as ad
from functional import seq

# import emission.analysis.classification.cleaning.location_smoothing as ls
import bson.json_util as bju

import emission.storage.decorations.location_queries as lq
import emission.storage.decorations.trip_queries as esdt
import emission.storage.decorations.place_queries as esdp
import emission.storage.decorations.stop_queries as esds
import emission.storage.decorations.section_queries as esdsc

import emission.storage.timeseries.abstract_timeseries as esta

import emission.core.wrapper.stop as ecws
import emission.core.wrapper.section as ecwsc

import emission.analysis.plotting.geojson.geojson_feature_converter as gfc
import emission.analysis.plotting.leaflet_osm.folium_geojson_plugin as fgjp

import emission.net.usercache.abstract_usercache as enua
import emission.net.api.usercache as enau

all_color_list = ['black', 'brown', 'blue', 'chocolate', 'cyan', 'fuschia', 'green', 'lime', 'magenta', 'navy', 'pink', 'purple', 'red', 'snow', 'yellow']
sel_color_list = ['black', 'blue', 'chocolate', 'cyan', 'fuschia', 'green', 'lime', 'magenta', 'pink', 'purple', 'red', 'yellow']

def df_to_string_list(df):
    """
    Convert the input df into a list of strings, suitable for using as popups in a map.
    This is a utility function.
    """
    # print "Converting df with size %s to string list" % df.shape[0]
    array_list = df.to_dict(orient='records')
    return [str(line) for line in array_list]

def get_maps_for_range(user_id, start_ts, end_ts):
    map_list = []
    geojson_list = gfc.get_geojson_for_ts(user_id, start_ts, end_ts)
    return get_maps_for_geojson_list(geojson_list)

def get_maps_for_usercache(user_id):
    data_to_phone = seq(enau.sync_server_to_phone(user_id))
    logging.debug("Before pipeline, trips to phone list has length %d" % len(data_to_phone.to_list()))
    logging.debug("keys are %s" % data_to_phone.map(lambda e: ad.AttrDict(e).metadata.key))
    trips_to_phone = data_to_phone.map(lambda e: ad.AttrDict(e))\
                                    .filter(lambda e: e.metadata.key.startswith("diary/trips")) \
                                    .map(lambda e: e.data)
    logging.debug("After pipeline, trips to phone list has length %d" % len(trips_to_phone.to_list()))
    # logging.debug("trips_to_phone = %s" % trips_to_phone)
    maps_for_day = []
    for day in trips_to_phone:
        maps_for_day.append(get_maps_for_geojson_list(day))
    return maps_for_day

def get_maps_for_geojson_list(trip_geojson_list):
    map_list = []
    for trip_doc in trip_geojson_list:
        # logging.debug(trip_doc)
        trip_geojson = ad.AttrDict(trip_doc)
        logging.debug("centering based on start = %s, end = %s " % (trip_geojson.features[0], trip_geojson.features[1]))
        flipped_midpoint = lambda(p1, p2): [(p1.coordinates[1] + p2.coordinates[1])/2,
                                            (p1.coordinates[0] + p2.coordinates[0])/2]

        curr_map = folium.Map(flipped_midpoint((trip_geojson.features[0].geometry,
                                                trip_geojson.features[1].geometry)))
        curr_plugin = fgjp.FoliumGeojsonPlugin(dict(trip_geojson))
        curr_map.add_plugin(curr_plugin)
        map_list.append(curr_map)
    return map_list

def flipped(coord):
    return (coord[1], coord[0])

def get_center_for_map(coords):
    # logging.debug(trip_geojson)
    midpoint = lambda(p1, p2): [(p1[0] + p2[0])/2,
                                (p1[1] + p2[1])/2]
    if len(coords) == 0:
        return None
    if len(coords) == 1:
        return flipped(coords)
    if len(coords) > 0:
        logging.debug("Getting midpoint of %s and %s" % (coords[0], coords[-1]))
        return flipped(midpoint((coords[0], coords[-1])))

def get_maps_for_geojson_unsectioned(feature_list):
    map_list = []
    for feature in feature_list:
        # logging.debug("Getting map for feature %s" % bju.dumps(feature))
        feature_coords = list(get_coords(feature))
        # feature_coords = list(gj.utils.coords(feature))
        curr_map = folium.Map(get_center_for_map(feature_coords))
        curr_plugin = fgjp.FoliumGeojsonPlugin(dict(feature))
        curr_map.add_plugin(curr_plugin)
        map_list.append(curr_map)
    return map_list

def get_coords(feature):
    # logging.debug("Getting coordinates for feature %s" % bju.dumps(feature))
    if feature["type"] == "FeatureCollection":
        retVal = []
        for f in feature["features"]:
            retVal.extend(get_coords(f))
        return retVal
    else:
        return gj.utils.coords(feature)

def get_maps_for_range_old(user_id, start_ts, end_ts):
    # First, get the timeline for that range.
    ts = esta.TimeSeries.get_time_series(user_id)
    trip_list = esdt.get_trips(user_id, enua.UserCache.TimeQuery("start_ts", start_ts, end_ts))
    # TODO: Should the timeline support random access as well?
    # If it did, we wouldn't need this additional map
    # I think that it would be good to support a doubly linked list, i.e. prev and next in addition
    # to the iteration interface
    place_list = esdp.get_places(user_id, enua.UserCache.TimeQuery("exit_ts", start_ts, end_ts))
    place_list = place_list + (esdp.get_places(user_id, enua.UserCache.TimeQuery("enter_ts", start_ts, end_ts)))
    place_map = dict([(p.get_id(), p) for p in place_list])
    map_list = []
    flipped_midpoint = lambda(p1, p2): [(p1.coordinates[1] + p2.coordinates[1])/2,
                                        (p1.coordinates[0] + p2.coordinates[0])/2]
    for i, trip in enumerate(trip_list):
        logging.debug("-" * 20 + trip.start_fmt_time + "=>" + trip.end_fmt_time
                      + "(" + str(trip.end_ts - trip.start_ts) + ")")
        if (len(esdt.get_sections_for_trip(user_id, trip.get_id())) == 0 and
            len(esdt.get_stops_for_trip(user_id, trip.get_id())) == 0):
            logging.debug("Skipping trip because it has no stops and no sections")
            continue

        start_point = gj.GeoJSON.to_instance(trip.start_loc)
        end_point = gj.GeoJSON.to_instance(trip.end_loc)
        curr_map = folium.Map(flipped_midpoint((start_point, end_point)))
        map_list.append(curr_map)
        logging.debug("About to display places %s and %s" % (trip.start_place, trip.end_place))
        update_place(curr_map, trip.start_place, place_map, marker_color='green')
        update_place(curr_map, trip.end_place, place_map, marker_color='red')
        # TODO: Should get_timeline_for_trip work on a trip_id or on a trip object
        # it seems stupid to convert trip object -> id -> trip object
        curr_trip_timeline = esdt.get_timeline_for_trip(user_id, trip.get_id())
        for i, trip_element in enumerate(curr_trip_timeline):
            # logging.debug("Examining element %s of type %s" % (trip_element, type(trip_element)))
            if type(trip_element) == ecws.Stop:
                time_query = esds.get_time_query_for_stop(trip_element.get_id())
                logging.debug("time_query for stop %s = %s" % (trip_element, time_query))
                stop_points_df = ts.get_data_df("background/filtered_location", time_query)
                # logging.debug("stop_points_df.head() = %s" % stop_points_df.head())
                if len(stop_points_df) > 0:
                    update_line(curr_map, stop_points_df, line_color = sel_color_list[-1],
                                popup="%s -> %s" % (trip_element.enter_fmt_time, trip_element.exit_fmt_time))
            else:
                assert(type(trip_element) == ecwsc.Section)
                time_query = esdsc.get_time_query_for_section(trip_element.get_id())
                logging.debug("time_query for section %s = %s" %
                              (trip_element, "[%s,%s,%s]" % (time_query.timeType, time_query.startTs, time_query.endTs)))
                section_points_df = ts.get_data_df("background/filtered_location", time_query)
                logging.debug("section_points_df.tail() = %s" % section_points_df.tail())
                if len(section_points_df) > 0:
                    update_line(curr_map, section_points_df, line_color = sel_color_list[trip_element.sensed_mode.value],
                                popup="%s (%s -> %s)" % (trip_element.sensed_mode, trip_element.start_fmt_time,
                                                         trip_element.end_fmt_time))
                else:
                    logging.warn("found no points for section %s" % trip_element)
    return map_list

def update_place(curr_map, place_id, place_map, marker_color='blue'):
    if place_id is not None and place_id in place_map:
        place = place_map[place_id]
        logging.debug("Retrieved place %s" % place)
        if hasattr(place, "location"):
            coords = copy.copy(place.location.coordinates)
            coords.reverse()
            logging.debug("Displaying place at %s" % coords)
            curr_map.simple_marker(location=coords, popup=str(place), marker_color=marker_color)
        else:
            logging.debug("starting place has no location, skipping")
    else:
        logging.warn("place not mapped because place_id = %s and place_id in place_map = %s" % (place_id, place_id in place_map))

def update_line(currMap, line_points, line_color = None, popup=None):
    currMap.div_markers(line_points[['latitude', 'longitude']].as_matrix().tolist(),
        df_to_string_list(line_points), marker_size=5)
    currMap.line(line_points[['latitude', 'longitude']].as_matrix().tolist(),
        line_color = line_color,
        popup = popup)

##########################
# Everything below this line is from the time when we were evaluating
# segmentation and can potentially be deleted. It is also likely to have bitrotted.
#  Let's hold off a bit on that until we have the replacement, though
##########################

def get_map_list(df, potential_splits):
    mapList = []
    potential_splits_list = list(potential_splits)
    for start, end in zip(potential_splits_list, potential_splits_list[1:]):
        trip = df[start:end]
        print "Considering trip from %s to %s because start = %d and end = %d" % (df.formatted_time.loc[start], df.formatted_time.loc[end], start, end)
        if end - start < 4:
            # If there are only 3 entries, that means that there is only one
            # point other than the start and the end, bail
            print "Ignoring trip from %s to %s because start = %d and end = %d" % (df.formatted_time.loc[start], df.formatted_time.loc[end], start, end)
            continue
        mapList.append(get_map(trip))
    return mapList

def get_map_list_after_segmentation(section_map, outlier_algo = None, filter_algo = None):
    mapList = []
    for trip, section_list in section_map:
        logging.debug("%s %s -> %s %s" % ("=" * 20, trip.start_time, trip.end_time, "=" * 20))
        trip_df = lq.get_points_for_section(trip)
        curr_map = folium.Map([trip_df.mLatitude.mean(), trip_df.mLongitude.mean()])
        last_section_end = None
        for (i, section) in enumerate(section_list):
            logging.debug("%s %s: %s -> %s %s" %
                ("-" * 20, i, section.start_time, section.end_time, "-" * 20))
            raw_section_df = trip_df[np.logical_and(trip_df.mTime >= section.start_ts,
                                                trip_df.mTime <= section.end_ts)]
            section_df = ls.filter_points(raw_section_df, outlier_algo, filter_algo)
            if section_df.shape[0] == 0:
                logging.info("Found empty df! skipping...")
                continue
            logging.debug("for section %s, section_df.shape = %s, formatted_time.head() = %s" %
                    (section, section_df.shape, section_df["formatted_time"].head()))
            update_map(curr_map, section_df, line_color = sel_color_list[section.activity.value],
                        popup = "%s" % (section.activity))
            if section_df.shape[0] > 0:
                curr_section_start = section_df.iloc[0]
                if i != 0 and last_section_end is not None:
                    # We want to join this to the previous section.
                    curr_map.line([[last_section_end.mLatitude, last_section_end.mLongitude],
                                   [curr_section_start.mLatitude, curr_section_start.mLongitude]],
                                   line_color = sel_color_list[-1],
                                   popup = "%s -> %s" % (section_list[i-1].activity, section.activity))
                last_section_end = section_df.iloc[-1]
        mapList.append(curr_map)
    return mapList

def get_map(section_points, line_color = None, popup=None):
    currMap = folium.Map([section_points.mLatitude.mean(), section_points.mLongitude.mean()])
    update_map(currMap, section_points, line_color, popup)
    return currMap

def update_map(currMap, section_points, line_color = None, popup=None):
    currMap.div_markers(section_points[['mLatitude', 'mLongitude']].as_matrix().tolist(),
        df_to_string_list(section_points), marker_size=5)
    currMap.line(section_points[['mLatitude', 'mLongitude']].as_matrix().tolist(),
        line_color = line_color,
        popup = popup)

def evaluate_filtering(section_list, outlier_algos, filtering_algos):
    """
    TODO: Is this the best place for this? If not, what is?
    It almost seems like we need to have a separate evaluation module that is
    separate from the plotting and the calculation modules.
    But then, what is the purpose of this module?
    """
    nCols = 2 + len(outlier_algos) * len(filtering_algos)
    nRows = len(section_list)

    map_list = []

    for section in section_list:
        curr_compare_list = []
        section_df = ls.get_section_points(section)
        curr_compare_list.append(get_map(section_df))
        curr_compare_list.append(get_map(ls.filter_points(section_df, None, None)))
        for (oa, fa) in itertools.product(outlier_algos, filtering_algos):
            curr_filtered_df = ls.filter_points(section_df, oa, fa)
            print ("After filtering with %s, %s, size is %s" % (oa, fa, curr_filtered_df.shape))
            if "activity" in section:
                curr_compare_list.append(get_map(curr_filtered_df,
                                            line_color = sel_color_list[section.activity.value],
                                                 popup = "%s" % (section.activity)))
            else:
                curr_compare_list.append(get_map(curr_filtered_df))
        assert(len(curr_compare_list) == nCols)
        map_list.append(curr_compare_list)
    assert(len(map_list) == nRows)
    return map_list