import sys

import numpy as np
from numpy.testing import assert_array_almost_equal
import matplotlib.pyplot as plt
from matplotlib.testing.decorators import image_comparison
import matplotlib.transforms as mtransforms


on_win = (sys.platform == 'win32')
on_mac = (sys.platform == 'darwin')


def velocity_field():
    Y, X = np.mgrid[-3:3:100j, -3:3:100j]
    U = -1 - X**2 + Y
    V = 1 + X - Y**2
    return X, Y, U, V


def swirl_velocity_field():
    x = np.linspace(-3., 3., 100)
    y = np.linspace(-3., 3., 100)
    X, Y = np.meshgrid(x, y)
    a = 0.1
    U = np.cos(a) * (-Y) - np.sin(a) * X
    V = np.sin(a) * (-Y) + np.cos(a) * X
    return x, y, U, V


@image_comparison(['streamplot_startpoints'], remove_text=True, style='mpl20')
def test_startpoints():
    X, Y, U, V = velocity_field()
    start_x = np.linspace(X.min(), X.max(), 10)
    start_y = np.linspace(Y.min(), Y.max(), 10)
    start_points = np.column_stack([start_x, start_y])
    plt.streamplot(X, Y, U, V, start_points=start_points)
    plt.plot(start_x, start_y, 'ok')


@image_comparison(['streamplot_colormap'],
                  tol=.04, remove_text=True, style='mpl20')
def test_colormap():
    X, Y, U, V = velocity_field()
    plt.streamplot(X, Y, U, V, color=U, density=0.6, linewidth=2,
                   cmap=plt.cm.autumn)
    plt.colorbar()


@image_comparison(['streamplot_linewidth'], remove_text=True, style='mpl20')
def test_linewidth():
    X, Y, U, V = velocity_field()
    speed = np.hypot(U, V)
    lw = 5 * speed / speed.max()
    # Compatibility for old test image
    df = 25 / 30
    ax = plt.figure().subplots()
    ax.set(xlim=(-3.0, 2.9999999999999947),
           ylim=(-3.0000000000000004, 2.9999999999999947))
    ax.streamplot(X, Y, U, V, density=[0.5 * df, 1. * df], color='k',
                  linewidth=lw)


@image_comparison(['streamplot_masks_and_nans'],
                  remove_text=True, style='mpl20', tol=0.04 if on_win else 0)
def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 40:60] = 1
    U[:20, :20] = np.nan
    U = np.ma.array(U, mask=mask)
    # Compatibility for old test image
    ax = plt.figure().subplots()
    ax.set(xlim=(-3.0, 2.9999999999999947),
           ylim=(-3.0000000000000004, 2.9999999999999947))
    with np.errstate(invalid='ignore'):
        ax.streamplot(X, Y, U, V, color=U, cmap=plt.cm.Blues)


@image_comparison(['streamplot_maxlength.png'],
                  remove_text=True, style='mpl20',
                  tol=0.002 if on_mac else 0)
def test_maxlength():
    x, y, U, V = swirl_velocity_field()
    ax = plt.figure().subplots()
    ax.streamplot(x, y, U, V, maxlength=10., start_points=[[0., 1.5]],
                  linewidth=2, density=2)
    assert ax.get_xlim()[-1] == ax.get_ylim()[-1] == 3
    # Compatibility for old test image
    ax.set(xlim=(None, 3.2555988021882305), ylim=(None, 3.078326760195413))


@image_comparison(['streamplot_direction.png'],
                  remove_text=True, style='mpl20')
def test_direction():
    x, y, U, V = swirl_velocity_field()
    plt.streamplot(x, y, U, V, integration_direction='backward',
                   maxlength=1.5, start_points=[[1.5, 0.]],
                   linewidth=2, density=2)


def test_streamplot_limits():
    ax = plt.axes()
    x = np.linspace(-5, 10, 20)
    y = np.linspace(-2, 4, 10)
    y, x = np.meshgrid(y, x)
    trans = mtransforms.Affine2D().translate(25, 32) + ax.transData
    plt.barbs(x, y, np.sin(x), np.cos(y), transform=trans)
    # The calculated bounds are approximately the bounds of the original data,
    # this is because the entire path is taken into account when updating the
    # datalim.
    assert_array_almost_equal(ax.dataLim.bounds, (20, 30, 15, 6),
                              decimal=1)