"""Plotting module for Sympy.

A plot is represented by the ``Plot`` class that contains a reference to the
backend and a list of the data series to be plotted. The data series are
instances of classes meant to simplify getting points and meshes from sympy
expressions. ``plot_backends`` is a dictionary with all the backends.

This module gives only the essential. For all the fancy stuff use directly
the backend. You can get the backend wrapper for every plot from the
``_backend`` attribute. Moreover the data series classes have various useful
methods like ``get_points``, ``get_segments``, ``get_meshes``, etc, that may
be useful if you wish to use another plotting library.

Especially if you need publication ready graphs and this module is not enough
for you - just get the ``_backend`` attribute and add whatever you want
directly to it. In the case of matplotlib (the common way to graph data in
python) just copy ``_backend.fig`` which is the figure and ``_backend.ax``
which is the axis and work on them as you would on any other matplotlib object.

Simplicity of code takes much greater importance than performance. Don't use it
if you care at all about performance. A new backend instance is initialized
every time you call ``show()`` and the old one is left to the garbage collector.
"""

from __future__ import print_function, division

import inspect
from collections import Callable
import warnings
import sys

from sympy import sympify, Expr, Tuple, Dummy, Symbol
from sympy.external import import_module
from sympy.core.compatibility import range
from sympy.utilities.decorator import doctest_depends_on
from sympy.utilities.iterables import is_sequence
from .experimental_lambdify import (vectorized_lambdify, lambdify)

# N.B.
# When changing the minimum module version for matplotlib, please change
# the same in the `SymPyDocTestFinder`` in `sympy/utilities/runtests.py`

# Backend specific imports - textplot
from sympy.plotting.textplot import textplot

# Global variable
# Set to False when running tests / doctests so that the plots don't show.
_show = True

def unset_show():
    global _show
    _show = False

##############################################################################
# The public interface
##############################################################################

def _arity(f):
    """
    Python 2 and 3 compatible version that do not raise a Deprecation warning.
    """
    if sys.version_info < (3,):
        return len(inspect.getargspec(f)[0])
    else:
       param = inspect.signature(f).parameters.values()
       return len([p for p in param if p.kind == p.POSITIONAL_OR_KEYWORD])

class Plot(object):
    """The central class of the plotting module.

    For interactive work the function ``plot`` is better suited.

    This class permits the plotting of sympy expressions using numerous
    backends (matplotlib, textplot, the old pyglet module for sympy, Google
    charts api, etc).

    The figure can contain an arbitrary number of plots of sympy expressions,
    lists of coordinates of points, etc. Plot has a private attribute _series that
    contains all data series to be plotted (expressions for lines or surfaces,
    lists of points, etc (all subclasses of BaseSeries)). Those data series are
    instances of classes not imported by ``from sympy import *``.

    The customization of the figure is on two levels. Global options that
    concern the figure as a whole (eg title, xlabel, scale, etc) and
    per-data series options (eg name) and aesthetics (eg. color, point shape,
    line type, etc.).

    The difference between options and aesthetics is that an aesthetic can be
    a function of the coordinates (or parameters in a parametric plot). The
    supported values for an aesthetic are:
    - None (the backend uses default values)
    - a constant
    - a function of one variable (the first coordinate or parameter)
    - a function of two variables (the first and second coordinate or
    parameters)
    - a function of three variables (only in nonparametric 3D plots)
    Their implementation depends on the backend so they may not work in some
    backends.

    If the plot is parametric and the arity of the aesthetic function permits
    it the aesthetic is calculated over parameters and not over coordinates.
    If the arity does not permit calculation over parameters the calculation is
    done over coordinates.

    Only cartesian coordinates are supported for the moment, but you can use
    the parametric plots to plot in polar, spherical and cylindrical
    coordinates.

    The arguments for the constructor Plot must be subclasses of BaseSeries.

    Any global option can be specified as a keyword argument.

    The global options for a figure are:

    - title : str
    - xlabel : str
    - ylabel : str
    - legend : bool
    - xscale : {'linear', 'log'}
    - yscale : {'linear', 'log'}
    - axis : bool
    - axis_center : tuple of two floats or {'center', 'auto'}
    - xlim : tuple of two floats
    - ylim : tuple of two floats
    - aspect_ratio : tuple of two floats or {'auto'}
    - autoscale : bool
    - margin : float in [0, 1]

    The per data series options and aesthetics are:
    There are none in the base series. See below for options for subclasses.

    Some data series support additional aesthetics or options:

    ListSeries, LineOver1DRangeSeries, Parametric2DLineSeries,
    Parametric3DLineSeries support the following:

    Aesthetics:

    - line_color : function which returns a float.

    options:

    - label : str
    - steps : bool
    - integers_only : bool

    SurfaceOver2DRangeSeries, ParametricSurfaceSeries support the following:

    aesthetics:

    - surface_color : function which returns a float.
    """

    def __init__(self, *args, **kwargs):
        super(Plot, self).__init__()

        #  Options for the graph as a whole.
        #  The possible values for each option are described in the docstring of
        # Plot. They are based purely on convention, no checking is done.
        self.title = None
        self.xlabel = None
        self.ylabel = None
        self.aspect_ratio = 'auto'
        self.xlim = None
        self.ylim = None
        self.axis_center = 'auto'
        self.axis = True
        self.xscale = 'linear'
        self.yscale = 'linear'
        self.legend = False
        self.autoscale = True
        self.margin = 0

        # Contains the data objects to be plotted. The backend should be smart
        # enough to iterate over this list.
        self._series = []
        self._series.extend(args)

        # The backend type. On every show() a new backend instance is created
        # in self._backend which is tightly coupled to the Plot instance
        # (thanks to the parent attribute of the backend).
        self.backend = DefaultBackend

        # The keyword arguments should only contain options for the plot.
        for key, val in kwargs.items():
            if hasattr(self, key):
                setattr(self, key, val)

    def show(self):
        # TODO move this to the backend (also for save)
        if hasattr(self, '_backend'):
            self._backend.close()
        self._backend = self.backend(self)
        self._backend.show()

    def save(self, path):
        if hasattr(self, '_backend'):
            self._backend.close()
        self._backend = self.backend(self)
        self._backend.save(path)

    def __str__(self):
        series_strs = [('[%d]: ' % i) + str(s)
                       for i, s in enumerate(self._series)]
        return 'Plot object containing:\n' + '\n'.join(series_strs)

    def __getitem__(self, index):
        return self._series[index]

    def __setitem__(self, index, *args):
        if len(args) == 1 and isinstance(args[0], BaseSeries):
            self._series[index] = args

    def __delitem__(self, index):
        del self._series[index]

    @doctest_depends_on(modules=('numpy', 'matplotlib',))
    def append(self, arg):
        """Adds an element from a plot's series to an existing plot.

        Examples
        ========

        Consider two ``Plot`` objects, ``p1`` and ``p2``. To add the
        second plot's first series object to the first, use the
        ``append`` method, like so:

        >>> from sympy import symbols
        >>> from sympy.plotting import plot
        >>> x = symbols('x')
        >>> p1 = plot(x*x)
        >>> p2 = plot(x)
        >>> p1.append(p2[0])
        >>> p1
        Plot object containing:
        [0]: cartesian line: x**2 for x over (-10.0, 10.0)
        [1]: cartesian line: x for x over (-10.0, 10.0)

        See Also
        ========
        extend

        """
        if isinstance(arg, BaseSeries):
            self._series.append(arg)
        else:
            raise TypeError('Must specify element of plot to append.')

    @doctest_depends_on(modules=('numpy', 'matplotlib',))
    def extend(self, arg):
        """Adds all series from another plot.

        Examples
        ========

        Consider two ``Plot`` objects, ``p1`` and ``p2``. To add the
        second plot to the first, use the ``extend`` method, like so:

        >>> from sympy import symbols
        >>> from sympy.plotting import plot
        >>> x = symbols('x')
        >>> p1 = plot(x*x)
        >>> p2 = plot(x)
        >>> p1.extend(p2)
        >>> p1
        Plot object containing:
        [0]: cartesian line: x**2 for x over (-10.0, 10.0)
        [1]: cartesian line: x for x over (-10.0, 10.0)

        """
        if isinstance(arg, Plot):
            self._series.extend(arg._series)
        elif is_sequence(arg):
            self._series.extend(arg)
        else:
            raise TypeError('Expecting Plot or sequence of BaseSeries')

##############################################################################
# Data Series
##############################################################################
#TODO more general way to calculate aesthetics (see get_color_array)

### The base class for all series
class BaseSeries(object):
    """Base class for the data objects containing stuff to be plotted.

    The backend should check if it supports the data series that it's given.
    (eg TextBackend supports only LineOver1DRange).
    It's the backend responsibility to know how to use the class of
    data series that it's given.

    Some data series classes are grouped (using a class attribute like is_2Dline)
    according to the api they present (based only on convention). The backend is
    not obliged to use that api (eg. The LineOver1DRange belongs to the
    is_2Dline group and presents the get_points method, but the
    TextBackend does not use the get_points method).
    """

    # Some flags follow. The rationale for using flags instead of checking base
    # classes is that setting multiple flags is simpler than multiple
    # inheritance.

    is_2Dline = False
    # Some of the backends expect:
    #  - get_points returning 1D np.arrays list_x, list_y
    #  - get_segments returning np.array (done in Line2DBaseSeries)
    #  - get_color_array returning 1D np.array (done in Line2DBaseSeries)
    # with the colors calculated at the points from get_points

    is_3Dline = False
    # Some of the backends expect:
    #  - get_points returning 1D np.arrays list_x, list_y, list_y
    #  - get_segments returning np.array (done in Line2DBaseSeries)
    #  - get_color_array returning 1D np.array (done in Line2DBaseSeries)
    # with the colors calculated at the points from get_points

    is_3Dsurface = False
    # Some of the backends expect:
    #   - get_meshes returning mesh_x, mesh_y, mesh_z (2D np.arrays)
    #   - get_points an alias for get_meshes

    is_contour = False
    # Some of the backends expect:
    #   - get_meshes returning mesh_x, mesh_y, mesh_z (2D np.arrays)
    #   - get_points an alias for get_meshes

    is_implicit = False
    # Some of the backends expect:
    #   - get_meshes returning mesh_x (1D array), mesh_y(1D array,
    #     mesh_z (2D np.arrays)
    #   - get_points an alias for get_meshes
    #Different from is_contour as the colormap in backend will be
    #different

    is_parametric = False
    # The calculation of aesthetics expects:
    #   - get_parameter_points returning one or two np.arrays (1D or 2D)
    # used for calculation aesthetics

    def __init__(self):
        super(BaseSeries, self).__init__()

    @property
    def is_3D(self):
        flags3D = [
            self.is_3Dline,
            self.is_3Dsurface
        ]
        return any(flags3D)

    @property
    def is_line(self):
        flagslines = [
            self.is_2Dline,
            self.is_3Dline
        ]
        return any(flagslines)

### 2D lines
class Line2DBaseSeries(BaseSeries):
    """A base class for 2D lines.

    - adding the label, steps and only_integers options
    - making is_2Dline true
    - defining get_segments and get_color_array
    """

    is_2Dline = True

    _dim = 2

    def __init__(self):
        super(Line2DBaseSeries, self).__init__()
        self.label = None
        self.steps = False
        self.only_integers = False
        self.line_color = None

    def get_segments(self):
        np = import_module('numpy')
        points = self.get_points()
        if self.steps is True:
            x = np.array((points[0], points[0])).T.flatten()[1:]
            y = np.array((points[1], points[1])).T.flatten()[:-1]
            points = (x, y)
        points = np.ma.array(points).T.reshape(-1, 1, self._dim)
        return np.ma.concatenate([points[:-1], points[1:]], axis=1)

    def get_color_array(self):
        np = import_module('numpy')
        c = self.line_color
        if hasattr(c, '__call__'):
            f = np.vectorize(c)
            arity = _arity(c)
            if arity == 1 and self.is_parametric:
                x = self.get_parameter_points()
                return f(centers_of_segments(x))
            else:
                variables = list(map(centers_of_segments, self.get_points()))
                if arity == 1:
                    return f(variables[0])
                elif arity == 2:
                    return f(*variables[:2])
                else:  # only if the line is 3D (otherwise raises an error)
                    return f(*variables)
        else:
            return c*np.ones(self.nb_of_points)

