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


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


def swirl_velocity_field():
    x = np.linspace(-3., 3., 200)
    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.png'], remove_text=True, style='mpl20',
                  tol=0.003)
def test_startpoints():
    # Test varying startpoints. Also tests a non-default num_arrows argument.
    X, Y, U, V = velocity_field()
    start_x, start_y = np.meshgrid(np.linspace(X.min(), X.max(), 5),
                                   np.linspace(Y.min(), Y.max(), 5))
    start_points = np.column_stack([start_x.ravel(), start_y.ravel()])
    plt.streamplot(X, Y, U, V, start_points=start_points, num_arrows=4)
    plt.plot(start_x, start_y, 'ok')


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


@image_comparison(['streamplot_linewidth.png'], remove_text=True, style='mpl20',
                  tol=0.03)
def test_linewidth():
    X, Y, U, V = velocity_field()
    speed = np.hypot(U, V)
    lw = 5 * speed / speed.max()
    ax = plt.figure().subplots()
    ax.streamplot(X, Y, U, V, density=[0.5, 1], color='k', linewidth=lw, num_arrows=2)


@image_comparison(['streamplot_masks_and_nans.png'],
                  remove_text=True, style='mpl20')
def test_masks_and_nans():
    X, Y, U, V = velocity_field()
    mask = np.zeros(U.shape, dtype=bool)
    mask[40:60, 80:120] = 1
    U[:20, :40] = np.nan
    U = np.ma.array(U, mask=mask)
    ax = plt.figure().subplots()
    with np.errstate(invalid='ignore'):
        ax.streamplot(X, Y, U, V, color=U, cmap="Blues")


@image_comparison(['streamplot_maxlength.png'],
                  remove_text=True, style='mpl20', tol=0.302)
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_maxlength_no_broken.png'],
                  remove_text=True, style='mpl20', tol=0.302)
def test_maxlength_no_broken():
    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, broken_streamlines=False)
    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', tol=0.073)
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)


@image_comparison(['streamplot_integration.png'], style='mpl20', tol=0.05)
def test_integration_options():
    # Linear potential flow over a lifting cylinder
    n = 50
    x, y = np.meshgrid(np.linspace(-2, 2, n), np.linspace(-3, 3, n))
    th = np.arctan2(y, x)
    r = np.sqrt(x**2 + y**2)
    vr = -np.cos(th) / r**2
    vt = -np.sin(th) / r**2 - 1 / r
    vx = vr * np.cos(th) - vt * np.sin(th) + 1.0
    vy = vr * np.sin(th) + vt * np.cos(th)

    # Seed points
    n_seed = 50
    seed_pts = np.column_stack((np.full(n_seed, -1.75), np.linspace(-2, 2, n_seed)))

    fig, axs = plt.subplots(3, 1, figsize=(6, 14))
    th_circ = np.linspace(0, 2 * np.pi, 100)
    for ax, max_val in zip(axs, [0.05, 1, 5]):
        ax_ins = ax.inset_axes([0.0, 0.7, 0.3, 0.35])
        for ax_curr, is_inset in zip([ax, ax_ins], [False, True]):
            ax_curr.streamplot(
                x,
                y,
                vx,
                vy,
                start_points=seed_pts,
                broken_streamlines=False,
                arrowsize=1e-10,
                linewidth=2 if is_inset else 0.6,
                color="k",
                integration_max_step_scale=max_val,
                integration_max_error_scale=max_val,
            )

            # Draw the cylinder
            ax_curr.fill(
                np.cos(th_circ),
                np.sin(th_circ),
                color="w",
                ec="k",
                lw=6 if is_inset else 2,
            )

            # Set axis properties
            ax_curr.set_aspect("equal")

        # Set axis limits and show zoomed region
        ax_ins.set_xlim(-1.2, -0.7)
        ax_ins.set_ylim(-0.8, -0.4)
        ax_ins.set_yticks(())
        ax_ins.set_xticks(())

        ax.set_ylim(-1.5, 1.5)
        ax.axis("off")
        ax.indicate_inset_zoom(ax_ins, ec="k")

    fig.tight_layout()


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)


def test_streamplot_grid():
    u = np.ones((2, 2))
    v = np.zeros((2, 2))

    # Test for same rows and columns
    x = np.array([[10, 20], [10, 30]])
    y = np.array([[10, 10], [20, 20]])

    with pytest.raises(ValueError, match="The rows of 'x' must be equal"):
        plt.streamplot(x, y, u, v)

    x = np.array([[10, 20], [10, 20]])
    y = np.array([[10, 10], [20, 30]])

    with pytest.raises(ValueError, match="The columns of 'y' must be equal"):
        plt.streamplot(x, y, u, v)

    x = np.array([[10, 20], [10, 20]])
    y = np.array([[10, 10], [20, 20]])
    plt.streamplot(x, y, u, v)

    # Test for maximum dimensions
    x = np.array([0, 10])
    y = np.array([[[0, 10]]])

    with pytest.raises(ValueError, match="'y' can have at maximum "
                                         "2 dimensions"):
        plt.streamplot(x, y, u, v)

    # Test for equal spacing
    u = np.ones((3, 3))
    v = np.zeros((3, 3))
    x = np.array([0, 10, 20])
    y = np.array([0, 10, 30])

    with pytest.raises(ValueError, match="'y' values must be equally spaced"):
        plt.streamplot(x, y, u, v)

    # Test for strictly increasing
    x = np.array([0, 20, 40])
    y = np.array([0, 20, 10])


def test_streamplot_integration_params():
    x = np.array([[10, 20], [10, 20]])
    y = np.array([[10, 10], [20, 20]])
    u = np.ones((2, 2))
    v = np.zeros((2, 2))

    err_str = "The value of integration_max_step_scale must be > 0, got -0.5"
    with pytest.raises(ValueError, match=err_str):
        plt.streamplot(x, y, u, v, integration_max_step_scale=-0.5)

    err_str = "The value of integration_max_error_scale must be > 0, got 0.0"
    with pytest.raises(ValueError, match=err_str):
        plt.streamplot(x, y, u, v, integration_max_error_scale=0.0)


def test_streamplot_inputs():  # test no exception occurs.
    # fully-masked
    plt.streamplot(np.arange(3), np.arange(3),
                   np.full((3, 3), np.nan), np.full((3, 3), np.nan),
                   color=np.random.rand(3, 3))
    # array-likes
    plt.streamplot(range(3), range(3),
                   np.random.rand(3, 3), np.random.rand(3, 3))