"""
Wrappers for :mod:`matplotlib` functions to facilitate plotting grids,
2D objects, etc.

This module loads all functions from :mod:`matplotlib.pyplot`, adds new
functions and overwrites some others (like :func:`~fatiando.vis.mpl.contour`,
:func:`~fatiando.vis.mpl.pcolor`, etc).

**Grids**

* :func:`~fatiando.vis.mpl.contour`
* :func:`~fatiando.vis.mpl.contourf`
* :func:`~fatiando.vis.mpl.pcolor`

Grids are automatically reshaped and interpolated if desired or necessary.

**2D objects**

* :func:`~fatiando.vis.mpl.points`
* :func:`~fatiando.vis.mpl.paths`
* :func:`~fatiando.vis.mpl.square`
* :func:`~fatiando.vis.mpl.squaremesh`
* :func:`~fatiando.vis.mpl.polygon`
* :func:`~fatiando.vis.mpl.layers`
* :func:`~fatiando.vis.mpl.seismic_image`
* :func:`~fatiando.vis.mpl.seismic_wiggle`

**Interactive**

* :func:`~fatiando.vis.mpl.draw_polygon`
* :func:`~fatiando.vis.mpl.draw_layers`
* :func:`~fatiando.vis.mpl.pick_points`

**Basemap (map projections)**

* :func:`~fatiando.vis.mpl.basemap`
* :func:`~fatiando.vis.mpl.draw_geolines`
* :func:`~fatiando.vis.mpl.draw_countries`
* :func:`~fatiando.vis.mpl.draw_coastlines`

**Auxiliary**

* :func:`~fatiando.vis.mpl.set_area`
* :func:`~fatiando.vis.mpl.m2km`

----

"""

import numpy
from matplotlib import pyplot, widgets
# Quick hack so that the docs can build using the mocks for readthedocs
# Ideal would be to log an error message saying that functions from pyplot
# were not imported
try:
    from matplotlib.pyplot import *
except:
    pass

import fatiando.gridder

# Dummy variable to laizy import the basemap toolkit (slow)
Basemap = None

def draw_polygon(area, axes, style='-', marker='o', color='k', width=2,
                 alpha=0.5, xy2ne=False):
    """
    Draw a polygon by clicking with the mouse.

    INSTRUCTIONS:

    * Left click to pick the edges of the polygon;
    * Draw edges CLOCKWISE;
    * Press 'e' to erase the last edge;
    * Right click to close the polygon;
    * Close the figure window to finish;

    Parameters:

    * area : list = [x1, x2, y1, y2]
        Borders of the area containing the polygon
    * axes : matplotlib Axes
        The figure to use for drawing the polygon.
        To get an Axes instace, just do::

            from matplotlib import pyplot
            axes = pyplot.figure().add_subplot(1,1,1)

        You can plot things to ``axes`` before calling this function so that
        they'll appear on the background.
    * style : str
        Line style (as in matplotlib.pyplot.plot)
    * marker : str
        Style of the point markers (as in matplotlib.pyplot.plot)
    * color : str
        Line color (as in matplotlib.pyplot.plot)
    * width : float
        The line width (as in matplotlib.pyplot.plot)
    * alpha : float
        Transparency of the fill of the polygon. 0 for transparent, 1 for
        opaque (fills the polygon once done drawing)
    * xy2ne : True or False
        If True, will exchange the x and y axis so that x points north.
        Use this when drawing on a map viewed from above. If the y-axis of the
        plot is supposed to be z (depth), then use ``xy2ne=False``.

    Returns:

    * edges : list of lists
        List of ``[x, y]`` pairs with the edges of the polygon

    """
    axes.set_title("Click to draw polygon. Right click when done.")
    if xy2ne:
        axes.set_xlim(area[2], area[3])
        axes.set_ylim(area[0], area[1])
    else:
        axes.set_xlim(area[0], area[1])
        axes.set_ylim(area[2], area[3])
    # start with an empty line
    line, = axes.plot([], [], marker=marker, linestyle=style, color=color,
                      linewidth=width)
    tmpline, = axes.plot([], [], marker=marker, linestyle=style, color=color,
                         linewidth=width)
    draw = axes.figure.canvas.draw
    x = []
    y = []
    plotx = []
    ploty = []
    # Hack because Python 2 doesn't like nonlocal variables that change value.
    # Lists it doesn't mind.
    picking = [True]

    def draw_guide(px, py):
        if len(x) != 0:
            tmpline.set_data([x[-1], px], [y[-1], py])

    def move(event):
        if event.inaxes != axes:
            return 0
        if picking[0]:
            draw_guide(event.xdata, event.ydata)
            draw()

    def pick(event):
        if event.inaxes != axes:
            return 0
        if event.button == 1 and picking[0]:
            x.append(event.xdata)
            y.append(event.ydata)
            plotx.append(event.xdata)
            ploty.append(event.ydata)
        if event.button == 3 or event.button == 2 and picking[0]:
            if len(x) < 3:
                axes.set_title("Need at least 3 points to make a polygon")
            else:
                picking[0] = False
                axes.set_title("Done! You can close the window now.")
                plotx.append(x[0])
                ploty.append(y[0])
                tmpline.set_data([], [])
                axes.fill(plotx, ploty, color=color, alpha=alpha)
        line.set_data(plotx, ploty)
        draw()

    def erase(event):
        if event.key == 'e' and picking[0]:
            x.pop()
            y.pop()
            plotx.pop()
            ploty.pop()
            line.set_data(plotx, ploty)
            draw_guide(event.xdata, event.ydata)
            draw()
    line.figure.canvas.mpl_connect('button_press_event', pick)
    line.figure.canvas.mpl_connect('key_press_event', erase)
    line.figure.canvas.mpl_connect('motion_notify_event', move)
    pyplot.show()
    if len(x) < 3:
        raise ValueError("Need at least 3 points to make a polygon")
    if xy2ne:
        verts = numpy.transpose([y, x])
    else:
        verts = numpy.transpose([x, y])
    return verts

def pick_points(area, axes, marker='o', color='k', size=8, xy2ne=False):
    """
    Get the coordinates of points by clicking with the mouse.

    INSTRUCTIONS:

    * Left click to pick the points;
    * Press 'e' to erase the last point picked;
    * Close the figure window to finish;

    Parameters:

    * area : list = [x1, x2, y1, y2]
        Borders of the area containing the points
    * axes : matplotlib Axes
        The figure to use for drawing the polygon.
        To get an Axes instace, just do::

            from matplotlib import pyplot
            axes = pyplot.figure().add_subplot(1,1,1)

        You can plot things to ``axes`` before calling this function so that
        they'll appear on the background.
    * marker : str
        Style of the point markers (as in matplotlib.pyplot.plot)
    * color : str
        Line color (as in matplotlib.pyplot.plot)
    * size : float
        Marker size (as in matplotlib.pyplot.plot)
    * xy2ne : True or False
        If True, will exchange the x and y axis so that x points north.
        Use this when drawing on a map viewed from above. If the y-axis of the
        plot is supposed to be z (depth), then use ``xy2ne=False``.

    Returns:

    * points : list of lists
        List of ``[x, y]`` coordinates of the points

    """
    axes.set_title("Click to pick points. Close window when done.")
    if xy2ne:
        axes.set_xlim(area[2], area[3])
        axes.set_ylim(area[0], area[1])
    else:
        axes.set_xlim(area[0], area[1])
        axes.set_ylim(area[2], area[3])
    # start with an empty set
    line, = axes.plot([], [], marker=marker, color=color, markersize=size)
    line.figure.canvas.draw()
    x = []
    y = []
    plotx = []
    ploty = []
    # Hack because Python 2 doesn't like nonlocal variables that change value.
    # Lists it doesn't mind.
    picking = [True]

    def pick(event):
        if event.inaxes != axes:
            return 0
        if event.button == 1 and picking[0]:
            x.append(event.xdata)
            y.append(event.ydata)
            plotx.append(event.xdata)
            ploty.append(event.ydata)
            line.set_color(color)
            line.set_marker(marker)
            line.set_markersize(size)
            line.set_linestyle('')
            line.set_data(plotx, ploty)
            line.figure.canvas.draw()

    def erase(event):
        if event.key == 'e' and picking[0]:
            x.pop()
            y.pop()
            plotx.pop()
            ploty.pop()
            line.set_data(plotx, ploty)
            line.figure.canvas.draw()
    line.figure.canvas.mpl_connect('button_press_event', pick)
    line.figure.canvas.mpl_connect('key_press_event', erase)
    pyplot.show()
    if xy2ne:
        points = numpy.transpose([y, x])
    else:
        points = numpy.transpose([x, y])
    return points

def draw_layers(area, axes, style='-', marker='o', color='k', width=2):
    """
    Draw series of horizontal layers by clicking with the mouse.

    The y-axis is assumed to be depth, the x-axis is the physical property of
    each layer.

    INSTRUCTIONS:

    * Click to make a new layer;
    * Press 'e' to erase the last layer;
    * Close the figure window to finish;

    Parameters:

    * area : list = [x1, x2, y1, y2]
        Borders of the area containing the polygon
    * axes : matplotlib Axes
        The figure to use for drawing the polygon.
        To get an Axes instace, just do::

            from matplotlib import pyplot
            axes = pyplot.figure().add_subplot(1,1,1)

        You can plot things to ``axes`` before calling this function so that
        they'll appear on the background.
    * style : str
        Line style (as in matplotlib.pyplot.plot)
    * marker : str
        Style of the point markers (as in matplotlib.pyplot.plot)
    * color : str
        Line color (as in matplotlib.pyplot.plot)
    * width : float
        The line width (as in matplotlib.pyplot.plot)

    Returns:

    * layers : list = [thickness, values]

        * thickness : list
            The thickness of each layer, in order of increasing depth
        * values : list
            The physical property value of each layer, in the same order

    """
    axes.set_title("Click to set a layer. Close the window when done.")
    axes.grid()
    vmin, vmax, zmin, zmax = area
    axes.set_xlim(vmin, vmax)
    axes.set_ylim(zmax, zmin)
    # start with an empty line
    line, = axes.plot([], [], marker=marker, linestyle=style,
                      color=color, linewidth=width)
    midv = 0.5 * (vmax + vmin)
    # this is the line that moves around with the mouse
    tmpline, = axes.plot([midv], [zmin], marker=marker, linestyle='--',
                         color=color, linewidth=width)
    # Make a proxy for drawing
    draw = axes.figure.canvas.draw
    depths = [zmin]
    values = []
    plotv = []
    plotz = []
    tmpz = [zmin]
    # Hack because Python 2 doesn't like nonlocal variables that change value.
    # Lists it doesn't mind.
    picking = [True]

    def draw_guide(v, z):
        if len(values) == 0:
            tmpline.set_data([v, v], [tmpz[0], z])
        else:
            if z > tmpz[0]:
                tmpline.set_data([values[-1], v, v], [tmpz[0], tmpz[0], z])
            else:
                tmpline.set_data([values[-1], v], [tmpz[0], tmpz[0]])

    def move(event):
        if event.inaxes != axes:
            return 0
        v, z = event.xdata, event.ydata
        if picking[0]:
            draw_guide(v, z)
            draw()

    def pick(event):
        if event.inaxes != axes:
            return 0
        if event.button == 1 and picking[0]:
            v, z = event.xdata, event.ydata
            if z > tmpz[0]:
                depths.append(z)
                values.append(v)
                plotz.extend([tmpz[0], z])
                plotv.extend([v, v])
                tmpz[0] = z
                line.set_data(plotv, plotz)
                draw()

    def erase(event):
        if picking[0] and len(values) > 0 and event.key == 'e':
            depths.pop()
            values.pop()
            tmpz[0] = depths[-1]
            plotv.pop()
            plotv.pop()
            plotz.pop()
            plotz.pop()
            line.set_data(plotv, plotz)
            draw_guide(event.xdata, event.ydata)
            draw()
    line.figure.canvas.mpl_connect('button_press_event', pick)
    line.figure.canvas.mpl_connect('key_press_event', erase)
    line.figure.canvas.mpl_connect('motion_notify_event', move)
    pyplot.show()
    thickness = [depths[i + 1] - depths[i] for i in xrange(len(depths) - 1)]
    return thickness, values