class List2DSeries(Line2DBaseSeries):
    """Representation for a line consisting of list of points."""

    def __init__(self, list_x, list_y):
        np = import_module('numpy')
        super(List2DSeries, self).__init__()
        self.list_x = np.array(list_x)
        self.list_y = np.array(list_y)
        self.label = 'list'

    def __str__(self):
        return 'list plot'

    def get_points(self):
        return (self.list_x, self.list_y)

class LineOver1DRangeSeries(Line2DBaseSeries):
    """Representation for a line consisting of a SymPy expression over a range."""

    def __init__(self, expr, var_start_end, **kwargs):
        super(LineOver1DRangeSeries, self).__init__()
        self.expr = sympify(expr)
        self.label = str(self.expr)
        self.var = sympify(var_start_end[0])
        self.start = float(var_start_end[1])
        self.end = float(var_start_end[2])
        self.nb_of_points = kwargs.get('nb_of_points', 300)
        self.adaptive = kwargs.get('adaptive', True)
        self.depth = kwargs.get('depth', 12)
        self.line_color = kwargs.get('line_color', None)

    def __str__(self):
        return 'cartesian line: %s for %s over %s' % (
            str(self.expr), str(self.var), str((self.start, self.end)))

    def get_segments(self):
        """
        Adaptively gets segments for plotting.

        The adaptive sampling is done by recursively checking if three
        points are almost collinear. If they are not collinear, then more
        points are added between those points.

        References
        ==========
        [1] Adaptive polygonal approximation of parametric curves,
            Luiz Henrique de Figueiredo.

        """
        if self.only_integers or not self.adaptive:
            return super(LineOver1DRangeSeries, self).get_segments()
        else:
            f = lambdify([self.var], self.expr)
            list_segments = []

            def sample(p, q, depth):
                """ Samples recursively if three points are almost collinear.
                For depth < 6, points are added irrespective of whether they
                satisfy the collinearity condition or not. The maximum depth
                allowed is 12.
                """
                np = import_module('numpy')
                #Randomly sample to avoid aliasing.
                random = 0.45 + np.random.rand() * 0.1
                xnew = p[0] + random * (q[0] - p[0])
                ynew = f(xnew)
                new_point = np.array([xnew, ynew])

                #Maximum depth
                if depth > self.depth:
                    list_segments.append([p, q])

                #Sample irrespective of whether the line is flat till the
                #depth of 6. We are not using linspace to avoid aliasing.
                elif depth < 6:
                    sample(p, new_point, depth + 1)
                    sample(new_point, q, depth + 1)

                #Sample ten points if complex values are encountered
                #at both ends. If there is a real value in between, then
                #sample those points further.
                elif p[1] is None and q[1] is None:
                    xarray = np.linspace(p[0], q[0], 10)
                    yarray = list(map(f, xarray))
                    if any(y is not None for y in yarray):
                        for i in range(len(yarray) - 1):
                            if yarray[i] is not None or yarray[i + 1] is not None:
                                sample([xarray[i], yarray[i]],
                                    [xarray[i + 1], yarray[i + 1]], depth + 1)

                #Sample further if one of the end points in None( i.e. a complex
                #value) or the three points are not almost collinear.
                elif (p[1] is None or q[1] is None or new_point[1] is None
                        or not flat(p, new_point, q)):
                    sample(p, new_point, depth + 1)
                    sample(new_point, q, depth + 1)
                else:
                    list_segments.append([p, q])

            f_start = f(self.start)
            f_end = f(self.end)
            sample([self.start, f_start], [self.end, f_end], 0)
            return list_segments

    def get_points(self):
        np = import_module('numpy')
        if self.only_integers is True:
            list_x = np.linspace(int(self.start), int(self.end),
                    num=int(self.end) - int(self.start) + 1)
        else:
            list_x = np.linspace(self.start, self.end, num=self.nb_of_points)
        f = vectorized_lambdify([self.var], self.expr)
        list_y = f(list_x)
        return (list_x, list_y)

class Parametric2DLineSeries(Line2DBaseSeries):
    """Representation for a line consisting of two parametric sympy expressions
    over a range."""

    is_parametric = True

    def __init__(self, expr_x, expr_y, var_start_end, **kwargs):
        super(Parametric2DLineSeries, self).__init__()
        self.expr_x = sympify(expr_x)
        self.expr_y = sympify(expr_y)
        self.label = "(%s, %s)" % (str(self.expr_x), str(self.expr_y))
        self.var = sympify(var_start_end[0])
        self.start = float(var_start_end[1])
        self.end = float(var_start_end[2])
        self.nb_of_points = kwargs.get('nb_of_points', 300)
        self.adaptive = kwargs.get('adaptive', True)
        self.depth = kwargs.get('depth', 12)
        self.line_color = kwargs.get('line_color', None)

    def __str__(self):
        return 'parametric cartesian line: (%s, %s) for %s over %s' % (
            str(self.expr_x), str(self.expr_y), str(self.var),
            str((self.start, self.end)))

    def get_parameter_points(self):
        np = import_module('numpy')
        return np.linspace(self.start, self.end, num=self.nb_of_points)

    def get_points(self):
        param = self.get_parameter_points()
        fx = vectorized_lambdify([self.var], self.expr_x)
        fy = vectorized_lambdify([self.var], self.expr_y)
        list_x = fx(param)
        list_y = fy(param)
        return (list_x, list_y)

    def get_segments(self):
        """
        Adaptively gets segments for plotting.

        The adaptive sampling is done by recursively checking if three
        points are almost collinear. If they are not collinear, then more
        points are added between those points.

        References
        ==========
        [1] Adaptive polygonal approximation of parametric curves,
            Luiz Henrique de Figueiredo.

        """
        if not self.adaptive:
            return super(Parametric2DLineSeries, self).get_segments()

        f_x = lambdify([self.var], self.expr_x)
        f_y = lambdify([self.var], self.expr_y)
        list_segments = []

        def sample(param_p, param_q, p, q, depth):
            """ Samples recursively if three points are almost collinear.
            For depth < 6, points are added irrespective of whether they
            satisfy the collinearity condition or not. The maximum depth
            allowed is 12.
            """
            #Randomly sample to avoid aliasing.
            np = import_module('numpy')
            random = 0.45 + np.random.rand() * 0.1
            param_new = param_p + random * (param_q - param_p)
            xnew = f_x(param_new)
            ynew = f_y(param_new)
            new_point = np.array([xnew, ynew])

            #Maximum depth
            if depth > self.depth:
                list_segments.append([p, q])

            #Sample irrespective of whether the line is flat till the
            #depth of 6. We are not using linspace to avoid aliasing.
            elif depth < 6:
                sample(param_p, param_new, p, new_point, depth + 1)
                sample(param_new, param_q, new_point, q, depth + 1)

            #Sample ten points if complex values are encountered
            #at both ends. If there is a real value in between, then
            #sample those points further.
            elif ((p[0] is None and q[1] is None) or
                    (p[1] is None and q[1] is None)):
                param_array = np.linspace(param_p, param_q, 10)
                x_array = list(map(f_x, param_array))
                y_array = list(map(f_y, param_array))
                if any(x is not None and y is not None
                        for x, y in zip(x_array, y_array)):
                    for i in range(len(y_array) - 1):
                        if ((x_array[i] is not None and y_array[i] is not None) or
                                (x_array[i + 1] is not None and y_array[i + 1] is not None)):
                            point_a = [x_array[i], y_array[i]]
                            point_b = [x_array[i + 1], y_array[i + 1]]
                            sample(param_array[i], param_array[i], point_a,
                                   point_b, depth + 1)

            #Sample further if one of the end points in None( ie a complex
            #value) or the three points are not almost collinear.
            elif (p[0] is None or p[1] is None
                    or q[1] is None or q[0] is None
                    or not flat(p, new_point, q)):
                sample(param_p, param_new, p, new_point, depth + 1)
                sample(param_new, param_q, new_point, q, depth + 1)
            else:
                list_segments.append([p, q])

        f_start_x = f_x(self.start)
        f_start_y = f_y(self.start)
        start = [f_start_x, f_start_y]
        f_end_x = f_x(self.end)
        f_end_y = f_y(self.end)
        end = [f_end_x, f_end_y]
        sample(self.start, self.end, start, end, 0)
        return list_segments

### 3D lines
class Line3DBaseSeries(Line2DBaseSeries):
    """A base class for 3D lines.

    Most of the stuff is derived from Line2DBaseSeries."""

    is_2Dline = False
    is_3Dline = True
    _dim = 3

    def __init__(self):
        super(Line3DBaseSeries, self).__init__()

class Parametric3DLineSeries(Line3DBaseSeries):
    """Representation for a 3D line consisting of two parametric sympy
    expressions and a range."""

    def __init__(self, expr_x, expr_y, expr_z, var_start_end, **kwargs):
        super(Parametric3DLineSeries, self).__init__()
        self.expr_x = sympify(expr_x)
        self.expr_y = sympify(expr_y)
        self.expr_z = sympify(expr_z)
        self.label = "(%s, %s)" % (str(self.expr_x), str(self.expr_y))
        self.var = sympify(var_start_end[0])
        self.start = float(var_start_end[1])
        self.end = float(var_start_end[2])
        self.nb_of_points = kwargs.get('nb_of_points', 300)
        self.line_color = kwargs.get('line_color', None)

    def __str__(self):
        return '3D parametric cartesian line: (%s, %s, %s) for %s over %s' % (
            str(self.expr_x), str(self.expr_y), str(self.expr_z),
            str(self.var), str((self.start, self.end)))

    def get_parameter_points(self):
        np = import_module('numpy')
        return np.linspace(self.start, self.end, num=self.nb_of_points)

    def get_points(self):
        param = self.get_parameter_points()
        fx = vectorized_lambdify([self.var], self.expr_x)
        fy = vectorized_lambdify([self.var], self.expr_y)
        fz = vectorized_lambdify([self.var], self.expr_z)
        list_x = fx(param)
        list_y = fy(param)
        list_z = fz(param)
        return (list_x, list_y, list_z)

### Surfaces
class SurfaceBaseSeries(BaseSeries):
    """A base class for 3D surfaces."""

    is_3Dsurface = True

    def __init__(self):
        super(SurfaceBaseSeries, self).__init__()
        self.surface_color = None

    def get_color_array(self):
        np = import_module('numpy')
        c = self.surface_color
        if isinstance(c, Callable):
            f = np.vectorize(c)
            arity = _arity(c)
            if self.is_parametric:
                variables = list(map(centers_of_faces, self.get_parameter_meshes()))
                if arity == 1:
                    return f(variables[0])
                elif arity == 2:
                    return f(*variables)
            variables = list(map(centers_of_faces, self.get_meshes()))
            if arity == 1:
                return f(variables[0])
            elif arity == 2:
                return f(*variables[:2])
            else:
                return f(*variables)
        else:
            return c*np.ones(self.nb_of_points)

class SurfaceOver2DRangeSeries(SurfaceBaseSeries):
    """Representation for a 3D surface consisting of a sympy expression and 2D
    range."""
    def __init__(self, expr, var_start_end_x, var_start_end_y, **kwargs):
        super(SurfaceOver2DRangeSeries, self).__init__()
        self.expr = sympify(expr)
        self.var_x = sympify(var_start_end_x[0])
        self.start_x = float(var_start_end_x[1])
        self.end_x = float(var_start_end_x[2])
        self.var_y = sympify(var_start_end_y[0])
        self.start_y = float(var_start_end_y[1])
        self.end_y = float(var_start_end_y[2])
        self.nb_of_points_x = kwargs.get('nb_of_points_x', 50)
        self.nb_of_points_y = kwargs.get('nb_of_points_y', 50)
        self.surface_color = kwargs.get('surface_color', None)

    def __str__(self):
        return ('cartesian surface: %s for'
                ' %s over %s and %s over %s') % (
                    str(self.expr),
                    str(self.var_x),
                    str((self.start_x, self.end_x)),
                    str(self.var_y),
                    str((self.start_y, self.end_y)))

    def get_meshes(self):
        np = import_module('numpy')
        mesh_x, mesh_y = np.meshgrid(np.linspace(self.start_x, self.end_x,
                                                 num=self.nb_of_points_x),
                                     np.linspace(self.start_y, self.end_y,
                                                 num=self.nb_of_points_y))
        f = vectorized_lambdify((self.var_x, self.var_y), self.expr)
        return (mesh_x, mesh_y, f(mesh_x, mesh_y))

class ParametricSurfaceSeries(SurfaceBaseSeries):
    """Representation for a 3D surface consisting of three parametric sympy
    expressions and a range."""

    is_parametric = True

    def __init__(
        self, expr_x, expr_y, expr_z, var_start_end_u, var_start_end_v,
            **kwargs):
        super(ParametricSurfaceSeries, self).__init__()
        self.expr_x = sympify(expr_x)
        self.expr_y = sympify(expr_y)
        self.expr_z = sympify(expr_z)
        self.var_u = sympify(var_start_end_u[0])
        self.start_u = float(var_start_end_u[1])
        self.end_u = float(var_start_end_u[2])
        self.var_v = sympify(var_start_end_v[0])
        self.start_v = float(var_start_end_v[1])
        self.end_v = float(var_start_end_v[2])
        self.nb_of_points_u = kwargs.get('nb_of_points_u', 50)
        self.nb_of_points_v = kwargs.get('nb_of_points_v', 50)
        self.surface_color = kwargs.get('surface_color', None)

    def __str__(self):
        return ('parametric cartesian surface: (%s, %s, %s) for'
                ' %s over %s and %s over %s') % (
                    str(self.expr_x),
                    str(self.expr_y),
                    str(self.expr_z),
                    str(self.var_u),
                    str((self.start_u, self.end_u)),
                    str(self.var_v),
                    str((self.start_v, self.end_v)))

    def get_parameter_meshes(self):
        np = import_module('numpy')
        return np.meshgrid(np.linspace(self.start_u, self.end_u,
                                       num=self.nb_of_points_u),
                           np.linspace(self.start_v, self.end_v,
                                       num=self.nb_of_points_v))

    def get_meshes(self):
        mesh_u, mesh_v = self.get_parameter_meshes()
        fx = vectorized_lambdify((self.var_u, self.var_v), self.expr_x)
        fy = vectorized_lambdify((self.var_u, self.var_v), self.expr_y)
        fz = vectorized_lambdify((self.var_u, self.var_v), self.expr_z)
        return (fx(mesh_u, mesh_v), fy(mesh_u, mesh_v), fz(mesh_u, mesh_v))

### Contours
class ContourSeries(BaseSeries):
    """Representation for a contour plot."""
    #The code is mostly repetition of SurfaceOver2DRange.
    #XXX: Presently not used in any of those functions.
    #XXX: Add contour plot and use this seties.

    is_contour = True

    def __init__(self, expr, var_start_end_x, var_start_end_y):
        super(ContourSeries, self).__init__()
        self.nb_of_points_x = 50
        self.nb_of_points_y = 50
        self.expr = sympify(expr)
        self.var_x = sympify(var_start_end_x[0])
        self.start_x = float(var_start_end_x[1])
        self.end_x = float(var_start_end_x[2])
        self.var_y = sympify(var_start_end_y[0])
        self.start_y = float(var_start_end_y[1])
        self.end_y = float(var_start_end_y[2])

        self.get_points = self.get_meshes

    def __str__(self):
        return ('contour: %s for '
                '%s over %s and %s over %s') % (
                    str(self.expr),
                    str(self.var_x),
                    str((self.start_x, self.end_x)),
                    str(self.var_y),
                    str((self.start_y, self.end_y)))

    def get_meshes(self):
        np = import_module('numpy')
        mesh_x, mesh_y = np.meshgrid(np.linspace(self.start_x, self.end_x,
                                                 num=self.nb_of_points_x),
                                     np.linspace(self.start_y, self.end_y,
                                                 num=self.nb_of_points_y))
        f = vectorized_lambdify((self.var_x, self.var_y), self.expr)
        return (mesh_x, mesh_y, f(mesh_x, mesh_y))

##############################################################################
# Backends
##############################################################################

class BaseBackend(object):
    def __init__(self, parent):
        super(BaseBackend, self).__init__()
        self.parent = parent

## don't have to check for the success of importing matplotlib in each case;
## we will only be using this backend if we can successfully import matploblib
class MatplotlibBackend(BaseBackend):
    def __init__(self, parent):
        super(MatplotlibBackend, self).__init__(parent)
        are_3D = [s.is_3D for s in self.parent._series]
        self.matplotlib = import_module('matplotlib',
            __import__kwargs={'fromlist': ['pyplot', 'cm', 'collections']},
            min_module_version='1.1.0', catch=(RuntimeError,))
        self.plt = self.matplotlib.pyplot
        self.cm = self.matplotlib.cm
        self.LineCollection = self.matplotlib.collections.LineCollection
        if any(are_3D) and not all(are_3D):
            raise ValueError('The matplotlib backend can not mix 2D and 3D.')
        elif not any(are_3D):
            self.fig = self.plt.figure()
            self.ax = self.fig.add_subplot(111)
            self.ax.spines['left'].set_position('zero')
            self.ax.spines['right'].set_color('none')
            self.ax.spines['bottom'].set_position('zero')
            self.ax.spines['top'].set_color('none')
            self.ax.spines['left'].set_smart_bounds(True)
            self.ax.spines['bottom'].set_smart_bounds(False)
            self.ax.xaxis.set_ticks_position('bottom')
            self.ax.yaxis.set_ticks_position('left')
        elif all(are_3D):
            ## mpl_toolkits.mplot3d is necessary for
            ##      projection='3d'
            mpl_toolkits = import_module('mpl_toolkits',
                                     __import__kwargs={'fromlist': ['mplot3d']})
            self.fig = self.plt.figure()
            self.ax = self.fig.add_subplot(111, projection='3d')

    def process_series(self):
        parent = self.parent

        for s in self.parent._series:
            # Create the collections
            if s.is_2Dline:
                collection = self.LineCollection(s.get_segments())
                self.ax.add_collection(collection)
            elif s.is_contour:
                self.ax.contour(*s.get_meshes())
            elif s.is_3Dline:
                # TODO too complicated, I blame matplotlib
                mpl_toolkits = import_module('mpl_toolkits',
                    __import__kwargs={'fromlist': ['mplot3d']})
                art3d = mpl_toolkits.mplot3d.art3d
                collection = art3d.Line3DCollection(s.get_segments())
                self.ax.add_collection(collection)
                x, y, z = s.get_points()
                self.ax.set_xlim((min(x), max(x)))
                self.ax.set_ylim((min(y), max(y)))
                self.ax.set_zlim((min(z), max(z)))
            elif s.is_3Dsurface:
                x, y, z = s.get_meshes()
                collection = self.ax.plot_surface(x, y, z, cmap=self.cm.jet,
                                                  rstride=1, cstride=1,
                                                  linewidth=0.1)
            elif s.is_implicit:
                #Smart bounds have to be set to False for implicit plots.
                self.ax.spines['left'].set_smart_bounds(False)
                self.ax.spines['bottom'].set_smart_bounds(False)
                points = s.get_raster()
                if len(points) == 2:
                    #interval math plotting
                    x, y = _matplotlib_list(points[0])
                    self.ax.fill(x, y, facecolor=s.line_color, edgecolor='None')
                else:
                    # use contourf or contour depending on whether it is
                    # an inequality or equality.
                    #XXX: ``contour`` plots multiple lines. Should be fixed.
                    ListedColormap = self.matplotlib.colors.ListedColormap
                    colormap = ListedColormap(["white", s.line_color])
                    xarray, yarray, zarray, plot_type = points
                    if plot_type == 'contour':
                        self.ax.contour(xarray, yarray, zarray,
                                contours=(0, 0), fill=False, cmap=colormap)
                    else:
                        self.ax.contourf(xarray, yarray, zarray, cmap=colormap)
            else:
                raise ValueError('The matplotlib backend supports only '
                                 'is_2Dline, is_3Dline, is_3Dsurface and '
                                 'is_contour objects.')

            # Customise the collections with the corresponding per-series
            # options.
            if hasattr(s, 'label'):
                collection.set_label(s.label)
            if s.is_line and s.line_color:
                if isinstance(s.line_color, (float, int)) or isinstance(s.line_color, Callable):
                    color_array = s.get_color_array()
                    collection.set_array(color_array)
                else:
                    collection.set_color(s.line_color)
            if s.is_3Dsurface and s.surface_color:
                if self.matplotlib.__version__ < "1.2.0":  # TODO in the distant future remove this check
                    warnings.warn('The version of matplotlib is too old to use surface coloring.')
                elif isinstance(s.surface_color, (float, int)) or isinstance(s.surface_color, Callable):
                    color_array = s.get_color_array()
                    color_array = color_array.reshape(color_array.size)
                    collection.set_array(color_array)
                else:
                    collection.set_color(s.surface_color)

        # Set global options.
        # TODO The 3D stuff
        # XXX The order of those is important.

        mpl_toolkits = import_module('mpl_toolkits',
            __import__kwargs={'fromlist': ['mplot3d']})
        Axes3D = mpl_toolkits.mplot3d.Axes3D
        if parent.xscale and not isinstance(self.ax, Axes3D):
            self.ax.set_xscale(parent.xscale)
        if parent.yscale and not isinstance(self.ax, Axes3D):
            self.ax.set_yscale(parent.yscale)
        if parent.xlim:
            self.ax.set_xlim(parent.xlim)
        else:
            if all(isinstance(s, LineOver1DRangeSeries) for s in parent._series):
                starts = [s.start for s in parent._series]
                ends = [s.end for s in parent._series]
                self.ax.set_xlim(min(starts), max(ends))
        if parent.ylim:
            self.ax.set_ylim(parent.ylim)
        if not isinstance(self.ax, Axes3D) or self.matplotlib.__version__ >= '1.2.0':  # XXX in the distant future remove this check
            self.ax.set_autoscale_on(parent.autoscale)
        if parent.axis_center:
            val = parent.axis_center
            if isinstance(self.ax, Axes3D):
                pass
            elif val == 'center':
                self.ax.spines['left'].set_position('center')
                self.ax.spines['bottom'].set_position('center')
            elif val == 'auto':
                xl, xh = self.ax.get_xlim()
                yl, yh = self.ax.get_ylim()
                pos_left = ('data', 0) if xl*xh <= 0 else 'center'
                pos_bottom = ('data', 0) if yl*yh <= 0 else 'center'
                self.ax.spines['left'].set_position(pos_left)
                self.ax.spines['bottom'].set_position(pos_bottom)
            else:
                self.ax.spines['left'].set_position(('data', val[0]))
                self.ax.spines['bottom'].set_position(('data', val[1]))
        if not parent.axis:
            self.ax.set_axis_off()
        if parent.legend:
            if self.ax.legend():
                self.ax.legend_.set_visible(parent.legend)
        if parent.margin:
            self.ax.set_xmargin(parent.margin)
            self.ax.set_ymargin(parent.margin)
        if parent.title:
            self.ax.set_title(parent.title)
        if parent.xlabel:
            self.ax.set_xlabel(parent.xlabel, position=(1, 0))
        if parent.ylabel:
            self.ax.set_ylabel(parent.ylabel, position=(0, 1))

    def show(self):
        self.process_series()
        #TODO after fixing https://github.com/ipython/ipython/issues/1255
        # you can uncomment the next line and remove the pyplot.show() call
        #self.fig.show()
        if _show:
            self.plt.show()

    def save(self, path):
        self.process_series()
        self.fig.savefig(path)

    def close(self):
        self.plt.close(self.fig)