def draw_geolines(area, dlon, dlat, basemap, linewidth=1):
    """
    Draw the parallels and meridians on a basemap plot.

    Parameters:

    * area : list
        ``[west, east, south, north]``, i.e., the area where the lines will
        be plotted
    * dlon, dlat : float
        The spacing between the lines in the longitude and latitude directions,
        respectively (in decimal degrees)
    * basemap : mpl_toolkits.basemap.Basemap
        The basemap used for plotting (see :func:`~fatiando.vis.mpl.basemap`)
    * linewidth : float
        The width of the lines

    """
    west, east, south, north = area
    meridians = basemap.drawmeridians(numpy.arange(west, east, dlon),
                                      labels=[0, 0, 0, 1], linewidth=linewidth)
    parallels = basemap.drawparallels(numpy.arange(south, north, dlat),
                                      labels=[1, 0, 0, 0], linewidth=linewidth)

def draw_countries(basemap, linewidth=1, style='dashed'):
    """
    Draw the country borders using the given basemap.

    Parameters:

    * basemap : mpl_toolkits.basemap.Basemap
        The basemap used for plotting (see :func:`~fatiando.vis.mpl.basemap`)
    * linewidth : float
        The width of the lines
    * style : str
        The style of the lines. Can be: 'solid', 'dashed', 'dashdot' or
        'dotted'

    """
    lines = basemap.drawcountries(linewidth=linewidth)
    lines.set_linestyles(style)

def draw_coastlines(basemap, linewidth=1, style='solid'):
    """
    Draw the coastlines using the given basemap.

    Parameters:

    * basemap : mpl_toolkits.basemap.Basemap
        The basemap used for plotting (see :func:`~fatiando.vis.mpl.basemap`)
    * linewidth : float
        The width of the lines
    * style : str
        The style of the lines. Can be: 'solid', 'dashed', 'dashdot' or
        'dotted'

    """
    lines = basemap.drawcoastlines(linewidth=linewidth)
    lines.set_linestyles(style)

def basemap(area, projection, resolution='c'):
    """
    Make a basemap to use when plotting with map projections.

    Uses the matplotlib basemap toolkit.

    Parameters:

    * area : list
        ``[west, east, south, north]``, i.e., the area of the data that is
        going to be plotted
    * projection : str
        The name of the projection you want to use. Choose from:

        * 'ortho': Orthographic
        * 'geos': Geostationary
        * 'robin': Robinson
        * 'cass': Cassini
        * 'merc': Mercator
        * 'poly': Polyconic
        * 'lcc': Lambert Conformal
        * 'stere': Stereographic

    * resolution : str
        The resolution for the coastlines. Can be 'c' for crude, 'l' for low,
        'i' for intermediate, 'h' for high

    Returns:

    * basemap : mpl_toolkits.basemap.Basemap
        The basemap

    """
    if projection not in ['ortho', 'aeqd', 'geos', 'robin', 'cass', 'merc',
                          'poly', 'lcc', 'stere']:
        raise ValueError("Unsuported projection '%s'" % (projection))
    global Basemap
    if Basemap is None:
        try:
            from mpl_toolkits.basemap import Basemap
        except ImportError:
            raise
    west, east, south, north = area
    lon_0 = 0.5 * (east + west)
    lat_0 = 0.5 * (north + south)
    if projection == 'ortho':
        bm = Basemap(projection=projection, lon_0=lon_0, lat_0=lat_0,
                     resolution=resolution)
    elif projection == 'geos' or projection == 'robin':
        bm = Basemap(projection=projection, lon_0=lon_0, resolution=resolution)
    elif (projection == 'cass' or
          projection == 'poly'):
        bm = Basemap(projection=projection, llcrnrlon=west, urcrnrlon=east,
                     llcrnrlat=south, urcrnrlat=north, lat_0=lat_0,
                     lon_0=lon_0, resolution=resolution)
    elif projection == 'merc':
        bm = Basemap(projection=projection, llcrnrlon=west, urcrnrlon=east,
                     llcrnrlat=south, urcrnrlat=north, lat_ts=lat_0,
                     resolution=resolution)
    elif projection == 'lcc':
        bm = Basemap(projection=projection, llcrnrlon=west, urcrnrlon=east,
                     llcrnrlat=south, urcrnrlat=north, lat_0=lat_0,
                     lon_0=lon_0, rsphere=(6378137.00, 6356752.3142),
                     lat_1=lat_0, resolution=resolution)
    elif projection == 'stere':
        bm = Basemap(projection=projection, llcrnrlon=west, urcrnrlon=east,
                     llcrnrlat=south, urcrnrlat=north, lat_0=lat_0,
                     lon_0=lon_0, lat_ts=lat_0, resolution=resolution)
    return bm

def m2km(axis=None):
    """
    Convert the x and y tick labels from meters to kilometers.

    Parameters:

    * axis : matplotlib axis instance
        The plot.

    .. tip:: Use ``fatiando.vis.gca()`` to get the current axis. Or the value
        returned by ``fatiando.vis.subplot`` or ``matplotlib.pyplot.subplot``.

    """
    if axis is None:
        axis = pyplot.gca()
    axis.set_xticklabels(['%g' % (0.001 * l) for l in axis.get_xticks()])
    axis.set_yticklabels(['%g' % (0.001 * l) for l in axis.get_yticks()])

def set_area(area):
    """
    Set the area of a Matplolib plot using xlim and ylim.

    Parameters:

    * area : list = [x1, x2, y1, y2]
        Coordinates of the top right and bottom left corners of the area

    """
    x1, x2, y1, y2 = area
    pyplot.xlim(x1, x2)
    pyplot.ylim(y1, y2)

def points(pts, style='.k', size=10, label=None, xy2ne=False):
    """
    Plot a list of points.

    Parameters:

    * pts : list of lists
        List of [x, y] pairs with the coordinates of the points
    * style : str
        String with the color and line style (as in matplotlib.pyplot.plot)
    * size : int
        Size of the plotted points
    * label : str
        If not None, then the string that will show in the legend
    * xy2ne : True or False
        If True, will exchange the x and y axis so that the x coordinates of
        the polygon are north. Use this when drawing on a map viewed from
        above. If the y-axis of the plot is supposed to be z (depth), then use
        ``xy2ne=False``.

    Returns:

    * axes : ``matplitlib.axes``
        The axes element of the plot

    """
    x, y = numpy.array(pts).T
    if xy2ne:
        x, y = y, x
    kwargs = {}
    if label is not None:
        kwargs['label'] = label
    return pyplot.plot(x, y, style, markersize=size, **kwargs)

def paths(pts1, pts2, style='-k', linewidth=1, label=None):
    """
    Plot paths between the two sets of points.

    Parameters:

    * pts1 : list of lists
        List of (x, y) pairs with the coordinates of the points
    * pts2 : list of lists
        List of (x, y) pairs with the coordinates of the points
    * style : str
        String with the color and line style (as in matplotlib.pyplot.plot)
    * linewidth : float
        The width of the lines representing the paths
    * label : str
        If not None, then the string that will show in the legend

    """
    kwargs = {'linewidth': linewidth}
    if label is not None:
        kwargs['label'] = label
    for p1, p2 in zip(pts1, pts2):
        pyplot.plot([p1[0], p2[0]], [p1[1], p2[1]], style, **kwargs)

def layers(thickness, values, style='-k', z0=0., linewidth=1, label=None,
           **kwargs):
    """
    Plot a series of layers and values associated to each layer.

    Parameters:

    * thickness : list
        The thickness of each layer in order of increasing depth
    * values : list
        The value associated with each layer in order of increasing
        depth
    * style : str
        String with the color and line style (as in matplotlib.pyplot.plot)
    * z0 : float
        The depth of the top of the first layer
    * linewidth : float
        Line width
    * label : str
        label associated with the square.

    Returns:

    * axes : ``matplitlib.axes``
        The axes element of the plot

    """
    if len(thickness) != len(values):
        raise ValueError("thickness and values must have same length")
    nlayers = len(thickness)
    interfaces = [z0 + sum(thickness[:i]) for i in xrange(nlayers + 1)]
    ys = [interfaces[0]]
    for y in interfaces[1:-1]:
        ys.append(y)
        ys.append(y)
    ys.append(interfaces[-1])
    xs = []
    for x in values:
        xs.append(x)
        xs.append(x)
    kwargs['linewidth'] = linewidth
    if label is not None:
        kwargs['label'] = label
    plot, = pyplot.plot(xs, ys, style, **kwargs)
    return plot

def square(area, style='-k', linewidth=1, fill=None, alpha=1., label=None,
           xy2ne=False):
    """
    Plot a square.

    Parameters:

    * area : list = [x1, x2, y1, y2]
        Borders of the square
    * style : str
        String with the color and line style (as in matplotlib.pyplot.plot)
    * linewidth : float
        Line width
    * fill : str
        A color string used to fill the square. If None, the square is not
        filled
    * alpha : float
        Transparency of the fill (1 >= alpha >= 0). 0 is transparent and 1 is
        opaque
    * label : str
        label associated with the square.
    * xy2ne : True or False
        If True, will exchange the x and y axis so that the x coordinates of
        the polygon are north. Use this when drawing on a map viewed from
        above. If the y-axis of the plot is supposed to be z (depth), then use
        ``xy2ne=False``.

    Returns:

    * axes : ``matplitlib.axes``
        The axes element of the plot

    """
    x1, x2, y1, y2 = area
    if xy2ne:
        x1, x2, y1, y2 = y1, y2, x1, x2
    xs = [x1, x1, x2, x2, x1]
    ys = [y1, y2, y2, y1, y1]
    kwargs = {'linewidth': linewidth}
    if label is not None:
        kwargs['label'] = label
    plot, = pyplot.plot(xs, ys, style, **kwargs)
    if fill is not None:
        pyplot.fill(xs, ys, color=fill, alpha=alpha)
    return plot