class TextBackend(BaseBackend):
    def __init__(self, parent):
        super(TextBackend, self).__init__(parent)

    def show(self):
        if len(self.parent._series) != 1:
            raise ValueError(
                'The TextBackend supports only one graph per Plot.')
        elif not isinstance(self.parent._series[0], LineOver1DRangeSeries):
            raise ValueError(
                'The TextBackend supports only expressions over a 1D range')
        else:
            ser = self.parent._series[0]
            textplot(ser.expr, ser.start, ser.end)

    def close(self):
        pass

class DefaultBackend(BaseBackend):
    def __new__(cls, parent):
        matplotlib = import_module('matplotlib', min_module_version='1.1.0', catch=(RuntimeError,))
        if matplotlib:
            return MatplotlibBackend(parent)
        else:
            return TextBackend(parent)

plot_backends = {
    'matplotlib': MatplotlibBackend,
    'text': TextBackend,
    'default': DefaultBackend
}

##############################################################################
# Finding the centers of line segments or mesh faces
##############################################################################

def centers_of_segments(array):
    np = import_module('numpy')
    return np.average(np.vstack((array[:-1], array[1:])), 0)

def centers_of_faces(array):
    np = import_module('numpy')
    return np.average(np.dstack((array[:-1, :-1],
                                 array[1:, :-1],
                                 array[:-1, 1: ],
                                 array[:-1, :-1],
                                 )), 2)

def flat(x, y, z, eps=1e-3):
    """Checks whether three points are almost collinear"""
    np = import_module('numpy')
    # Workaround plotting piecewise (#8577):
    #   workaround for `lambdify` in `.experimental_lambdify` fails
    #   to return numerical values in some cases. Lower-level fix
    #   in `lambdify` is possible.
    vector_a = (x - y).astype(np.float)
    vector_b = (z - y).astype(np.float)
    dot_product = np.dot(vector_a, vector_b)
    vector_a_norm = np.linalg.norm(vector_a)
    vector_b_norm = np.linalg.norm(vector_b)
    cos_theta = dot_product / (vector_a_norm * vector_b_norm)
    return abs(cos_theta + 1) < eps

def _matplotlib_list(interval_list):
    """
    Returns lists for matplotlib ``fill`` command from a list of bounding
    rectangular intervals
    """
    xlist = []
    ylist = []
    if len(interval_list):
        for intervals in interval_list:
            intervalx = intervals[0]
            intervaly = intervals[1]
            xlist.extend([intervalx.start, intervalx.start,
                          intervalx.end, intervalx.end, None])
            ylist.extend([intervaly.start, intervaly.end,
                          intervaly.end, intervaly.start, None])
    else:
        #XXX Ugly hack. Matplotlib does not accept empty lists for ``fill``
        xlist.extend([None, None, None, None])
        ylist.extend([None, None, None, None])
    return xlist, ylist

####New API for plotting module ####

# TODO: Add color arrays for plots.
# TODO: Add more plotting options for 3d plots.
# TODO: Adaptive sampling for 3D plots.

@doctest_depends_on(modules=('numpy', 'matplotlib',))
def plot(*args, **kwargs):
    """
    Plots a function of a single variable and returns an instance of
    the ``Plot`` class (also, see the description of the
    ``show`` keyword argument below).

    The plotting uses an adaptive algorithm which samples recursively to
    accurately plot the plot. The adaptive algorithm uses a random point near
    the midpoint of two points that has to be further sampled. Hence the same
    plots can appear slightly different.

    Usage
    =====

    Single Plot

    ``plot(expr, range, **kwargs)``

    If the range is not specified, then a default range of (-10, 10) is used.

    Multiple plots with same range.

    ``plot(expr1, expr2, ..., range, **kwargs)``

    If the range is not specified, then a default range of (-10, 10) is used.

    Multiple plots with different ranges.

    ``plot((expr1, range), (expr2, range), ..., **kwargs)``

    Range has to be specified for every expression.

    Default range may change in the future if a more advanced default range
    detection algorithm is implemented.

    Arguments
    =========

    ``expr`` : Expression representing the function of single variable

    ``range``: (x, 0, 5), A 3-tuple denoting the range of the free variable.

    Keyword Arguments
    =================

    Arguments for ``plot`` function:

    ``show``: Boolean. The default value is set to ``True``. Set show to
    ``False`` and the function will not display the plot. The returned
    instance of the ``Plot`` class can then be used to save or display
    the plot by calling the ``save()`` and ``show()`` methods
    respectively.

    Arguments for ``LineOver1DRangeSeries`` class:

    ``adaptive``: Boolean. The default value is set to True. Set adaptive to False and
    specify ``nb_of_points`` if uniform sampling is required.

    ``depth``: int Recursion depth of the adaptive algorithm. A depth of value ``n``
    samples a maximum of `2^{n}` points.

    ``nb_of_points``: int. Used when the ``adaptive`` is set to False. The function
    is uniformly sampled at ``nb_of_points`` number of points.

    Aesthetics options:

    ``line_color``: float. Specifies the color for the plot.
    See ``Plot`` to see how to set color for the plots.

    If there are multiple plots, then the same series series are applied to
    all the plots. If you want to set these options separately, you can index
    the ``Plot`` object returned and set it.

    Arguments for ``Plot`` class:

    ``title`` : str. Title of the plot. It is set to the latex representation of
    the expression, if the plot has only one expression.

    ``xlabel`` : str. Label for the x-axis.

    ``ylabel`` : str. Label for the y-axis.

    ``xscale``: {'linear', 'log'} Sets the scaling of the x-axis.

    ``yscale``: {'linear', 'log'} Sets the scaling if the y-axis.

    ``axis_center``: tuple of two floats denoting the coordinates of the center or
    {'center', 'auto'}

    ``xlim`` : tuple of two floats, denoting the x-axis limits.

    ``ylim`` : tuple of two floats, denoting the y-axis limits.

    Examples
    ========

    >>> from sympy import symbols
    >>> from sympy.plotting import plot
    >>> x = symbols('x')

    Single Plot

    >>> plot(x**2, (x, -5, 5))
    Plot object containing:
    [0]: cartesian line: x**2 for x over (-5.0, 5.0)

    Multiple plots with single range.

    >>> plot(x, x**2, x**3, (x, -5, 5))
    Plot object containing:
    [0]: cartesian line: x for x over (-5.0, 5.0)
    [1]: cartesian line: x**2 for x over (-5.0, 5.0)
    [2]: cartesian line: x**3 for x over (-5.0, 5.0)

    Multiple plots with different ranges.

    >>> plot((x**2, (x, -6, 6)), (x, (x, -5, 5)))
    Plot object containing:
    [0]: cartesian line: x**2 for x over (-6.0, 6.0)
    [1]: cartesian line: x for x over (-5.0, 5.0)

    No adaptive sampling.

    >>> plot(x**2, adaptive=False, nb_of_points=400)
    Plot object containing:
    [0]: cartesian line: x**2 for x over (-10.0, 10.0)

    See Also
    ========

    Plot, LineOver1DRangeSeries.

    """
    args = list(map(sympify, args))
    free = set()
    for a in args:
        if isinstance(a, Expr):
            free |= a.free_symbols
            if len(free) > 1:
                raise ValueError(
                    'The same variable should be used in all '
                    'univariate expressions being plotted.')
    x = free.pop() if free else Symbol('x')
    kwargs.setdefault('xlabel', x.name)
    kwargs.setdefault('ylabel', 'f(%s)' % x.name)
    show = kwargs.pop('show', True)
    series = []
    plot_expr = check_arguments(args, 1, 1)
    series = [LineOver1DRangeSeries(*arg, **kwargs) for arg in plot_expr]

    plots = Plot(*series, **kwargs)
    if show:
        plots.show()
    return plots

@doctest_depends_on(modules=('numpy', 'matplotlib',))
def plot_parametric(*args, **kwargs):
    """
    Plots a 2D parametric plot.

    The plotting uses an adaptive algorithm which samples recursively to
    accurately plot the plot. The adaptive algorithm uses a random point near
    the midpoint of two points that has to be further sampled. Hence the same
    plots can appear slightly different.

    Usage
    =====

    Single plot.

    ``plot_parametric(expr_x, expr_y, range, **kwargs)``

    If the range is not specified, then a default range of (-10, 10) is used.

    Multiple plots with same range.

    ``plot_parametric((expr1_x, expr1_y), (expr2_x, expr2_y), range, **kwargs)``

    If the range is not specified, then a default range of (-10, 10) is used.

    Multiple plots with different ranges.

    ``plot_parametric((expr_x, expr_y, range), ..., **kwargs)``

    Range has to be specified for every expression.

    Default range may change in the future if a more advanced default range
    detection algorithm is implemented.

    Arguments
    =========

    ``expr_x`` : Expression representing the function along x.

    ``expr_y`` : Expression representing the function along y.

    ``range``: (u, 0, 5), A 3-tuple denoting the range of the parameter
    variable.

    Keyword Arguments
    =================

    Arguments for ``Parametric2DLineSeries`` class:

    ``adaptive``: Boolean. The default value is set to True. Set adaptive to
    False and specify ``nb_of_points`` if uniform sampling is required.

    ``depth``: int Recursion depth of the adaptive algorithm. A depth of
    value ``n`` samples a maximum of `2^{n}` points.

    ``nb_of_points``: int. Used when the ``adaptive`` is set to False. The
    function is uniformly sampled at ``nb_of_points`` number of points.

    Aesthetics
    ----------

    ``line_color``: function which returns a float. Specifies the color for the
    plot. See ``sympy.plotting.Plot`` for more details.

    If there are multiple plots, then the same Series arguments are applied to
    all the plots. If you want to set these options separately, you can index
    the returned ``Plot`` object and set it.

    Arguments for ``Plot`` class:

    ``xlabel`` : str. Label for the x-axis.

    ``ylabel`` : str. Label for the y-axis.

    ``xscale``: {'linear', 'log'} Sets the scaling of the x-axis.

    ``yscale``: {'linear', 'log'} Sets the scaling if the y-axis.

    ``axis_center``: tuple of two floats denoting the coordinates of the center
    or {'center', 'auto'}

    ``xlim`` : tuple of two floats, denoting the x-axis limits.

    ``ylim`` : tuple of two floats, denoting the y-axis limits.

    Examples
    ========

    >>> from sympy import symbols, cos, sin
    >>> from sympy.plotting import plot_parametric
    >>> u = symbols('u')

    Single Parametric plot

    >>> plot_parametric(cos(u), sin(u), (u, -5, 5))
    Plot object containing:
    [0]: parametric cartesian line: (cos(u), sin(u)) for u over (-5.0, 5.0)

    Multiple parametric plot with single range.

    >>> plot_parametric((cos(u), sin(u)), (u, cos(u)))
    Plot object containing:
    [0]: parametric cartesian line: (cos(u), sin(u)) for u over (-10.0, 10.0)
    [1]: parametric cartesian line: (u, cos(u)) for u over (-10.0, 10.0)

    Multiple parametric plots.

    >>> plot_parametric((cos(u), sin(u), (u, -5, 5)),
    ...     (cos(u), u, (u, -5, 5)))
    Plot object containing:
    [0]: parametric cartesian line: (cos(u), sin(u)) for u over (-5.0, 5.0)
    [1]: parametric cartesian line: (cos(u), u) for u over (-5.0, 5.0)

    See Also
    ========
    Plot, Parametric2DLineSeries

    """
    args = list(map(sympify, args))
    show = kwargs.pop('show', True)
    series = []
    plot_expr = check_arguments(args, 2, 1)
    series = [Parametric2DLineSeries(*arg, **kwargs) for arg in plot_expr]
    plots = Plot(*series, **kwargs)
    if show:
        plots.show()
    return plots

@doctest_depends_on(modules=('numpy', 'matplotlib',))
def plot3d_parametric_line(*args, **kwargs):
    """
    Plots a 3D parametric line plot.

    Usage
    =====

    Single plot:

    ``plot3d_parametric_line(expr_x, expr_y, expr_z, range, **kwargs)``

    If the range is not specified, then a default range of (-10, 10) is used.

    Multiple plots.

    ``plot3d_parametric_line((expr_x, expr_y, expr_z, range), ..., **kwargs)``

    Ranges have to be specified for every expression.

    Default range may change in the future if a more advanced default range
    detection algorithm is implemented.

    Arguments
    =========

    ``expr_x`` : Expression representing the function along x.

    ``expr_y`` : Expression representing the function along y.

    ``expr_z`` : Expression representing the function along z.

    ``range``: ``(u, 0, 5)``, A 3-tuple denoting the range of the parameter
    variable.

    Keyword Arguments
    =================

    Arguments for ``Parametric3DLineSeries`` class.

    ``nb_of_points``: The range is uniformly sampled at ``nb_of_points``
    number of points.

    Aesthetics:

    ``line_color``: function which returns a float. Specifies the color for the
    plot. See ``sympy.plotting.Plot`` for more details.

    If there are multiple plots, then the same series arguments are applied to
    all the plots. If you want to set these options separately, you can index
    the returned ``Plot`` object and set it.

    Arguments for ``Plot`` class.

    ``title`` : str. Title of the plot.

    Examples
    ========

    >>> from sympy import symbols, cos, sin
    >>> from sympy.plotting import plot3d_parametric_line
    >>> u = symbols('u')

    Single plot.

    >>> plot3d_parametric_line(cos(u), sin(u), u, (u, -5, 5))
    Plot object containing:
    [0]: 3D parametric cartesian line: (cos(u), sin(u), u) for u over (-5.0, 5.0)

    Multiple plots.

    >>> plot3d_parametric_line((cos(u), sin(u), u, (u, -5, 5)),
    ...     (sin(u), u**2, u, (u, -5, 5)))
    Plot object containing:
    [0]: 3D parametric cartesian line: (cos(u), sin(u), u) for u over (-5.0, 5.0)
    [1]: 3D parametric cartesian line: (sin(u), u**2, u) for u over (-5.0, 5.0)

    See Also
    ========

    Plot, Parametric3DLineSeries

    """
    args = list(map(sympify, args))
    show = kwargs.pop('show', True)
    series = []
    plot_expr = check_arguments(args, 3, 1)
    series = [Parametric3DLineSeries(*arg, **kwargs) for arg in plot_expr]
    plots = Plot(*series, **kwargs)
    if show:
        plots.show()
    return plots

@doctest_depends_on(modules=('numpy', 'matplotlib',))
def plot3d(*args, **kwargs):
    """
    Plots a 3D surface plot.

    Usage
    =====

    Single plot

    ``plot3d(expr, range_x, range_y, **kwargs)``

    If the ranges are not specified, then a default range of (-10, 10) is used.

    Multiple plot with the same range.

    ``plot3d(expr1, expr2, range_x, range_y, **kwargs)``

    If the ranges are not specified, then a default range of (-10, 10) is used.

    Multiple plots with different ranges.

    ``plot3d((expr1, range_x, range_y), (expr2, range_x, range_y), ..., **kwargs)``

    Ranges have to be specified for every expression.

    Default range may change in the future if a more advanced default range
    detection algorithm is implemented.

    Arguments
    =========

    ``expr`` : Expression representing the function along x.

    ``range_x``: (x, 0, 5), A 3-tuple denoting the range of the x
    variable.

    ``range_y``: (y, 0, 5), A 3-tuple denoting the range of the y
     variable.

    Keyword Arguments
    =================

    Arguments for ``SurfaceOver2DRangeSeries`` class:

    ``nb_of_points_x``: int. The x range is sampled uniformly at
    ``nb_of_points_x`` of points.

    ``nb_of_points_y``: int. The y range is sampled uniformly at
    ``nb_of_points_y`` of points.

    Aesthetics:

    ``surface_color``: Function which returns a float. Specifies the color for
    the surface of the plot. See ``sympy.plotting.Plot`` for more details.

    If there are multiple plots, then the same series arguments are applied to
    all the plots. If you want to set these options separately, you can index
    the returned ``Plot`` object and set it.

    Arguments for ``Plot`` class:

    ``title`` : str. Title of the plot.

    Examples
    ========

    >>> from sympy import symbols
    >>> from sympy.plotting import plot3d
    >>> x, y = symbols('x y')

    Single plot

    >>> plot3d(x*y, (x, -5, 5), (y, -5, 5))
    Plot object containing:
    [0]: cartesian surface: x*y for x over (-5.0, 5.0) and y over (-5.0, 5.0)

    Multiple plots with same range

    >>> plot3d(x*y, -x*y, (x, -5, 5), (y, -5, 5))
    Plot object containing:
    [0]: cartesian surface: x*y for x over (-5.0, 5.0) and y over (-5.0, 5.0)
    [1]: cartesian surface: -x*y for x over (-5.0, 5.0) and y over (-5.0, 5.0)

    Multiple plots with different ranges.

    >>> plot3d((x**2 + y**2, (x, -5, 5), (y, -5, 5)),
    ...     (x*y, (x, -3, 3), (y, -3, 3)))
    Plot object containing:
    [0]: cartesian surface: x**2 + y**2 for x over (-5.0, 5.0) and y over (-5.0, 5.0)
    [1]: cartesian surface: x*y for x over (-3.0, 3.0) and y over (-3.0, 3.0)

    See Also
    ========
    Plot, SurfaceOver2DRangeSeries

    """

    args = list(map(sympify, args))
    show = kwargs.pop('show', True)
    series = []
    plot_expr = check_arguments(args, 1, 2)
    series = [SurfaceOver2DRangeSeries(*arg, **kwargs) for arg in plot_expr]
    plots = Plot(*series, **kwargs)
    if show:
        plots.show()
    return plots

@doctest_depends_on(modules=('numpy', 'matplotlib',))
def plot3d_parametric_surface(*args, **kwargs):
    """
    Plots a 3D parametric surface plot.

    Usage
    =====

    Single plot.

    ``plot3d_parametric_surface(expr_x, expr_y, expr_z, range_u, range_v, **kwargs)``

    If the ranges is not specified, then a default range of (-10, 10) is used.

    Multiple plots.

    ``plot3d_parametric_surface((expr_x, expr_y, expr_z, range_u, range_v), ..., **kwargs)``

    Ranges have to be specified for every expression.

    Default range may change in the future if a more advanced default range
    detection algorithm is implemented.

    Arguments
    =========

    ``expr_x``: Expression representing the function along ``x``.

    ``expr_y``: Expression representing the function along ``y``.

    ``expr_z``: Expression representing the function along ``z``.

    ``range_u``: ``(u, 0, 5)``,  A 3-tuple denoting the range of the ``u``
    variable.

    ``range_v``: ``(v, 0, 5)``,  A 3-tuple denoting the range of the v
    variable.

    Keyword Arguments
    =================

    Arguments for ``ParametricSurfaceSeries`` class:

    ``nb_of_points_u``: int. The ``u`` range is sampled uniformly at
    ``nb_of_points_v`` of points

    ``nb_of_points_y``: int. The ``v`` range is sampled uniformly at
    ``nb_of_points_y`` of points

    Aesthetics:

    ``surface_color``: Function which returns a float. Specifies the color for
    the surface of the plot. See ``sympy.plotting.Plot`` for more details.

    If there are multiple plots, then the same series arguments are applied for
    all the plots. If you want to set these options separately, you can index
    the returned ``Plot`` object and set it.

    Arguments for ``Plot`` class:

    ``title`` : str. Title of the plot.

    Examples
    ========

    >>> from sympy import symbols, cos, sin
    >>> from sympy.plotting import plot3d_parametric_surface
    >>> u, v = symbols('u v')

    Single plot.

    >>> plot3d_parametric_surface(cos(u + v), sin(u - v), u - v,
    ...     (u, -5, 5), (v, -5, 5))
    Plot object containing:
    [0]: parametric cartesian surface: (cos(u + v), sin(u - v), u - v) for u over (-5.0, 5.0) and v over (-5.0, 5.0)

    See Also
    ========
    Plot, ParametricSurfaceSeries

    """

    args = list(map(sympify, args))
    show = kwargs.pop('show', True)
    series = []
    plot_expr = check_arguments(args, 3, 2)
    series = [ParametricSurfaceSeries(*arg, **kwargs) for arg in plot_expr]
    plots = Plot(*series, **kwargs)
    if show:
        plots.show()
    return plots

def check_arguments(args, expr_len, nb_of_free_symbols):
    """
    Checks the arguments and converts into tuples of the
    form (exprs, ranges)

    Examples
    ========

    >>> from sympy import plot, cos, sin, symbols
    >>> from sympy.plotting.plot import check_arguments
    >>> x = symbols('x')
    >>> check_arguments([cos(x), sin(x)], 2, 1)
        [(cos(x), sin(x), (x, -10, 10))]

    >>> check_arguments([x, x**2], 1, 1)
        [(x, (x, -10, 10)), (x**2, (x, -10, 10))]
    """
    if expr_len > 1 and isinstance(args[0], Expr):
        # Multiple expressions same range.
        # The arguments are tuples when the expression length is
        # greater than 1.
        if len(args) < expr_len:
            raise ValueError("len(args) should not be less than expr_len")
        for i in range(len(args)):
            if isinstance(args[i], Tuple):
                break
        else:
            i = len(args) + 1

        exprs = Tuple(*args[:i])
        free_symbols = list(set().union(*[e.free_symbols for e in exprs]))
        if len(args) == expr_len + nb_of_free_symbols:
            #Ranges given
            plots = [exprs + Tuple(*args[expr_len:])]
        else:
            default_range = Tuple(-10, 10)
            ranges = []
            for symbol in free_symbols:
                ranges.append(Tuple(symbol) + default_range)

            for i in range(len(free_symbols) - nb_of_free_symbols):
                ranges.append(Tuple(Dummy()) + default_range)
            plots = [exprs + Tuple(*ranges)]
        return plots

    if isinstance(args[0], Expr) or (isinstance(args[0], Tuple) and
                                     len(args[0]) == expr_len and
                                     expr_len != 3):
        # Cannot handle expressions with number of expression = 3. It is
        # not possible to differentiate between expressions and ranges.
        #Series of plots with same range
        for i in range(len(args)):
            if isinstance(args[i], Tuple) and len(args[i]) != expr_len:
                break
            if not isinstance(args[i], Tuple):
                args[i] = Tuple(args[i])
        else:
            i = len(args) + 1

        exprs = args[:i]
        assert all(isinstance(e, Expr) for expr in exprs for e in expr)
        free_symbols = list(set().union(*[e.free_symbols for expr in exprs
                                        for e in expr]))

        if len(free_symbols) > nb_of_free_symbols:
            raise ValueError("The number of free_symbols in the expression "
                             "is greater than %d" % nb_of_free_symbols)
        if len(args) == i + nb_of_free_symbols and isinstance(args[i], Tuple):
            ranges = Tuple(*[range_expr for range_expr in args[
                           i:i + nb_of_free_symbols]])
            plots = [expr + ranges for expr in exprs]
            return plots
        else:
            #Use default ranges.
            default_range = Tuple(-10, 10)
            ranges = []
            for symbol in free_symbols:
                ranges.append(Tuple(symbol) + default_range)

            for i in range(len(free_symbols) - nb_of_free_symbols):
                ranges.append(Tuple(Dummy()) + default_range)
            ranges = Tuple(*ranges)
            plots = [expr + ranges for expr in exprs]
            return plots

    elif isinstance(args[0], Tuple) and len(args[0]) == expr_len + nb_of_free_symbols:
        #Multiple plots with different ranges.
        for arg in args:
            for i in range(expr_len):
                if not isinstance(arg[i], Expr):
                    raise ValueError("Expected an expression, given %s" %
                                     str(arg[i]))
            for i in range(nb_of_free_symbols):
                if not len(arg[i + expr_len]) == 3:
                    raise ValueError("The ranges should be a tuple of "
                                     "length 3, got %s" % str(arg[i + expr_len]))
        return args


"""
=========================
Multi-dimensional scaling
=========================

An illustration of the metric and non-metric MDS on generated noisy data.

The reconstructed points using the metric MDS and non metric MDS are slightly
shifted to avoid overlapping.
"""