def squaremesh(mesh, prop, cmap=pyplot.cm.jet, vmin=None, vmax=None):
    """
    Make a pseudo-color plot of a mesh of squares

    Parameters:

    * mesh : :class:`fatiando.mesher.SquareMesh` or compatible
        The mesh (a compatible mesh must implement the methods ``get_xs`` and
        ``get_ys``)
    * prop : str
        The physical property of the squares to use as the color scale.
    * cmap : colormap
        Color map to be used. (see pyplot.cm module)
    * vmin, vmax : float
        Saturation values of the colorbar.

    Returns:

    * axes : ``matplitlib.axes``
        The axes element of the plot

    """
    if prop not in mesh.props:
        raise ValueError("Can't plot because 'mesh' doesn't have property '%s'"
                         % (prop))
    xs = mesh.get_xs()
    ys = mesh.get_ys()
    X, Y = numpy.meshgrid(xs, ys)
    V = numpy.reshape(mesh.props[prop], mesh.shape)
    plot = pyplot.pcolor(X, Y, V, cmap=cmap, vmin=vmin, vmax=vmax, picker=True)
    pyplot.xlim(xs.min(), xs.max())
    pyplot.ylim(ys.min(), ys.max())
    return plot

def polygon(polygon, style='-k', linewidth=1, fill=None, alpha=1., label=None,
            xy2ne=False, linealpha=1.):
    """
    Plot a polygon.

    Parameters:

    * polygon : :class:`fatiando.mesher.Polygon`
        The polygon
    * style : str
        Color and line style string (as in matplotlib.pyplot.plot)
    * linewidth : float
        Line width
    * fill : str
        A color string used to fill the polygon. If None, the polygon is not
        filled
    * alpha : float
        Transparency of the fill (1 >= alpha >= 0). 0 is transparent and 1 is
        opaque
    * linealpha : float
        Transparency of the line (1 >= alpha >= 0). 0 is transparent and 1 is
        opaque
    * label : str
        String with the label identifying the polygon in the legend
    * xy2ne : True or False
        If True, will exchange the x and y axis so that the x coordinates of
        the polygon are north. Use this when drawing on a map viewed from
        above. If the y-axis of the plot is supposed to be z (depth), then use
        ``xy2ne=False``.

    Returns:

    * lines : matplotlib Line object
        Line corresponding to the polygon plotted

    """
    if xy2ne:
        tmpx = [y for y in polygon.y]
        tmpx.append(polygon.y[0])
        tmpy = [x for x in polygon.x]
        tmpy.append(polygon.x[0])
    else:
        tmpx = [x for x in polygon.x]
        tmpx.append(polygon.x[0])
        tmpy = [y for y in polygon.y]
        tmpy.append(polygon.y[0])
    kwargs = {'linewidth': linewidth, 'alpha': linealpha}
    if label is not None:
        kwargs['label'] = label
    line, = pyplot.plot(tmpx, tmpy, style, **kwargs)
    if fill is not None:
        pyplot.fill(tmpx, tmpy, color=fill, alpha=alpha)
    return line

def contour(x, y, v, shape, levels, interp=False, extrapolate=False, color='k',
            label=None, clabel=True, style='solid', linewidth=1.0,
            basemap=None):
    """
    Make a contour plot of the data.

    Parameters:

    * x, y : array
        Arrays with the x and y coordinates of the grid points. If the data is
        on a regular grid, then assume x varies first (ie, inner loop), then y.
    * v : array
        The scalar value assigned to the grid points.
    * shape : tuple = (ny, nx)
        Shape of the regular grid.
        If interpolation is not False, then will use *shape* to grid the data.
    * levels : int or list
        Number of contours to use or a list with the contour values.
    * interp : True or False
        Wether or not to interpolate before trying to plot. If data is not on
        regular grid, set to True!
    * extrapolate : True or False
        Wether or not to extrapolate the data when interp=True
    * color : str
        Color of the contour lines.
    * label : str
        String with the label of the contour that would show in a legend.
    * clabel : True or False
        Wether or not to print the numerical value of the contour lines
    * style : str
        The style of the contour lines. Can be ``'dashed'``, ``'solid'`` or
        ``'mixed'`` (solid lines for positive contours and dashed for negative)
    * linewidth : float
        Width of the contour lines
    * basemap : mpl_toolkits.basemap.Basemap
        If not None, will use this basemap for plotting with a map projection
        (see :func:`~fatiando.vis.mpl.basemap` for creating basemaps)

    Returns:

    * levels : list
        List with the values of the contour levels

    """
    if style not in ['solid', 'dashed', 'mixed']:
        raise ValueError("Invalid contour style %s" % (style))
    if x.shape != y.shape != v.shape:
        raise ValueError("Input arrays x, y, and v must have same shape!")
    if interp:
        x, y, v = fatiando.gridder.interp(x, y, v, shape,
                                          extrapolate=extrapolate)
    X = numpy.reshape(x, shape)
    Y = numpy.reshape(y, shape)
    V = numpy.reshape(v, shape)
    kwargs = dict(colors=color, picker=True)
    if basemap is None:
        ct_data = pyplot.contour(X, Y, V, levels, **kwargs)
        pyplot.xlim(X.min(), X.max())
        pyplot.ylim(Y.min(), Y.max())
    else:
        lon, lat = basemap(X, Y)
        ct_data = basemap.contour(lon, lat, V, levels, **kwargs)
    if clabel:
        ct_data.clabel(fmt='%g')
    if label is not None:
        ct_data.collections[0].set_label(label)
    if style != 'mixed':
        for c in ct_data.collections:
            c.set_linestyle(style)
    for c in ct_data.collections:
        c.set_linewidth(linewidth)
    return ct_data.levels

def contourf(x, y, v, shape, levels, interp=False, extrapolate=False,
             vmin=None, vmax=None, cmap=pyplot.cm.jet, basemap=None):
    """
    Make a filled contour plot of the data.

    Parameters:

    * x, y : array
        Arrays with the x and y coordinates of the grid points. If the data is
        on a regular grid, then assume x varies first (ie, inner loop), then y.
    * v : array
        The scalar value assigned to the grid points.
    * shape : tuple = (ny, nx)
        Shape of the regular grid.
        If interpolation is not False, then will use *shape* to grid the data.
    * levels : int or list
        Number of contours to use or a list with the contour values.
    * interp : True or False
        Wether or not to interpolate before trying to plot. If data is not on
        regular grid, set to True!
    * extrapolate : True or False
        Wether or not to extrapolate the data when interp=True
    * vmin, vmax
        Saturation values of the colorbar. If provided, will overwrite what is
        set by *levels*.
    * cmap : colormap
        Color map to be used. (see pyplot.cm module)
    * basemap : mpl_toolkits.basemap.Basemap
        If not None, will use this basemap for plotting with a map projection
        (see :func:`~fatiando.vis.mpl.basemap` for creating basemaps)

    Returns:

    * levels : list
        List with the values of the contour levels

    """
    if x.shape != y.shape != v.shape:
        raise ValueError("Input arrays x, y, and v must have same shape!")
    if interp:
        x, y, v = fatiando.gridder.interp(x, y, v, shape,
                                          extrapolate=extrapolate)
    X = numpy.reshape(x, shape)
    Y = numpy.reshape(y, shape)
    V = numpy.reshape(v, shape)
    kwargs = dict(vmin=vmin, vmax=vmax, cmap=cmap, picker=True)
    if basemap is None:
        ct_data = pyplot.contourf(X, Y, V, levels, **kwargs)
        pyplot.xlim(X.min(), X.max())
        pyplot.ylim(Y.min(), Y.max())
    else:
        lon, lat = basemap(X, Y)
        ct_data = basemap.contourf(lon, lat, V, levels, **kwargs)
    return ct_data.levels

def pcolor(x, y, v, shape, interp=False, extrapolate=False, cmap=pyplot.cm.jet,
           vmin=None, vmax=None, basemap=None):
    """
    Make a pseudo-color plot of the data.

    Parameters:

    * x, y : array
        Arrays with the x and y coordinates of the grid points. If the data is
        on a regular grid, then assume x varies first (ie, inner loop), then y.
    * v : array
        The scalar value assigned to the grid points.
    * shape : tuple = (ny, nx)
        Shape of the regular grid.
        If interpolation is not False, then will use *shape* to grid the data.
    * interp : True or False
        Wether or not to interpolate before trying to plot. If data is not on
        regular grid, set to True!
    * extrapolate : True or False
        Wether or not to extrapolate the data when interp=True
    * cmap : colormap
        Color map to be used. (see pyplot.cm module)
    * vmin, vmax
        Saturation values of the colorbar.
    * basemap : mpl_toolkits.basemap.Basemap
        If not None, will use this basemap for plotting with a map projection
        (see :func:`~fatiando.vis.mpl.basemap` for creating basemaps)

    Returns:

    * axes : ``matplitlib.axes``
        The axes element of the plot

    """
    if x.shape != y.shape != v.shape:
        raise ValueError("Input arrays x, y, and v must have same shape!")
    if vmin is None:
        vmin = v.min()
    if vmax is None:
        vmax = v.max()
    if interp:
        x, y, v = fatiando.gridder.interp(x, y, v, shape,
                                          extrapolate=extrapolate)
    X = numpy.reshape(x, shape)
    Y = numpy.reshape(y, shape)
    V = numpy.reshape(v, shape)
    if basemap is None:
        plot = pyplot.pcolor(X, Y, V, cmap=cmap, vmin=vmin, vmax=vmax,
                             picker=True)
        pyplot.xlim(X.min(), X.max())
        pyplot.ylim(Y.min(), Y.max())
    else:
        lon, lat = basemap(X, Y)
        plot = basemap.pcolor(lon, lat, V, cmap=cmap, vmin=vmin, vmax=vmax,
                              picker=True)
    return plot

def seismic_wiggle(section, dt=0.004, ranges=None, scale=1.,
                   color='k', normalize=False):
    """
    Plot a seismic section (numpy 2D array matrix) as wiggles.

    Parameters:

    * section :  2D array
        matrix of traces (first dimension time, second dimension traces)
    * dt : float
        sample rate in seconds (default 4 ms)
    * ranges : (x1, x2)
        min and max horizontal values (default trace number)
    * scale : float
        scale factor multiplied by the section values before plotting
    * color : tuple of strings
        Color for filling the wiggle, positive  and negative lobes.
    * normalize :
        True to normalizes all trace in the section using global max/min
        data will be in the range (-0.5, 0.5) zero centered

    .. warning::
        Slow for more than 200 traces, in this case decimate your
        data or use ``seismic_image``.

    """
    npts, ntraces = section.shape  # time/traces
    if ntraces < 1:
        raise IndexError("Nothing to plot")
    if npts < 1:
        raise IndexError("Nothing to plot")
    t = numpy.linspace(0, dt*npts, npts)
    amp = 1.  # normalization factor
    gmin = 0.  # global minimum
    toffset = 0.  # offset in time to make 0 centered
    if normalize:
        gmax = section.max()
        gmin = section.min()
        amp = (gmax-gmin)
        toffset = 0.5
    pyplot.ylim(max(t), 0)
    if ranges is None:
        ranges = (0, ntraces)
    x0, x1 = ranges
    # horizontal increment
    dx = float((x1-x0)/ntraces)
    pyplot.xlim(x0, x1)
    for i, trace in enumerate(section.transpose()):
        tr = (((trace-gmin)/amp)-toffset)*scale*dx
        x = x0+i*dx  # x positon for this trace
        pyplot.plot(x+tr, t, 'k')
        pyplot.fill_betweenx(t, x+tr, x, tr > 0, color=color)