# Author: Nelle Varoquaux <nelle.varoquaux@gmail.com>
# Licence: BSD

print(__doc__)
import numpy as np

from matplotlib import pyplot as plt
from matplotlib.collections import LineCollection

from sklearn import manifold
from sklearn.metrics import euclidean_distances
from sklearn.decomposition import PCA

n_samples = 20
seed = np.random.RandomState(seed=3)
X_true = seed.randint(0, 20, 2 * n_samples).astype(np.float)
X_true = X_true.reshape((n_samples, 2))
# Center the data
X_true -= X_true.mean()

similarities = euclidean_distances(X_true)

# Add noise to the similarities
noise = np.random.rand(n_samples, n_samples)
noise = noise + noise.T
noise[np.arange(noise.shape[0]), np.arange(noise.shape[0])] = 0
similarities += noise

mds = manifold.MDS(n_components=2, max_iter=3000, eps=1e-9, random_state=seed,
                   dissimilarity="precomputed", n_jobs=1)
pos = mds.fit(similarities).embedding_

nmds = manifold.MDS(n_components=2, metric=False, max_iter=3000, eps=1e-12,
                    dissimilarity="precomputed", random_state=seed, n_jobs=1,
                    n_init=1)
npos = nmds.fit_transform(similarities, init=pos)

# Rescale the data
pos *= np.sqrt((X_true ** 2).sum()) / np.sqrt((pos ** 2).sum())
npos *= np.sqrt((X_true ** 2).sum()) / np.sqrt((npos ** 2).sum())

# Rotate the data
clf = PCA(n_components=2)
X_true = clf.fit_transform(X_true)

pos = clf.fit_transform(pos)

npos = clf.fit_transform(npos)

fig = plt.figure(1)
ax = plt.axes([0., 0., 1., 1.])

plt.scatter(X_true[:, 0], X_true[:, 1], c='r', s=20)
plt.scatter(pos[:, 0], pos[:, 1], s=20, c='g')
plt.scatter(npos[:, 0], npos[:, 1], s=20, c='b')
plt.legend(('True position', 'MDS', 'NMDS'), loc='best')

similarities = similarities.max() / similarities * 100
similarities[np.isinf(similarities)] = 0

# Plot the edges
start_idx, end_idx = np.where(pos)
#a sequence of (*line0*, *line1*, *line2*), where::
#            linen = (x0, y0), (x1, y1), ... (xm, ym)
segments = [[X_true[i, :], X_true[j, :]]
            for i in range(len(pos)) for j in range(len(pos))]
values = np.abs(similarities)
lc = LineCollection(segments,
                    zorder=0, cmap=plt.cm.hot_r,
                    norm=plt.Normalize(0, values.max()))
lc.set_array(similarities.flatten())
lc.set_linewidths(0.5 * np.ones(len(segments)))
ax.add_collection(lc)

plt.show()


#!/usr/bin/env python
import numpy
import sncosmo
import scipy.optimize

import matplotlib.pyplot as plt

model=sncosmo.Model(source='salt2-extended')

def f(t ,rlim):
    # print t, model.bandflux('desr',t, zp = rlim, zpsys='ab')
    return model.bandflux('desr',t, zp = rlim, zpsys='ab')-1.

def controlTime(z,rlim):
    model.set(z=z, t0=55000.)
    model.set_source_peakabsmag(absmag=-19.3,band='bessellb',magsys='ab')
    pre = scipy.optimize.fsolve(f, 55000.-15*(1+z) ,args=(rlim),xtol=1e-8)
    post = scipy.optimize.fsolve(f, 55000.+20*(1+z) ,args=(rlim),xtol=1e-8)
    return max(post[0]-pre[0],0)
    # print scipy.optimize.fsolve(f, 55000.+40,args=(rlim),factor=1.,xtol=1e-8)

def plot():

    lmag = numpy.arange(19.5,21.6,0.5)
    zs = numpy.arange(0.02, 0.2501,0.02)

    ans = []
    for lm in lmag:
        ans_=[]
        for z in zs:
            ans_.append(controlTime(z,lm))
        ans.append(ans_)

    for lm, ct in zip(lmag, ans):
        plt.plot(zs, ct, label = '$r_{{lim}} = {}$'.format(str(lm)))

    plt.xlabel(r'$z$')
    plt.ylabel(r'control time (days)')
    plt.legend()
    plt.show()


"""
=============================================
Comparison of kernel ridge regression and SVR
=============================================

Both kernel ridge regression (KRR) and SVR learn a non-linear function by
employing the kernel trick, i.e., they learn a linear function in the space
induced by the respective kernel which corresponds to a non-linear function in
the original space. They differ in the loss functions (ridge versus
epsilon-insensitive loss). In contrast to SVR, fitting a KRR can be done in
closed-form and is typically faster for medium-sized datasets. On the other
hand, the learned model is non-sparse and thus slower than SVR at
prediction-time.

This example illustrates both methods on an artificial dataset, which
consists of a sinusoidal target function and strong noise added to every fifth
datapoint. The first figure compares the learned model of KRR and SVR when both
complexity/regularization and bandwidth of the RBF kernel are optimized using
grid-search. The learned functions are very similar; however, fitting KRR is
approx. seven times faster than fitting SVR (both with grid-search). However,
prediction of 100000 target values is more than tree times faster with SVR
since it has learned a sparse model using only approx. 1/3 of the 100 training
datapoints as support vectors.

The next figure compares the time for fitting and prediction of KRR and SVR for
different sizes of the training set. Fitting KRR is faster than SVR for medium-
sized training sets (less than 1000 samples); however, for larger training sets
SVR scales better. With regard to prediction time, SVR is faster than
KRR for all sizes of the training set because of the learned sparse
solution. Note that the degree of sparsity and thus the prediction time depends
on the parameters epsilon and C of the SVR.
"""

# Authors: Jan Hendrik Metzen <jhm@informatik.uni-bremen.de>
# License: BSD 3 clause

from __future__ import division
import time

import numpy as np

from sklearn.svm import SVR
from sklearn.grid_search import GridSearchCV
from sklearn.learning_curve import learning_curve
from sklearn.kernel_ridge import KernelRidge
import matplotlib.pyplot as plt

rng = np.random.RandomState(0)

#############################################################################
# Generate sample data
X = 5 * rng.rand(10000, 1)
y = np.sin(X).ravel()

# Add noise to targets
y[::5] += 3 * (0.5 - rng.rand(X.shape[0]/5))

X_plot = np.linspace(0, 5, 100000)[:, None]

#############################################################################
# Fit regression model
train_size = 100
svr = GridSearchCV(SVR(kernel='rbf', gamma=0.1), cv=5,
                   param_grid={"C": [1e0, 1e1, 1e2, 1e3],
                               "gamma": np.logspace(-2, 2, 5)})

kr = GridSearchCV(KernelRidge(kernel='rbf', gamma=0.1), cv=5,
                  param_grid={"alpha": [1e0, 0.1, 1e-2, 1e-3],
                              "gamma": np.logspace(-2, 2, 5)})

t0 = time.time()
svr.fit(X[:train_size], y[:train_size])
svr_fit = time.time() - t0
print("SVR complexity and bandwidth selected and model fitted in %.3f s"
      % svr_fit)

t0 = time.time()
kr.fit(X[:train_size], y[:train_size])
kr_fit = time.time() - t0
print("KRR complexity and bandwidth selected and model fitted in %.3f s"
      % kr_fit)

sv_ratio = svr.best_estimator_.support_.shape[0] / train_size
print("Support vector ratio: %.3f" % sv_ratio)

t0 = time.time()
y_svr = svr.predict(X_plot)
svr_predict = time.time() - t0
print("SVR prediction for %d inputs in %.3f s"
      % (X_plot.shape[0], svr_predict))

t0 = time.time()
y_kr = kr.predict(X_plot)
kr_predict = time.time() - t0
print("KRR prediction for %d inputs in %.3f s"
      % (X_plot.shape[0], kr_predict))

#############################################################################
# look at the results
sv_ind = svr.best_estimator_.support_
plt.scatter(X[sv_ind], y[sv_ind], c='r', s=50, label='SVR support vectors')
plt.scatter(X[:100], y[:100], c='k', label='data')
plt.hold('on')
plt.plot(X_plot, y_svr, c='r',
         label='SVR (fit: %.3fs, predict: %.3fs)' % (svr_fit, svr_predict))
plt.plot(X_plot, y_kr, c='g',
         label='KRR (fit: %.3fs, predict: %.3fs)' % (kr_fit, kr_predict))
plt.xlabel('data')
plt.ylabel('target')
plt.title('SVR versus Kernel Ridge')
plt.legend()

# Visualize training and prediction time
plt.figure()

# Generate sample data
X = 5 * rng.rand(10000, 1)
y = np.sin(X).ravel()
y[::5] += 3 * (0.5 - rng.rand(X.shape[0]/5))
sizes = np.logspace(1, 4, 7)
for name, estimator in {"KRR": KernelRidge(kernel='rbf', alpha=0.1,
                                           gamma=10),
                        "SVR": SVR(kernel='rbf', C=1e1, gamma=10)}.items():
    train_time = []
    test_time = []
    for train_test_size in sizes:
        t0 = time.time()
        estimator.fit(X[:train_test_size], y[:train_test_size])
        train_time.append(time.time() - t0)

        t0 = time.time()
        estimator.predict(X_plot[:1000])
        test_time.append(time.time() - t0)

    plt.plot(sizes, train_time, 'o-', color="r" if name == "SVR" else "g",
             label="%s (train)" % name)
    plt.plot(sizes, test_time, 'o--', color="r" if name == "SVR" else "g",
             label="%s (test)" % name)

plt.xscale("log")
plt.yscale("log")
plt.xlabel("Train size")
plt.ylabel("Time (seconds)")
plt.title('Execution Time')
plt.legend(loc="best")

# Visualize learning curves
plt.figure()

svr = SVR(kernel='rbf', C=1e1, gamma=0.1)
kr = KernelRidge(kernel='rbf', alpha=0.1, gamma=0.1)
train_sizes, train_scores_svr, test_scores_svr = \
    learning_curve(svr, X[:100], y[:100], train_sizes=np.linspace(0.1, 1, 10),
                   scoring="mean_squared_error", cv=10)
train_sizes_abs, train_scores_kr, test_scores_kr = \
    learning_curve(kr, X[:100], y[:100], train_sizes=np.linspace(0.1, 1, 10),
                   scoring="mean_squared_error", cv=10)

plt.plot(train_sizes, test_scores_svr.mean(1), 'o-', color="r",
         label="SVR")
plt.plot(train_sizes, test_scores_kr.mean(1), 'o-', color="g",
         label="KRR")
plt.xlabel("Train size")
plt.ylabel("Mean Squared Error")
plt.title('Learning curves')
plt.legend(loc="best")

plt.show()


from __future__ import print_function
## 3D Lattice Boltzmann (BGK) model of a fluid.
## D3Q19 model. At each timestep, particle densities propagate
## outwards in the directions indicated in the figure. An
## equivalent 'equilibrium' density is found, and the densities
## relax towards that state, in a proportion governed by omega.
##               Iain Haslam, March 2006.
import util
if util.Benchmark().bohrium:
    import bohrium as np
else:
    import numpy as np