def seismic_image(section, dt=0.004, ranges=None, cmap=pyplot.cm.gray,
                  aspect=None, vmin=None, vmax=None):
    """
    Plot a seismic section (numpy 2D array matrix) as an image.

    Parameters:

    * section :  2D array
        matrix of traces (first dimension time, second dimension traces)
    * dt : float
        sample rate in seconds (default 4 ms)
    * ranges : (x1, x2)
        min and max horizontal values (default trace number)
    * cmap : colormap
        color map to be used. (see pyplot.cm module)
    * aspect : float
        matplotlib imshow aspect parameter, ratio between axes
    * vmin, vmax : float
        min and max values for imshow

    """
    npts, maxtraces = section.shape  # time/traces
    if maxtraces < 1:
        raise IndexError("Nothing to plot")
    if npts < 1:
        raise IndexError("Nothing to plot")
    t = numpy.linspace(0, dt*npts, npts)
    data = section
    if ranges is None:
        ranges = (0, maxtraces)
    x0, x1 = ranges
    extent = (x0, x1, t[-1:], t[0])
    if aspect is None:  # guarantee a rectangular picture
        aspect = numpy.round((x1-x0)/numpy.max(t))
        aspect -= aspect*0.2
    pyplot.imshow(data, aspect=aspect, cmap=cmap, origin='upper',
                  extent=extent, vmin=vmin, vmax=vmax)


import scipy as sp
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# animation of the probability density of the wavefunction over the course
# of time
def probabilityDensity(x, t, V, psi):

    # convert to the probability density
    Nt = len(t)
    rho = sp.real(sp.conjugate(psi)*psi)

    # set the first frame properties and grab the line handles
    fig, ax = plt.subplots()

    line1, line2, line3, line4 = ax.plot(x, rho[:,1], 'k',
            x, sp.real(psi[:,1]), 'b:',
            x, sp.imag(psi[:,1]), 'r:',
            x, V, 'm--',
            linewidth=2.0)

    ax.set_xlabel("Position")
    ax.set_ylabel("Probability Density")
    ax.set_ylim([-rho.max(), rho.max()])
    ax.set_xlim([min(x), max(x)])

    # the animation function, to be called repeatedly
    def animate(i):
        # set the new data each frame
        line1.set_ydata(rho[:,i])
        line2.set_ydata(sp.real(psi[:,i]))
        line3.set_ydata(sp.imag(psi[:,i]))
        return line1, line2, line3

    # the initialization function, useful when blit=True
    def init():
        line1.set_ydata(sp.ma.array(x, mask=True))
        line2.set_ydata(sp.ma.array(x, mask=True))
        line3.set_ydata(sp.ma.array(x, mask=True))
        return line1, line2, line3

    # perform the animation
    ani = animation.FuncAnimation(fig, animate, sp.arange(1,Nt),
            init_func=init, interval=25, blit=True)
    plt.show()


#!/usr/bin/env python
"""
Test XML backend
"""
import os
import sys
import time
import unittest

import pymoca.parser as mo_parser
from pymoca.backends.xml import analysis, generator, sim_scipy
from pymoca.backends.xml import parser as xml_parser
# get matplotlib from analysis, since logic for plotting
# without display already handled there
from pymoca.backends.xml.analysis import plt

TEST_DIR = os.path.dirname(os.path.realpath(__file__))
MODEL_DIR = os.path.join(TEST_DIR, 'models')
GENERATED_DIR = os.path.join(TEST_DIR, 'generated')

class XmlTest(unittest.TestCase):
    """
    Xml tests
    """

    def setUp(self):
        pass

    def tearDown(self):
        pass

    @staticmethod
    def flush():
        sys.stdout.flush()
        sys.stdout.flush()
        time.sleep(0.01)

    def test_noise(self):

        # compile to ModelicaXML
        with open(os.path.join(MODEL_DIR, 'Noise.mo'), 'r') as f:
            txt = f.read()
        ast_tree = mo_parser.parse(txt)
        model_xml = generator.generate(ast_tree, 'Noise')

        # save xml model to disk
        with open(os.path.join(GENERATED_DIR, 'Noise.xml'), 'w') as f:
            f.write(model_xml)

        # load xml model
        model = xml_parser.parse(model_xml, verbose=False)
        print(model)

        # convert to ode
        model_ode = model.to_ode()  # type: model.HybridOde
        print(model_ode)

        # simulate
        data = sim_scipy.sim(model_ode, {'tf': 1, 'dt': 0.001, 'verbose': True})

        # plot
        analysis.plot(data, fields=['x', 'm'])
        plt.draw()
        plt.pause(0.1)
        plt.close()

    def test_simple_circuit(self):

        # compile to ModelicaXML
        with open(os.path.join(MODEL_DIR, 'SimpleCircuit.mo'), 'r') as f:
            txt = f.read()
        ast_tree = mo_parser.parse(txt)
        model_xml = generator.generate(ast_tree, 'SimpleCircuit')

        # save xml model to disk
        with open(os.path.join(GENERATED_DIR, 'SimpleCircuit.xml'), 'w') as f:
            f.write(model_xml)

        # load xml model
        model = xml_parser.parse(model_xml, verbose=False)
        print(model)

        # convert to ode
        model_ode = model.to_ode()  # type: model.HybridOde
        print(model_ode)

        # simulate
        data = sim_scipy.sim(model_ode, {'tf': 1, 'dt': 0.001, 'verbose': True})

        # plot
        analysis.plot(data, fields=['x', 'c', 'm'])
        plt.draw()
        plt.pause(0.1)
        plt.close()

    def test_bouncing_ball(self):

        # generate
        with open(os.path.join(MODEL_DIR, 'BouncingBall.mo'), 'r') as f:
            txt = f.read()
        ast_tree = mo_parser.parse(txt)
        generator.generate(ast_tree, 'BouncingBall')

        # parse
        example_file = os.path.join(
            MODEL_DIR, 'bouncing-ball.xml')
        model = xml_parser.parse_file(example_file, verbose=False)
        print(model)

        # convert to ode
        model_ode = model.to_ode()  # type: model.HybridOde
        model_ode.prop['x']['start'] = 1
        print(model_ode)

        # simulate
        data = sim_scipy.sim(model_ode, {'tf': 3.5, 'dt': 0.01, 'verbose': True})

        # plot
        analysis.plot(data, linewidth=0.5, marker='.', markersize=0.5)
        plt.draw()
        plt.pause(0.1)
        plt.close()

        # simulate in soft real-time
        do_realtime = False
        if do_realtime:
            print('\nsoft-realtime simulation')
            time_start = time.time()

            def realtime_callback(t, x, y, m, p, c):
                t_real = time.time() - time_start
                lag = t_real - t
                if abs(lag) > 0.1:
                    print("real: {:10f} > sim: {:10f}, lag: {:10f}".format(t_real, t, lag))
                elif lag < 0:
                    time.sleep(-lag)

            data = sim_scipy.sim(model_ode, {'tf': 3.5, 'dt': 0.01, 'verbose': True},
                                 user_callback=realtime_callback)

        # plt.gca().set_ylim(-2, 2)
        self.flush()


# Author: Travis Oliphant
# 2003
#
# Feb. 2010: Updated by Warren Weckesser:
#   Rewrote much of chirp()
#   Added sweep_poly()
from __future__ import division, print_function, absolute_import

import numpy as np
from numpy import asarray, zeros, place, nan, mod, pi, extract, log, sqrt, \
    exp, cos, sin, polyval, polyint

__all__ = ['sawtooth', 'square', 'gausspulse', 'chirp', 'sweep_poly']

def sawtooth(t, width=1):
    """
    Return a periodic sawtooth or triangle waveform.

    The sawtooth waveform has a period ``2*pi``, rises from -1 to 1 on the
    interval 0 to ``width*2*pi``, then drops from 1 to -1 on the interval
    ``width*2*pi`` to ``2*pi``. `width` must be in the interval [0, 1].

    Note that this is not band-limited.  It produces an infinite number
    of harmonics, which are aliased back and forth across the frequency
    spectrum.

    Parameters
    ----------
    t : array_like
        Time.
    width : array_like, optional
        Width of the rising ramp as a proportion of the total cycle.
        Default is 1, producing a rising ramp, while 0 produces a falling
        ramp.  `width` = 0.5 produces a triangle wave.
        If an array, causes wave shape to change over time, and must be the
        same length as t.

    Returns
    -------
    y : ndarray
        Output array containing the sawtooth waveform.

    Examples
    --------
    A 5 Hz waveform sampled at 500 Hz for 1 second:

    >>> from scipy import signal
    >>> import matplotlib.pyplot as plt
    >>> t = np.linspace(0, 1, 500)
    >>> plt.plot(t, signal.sawtooth(2 * np.pi * 5 * t))

    """
    t, w = asarray(t), asarray(width)
    w = asarray(w + (t - t))
    t = asarray(t + (w - w))
    if t.dtype.char in ['fFdD']:
        ytype = t.dtype.char
    else:
        ytype = 'd'
    y = zeros(t.shape, ytype)

    # width must be between 0 and 1 inclusive
    mask1 = (w > 1) | (w < 0)
    place(y, mask1, nan)

    # take t modulo 2*pi
    tmod = mod(t, 2 * pi)

    # on the interval 0 to width*2*pi function is
    #  tmod / (pi*w) - 1
    mask2 = (1 - mask1) & (tmod < w * 2 * pi)
    tsub = extract(mask2, tmod)
    wsub = extract(mask2, w)
    place(y, mask2, tsub / (pi * wsub) - 1)

    # on the interval width*2*pi to 2*pi function is
    #  (pi*(w+1)-tmod) / (pi*(1-w))

    mask3 = (1 - mask1) & (1 - mask2)
    tsub = extract(mask3, tmod)
    wsub = extract(mask3, w)
    place(y, mask3, (pi * (wsub + 1) - tsub) / (pi * (1 - wsub)))
    return y

def square(t, duty=0.5):
    """
    Return a periodic square-wave waveform.

    The square wave has a period ``2*pi``, has value +1 from 0 to
    ``2*pi*duty`` and -1 from ``2*pi*duty`` to ``2*pi``. `duty` must be in
    the interval [0,1].

    Note that this is not band-limited.  It produces an infinite number
    of harmonics, which are aliased back and forth across the frequency
    spectrum.

    Parameters
    ----------
    t : array_like
        The input time array.
    duty : array_like, optional
        Duty cycle.  Default is 0.5 (50% duty cycle).
        If an array, causes wave shape to change over time, and must be the
        same length as t.

    Returns
    -------
    y : ndarray
        Output array containing the square waveform.

    Examples
    --------
    A 5 Hz waveform sampled at 500 Hz for 1 second:

    >>> from scipy import signal
    >>> import matplotlib.pyplot as plt
    >>> t = np.linspace(0, 1, 500, endpoint=False)
    >>> plt.plot(t, signal.square(2 * np.pi * 5 * t))
    >>> plt.ylim(-2, 2)

    A pulse-width modulated sine wave:

    >>> plt.figure()
    >>> sig = np.sin(2 * np.pi * t)
    >>> pwm = signal.square(2 * np.pi * 30 * t, duty=(sig + 1)/2)
    >>> plt.subplot(2, 1, 1)
    >>> plt.plot(t, sig)
    >>> plt.subplot(2, 1, 2)
    >>> plt.plot(t, pwm)
    >>> plt.ylim(-1.5, 1.5)

    """
    t, w = asarray(t), asarray(duty)
    w = asarray(w + (t - t))
    t = asarray(t + (w - w))
    if t.dtype.char in ['fFdD']:
        ytype = t.dtype.char
    else:
        ytype = 'd'

    y = zeros(t.shape, ytype)

    # width must be between 0 and 1 inclusive
    mask1 = (w > 1) | (w < 0)
    place(y, mask1, nan)

    # on the interval 0 to duty*2*pi function is 1
    tmod = mod(t, 2 * pi)
    mask2 = (1 - mask1) & (tmod < w * 2 * pi)
    place(y, mask2, 1)

    # on the interval duty*2*pi to 2*pi function is
    #  (pi*(w+1)-tmod) / (pi*(1-w))
    mask3 = (1 - mask1) & (1 - mask2)
    place(y, mask3, -1)
    return y

def gausspulse(t, fc=1000, bw=0.5, bwr=-6, tpr=-60, retquad=False,
               retenv=False):
    """
    Return a Gaussian modulated sinusoid:

        ``exp(-a t^2) exp(1j*2*pi*fc*t).``

    If `retquad` is True, then return the real and imaginary parts
    (in-phase and quadrature).
    If `retenv` is True, then return the envelope (unmodulated signal).
    Otherwise, return the real part of the modulated sinusoid.

    Parameters
    ----------
    t : ndarray or the string 'cutoff'
        Input array.
    fc : int, optional
        Center frequency (e.g. Hz).  Default is 1000.
    bw : float, optional
        Fractional bandwidth in frequency domain of pulse (e.g. Hz).
        Default is 0.5.
    bwr : float, optional
        Reference level at which fractional bandwidth is calculated (dB).
        Default is -6.
    tpr : float, optional
        If `t` is 'cutoff', then the function returns the cutoff
        time for when the pulse amplitude falls below `tpr` (in dB).
        Default is -60.
    retquad : bool, optional
        If True, return the quadrature (imaginary) as well as the real part
        of the signal.  Default is False.
    retenv : bool, optional
        If True, return the envelope of the signal.  Default is False.

    Returns
    -------
    yI : ndarray
        Real part of signal.  Always returned.
    yQ : ndarray
        Imaginary part of signal.  Only returned if `retquad` is True.
    yenv : ndarray
        Envelope of signal.  Only returned if `retenv` is True.

    See Also
    --------
    scipy.signal.morlet

    Examples
    --------
    Plot real component, imaginary component, and envelope for a 5 Hz pulse,
    sampled at 100 Hz for 2 seconds:

    >>> from scipy import signal
    >>> import matplotlib.pyplot as plt
    >>> t = np.linspace(-1, 1, 2 * 100, endpoint=False)
    >>> i, q, e = signal.gausspulse(t, fc=5, retquad=True, retenv=True)
    >>> plt.plot(t, i, t, q, t, e, '--')

    """
    if fc < 0:
        raise ValueError("Center frequency (fc=%.2f) must be >=0." % fc)
    if bw <= 0:
        raise ValueError("Fractional bandwidth (bw=%.2f) must be > 0." % bw)
    if bwr >= 0:
        raise ValueError("Reference level for bandwidth (bwr=%.2f) must "
                         "be < 0 dB" % bwr)

    # exp(-a t^2) <->  sqrt(pi/a) exp(-pi^2/a * f^2)  = g(f)

    ref = pow(10.0, bwr / 20.0)
    # fdel = fc*bw/2:  g(fdel) = ref --- solve this for a
    #
    # pi^2/a * fc^2 * bw^2 /4=-log(ref)
    a = -(pi * fc * bw) ** 2 / (4.0 * log(ref))

    if t == 'cutoff':  # compute cut_off point
        #  Solve exp(-a tc**2) = tref  for tc
        #   tc = sqrt(-log(tref) / a) where tref = 10^(tpr/20)
        if tpr >= 0:
            raise ValueError("Reference level for time cutoff must be < 0 dB")
        tref = pow(10.0, tpr / 20.0)
        return sqrt(-log(tref) / a)

    yenv = exp(-a * t * t)
    yI = yenv * cos(2 * pi * fc * t)
    yQ = yenv * sin(2 * pi * fc * t)
    if not retquad and not retenv:
        return yI
    if not retquad and retenv:
        return yI, yenv
    if retquad and not retenv:
        return yI, yQ
    if retquad and retenv:
        return yI, yQ, yenv

def chirp(t, f0, t1, f1, method='linear', phi=0, vertex_zero=True):
    """Frequency-swept cosine generator.

    In the following, 'Hz' should be interpreted as 'cycles per unit';
    there is no requirement here that the unit is one second.  The
    important distinction is that the units of rotation are cycles, not
    radians. Likewise, `t` could be a measurement of space instead of time.

    Parameters
    ----------
    t : array_like
        Times at which to evaluate the waveform.
    f0 : float
        Frequency (e.g. Hz) at time t=0.
    t1 : float
        Time at which `f1` is specified.
    f1 : float
        Frequency (e.g. Hz) of the waveform at time `t1`.
    method : {'linear', 'quadratic', 'logarithmic', 'hyperbolic'}, optional
        Kind of frequency sweep.  If not given, `linear` is assumed.  See
        Notes below for more details.
    phi : float, optional
        Phase offset, in degrees. Default is 0.
    vertex_zero : bool, optional
        This parameter is only used when `method` is 'quadratic'.
        It determines whether the vertex of the parabola that is the graph
        of the frequency is at t=0 or t=t1.

    Returns
    -------
    y : ndarray
        A numpy array containing the signal evaluated at `t` with the
        requested time-varying frequency.  More precisely, the function
        returns ``cos(phase + (pi/180)*phi)`` where `phase` is the integral
        (from 0 to `t`) of ``2*pi*f(t)``. ``f(t)`` is defined below.

    See Also
    --------
    sweep_poly

    Notes
    -----
    There are four options for the `method`.  The following formulas give
    the instantaneous frequency (in Hz) of the signal generated by
    `chirp()`.  For convenience, the shorter names shown below may also be
    used.

    linear, lin, li:

        ``f(t) = f0 + (f1 - f0) * t / t1``

    quadratic, quad, q:

        The graph of the frequency f(t) is a parabola through (0, f0) and
        (t1, f1).  By default, the vertex of the parabola is at (0, f0).
        If `vertex_zero` is False, then the vertex is at (t1, f1).  The
        formula is:

        if vertex_zero is True:

            ``f(t) = f0 + (f1 - f0) * t**2 / t1**2``

        else:

            ``f(t) = f1 - (f1 - f0) * (t1 - t)**2 / t1**2``

        To use a more general quadratic function, or an arbitrary
        polynomial, use the function `scipy.signal.waveforms.sweep_poly`.

    logarithmic, log, lo:

        ``f(t) = f0 * (f1/f0)**(t/t1)``

        f0 and f1 must be nonzero and have the same sign.

        This signal is also known as a geometric or exponential chirp.

    hyperbolic, hyp:

        ``f(t) = f0*f1*t1 / ((f0 - f1)*t + f1*t1)``

        f0 and f1 must be nonzero.

    """
    # 'phase' is computed in _chirp_phase, to make testing easier.
    phase = _chirp_phase(t, f0, t1, f1, method, vertex_zero)
    # Convert  phi to radians.
    phi *= pi / 180
    return cos(phase + phi)

def _chirp_phase(t, f0, t1, f1, method='linear', vertex_zero=True):
    """
    Calculate the phase used by chirp_phase to generate its output.

    See `chirp_phase` for a description of the arguments.

    """
    t = asarray(t)
    f0 = float(f0)
    t1 = float(t1)
    f1 = float(f1)
    if method in ['linear', 'lin', 'li']:
        beta = (f1 - f0) / t1
        phase = 2 * pi * (f0 * t + 0.5 * beta * t * t)

    elif method in ['quadratic', 'quad', 'q']:
        beta = (f1 - f0) / (t1 ** 2)
        if vertex_zero:
            phase = 2 * pi * (f0 * t + beta * t ** 3 / 3)
        else:
            phase = 2 * pi * (f1 * t + beta * ((t1 - t) ** 3 - t1 ** 3) / 3)

    elif method in ['logarithmic', 'log', 'lo']:
        if f0 * f1 <= 0.0:
            raise ValueError("For a logarithmic chirp, f0 and f1 must be "
                             "nonzero and have the same sign.")
        if f0 == f1:
            phase = 2 * pi * f0 * t
        else:
            beta = t1 / log(f1 / f0)
            phase = 2 * pi * beta * f0 * (pow(f1 / f0, t / t1) - 1.0)

    elif method in ['hyperbolic', 'hyp']:
        if f0 == 0 or f1 == 0:
            raise ValueError("For a hyperbolic chirp, f0 and f1 must be "
                             "nonzero.")
        if f0 == f1:
            # Degenerate case: constant frequency.
            phase = 2 * pi * f0 * t
        else:
            # Singular point: the instantaneous frequency blows up
            # when t == sing.
            sing = -f1 * t1 / (f0 - f1)
            phase = 2 * pi * (-sing * f0) * log(np.abs(1 - t/sing))

    else:
        raise ValueError("method must be 'linear', 'quadratic', 'logarithmic',"
                         " or 'hyperbolic', but a value of %r was given."
                         % method)

    return phase

def sweep_poly(t, poly, phi=0):
    """
    Frequency-swept cosine generator, with a time-dependent frequency.

    This function generates a sinusoidal function whose instantaneous
    frequency varies with time.  The frequency at time `t` is given by
    the polynomial `poly`.

    Parameters
    ----------
    t : ndarray
        Times at which to evaluate the waveform.
    poly : 1-D array_like or instance of numpy.poly1d
        The desired frequency expressed as a polynomial.  If `poly` is
        a list or ndarray of length n, then the elements of `poly` are
        the coefficients of the polynomial, and the instantaneous
        frequency is

          ``f(t) = poly[0]*t**(n-1) + poly[1]*t**(n-2) + ... + poly[n-1]``

        If `poly` is an instance of numpy.poly1d, then the
        instantaneous frequency is

          ``f(t) = poly(t)``

    phi : float, optional
        Phase offset, in degrees, Default: 0.

    Returns
    -------
    sweep_poly : ndarray
        A numpy array containing the signal evaluated at `t` with the
        requested time-varying frequency.  More precisely, the function
        returns ``cos(phase + (pi/180)*phi)``, where `phase` is the integral
        (from 0 to t) of ``2 * pi * f(t)``; ``f(t)`` is defined above.

    See Also
    --------
    chirp

    Notes
    -----
    .. versionadded:: 0.8.0

    If `poly` is a list or ndarray of length `n`, then the elements of
    `poly` are the coefficients of the polynomial, and the instantaneous
    frequency is:

        ``f(t) = poly[0]*t**(n-1) + poly[1]*t**(n-2) + ... + poly[n-1]``

    If `poly` is an instance of `numpy.poly1d`, then the instantaneous
    frequency is:

          ``f(t) = poly(t)``

    Finally, the output `s` is:

        ``cos(phase + (pi/180)*phi)``

    where `phase` is the integral from 0 to `t` of ``2 * pi * f(t)``,
    ``f(t)`` as defined above.

    """
    # 'phase' is computed in _sweep_poly_phase, to make testing easier.
    phase = _sweep_poly_phase(t, poly)
    # Convert to radians.
    phi *= pi / 180
    return cos(phase + phi)