def main():
    B = util.Benchmark()

    nx      = B.size[0]
    ny      = B.size[1]
    nz      = B.size[2]
    ITER    = B.size[3]

    NO_OBST = 1
    omega   = 1.0
    density = 1.0
    deltaU  = 1e-7
    t1      = 1/3.0
    t2      = 1/18.0
    t3      = 1/36.0

    B.start()
    F       = np.ones((19, nx, ny, nz), dtype=np.float64)
    F[:]    = density/19.0
    FEQ     = np.ones((19, nx, ny, nz), dtype=np.float64)
    FEQ[:]  = density/19.0
    T       = np.zeros((19, nx, ny, nz), dtype=np.float64)

    #Create the scenery.
    BOUND   = np.zeros((nx, ny, nz), dtype=np.float64)
    BOUNDi  = np.ones((nx, ny, nz), dtype=np.float64)
    """
    if not NO_OBST:
        for i in xrange(nx):
            for j in xrange(ny):
                for k in xrange(nz):
                    if ((i-4)**2+(j-5)**2+(k-6)**2) < 6:
                        BOUND[i,j,k] += 1.0
                        BOUNDi[i,j,k] += 0.0
    BOUND[:,0,:]    += 1.0
    BOUNDi[:,0,:]   *= 0.0
    """

    if util.Benchmark().bohrium:
        np.flush()
    for ts in xrange(0, ITER):

        ##Propagate / Streaming step
        T[:] = F
        #nearest-neighbours
        F[1,:,:,0]   = T[1,:,:,-1]
        F[1,:,:,1:]  = T[1,:,:,:-1]
        F[2,:,:,:-1] = T[2,:,:,1:]
        F[2,:,:,-1]  = T[2,:,:,0]
        F[3,:,0,:]   = T[3,:,-1,:]
        F[3,:,1:,:]  = T[3,:,:-1,:]
        F[4,:,:-1,:] = T[4,:,1:,:]
        F[4,:,-1,:]  = T[4,:,0,:]
        F[5,0,:,:]   = T[5,-1,:,:]
        F[5,1:,:,:]  = T[5,:-1,:,:]
        F[6,:-1,:,:] = T[6,1:,:,:]
        F[6,-1,:,:]  = T[6,0,:,:]
        #next-nearest neighbours
        F[7,0 ,0 ,:] = T[7,-1 , -1,:]
        F[7,0 ,1:,:] = T[7,-1 ,:-1,:]
        F[7,1:,0 ,:] = T[7,:-1, -1,:]
        F[7,1:,1:,:] = T[7,:-1,:-1,:]

        F[8,0 ,:-1,:] = T[8,-1 ,1:,:]
        F[8,0 , -1,:] = T[8,-1 ,0 ,:]
        F[8,1:,:-1,:] = T[8,:-1,1:,:]
        F[8,1:, -1,:] = T[8,:-1,0 ,:]

        F[9,:-1,0 ,:] = T[9,1:, -1,:]
        F[9,:-1,1:,:] = T[9,1:,:-1,:]
        F[9,-1 ,0 ,:] = T[9,0 ,  0,:]
        F[9,-1 ,1:,:] = T[9,0 ,:-1,:]

        F[10,:-1,:-1,:] = T[10,1:,1:,:]
        F[10,:-1, -1,:] = T[10,1:,0 ,:]
        F[10,-1 ,:-1,:] = T[10,0 ,1:,:]
        F[10,-1 , -1,:] = T[10,0 ,0 ,:]

        F[11,0 ,:,0 ] = T[11,0  ,:, -1]
        F[11,0 ,:,1:] = T[11,0  ,:,:-1]
        F[11,1:,:,0 ] = T[11,:-1,:, -1]
        F[11,1:,:,1:] = T[11,:-1,:,:-1]

        F[12,0 ,:,:-1] = T[12, -1,:,1:]
        F[12,0 ,:, -1] = T[12, -1,:,0 ]
        F[12,1:,:,:-1] = T[12,:-1,:,1:]
        F[12,1:,:, -1] = T[12,:-1,:,0 ]

        F[13,:-1,:,0 ] = T[13,1:,:, -1]
        F[13,:-1,:,1:] = T[13,1:,:,:-1]
        F[13, -1,:,0 ] = T[13,0 ,:, -1]
        F[13, -1,:,1:] = T[13,0 ,:,:-1]

        F[14,:-1,:,:-1] = T[14,1:,:,1:]
        F[14,:-1,:, -1] = T[14,1:,:,0 ]
        F[14,-1 ,:,:-1] = T[14,0 ,:,1:]
        F[14,-1 ,:, -1] = T[14,0 ,:,0 ]

        F[15,:,0 ,0 ] = T[15,:, -1, -1]
        F[15,:,0 ,1:] = T[15,:, -1,:-1]
        F[15,:,1:,0 ] = T[15,:,:-1, -1]
        F[15,:,1:,1:] = T[15,:,:-1,:-1]

        F[16,:,0 ,:-1] = T[16,:, -1,1:]
        F[16,:,0 , -1] = T[16,:, -1,0 ]
        F[16,:,1:,:-1] = T[16,:,:-1,1:]
        F[16,:,1:, -1] = T[16,:,:-1,0 ]

        F[17,:,:-1,0 ] = T[17,:,1:, -1]
        F[17,:,:-1,1:] = T[17,:,1:,:-1]
        F[17,:, -1,0 ] = T[17,:,0 , -1]
        F[17,:, -1,1:] = T[17,:,0 ,:-1]

        F[18,:,:-1,:-1] = T[18,:,1:,1:]
        F[18,:,:-1, -1] = T[18,:,1:,0 ]
        F[18,:,-1 ,:-1] = T[18,:,0 ,1:]
        F[18,:,-1 , -1] = T[18,:,0 ,0 ]
        #Densities bouncing back at next timestep
        BB = np.empty(F.shape)
        T[:] = F
        T[1:,:,:,:] *= BOUND[np.newaxis,:,:,:]
        BB[2 ,:,:,:] += T[1 ,:,:,:]
        BB[1 ,:,:,:] += T[2 ,:,:,:]
        BB[4 ,:,:,:] += T[3 ,:,:,:]
        BB[3 ,:,:,:] += T[4 ,:,:,:]
        BB[6 ,:,:,:] += T[5 ,:,:,:]
        BB[5 ,:,:,:] += T[6 ,:,:,:]
        BB[10,:,:,:] += T[7 ,:,:,:]
        BB[9 ,:,:,:] += T[8 ,:,:,:]
        BB[8 ,:,:,:] += T[9 ,:,:,:]
        BB[7 ,:,:,:] += T[10,:,:,:]
        BB[14,:,:,:] += T[11,:,:,:]
        BB[13,:,:,:] += T[12,:,:,:]
        BB[12,:,:,:] += T[13,:,:,:]
        BB[11,:,:,:] += T[14,:,:,:]
        BB[18,:,:,:] += T[15,:,:,:]
        BB[17,:,:,:] += T[16,:,:,:]
        BB[16,:,:,:] += T[17,:,:,:]
        BB[15,:,:,:] += T[18,:,:,:]

        # Relax calculate equilibrium state (FEQ) with equivalent speed and density to F
        DENSITY = np.add.reduce(F)

        #UX = F[5,:,:,:].copy()
        UX = np.ones(F[5,:,:,:].shape, dtype=np.float64)
        UX[:,:,:] = F[5,:,:,:]

        UX += F[7,:,:,:]
        UX += F[8,:,:,:]
        UX += F[11,:,:,:]
        UX += F[12,:,:,:]
        UX -= F[6,:,:,:]
        UX -= F[9,:,:,:]
        UX -= F[10,:,:,:]
        UX -= F[13,:,:,:]
        UX -= F[14,:,:,:]
        UX /=DENSITY

        #UY = F[3,:,:,:].copy()
        UY = np.ones(F[3,:,:,:].shape, dtype=np.float64)
        UY[:,:,:] = F[3,:,:,:]

        UY += F[7,:,:,:]
        UY += F[9,:,:,:]
        UY += F[15,:,:,:]
        UY += F[16,:,:,:]
        UY -= F[4,:,:,:]
        UY -= F[8,:,:,:]
        UY -= F[10,:,:,:]
        UY -= F[17,:,:,:]
        UY -= F[18,:,:,:]
        UY /=DENSITY

        #UZ = F[1,:,:,:].copy()
        UZ = np.ones(F[1,:,:,:].shape, dtype=np.float64)
        UZ[:,:,:] = F[1,:,:,:]

        UZ += F[11,:,:,:]
        UZ += F[13,:,:,:]
        UZ += F[15,:,:,:]
        UZ += F[17,:,:,:]
        UZ -= F[2,:,:,:]
        UZ -= F[12,:,:,:]
        UZ -= F[14,:,:,:]
        UZ -= F[16,:,:,:]
        UZ -= F[18,:,:,:]
        UZ /=DENSITY

        UX[0,:,:] += deltaU #Increase inlet pressure
                            #Set bourderies to zero.
        UX[:,:,:] *= BOUNDi
        UY[:,:,:] *= BOUNDi
        UZ[:,:,:] *= BOUNDi
        DENSITY[:,:,:] *= BOUNDi

        U_SQU = UX**2 + UY**2 + UZ**2

        # Calculate equilibrium distribution: stationary
        FEQ[0,:,:,:] = (t1*DENSITY)*(1.0-3.0*U_SQU/2.0)
        # nearest-neighbours
        T1 = 3.0/2.0*U_SQU
        tDENSITY = t2*DENSITY
        FEQ[1,:,:,:]=tDENSITY*(1.0 + 3.0*UZ + 9.0/2.0*UZ**2 - T1)
        FEQ[2,:,:,:]=tDENSITY*(1.0 - 3.0*UZ + 9.0/2.0*UZ**2 - T1)
        FEQ[3,:,:,:]=tDENSITY*(1.0 + 3.0*UY + 9.0/2.0*UY**2 - T1)
        FEQ[4,:,:,:]=tDENSITY*(1.0 - 3.0*UY + 9.0/2.0*UY**2 - T1)
        FEQ[5,:,:,:]=tDENSITY*(1.0 + 3.0*UX + 9.0/2.0*UX**2 - T1)
        FEQ[6,:,:,:]=tDENSITY*(1.0 - 3.0*UX + 9.0/2.0*UX**2 - T1)
        # next-nearest neighbours
        T1 = 3.0*U_SQU/2.0
        tDENSITY = t3*DENSITY
        U8 = UX+UY
        FEQ[7,:,:,:] =tDENSITY*(1.0 + 3.0*U8  + 9.0/2.0*(U8)**2  - T1)
        U9 = UX-UY
        FEQ[8,:,:,:] =tDENSITY*(1.0 + 3.0*U9  + 9.0/2.0*(U9)**2  - T1)
        U10 = -UX+UY
        FEQ[9,:,:,:] =tDENSITY*(1.0 + 3.0*U10 + 9.0/2.0*(U10)**2 - T1)
        U8 *= -1.0
        FEQ[10,:,:,:]=tDENSITY*(1.0 + 3.0*U8 + 9.0/2.0*(U8)**2 - T1)
        U12 = UX+UZ
        FEQ[11,:,:,:]=tDENSITY*(1.0 + 3.0*U12 + 9.0/2.0*(U12)**2 - T1)
        U12 *= 1.0
        FEQ[14,:,:,:]=tDENSITY*(1.0 + 3.0*U12 + 9.0/2.0*(U12)**2 - T1)
        U13 = UX-UZ
        FEQ[12,:,:,:]=tDENSITY*(1.0 + 3.0*U13 + 9.0/2.0*(U13)**2 - T1)
        U13 *= -1.0
        FEQ[13,:,:,:]=tDENSITY*(1.0 + 3.0*U13 + 9.0/2.0*(U13)**2 - T1)
        U16 = UY+UZ
        FEQ[15,:,:,:]=tDENSITY*(1.0 + 3.0*U16 + 9.0/2.0*(U16)**2 - T1)
        U17 = UY-UZ
        FEQ[16,:,:,:]=tDENSITY*(1.0 + 3.0*U17 + 9.0/2.0*(U17)**2 - T1)
        U17 *= -1.0
        FEQ[17,:,:,:]=tDENSITY*(1.0 + 3.0*U17 + 9.0/2.0*(U17)**2 - T1)
        U16 *= -1.0
        FEQ[18,:,:,:]=tDENSITY*(1.0 + 3.0*U16 + 9.0/2.0*(U16)**2 - T1)
        F *= (1.0-omega)
        F += omega * FEQ

        #Densities bouncing back at next timestep
        F[1:,:,:,:] *= BOUNDi[np.newaxis,:,:,:]
        F[1:,:,:,:] += BB[1:,:,:,:]

        del BB
        del T1
        del UX, UY, UZ
        del U_SQU
        del DENSITY, tDENSITY
        del U8, U9, U10, U12, U13, U16, U17
        if util.Benchmark().bohrium:
            np.flush()

    B.stop()
    B.pprint()

    if B.outputfn:
        B.tofile(B.outputfn, {'res': UX})

    """
    import matplotlib.pyplot as plt
    UX *= -1
    plt.hold(True)
    plt.quiver(UY[:,:,4],UX[:,:,4], pivot='middle')
    plt.imshow(BOUND[:,:,4])
    plt.show()

    """