def _sweep_poly_phase(t, poly):
    """
    Calculate the phase used by sweep_poly to generate its output.

    See `sweep_poly` for a description of the arguments.

    """
    # polyint handles lists, ndarrays and instances of poly1d automatically.
    intpoly = polyint(poly)
    phase = 2 * pi * polyval(intpoly, t)
    return phase


"""
==================
Find ECG artifacts
==================

Locate QRS component of ECG.

"""
# Authors: Alexandre Gramfort <alexandre.gramfort@telecom-paristech.fr>
#
# License: BSD (3-clause)

import numpy as np
import matplotlib.pyplot as plt

import mne
from mne import io
from mne.datasets import sample

print(__doc__)

data_path = sample.data_path()

###############################################################################
# Set parameters
raw_fname = data_path + '/MEG/sample/sample_audvis_raw.fif'

# Setup for reading the raw data
raw = io.Raw(raw_fname)

event_id = 999
ecg_events, _, _ = mne.preprocessing.find_ecg_events(raw, event_id,
                                                     ch_name='MEG 1531')

# Read epochs
picks = mne.pick_types(raw.info, meg=False, eeg=False, stim=False, eog=False,
                       include=['MEG 1531'], exclude='bads')
tmin, tmax = -0.1, 0.1
epochs = mne.Epochs(raw, ecg_events, event_id, tmin, tmax, picks=picks,
                    proj=False)
data = epochs.get_data()

print("Number of detected ECG artifacts : %d" % len(data))

###############################################################################
# Plot ECG artifacts
plt.plot(1e3 * epochs.times, np.squeeze(data).T)
plt.xlabel('Times (ms)')
plt.ylabel('ECG')
plt.show()


"""Compatibility fixes for older version of python, numpy and scipy

If you add content to this file, please give the version of the package
at which the fixe is no longer needed.

# XXX : copied from scikit-learn

"""
# Authors: Emmanuelle Gouillart <emmanuelle.gouillart@normalesup.org>
#          Gael Varoquaux <gael.varoquaux@normalesup.org>
#          Fabian Pedregosa <fpedregosa@acm.org>
#          Lars Buitinck <L.J.Buitinck@uva.nl>
# License: BSD

from __future__ import division
import collections
from operator import itemgetter
import inspect

import warnings
import numpy as np
import scipy
from scipy import linalg, sparse
from math import ceil, log
from numpy.fft import irfft
from distutils.version import LooseVersion
from functools import partial
from .externals import six
from .externals.six.moves import copyreg, xrange
from gzip import GzipFile

###############################################################################
# Misc

class gzip_open(GzipFile):  # python2.6 doesn't have context managing

    def __enter__(self):
        if hasattr(GzipFile, '__enter__'):
            return GzipFile.__enter__(self)
        else:
            return self

    def __exit__(self, exc_type, exc_value, traceback):
        if hasattr(GzipFile, '__exit__'):
            return GzipFile.__exit__(self, exc_type, exc_value, traceback)
        else:
            return self.close()

class _Counter(collections.defaultdict):
    """Partial replacement for Python 2.7 collections.Counter."""
    def __init__(self, iterable=(), **kwargs):
        super(_Counter, self).__init__(int, **kwargs)
        self.update(iterable)

    def most_common(self):
        return sorted(six.iteritems(self), key=itemgetter(1), reverse=True)

    def update(self, other):
        """Adds counts for elements in other"""
        if isinstance(other, self.__class__):
            for x, n in six.iteritems(other):
                self[x] += n
        else:
            for x in other:
                self[x] += 1

try:
    Counter = collections.Counter
except AttributeError:
    Counter = _Counter

def _unique(ar, return_index=False, return_inverse=False):
    """A replacement for the np.unique that appeared in numpy 1.4.

    While np.unique existed long before, keyword return_inverse was
    only added in 1.4.
    """
    try:
        ar = ar.flatten()
    except AttributeError:
        if not return_inverse and not return_index:
            items = sorted(set(ar))
            return np.asarray(items)
        else:
            ar = np.asarray(ar).flatten()

    if ar.size == 0:
        if return_inverse and return_index:
            return ar, np.empty(0, np.bool), np.empty(0, np.bool)
        elif return_inverse or return_index:
            return ar, np.empty(0, np.bool)
        else:
            return ar

    if return_inverse or return_index:
        perm = ar.argsort()
        aux = ar[perm]
        flag = np.concatenate(([True], aux[1:] != aux[:-1]))
        if return_inverse:
            iflag = np.cumsum(flag) - 1
            iperm = perm.argsort()
            if return_index:
                return aux[flag], perm[flag], iflag[iperm]
            else:
                return aux[flag], iflag[iperm]
        else:
            return aux[flag], perm[flag]

    else:
        ar.sort()
        flag = np.concatenate(([True], ar[1:] != ar[:-1]))
        return ar[flag]

if LooseVersion(np.__version__) < LooseVersion('1.5'):
    unique = _unique
else:
    unique = np.unique

def _bincount(X, weights=None, minlength=None):
    """Replacing np.bincount in numpy < 1.6 to provide minlength."""
    result = np.bincount(X, weights)
    if minlength is None or len(result) >= minlength:
        return result
    out = np.zeros(minlength, np.int)
    out[:len(result)] = result
    return out

if LooseVersion(np.__version__) < LooseVersion('1.6'):
    bincount = _bincount
else:
    bincount = np.bincount

def _copysign(x1, x2):
    """Slow replacement for np.copysign, which was introduced in numpy 1.4"""
    return np.abs(x1) * np.sign(x2)

if not hasattr(np, 'copysign'):
    copysign = _copysign
else:
    copysign = np.copysign

def _in1d(ar1, ar2, assume_unique=False, invert=False):
    """Replacement for in1d that is provided for numpy >= 1.4"""
    # Ravel both arrays, behavior for the first array could be different
    ar1 = np.asarray(ar1).ravel()
    ar2 = np.asarray(ar2).ravel()

    # This code is significantly faster when the condition is satisfied.
    if len(ar2) < 10 * len(ar1) ** 0.145:
        if invert:
            mask = np.ones(len(ar1), dtype=np.bool)
            for a in ar2:
                mask &= (ar1 != a)
        else:
            mask = np.zeros(len(ar1), dtype=np.bool)
            for a in ar2:
                mask |= (ar1 == a)
        return mask

    # Otherwise use sorting
    if not assume_unique:
        ar1, rev_idx = unique(ar1, return_inverse=True)
        ar2 = np.unique(ar2)

    ar = np.concatenate((ar1, ar2))
    # We need this to be a stable sort, so always use 'mergesort'
    # here. The values from the first array should always come before
    # the values from the second array.
    order = ar.argsort(kind='mergesort')
    sar = ar[order]
    if invert:
        bool_ar = (sar[1:] != sar[:-1])
    else:
        bool_ar = (sar[1:] == sar[:-1])
    flag = np.concatenate((bool_ar, [invert]))
    indx = order.argsort(kind='mergesort')[:len(ar1)]

    if assume_unique:
        return flag[indx]
    else:
        return flag[indx][rev_idx]

if not hasattr(np, 'in1d') or LooseVersion(np.__version__) < '1.8':
    in1d = _in1d
else:
    in1d = np.in1d

def _digitize(x, bins, right=False):
    """Replacement for digitize with right kwarg (numpy < 1.7).

    Notes
    -----
    This fix is only meant for integer arrays. If ``right==True`` but either
    ``x`` or ``bins`` are of a different type, a NotImplementedError will be
    raised.
    """
    if right:
        x = np.asarray(x)
        bins = np.asarray(bins)
        if (x.dtype.kind not in 'ui') or (bins.dtype.kind not in 'ui'):
            raise NotImplementedError("Only implemented for integer input")
        return np.digitize(x - 1e-5, bins)
    else:
        return np.digitize(x, bins)

if LooseVersion(np.__version__) < LooseVersion('1.7'):
    digitize = _digitize
else:
    digitize = np.digitize

def _tril_indices(n, k=0):
    """Replacement for tril_indices that is provided for numpy >= 1.4"""
    mask = np.greater_equal(np.subtract.outer(np.arange(n), np.arange(n)), -k)
    indices = np.where(mask)

    return indices

if not hasattr(np, 'tril_indices'):
    tril_indices = _tril_indices
else:
    tril_indices = np.tril_indices

def _unravel_index(indices, dims):
    """Add support for multiple indices in unravel_index that is provided
    for numpy >= 1.4"""
    indices_arr = np.asarray(indices)
    if indices_arr.size == 1:
        return np.unravel_index(indices, dims)
    else:
        if indices_arr.ndim != 1:
            raise ValueError('indices should be one dimensional')

        ndims = len(dims)
        unraveled_coords = np.empty((indices_arr.size, ndims), dtype=np.int)
        for coord, idx in zip(unraveled_coords, indices_arr):
            coord[:] = np.unravel_index(idx, dims)
        return tuple(unraveled_coords.T)

if LooseVersion(np.__version__) < LooseVersion('1.4'):
    unravel_index = _unravel_index
else:
    unravel_index = np.unravel_index

def _qr_economic_old(A, **kwargs):
    """
    Compat function for the QR-decomposition in economic mode
    Scipy 0.9 changed the keyword econ=True to mode='economic'
    """
    with warnings.catch_warnings(record=True):
        return linalg.qr(A, econ=True, **kwargs)

def _qr_economic_new(A, **kwargs):
    return linalg.qr(A, mode='economic', **kwargs)

if LooseVersion(scipy.__version__) < LooseVersion('0.9'):
    qr_economic = _qr_economic_old
else:
    qr_economic = _qr_economic_new

def savemat(file_name, mdict, oned_as="column", **kwargs):
    """MATLAB-format output routine that is compatible with SciPy 0.7's.

    0.7.2 (or .1?) added the oned_as keyword arg with 'column' as the default
    value. It issues a warning if this is not provided, stating that "This will
    change to 'row' in future versions."
    """
    import scipy.io
    try:
        return scipy.io.savemat(file_name, mdict, oned_as=oned_as, **kwargs)
    except TypeError:
        return scipy.io.savemat(file_name, mdict, **kwargs)

if hasattr(np, 'count_nonzero'):
    from numpy import count_nonzero
else:
    def count_nonzero(X):
        return len(np.flatnonzero(X))

# little dance to see if np.copy has an 'order' keyword argument
if 'order' in inspect.getargspec(np.copy)[0]:
    def safe_copy(X):
        # Copy, but keep the order
        return np.copy(X, order='K')
else:
    # Before an 'order' argument was introduced, numpy wouldn't muck with
    # the ordering
    safe_copy = np.copy

def _meshgrid(*xi, **kwargs):
    """
    Return coordinate matrices from coordinate vectors.
    Make N-D coordinate arrays for vectorized evaluations of
    N-D scalar/vector fields over N-D grids, given
    one-dimensional coordinate arrays x1, x2,..., xn.
    .. versionchanged:: 1.9
       1-D and 0-D cases are allowed.
    Parameters
    ----------
    x1, x2,..., xn : array_like
        1-D arrays representing the coordinates of a grid.
    indexing : {'xy', 'ij'}, optional
        Cartesian ('xy', default) or matrix ('ij') indexing of output.
        See Notes for more details.
        .. versionadded:: 1.7.0
    sparse : bool, optional
        If True a sparse grid is returned in order to conserve memory.
        Default is False.
        .. versionadded:: 1.7.0
    copy : bool, optional
        If False, a view into the original arrays are returned in order to
        conserve memory.  Default is True.  Please note that
        ``sparse=False, copy=False`` will likely return non-contiguous
        arrays.  Furthermore, more than one element of a broadcast array
        may refer to a single memory location.  If you need to write to the
        arrays, make copies first.
        .. versionadded:: 1.7.0
    Returns
    -------
    X1, X2,..., XN : ndarray
        For vectors `x1`, `x2`,..., 'xn' with lengths ``Ni=len(xi)`` ,
        return ``(N1, N2, N3,...Nn)`` shaped arrays if indexing='ij'
        or ``(N2, N1, N3,...Nn)`` shaped arrays if indexing='xy'
        with the elements of `xi` repeated to fill the matrix along
        the first dimension for `x1`, the second for `x2` and so on.
    """
    ndim = len(xi)

    copy_ = kwargs.pop('copy', True)
    sparse = kwargs.pop('sparse', False)
    indexing = kwargs.pop('indexing', 'xy')

    if kwargs:
        raise TypeError("meshgrid() got an unexpected keyword argument '%s'"
                        % (list(kwargs)[0],))

    if indexing not in ['xy', 'ij']:
        raise ValueError(
            "Valid values for `indexing` are 'xy' and 'ij'.")

    s0 = (1,) * ndim
    output = [np.asanyarray(x).reshape(s0[:i] + (-1,) + s0[i + 1::])
              for i, x in enumerate(xi)]

    shape = [x.size for x in output]

    if indexing == 'xy' and ndim > 1:
        # switch first and second axis
        output[0].shape = (1, -1) + (1,) * (ndim - 2)
        output[1].shape = (-1, 1) + (1,) * (ndim - 2)
        shape[0], shape[1] = shape[1], shape[0]

    if sparse:
        if copy_:
            return [x.copy() for x in output]
        else:
            return output
    else:
        # Return the full N-D matrix (not only the 1-D vector)
        if copy_:
            mult_fact = np.ones(shape, dtype=int)
            return [x * mult_fact for x in output]
        else:
            return np.broadcast_arrays(*output)

if LooseVersion(np.__version__) < LooseVersion('1.7'):
    meshgrid = _meshgrid
else:
    meshgrid = np.meshgrid

###############################################################################
# Back porting firwin2 for older scipy

# Original version of firwin2 from scipy ticket #457, submitted by "tash".
#
# Rewritten by Warren Weckesser, 2010.

def _firwin2(numtaps, freq, gain, nfreqs=None, window='hamming', nyq=1.0):
    """FIR filter design using the window method.

    From the given frequencies `freq` and corresponding gains `gain`,
    this function constructs an FIR filter with linear phase and
    (approximately) the given frequency response.

    Parameters
    ----------
    numtaps : int
        The number of taps in the FIR filter.  `numtaps` must be less than
        `nfreqs`.  If the gain at the Nyquist rate, `gain[-1]`, is not 0,
        then `numtaps` must be odd.

    freq : array-like, 1D
        The frequency sampling points. Typically 0.0 to 1.0 with 1.0 being
        Nyquist.  The Nyquist frequency can be redefined with the argument
        `nyq`.

        The values in `freq` must be nondecreasing.  A value can be repeated
        once to implement a discontinuity.  The first value in `freq` must
        be 0, and the last value must be `nyq`.

    gain : array-like
        The filter gains at the frequency sampling points.

    nfreqs : int, optional
        The size of the interpolation mesh used to construct the filter.
        For most efficient behavior, this should be a power of 2 plus 1
        (e.g, 129, 257, etc).  The default is one more than the smallest
        power of 2 that is not less than `numtaps`.  `nfreqs` must be greater
        than `numtaps`.

    window : string or (string, float) or float, or None, optional
        Window function to use. Default is "hamming".  See
        `scipy.signal.get_window` for the complete list of possible values.
        If None, no window function is applied.

    nyq : float
        Nyquist frequency.  Each frequency in `freq` must be between 0 and
        `nyq` (inclusive).

    Returns
    -------
    taps : numpy 1D array of length `numtaps`
        The filter coefficients of the FIR filter.

    Examples
    --------
    A lowpass FIR filter with a response that is 1 on [0.0, 0.5], and
    that decreases linearly on [0.5, 1.0] from 1 to 0:

    >>> taps = firwin2(150, [0.0, 0.5, 1.0], [1.0, 1.0, 0.0])  # doctest: +SKIP
    >>> print(taps[72:78])  # doctest: +SKIP
    [-0.02286961 -0.06362756  0.57310236  0.57310236 -0.06362756 -0.02286961]

    See also
    --------
    scipy.signal.firwin

    Notes
    -----

    From the given set of frequencies and gains, the desired response is
    constructed in the frequency domain.  The inverse FFT is applied to the
    desired response to create the associated convolution kernel, and the
    first `numtaps` coefficients of this kernel, scaled by `window`, are
    returned.

    The FIR filter will have linear phase.  The filter is Type I if `numtaps`
    is odd and Type II if `numtaps` is even.  Because Type II filters always
    have a zero at the Nyquist frequency, `numtaps` must be odd if `gain[-1]`
    is not zero.

    .. versionadded:: 0.9.0

    References
    ----------
    .. [1] Oppenheim, A. V. and Schafer, R. W., "Discrete-Time Signal
       Processing", Prentice-Hall, Englewood Cliffs, New Jersey (1989).
       (See, for example, Section 7.4.)

    .. [2] Smith, Steven W., "The Scientist and Engineer's Guide to Digital
       Signal Processing", Ch. 17. http://www.dspguide.com/ch17/1.htm

    """

    if len(freq) != len(gain):
        raise ValueError('freq and gain must be of same length.')

    if nfreqs is not None and numtaps >= nfreqs:
        raise ValueError('ntaps must be less than nfreqs, but firwin2 was '
                         'called with ntaps=%d and nfreqs=%s'
                         % (numtaps, nfreqs))

    if freq[0] != 0 or freq[-1] != nyq:
        raise ValueError('freq must start with 0 and end with `nyq`.')
    d = np.diff(freq)
    if (d < 0).any():
        raise ValueError('The values in freq must be nondecreasing.')
    d2 = d[:-1] + d[1:]
    if (d2 == 0).any():
        raise ValueError('A value in freq must not occur more than twice.')

    if numtaps % 2 == 0 and gain[-1] != 0.0:
        raise ValueError("A filter with an even number of coefficients must "
                         "have zero gain at the Nyquist rate.")

    if nfreqs is None:
        nfreqs = 1 + 2 ** int(ceil(log(numtaps, 2)))

    # Tweak any repeated values in freq so that interp works.
    eps = np.finfo(float).eps
    for k in range(len(freq)):
        if k < len(freq) - 1 and freq[k] == freq[k + 1]:
            freq[k] = freq[k] - eps
            freq[k + 1] = freq[k + 1] + eps

    # Linearly interpolate the desired response on a uniform mesh `x`.
    x = np.linspace(0.0, nyq, nfreqs)
    fx = np.interp(x, freq, gain)

    # Adjust the phases of the coefficients so that the first `ntaps` of the
    # inverse FFT are the desired filter coefficients.
    shift = np.exp(-(numtaps - 1) / 2. * 1.j * np.pi * x / nyq)
    fx2 = fx * shift

    # Use irfft to compute the inverse FFT.
    out_full = irfft(fx2)

    if window is not None:
        # Create the window to apply to the filter coefficients.
        from scipy.signal.signaltools import get_window
        wind = get_window(window, numtaps, fftbins=False)
    else:
        wind = 1

    # Keep only the first `numtaps` coefficients in `out`, and multiply by
    # the window.
    out = out_full[:numtaps] * wind

    return out

def get_firwin2():
    """Helper to get firwin2"""
    try:
        from scipy.signal import firwin2
    except ImportError:
        firwin2 = _firwin2
    return firwin2

def _filtfilt(*args, **kwargs):
    """wrap filtfilt, excluding padding arguments"""
    from scipy.signal import filtfilt
    # cut out filter args
    if len(args) > 4:
        args = args[:4]
    if 'padlen' in kwargs:
        del kwargs['padlen']
    return filtfilt(*args, **kwargs)

def get_filtfilt():
    """Helper to get filtfilt from scipy"""
    from scipy.signal import filtfilt

    if 'padlen' in inspect.getargspec(filtfilt)[0]:
        return filtfilt

    return _filtfilt

def _get_argrelmax():
    try:
        from scipy.signal import argrelmax
    except ImportError:
        argrelmax = _argrelmax
    return argrelmax

def _argrelmax(data, axis=0, order=1, mode='clip'):
    """Calculate the relative maxima of `data`.

    Parameters
    ----------
    data : ndarray
        Array in which to find the relative maxima.
    axis : int, optional
        Axis over which to select from `data`.  Default is 0.
    order : int, optional
        How many points on each side to use for the comparison
        to consider ``comparator(n, n+x)`` to be True.
    mode : str, optional
        How the edges of the vector are treated.
        Available options are 'wrap' (wrap around) or 'clip' (treat overflow
        as the same as the last (or first) element).
        Default 'clip'.  See `numpy.take`.

    Returns
    -------
    extrema : tuple of ndarrays
        Indices of the maxima in arrays of integers.  ``extrema[k]`` is
        the array of indices of axis `k` of `data`.  Note that the
        return value is a tuple even when `data` is one-dimensional.
    """
    comparator = np.greater
    if((int(order) != order) or (order < 1)):
        raise ValueError('Order must be an int >= 1')
    datalen = data.shape[axis]
    locs = np.arange(0, datalen)
    results = np.ones(data.shape, dtype=bool)
    main = data.take(locs, axis=axis, mode=mode)
    for shift in xrange(1, order + 1):
        plus = data.take(locs + shift, axis=axis, mode=mode)
        minus = data.take(locs - shift, axis=axis, mode=mode)
        results &= comparator(main, plus)
        results &= comparator(main, minus)
        if(~results.any()):
            return results
    return np.where(results)

###############################################################################
# Back porting matrix_rank for numpy < 1.7