if __name__ == "__main__":
    main()


#
# Copyright 2013 Quantopian, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pandas as pd
import pytz

from datetime import datetime
from dateutil import rrule
from functools import partial

start = pd.Timestamp('1990-01-01', tz='UTC')
end_base = pd.Timestamp('today', tz='UTC')
# Give an aggressive buffer for logic that needs to use the next trading
# day or minute.
end = end_base + pd.Timedelta(days=365)

def canonicalize_datetime(dt):
    # Strip out any HHMMSS or timezone info in the user's datetime, so that
    # all the datetimes we return will be 00:00:00 UTC.
    return datetime(dt.year, dt.month, dt.day, tzinfo=pytz.utc)

def get_non_trading_days(start, end):
    non_trading_rules = []

    start = canonicalize_datetime(start)
    end = canonicalize_datetime(end)

    weekends = rrule.rrule(
        rrule.YEARLY,
        byweekday=(rrule.SA, rrule.SU),
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(weekends)

    new_years = rrule.rrule(
        rrule.MONTHLY,
        byyearday=1,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(new_years)

    new_years_sunday = rrule.rrule(
        rrule.MONTHLY,
        byyearday=2,
        byweekday=rrule.MO,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(new_years_sunday)

    mlk_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=1,
        byweekday=(rrule.MO(+3)),
        cache=True,
        dtstart=datetime(1998, 1, 1, tzinfo=pytz.utc),
        until=end
    )
    non_trading_rules.append(mlk_day)

    presidents_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=2,
        byweekday=(rrule.MO(3)),
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(presidents_day)

    good_friday = rrule.rrule(
        rrule.DAILY,
        byeaster=-2,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(good_friday)

    memorial_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=5,
        byweekday=(rrule.MO(-1)),
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(memorial_day)

    july_4th = rrule.rrule(
        rrule.MONTHLY,
        bymonth=7,
        bymonthday=4,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(july_4th)

    july_4th_sunday = rrule.rrule(
        rrule.MONTHLY,
        bymonth=7,
        bymonthday=5,
        byweekday=rrule.MO,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(july_4th_sunday)

    july_4th_saturday = rrule.rrule(
        rrule.MONTHLY,
        bymonth=7,
        bymonthday=3,
        byweekday=rrule.FR,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(july_4th_saturday)

    labor_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=9,
        byweekday=(rrule.MO(1)),
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(labor_day)

    thanksgiving = rrule.rrule(
        rrule.MONTHLY,
        bymonth=11,
        byweekday=(rrule.TH(4)),
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(thanksgiving)

    christmas = rrule.rrule(
        rrule.MONTHLY,
        bymonth=12,
        bymonthday=25,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(christmas)

    christmas_sunday = rrule.rrule(
        rrule.MONTHLY,
        bymonth=12,
        bymonthday=26,
        byweekday=rrule.MO,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(christmas_sunday)

    # If Christmas is a Saturday then 24th, a Friday is observed.
    christmas_saturday = rrule.rrule(
        rrule.MONTHLY,
        bymonth=12,
        bymonthday=24,
        byweekday=rrule.FR,
        cache=True,
        dtstart=start,
        until=end
    )
    non_trading_rules.append(christmas_saturday)

    non_trading_ruleset = rrule.rruleset()

    for rule in non_trading_rules:
        non_trading_ruleset.rrule(rule)

    non_trading_days = non_trading_ruleset.between(start, end, inc=True)

    # Add September 11th closings
    # http://en.wikipedia.org/wiki/Aftermath_of_the_September_11_attacks
    # Due to the terrorist attacks, the stock market did not open on 9/11/2001
    # It did not open again until 9/17/2001.
    #
    #    September 2001
    # Su Mo Tu We Th Fr Sa
    #                    1
    #  2  3  4  5  6  7  8
    #  9 10 11 12 13 14 15
    # 16 17 18 19 20 21 22
    # 23 24 25 26 27 28 29
    # 30

    for day_num in range(11, 17):
        non_trading_days.append(
            datetime(2001, 9, day_num, tzinfo=pytz.utc))

    # Add closings due to Hurricane Sandy in 2012
    # http://en.wikipedia.org/wiki/Hurricane_sandy
    #
    # The stock exchange was closed due to Hurricane Sandy's
    # impact on New York.
    # It closed on 10/29 and 10/30, reopening on 10/31
    #     October 2012
    # Su Mo Tu We Th Fr Sa
    #     1  2  3  4  5  6
    #  7  8  9 10 11 12 13
    # 14 15 16 17 18 19 20
    # 21 22 23 24 25 26 27
    # 28 29 30 31

    for day_num in range(29, 31):
        non_trading_days.append(
            datetime(2012, 10, day_num, tzinfo=pytz.utc))

    # Misc closings from NYSE listing.
    # http://www.nyse.com/pdfs/closings.pdf
    #
    # National Days of Mourning
    # - President Richard Nixon
    non_trading_days.append(datetime(1994, 4, 27, tzinfo=pytz.utc))
    # - President Ronald W. Reagan - June 11, 2004
    non_trading_days.append(datetime(2004, 6, 11, tzinfo=pytz.utc))
    # - President Gerald R. Ford - Jan 2, 2007
    non_trading_days.append(datetime(2007, 1, 2, tzinfo=pytz.utc))

    non_trading_days.sort()
    return pd.DatetimeIndex(non_trading_days)

non_trading_days = get_non_trading_days(start, end)
trading_day = pd.tseries.offsets.CDay(holidays=non_trading_days)

def get_trading_days(start, end, trading_day=trading_day):
    return pd.date_range(start=start.date(),
                         end=end.date(),
                         freq=trading_day).tz_localize('UTC')

trading_days = get_trading_days(start, end)

def get_early_closes(start, end):
    # 1:00 PM close rules based on
    # http://quant.stackexchange.com/questions/4083/nyse-early-close-rules-july-4th-and-dec-25th # noqa
    # and verified against http://www.nyse.com/pdfs/closings.pdf

    # These rules are valid starting in 1993

    start = canonicalize_datetime(start)
    end = canonicalize_datetime(end)

    start = max(start, datetime(1993, 1, 1, tzinfo=pytz.utc))
    end = max(end, datetime(1993, 1, 1, tzinfo=pytz.utc))

    # Not included here are early closes prior to 1993
    # or unplanned early closes

    early_close_rules = []

    day_after_thanksgiving = rrule.rrule(
        rrule.MONTHLY,
        bymonth=11,
        # 4th Friday isn't correct if month starts on Friday, so restrict to
        # day range:
        byweekday=(rrule.FR),
        bymonthday=range(23, 30),
        cache=True,
        dtstart=start,
        until=end
    )
    early_close_rules.append(day_after_thanksgiving)

    christmas_eve = rrule.rrule(
        rrule.MONTHLY,
        bymonth=12,
        bymonthday=24,
        byweekday=(rrule.MO, rrule.TU, rrule.WE, rrule.TH),
        cache=True,
        dtstart=start,
        until=end
    )
    early_close_rules.append(christmas_eve)

    friday_after_christmas = rrule.rrule(
        rrule.MONTHLY,
        bymonth=12,
        bymonthday=26,
        byweekday=rrule.FR,
        cache=True,
        dtstart=start,
        # valid 1993-2007
        until=min(end, datetime(2007, 12, 31, tzinfo=pytz.utc))
    )
    early_close_rules.append(friday_after_christmas)

    day_before_independence_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=7,
        bymonthday=3,
        byweekday=(rrule.MO, rrule.TU, rrule.TH),
        cache=True,
        dtstart=start,
        until=end
    )
    early_close_rules.append(day_before_independence_day)

    day_after_independence_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=7,
        bymonthday=5,
        byweekday=rrule.FR,
        cache=True,
        dtstart=start,
        # starting in 2013: wednesday before independence day
        until=min(end, datetime(2012, 12, 31, tzinfo=pytz.utc))
    )
    early_close_rules.append(day_after_independence_day)

    wednesday_before_independence_day = rrule.rrule(
        rrule.MONTHLY,
        bymonth=7,
        bymonthday=3,
        byweekday=rrule.WE,
        cache=True,
        # starting in 2013
        dtstart=max(start, datetime(2013, 1, 1, tzinfo=pytz.utc)),
        until=max(end, datetime(2013, 1, 1, tzinfo=pytz.utc))
    )
    early_close_rules.append(wednesday_before_independence_day)

    early_close_ruleset = rrule.rruleset()

    for rule in early_close_rules:
        early_close_ruleset.rrule(rule)
    early_closes = early_close_ruleset.between(start, end, inc=True)

    # Misc early closings from NYSE listing.
    # http://www.nyse.com/pdfs/closings.pdf
    #
    # New Year's Eve
    nye_1999 = datetime(1999, 12, 31, tzinfo=pytz.utc)
    if start <= nye_1999 and nye_1999 <= end:
        early_closes.append(nye_1999)

    early_closes.sort()
    return pd.DatetimeIndex(early_closes)

early_closes = get_early_closes(start, end)

def get_open_and_close(day, early_closes):
    market_open = pd.Timestamp(
        datetime(
            year=day.year,
            month=day.month,
            day=day.day,
            hour=9,
            minute=31),
        tz='US/Eastern').tz_convert('UTC')
    # 1 PM if early close, 4 PM otherwise
    close_hour = 13 if day in early_closes else 16
    market_close = pd.Timestamp(
        datetime(
            year=day.year,
            month=day.month,
            day=day.day,
            hour=close_hour),
        tz='US/Eastern').tz_convert('UTC')

    return market_open, market_close

def get_open_and_closes(trading_days, early_closes, get_open_and_close):
    open_and_closes = pd.DataFrame(index=trading_days,
                                   columns=('market_open', 'market_close'))

    get_o_and_c = partial(get_open_and_close, early_closes=early_closes)

    open_and_closes['market_open'], open_and_closes['market_close'] = \
        zip(*open_and_closes.index.map(get_o_and_c))

    return open_and_closes

open_and_closes = get_open_and_closes(trading_days, early_closes,
                                      get_open_and_close)


"""
===================
Isotonic Regression
===================

An illustration of the isotonic regression on generated data. The
isotonic regression finds a non-decreasing approximation of a function
while minimizing the mean squared error on the training data. The benefit
of such a model is that it does not assume any form for the target
function such as linearity. For comparison a linear regression is also
presented.

"""
print(__doc__)

# Author: Nelle Varoquaux <nelle.varoquaux@gmail.com>
#         Alexandre Gramfort <alexandre.gramfort@inria.fr>
# Licence: BSD

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection

from sklearn.linear_model import LinearRegression
from sklearn.isotonic import IsotonicRegression
from sklearn.utils import check_random_state

n = 100
x = np.arange(n)
rs = check_random_state(0)
y = rs.randint(-50, 50, size=(n,)) + 50. * np.log(1 + np.arange(n))

###############################################################################
# Fit IsotonicRegression and LinearRegression models

ir = IsotonicRegression()

y_ = ir.fit_transform(x, y)

lr = LinearRegression()
lr.fit(x[:, np.newaxis], y)  # x needs to be 2d for LinearRegression

###############################################################################
# plot result

segments = [[[i, y[i]], [i, y_[i]]] for i in range(n)]
lc = LineCollection(segments, zorder=0)
lc.set_array(np.ones(len(y)))
lc.set_linewidths(0.5 * np.ones(n))

fig = plt.figure()
plt.plot(x, y, 'r.', markersize=12)
plt.plot(x, y_, 'g.-', markersize=12)
plt.plot(x, lr.predict(x[:, np.newaxis]), 'b-')
plt.gca().add_collection(lc)
plt.legend(('Data', 'Isotonic Fit', 'Linear Fit'), loc='lower right')
plt.title('Isotonic regression')
plt.show()