def _matrix_rank(M, tol=None):
    """ Return matrix rank of array using SVD method

    Rank of the array is the number of SVD singular values of the array that
    are greater than `tol`.

    Parameters
    ----------
    M : {(M,), (M, N)} array_like
        array of <=2 dimensions
    tol : {None, float}, optional
       threshold below which SVD values are considered zero. If `tol` is
       None, and ``S`` is an array with singular values for `M`, and
       ``eps`` is the epsilon value for datatype of ``S``, then `tol` is
       set to ``S.max() * max(M.shape) * eps``.

    Notes
    -----
    The default threshold to detect rank deficiency is a test on the magnitude
    of the singular values of `M`. By default, we identify singular values less
    than ``S.max() * max(M.shape) * eps`` as indicating rank deficiency (with
    the symbols defined above). This is the algorithm MATLAB uses [1]. It also
    appears in *Numerical recipes* in the discussion of SVD solutions for
    linear least squares [2].

    This default threshold is designed to detect rank deficiency accounting
    for the numerical errors of the SVD computation. Imagine that there is a
    column in `M` that is an exact (in floating point) linear combination of
    other columns in `M`. Computing the SVD on `M` will not produce a
    singular value exactly equal to 0 in general: any difference of the
    smallest SVD value from 0 will be caused by numerical imprecision in the
    calculation of the SVD. Our threshold for small SVD values takes this
    numerical imprecision into account, and the default threshold will detect
    such numerical rank deficiency. The threshold may declare a matrix `M`
    rank deficient even if the linear combination of some columns of `M` is
    not exactly equal to another column of `M` but only numerically very
    close to another column of `M`.

    We chose our default threshold because it is in wide use. Other
    thresholds are possible. For example, elsewhere in the 2007 edition of
    *Numerical recipes* there is an alternative threshold of ``S.max() *
    np.finfo(M.dtype).eps / 2. * np.sqrt(m + n + 1.)``. The authors describe
    this threshold as being based on "expected roundoff error" (p 71).

    The thresholds above deal with floating point roundoff error in the
    calculation of the SVD. However, you may have more information about the
    sources of error in `M` that would make you consider other tolerance
    values to detect *effective* rank deficiency. The most useful measure of
    the tolerance depends on the operations you intend to use on your matrix.
    For example, if your data come from uncertain measurements with
    uncertainties greater than floating point epsilon, choosing a tolerance
    near that uncertainty may be preferable. The tolerance may be absolute if
    the uncertainties are absolute rather than relative.

    References
    ----------
    .. [1] MATLAB reference documention, "Rank"
           http://www.mathworks.com/help/techdoc/ref/rank.html
    .. [2] W. H. Press, S. A. Teukolsky, W. T. Vetterling and B. P. Flannery,
           "Numerical Recipes (3rd edition)", Cambridge University Press, 2007,
           page 795.

    Examples
    --------
    >>> from numpy.linalg import matrix_rank
    >>> matrix_rank(np.eye(4)) # Full rank matrix
    4
    >>> I=np.eye(4); I[-1,-1] = 0. # rank deficient matrix
    >>> matrix_rank(I)
    3
    >>> matrix_rank(np.ones((4,))) # 1 dimension - rank 1 unless all 0
    1
    >>> matrix_rank(np.zeros((4,)))
    0
    """
    M = np.asarray(M)
    if M.ndim > 2:
        raise TypeError('array should have 2 or fewer dimensions')
    if M.ndim < 2:
        return np.int(not all(M == 0))
    S = np.linalg.svd(M, compute_uv=False)
    if tol is None:
        tol = S.max() * np.max(M.shape) * np.finfo(S.dtype).eps
    return np.sum(S > tol)

if LooseVersion(np.__version__) > '1.7.1':
    from numpy.linalg import matrix_rank
else:
    matrix_rank = _matrix_rank

def _reconstruct_partial(func, args, kwargs):
    """Helper to pickle partial functions"""
    return partial(func, *args, **(kwargs or {}))

def _reduce_partial(p):
    """Helper to pickle partial functions"""
    return _reconstruct_partial, (p.func, p.args, p.keywords)

# This adds pickling functionality to older Python 2.6
# Please always import partial from here.
copyreg.pickle(partial, _reduce_partial)

def normalize_colors(vmin, vmax, clip=False):
    """Helper to handle matplotlib API"""
    import matplotlib.pyplot as plt
    try:
        return plt.Normalize(vmin, vmax, clip=clip)
    except AttributeError:
        return plt.normalize(vmin, vmax, clip=clip)

def assert_true(expr, msg='False is not True'):
    """Fake assert_true without message"""
    if not expr:
        raise AssertionError(msg)

def assert_is(expr1, expr2, msg=None):
    """Fake assert_is without message"""
    assert_true(expr2 is expr2, msg)

def assert_is_not(expr1, expr2, msg=None):
    """Fake assert_is_not without message"""
    assert_true(expr1 is not expr2, msg)

def _sparse_block_diag(mats, format=None, dtype=None):
    """An implementation of scipy.sparse.block_diag since old versions of
    scipy don't have it. Forms a sparse matrix by stacking matrices in block
    diagonal form.

    Parameters
    ----------
    mats : list of matrices
        Input matrices.
    format : str, optional
        The sparse format of the result (e.g. "csr"). If not given, the
        matrix is returned in "coo" format.
    dtype : dtype specifier, optional
        The data-type of the output matrix. If not given, the dtype is
        determined from that of blocks.

    Returns
    -------
    res : sparse matrix
    """
    nmat = len(mats)
    rows = []
    for ia, a in enumerate(mats):
        row = [None] * nmat
        row[ia] = a
        rows.append(row)
    return sparse.bmat(rows, format=format, dtype=dtype)

try:
    from scipy.sparse import block_diag as sparse_block_diag
except Exception:
    sparse_block_diag = _sparse_block_diag

def _isclose(a, b, rtol=1.e-5, atol=1.e-8, equal_nan=False):
    """
    Returns a boolean array where two arrays are element-wise equal within a
    tolerance.

    The tolerance values are positive, typically very small numbers.  The
    relative difference (`rtol` * abs(`b`)) and the absolute difference
    `atol` are added together to compare against the absolute difference
    between `a` and `b`.

    Parameters
    ----------
    a, b : array_like
        Input arrays to compare.
    rtol : float
        The relative tolerance parameter (see Notes).
    atol : float
        The absolute tolerance parameter (see Notes).
    equal_nan : bool
        Whether to compare NaN's as equal.  If True, NaN's in `a` will be
        considered equal to NaN's in `b` in the output array.

    Returns
    -------
    y : array_like
        Returns a boolean array of where `a` and `b` are equal within the
        given tolerance. If both `a` and `b` are scalars, returns a single
        boolean value.

    See Also
    --------
    allclose

    Notes
    -----
    .. versionadded:: 1.7.0

    For finite values, isclose uses the following equation to test whether
    two floating point values are equivalent.

     absolute(`a` - `b`) <= (`atol` + `rtol` * absolute(`b`))

    The above equation is not symmetric in `a` and `b`, so that
    `isclose(a, b)` might be different from `isclose(b, a)` in
    some rare cases.

    Examples
    --------
    >>> isclose([1e10,1e-7], [1.00001e10,1e-8])
    array([ True, False], dtype=bool)
    >>> isclose([1e10,1e-8], [1.00001e10,1e-9])
    array([ True,  True], dtype=bool)
    >>> isclose([1e10,1e-8], [1.0001e10,1e-9])
    array([False,  True], dtype=bool)
    >>> isclose([1.0, np.nan], [1.0, np.nan])
    array([ True, False], dtype=bool)
    >>> isclose([1.0, np.nan], [1.0, np.nan], equal_nan=True)
    array([ True,  True], dtype=bool)
    """
    def within_tol(x, y, atol, rtol):
        with np.errstate(invalid='ignore'):
            result = np.less_equal(abs(x - y), atol + rtol * abs(y))
        if np.isscalar(a) and np.isscalar(b):
            result = bool(result)
        return result

    x = np.array(a, copy=False, subok=True, ndmin=1)
    y = np.array(b, copy=False, subok=True, ndmin=1)

    # Make sure y is an inexact type to avoid bad behavior on abs(MIN_INT).
    # This will cause casting of x later. Also, make sure to allow subclasses
    # (e.g., for numpy.ma).
    dt = np.core.multiarray.result_type(y, 1.)
    y = np.array(y, dtype=dt, copy=False, subok=True)

    xfin = np.isfinite(x)
    yfin = np.isfinite(y)
    if np.all(xfin) and np.all(yfin):
        return within_tol(x, y, atol, rtol)
    else:
        finite = xfin & yfin
        cond = np.zeros_like(finite, subok=True)
        # Because we're using boolean indexing, x & y must be the same shape.
        # Ideally, we'd just do x, y = broadcast_arrays(x, y). It's in
        # lib.stride_tricks, though, so we can't import it here.
        x = x * np.ones_like(cond)
        y = y * np.ones_like(cond)
        # Avoid subtraction with infinite/nan values...
        cond[finite] = within_tol(x[finite], y[finite], atol, rtol)
        # Check for equality of infinite values...
        cond[~finite] = (x[~finite] == y[~finite])
        if equal_nan:
            # Make NaN == NaN
            both_nan = np.isnan(x) & np.isnan(y)
            cond[both_nan] = both_nan[both_nan]
        return cond

if LooseVersion(np.__version__) < LooseVersion('1.7'):
    isclose = _isclose
else:
    isclose = np.isclose


"""
===========================================
Robust linear model estimation using RANSAC
===========================================

In this example we see how to robustly fit a linear model to faulty data using
the RANSAC algorithm.

"""
import numpy as np
from matplotlib import pyplot as plt

from sklearn import linear_model, datasets

n_samples = 1000
n_outliers = 50

X, y, coef = datasets.make_regression(n_samples=n_samples, n_features=1,
                                      n_informative=1, noise=10,
                                      coef=True, random_state=0)

# Add outlier data
np.random.seed(0)
X[:n_outliers] = 3 + 0.5 * np.random.normal(size=(n_outliers, 1))
y[:n_outliers] = -3 + 10 * np.random.normal(size=n_outliers)

# Fit line using all data
model = linear_model.LinearRegression()
model.fit(X, y)

# Robustly fit linear model with RANSAC algorithm
model_ransac = linear_model.RANSACRegressor(linear_model.LinearRegression())
model_ransac.fit(X, y)
inlier_mask = model_ransac.inlier_mask_
outlier_mask = np.logical_not(inlier_mask)

# Predict data of estimated models
line_X = np.arange(-5, 5)
line_y = model.predict(line_X[:, np.newaxis])
line_y_ransac = model_ransac.predict(line_X[:, np.newaxis])

# Compare estimated coefficients
print("Estimated coefficients (true, normal, RANSAC):")
print(coef, model.coef_, model_ransac.estimator_.coef_)

plt.plot(X[inlier_mask], y[inlier_mask], '.g', label='Inliers')
plt.plot(X[outlier_mask], y[outlier_mask], '.r', label='Outliers')
plt.plot(line_X, line_y, '-k', label='Linear regressor')
plt.plot(line_X, line_y_ransac, '-b', label='RANSAC regressor')
plt.legend(loc='lower right')
plt.show()