import io
from types import SimpleNamespace

import numpy as np
from numpy.testing import assert_array_equal, assert_array_almost_equal
import pytest

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.collections as mcollections
import matplotlib.colors as mcolors
import matplotlib.transforms as mtransforms
from matplotlib.collections import (Collection, LineCollection,
                                    EventCollection, PolyCollection)
from matplotlib.testing.decorators import check_figures_equal, image_comparison
from matplotlib._api.deprecation import MatplotlibDeprecationWarning


def generate_EventCollection_plot():
    """Generate the initial collection and plot it."""
    positions = np.array([0., 1., 2., 3., 5., 8., 13., 21.])
    extra_positions = np.array([34., 55., 89.])
    orientation = 'horizontal'
    lineoffset = 1
    linelength = .5
    linewidth = 2
    color = [1, 0, 0, 1]
    linestyle = 'solid'
    antialiased = True

    coll = EventCollection(positions,
                           orientation=orientation,
                           lineoffset=lineoffset,
                           linelength=linelength,
                           linewidth=linewidth,
                           color=color,
                           linestyle=linestyle,
                           antialiased=antialiased
                           )

    fig, ax = plt.subplots()
    ax.add_collection(coll)
    ax.set_title('EventCollection: default')
    props = {'positions': positions,
             'extra_positions': extra_positions,
             'orientation': orientation,
             'lineoffset': lineoffset,
             'linelength': linelength,
             'linewidth': linewidth,
             'color': color,
             'linestyle': linestyle,
             'antialiased': antialiased
             }
    ax.set_xlim(-1, 22)
    ax.set_ylim(0, 2)
    return ax, coll, props


@image_comparison(['EventCollection_plot__default'])
def test__EventCollection__get_props():
    _, coll, props = generate_EventCollection_plot()
    # check that the default segments have the correct coordinates
    check_segments(coll,
                   props['positions'],
                   props['linelength'],
                   props['lineoffset'],
                   props['orientation'])
    # check that the default positions match the input positions
    np.testing.assert_array_equal(props['positions'], coll.get_positions())
    # check that the default orientation matches the input orientation
    assert props['orientation'] == coll.get_orientation()
    # check that the default orientation matches the input orientation
    assert coll.is_horizontal()
    # check that the default linelength matches the input linelength
    assert props['linelength'] == coll.get_linelength()
    # check that the default lineoffset matches the input lineoffset
    assert props['lineoffset'] == coll.get_lineoffset()
    # check that the default linestyle matches the input linestyle
    assert coll.get_linestyle() == [(0, None)]
    # check that the default color matches the input color
    for color in [coll.get_color(), *coll.get_colors()]:
        np.testing.assert_array_equal(color, props['color'])


@image_comparison(['EventCollection_plot__set_positions'])
def test__EventCollection__set_positions():
    splt, coll, props = generate_EventCollection_plot()
    new_positions = np.hstack([props['positions'], props['extra_positions']])
    coll.set_positions(new_positions)
    np.testing.assert_array_equal(new_positions, coll.get_positions())
    check_segments(coll, new_positions,
                   props['linelength'],
                   props['lineoffset'],
                   props['orientation'])
    splt.set_title('EventCollection: set_positions')
    splt.set_xlim(-1, 90)


@image_comparison(['EventCollection_plot__add_positions'])
def test__EventCollection__add_positions():
    splt, coll, props = generate_EventCollection_plot()
    new_positions = np.hstack([props['positions'],
                               props['extra_positions'][0]])
    coll.switch_orientation()  # Test adding in the vertical orientation, too.
    coll.add_positions(props['extra_positions'][0])
    coll.switch_orientation()
    np.testing.assert_array_equal(new_positions, coll.get_positions())
    check_segments(coll,
                   new_positions,
                   props['linelength'],
                   props['lineoffset'],
                   props['orientation'])
    splt.set_title('EventCollection: add_positions')
    splt.set_xlim(-1, 35)


@image_comparison(['EventCollection_plot__append_positions'])
def test__EventCollection__append_positions():
    splt, coll, props = generate_EventCollection_plot()
    new_positions = np.hstack([props['positions'],
                               props['extra_positions'][2]])
    coll.append_positions(props['extra_positions'][2])
    np.testing.assert_array_equal(new_positions, coll.get_positions())
    check_segments(coll,
                   new_positions,
                   props['linelength'],
                   props['lineoffset'],
                   props['orientation'])
    splt.set_title('EventCollection: append_positions')
    splt.set_xlim(-1, 90)


@image_comparison(['EventCollection_plot__extend_positions'])
def test__EventCollection__extend_positions():
    splt, coll, props = generate_EventCollection_plot()
    new_positions = np.hstack([props['positions'],
                               props['extra_positions'][1:]])
    coll.extend_positions(props['extra_positions'][1:])
    np.testing.assert_array_equal(new_positions, coll.get_positions())
    check_segments(coll,
                   new_positions,
                   props['linelength'],
                   props['lineoffset'],
                   props['orientation'])
    splt.set_title('EventCollection: extend_positions')
    splt.set_xlim(-1, 90)


@image_comparison(['EventCollection_plot__switch_orientation'])
def test__EventCollection__switch_orientation():
    splt, coll, props = generate_EventCollection_plot()
    new_orientation = 'vertical'
    coll.switch_orientation()
    assert new_orientation == coll.get_orientation()
    assert not coll.is_horizontal()
    new_positions = coll.get_positions()
    check_segments(coll,
                   new_positions,
                   props['linelength'],
                   props['lineoffset'], new_orientation)
    splt.set_title('EventCollection: switch_orientation')
    splt.set_ylim(-1, 22)
    splt.set_xlim(0, 2)


@image_comparison(['EventCollection_plot__switch_orientation__2x'])
def test__EventCollection__switch_orientation_2x():
    """
    Check that calling switch_orientation twice sets the orientation back to
    the default.
    """
    splt, coll, props = generate_EventCollection_plot()
    coll.switch_orientation()
    coll.switch_orientation()
    new_positions = coll.get_positions()
    assert props['orientation'] == coll.get_orientation()
    assert coll.is_horizontal()
    np.testing.assert_array_equal(props['positions'], new_positions)
    check_segments(coll,
                   new_positions,
                   props['linelength'],
                   props['lineoffset'],
                   props['orientation'])
    splt.set_title('EventCollection: switch_orientation 2x')


@image_comparison(['EventCollection_plot__set_orientation'])
def test__EventCollection__set_orientation():
    splt, coll, props = generate_EventCollection_plot()
    new_orientation = 'vertical'
    coll.set_orientation(new_orientation)
    assert new_orientation == coll.get_orientation()
    assert not coll.is_horizontal()
    check_segments(coll,
                   props['positions'],
                   props['linelength'],
                   props['lineoffset'],
                   new_orientation)
    splt.set_title('EventCollection: set_orientation')
    splt.set_ylim(-1, 22)
    splt.set_xlim(0, 2)


@image_comparison(['EventCollection_plot__set_linelength'])
def test__EventCollection__set_linelength():
    splt, coll, props = generate_EventCollection_plot()
    new_linelength = 15
    coll.set_linelength(new_linelength)
    assert new_linelength == coll.get_linelength()
    check_segments(coll,
                   props['positions'],
                   new_linelength,
                   props['lineoffset'],
                   props['orientation'])
    splt.set_title('EventCollection: set_linelength')
    splt.set_ylim(-20, 20)


@image_comparison(['EventCollection_plot__set_lineoffset'])
def test__EventCollection__set_lineoffset():
    splt, coll, props = generate_EventCollection_plot()
    new_lineoffset = -5.
    coll.set_lineoffset(new_lineoffset)
    assert new_lineoffset == coll.get_lineoffset()
    check_segments(coll,
                   props['positions'],
                   props['linelength'],
                   new_lineoffset,
                   props['orientation'])
    splt.set_title('EventCollection: set_lineoffset')
    splt.set_ylim(-6, -4)


@image_comparison([
    'EventCollection_plot__set_linestyle',
    'EventCollection_plot__set_linestyle',
    'EventCollection_plot__set_linewidth',
])
def test__EventCollection__set_prop():
    for prop, value, expected in [
            ('linestyle', 'dashed', [(0, (6.0, 6.0))]),
            ('linestyle', (0, (6., 6.)), [(0, (6.0, 6.0))]),
            ('linewidth', 5, 5),
    ]:
        splt, coll, _ = generate_EventCollection_plot()
        coll.set(**{prop: value})
        assert plt.getp(coll, prop) == expected
        splt.set_title(f'EventCollection: set_{prop}')


@image_comparison(['EventCollection_plot__set_color'])
def test__EventCollection__set_color():
    splt, coll, _ = generate_EventCollection_plot()
    new_color = np.array([0, 1, 1, 1])
    coll.set_color(new_color)
    for color in [coll.get_color(), *coll.get_colors()]:
        np.testing.assert_array_equal(color, new_color)
    splt.set_title('EventCollection: set_color')


def check_segments(coll, positions, linelength, lineoffset, orientation):
    """
    Test helper checking that all values in the segment are correct, given a
    particular set of inputs.
    """
    segments = coll.get_segments()
    if (orientation.lower() == 'horizontal'
            or orientation.lower() == 'none' or orientation is None):
        # if horizontal, the position in is in the y-axis
        pos1 = 1
        pos2 = 0
    elif orientation.lower() == 'vertical':
        # if vertical, the position in is in the x-axis
        pos1 = 0
        pos2 = 1
    else:
        raise ValueError("orientation must be 'horizontal' or 'vertical'")

    # test to make sure each segment is correct
    for i, segment in enumerate(segments):
        assert segment[0, pos1] == lineoffset + linelength / 2
        assert segment[1, pos1] == lineoffset - linelength / 2
        assert segment[0, pos2] == positions[i]
        assert segment[1, pos2] == positions[i]


def test_null_collection_datalim():
    col = mcollections.PathCollection([])
    col_data_lim = col.get_datalim(mtransforms.IdentityTransform())
    assert_array_equal(col_data_lim.get_points(),
                       mtransforms.Bbox.null().get_points())


def test_add_collection():
    # Test if data limits are unchanged by adding an empty collection.
    # GitHub issue #1490, pull #1497.
    plt.figure()
    ax = plt.axes()
    coll = ax.scatter([0, 1], [0, 1])
    ax.add_collection(coll)
    bounds = ax.dataLim.bounds
    coll = ax.scatter([], [])
    assert ax.dataLim.bounds == bounds


@pytest.mark.style('mpl20')
@check_figures_equal(extensions=['png'])
def test_collection_log_datalim(fig_test, fig_ref):
    # Data limits should respect the minimum x/y when using log scale.
    x_vals = [4.38462e-6, 5.54929e-6, 7.02332e-6, 8.88889e-6, 1.12500e-5,
              1.42383e-5, 1.80203e-5, 2.28070e-5, 2.88651e-5, 3.65324e-5,
              4.62363e-5, 5.85178e-5, 7.40616e-5, 9.37342e-5, 1.18632e-4]
    y_vals = [0.0, 0.1, 0.182, 0.332, 0.604, 1.1, 2.0, 3.64, 6.64, 12.1, 22.0,
              39.6, 71.3]

    x, y = np.meshgrid(x_vals, y_vals)
    x = x.flatten()
    y = y.flatten()

    ax_test = fig_test.subplots()
    ax_test.set_xscale('log')
    ax_test.set_yscale('log')
    ax_test.margins = 0
    ax_test.scatter(x, y)

    ax_ref = fig_ref.subplots()
    ax_ref.set_xscale('log')
    ax_ref.set_yscale('log')
    ax_ref.plot(x, y, marker="o", ls="")


def test_quiver_limits():
    ax = plt.axes()
    x, y = np.arange(8), np.arange(10)
    u = v = np.linspace(0, 10, 80).reshape(10, 8)
    q = plt.quiver(x, y, u, v)
    assert q.get_datalim(ax.transData).bounds == (0., 0., 7., 9.)

    plt.figure()
    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.quiver(x, y, np.sin(x), np.cos(y), transform=trans)
    assert ax.dataLim.bounds == (20.0, 30.0, 15.0, 6.0)


def test_barb_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)


@image_comparison(['EllipseCollection_test_image.png'], remove_text=True)
def test_EllipseCollection():
    # Test basic functionality
    fig, ax = plt.subplots()
    x = np.arange(4)
    y = np.arange(3)
    X, Y = np.meshgrid(x, y)
    XY = np.vstack((X.ravel(), Y.ravel())).T

    ww = X / x[-1]
    hh = Y / y[-1]
    aa = np.ones_like(ww) * 20  # first axis is 20 degrees CCW from x axis

    ec = mcollections.EllipseCollection(ww, hh, aa,
                                        units='x',
                                        offsets=XY,
                                        transOffset=ax.transData,
                                        facecolors='none')
    ax.add_collection(ec)
    ax.autoscale_view()


@image_comparison(['polycollection_close.png'], remove_text=True)
def test_polycollection_close():
    from mpl_toolkits.mplot3d import Axes3D

    vertsQuad = [
        [[0., 0.], [0., 1.], [1., 1.], [1., 0.]],
        [[0., 1.], [2., 3.], [2., 2.], [1., 1.]],
        [[2., 2.], [2., 3.], [4., 1.], [3., 1.]],
        [[3., 0.], [3., 1.], [4., 1.], [4., 0.]]]

    fig = plt.figure()
    ax = fig.add_axes(Axes3D(fig, auto_add_to_figure=False))

    colors = ['r', 'g', 'b', 'y', 'k']
    zpos = list(range(5))

    poly = mcollections.PolyCollection(
        vertsQuad * len(zpos), linewidth=0.25)
    poly.set_alpha(0.7)

    # need to have a z-value for *each* polygon = element!
    zs = []
    cs = []
    for z, c in zip(zpos, colors):
        zs.extend([z] * len(vertsQuad))
        cs.extend([c] * len(vertsQuad))

    poly.set_color(cs)

    ax.add_collection3d(poly, zs=zs, zdir='y')

    # axis limit settings:
    ax.set_xlim3d(0, 4)
    ax.set_zlim3d(0, 3)
    ax.set_ylim3d(0, 4)


@image_comparison(['regularpolycollection_rotate.png'], remove_text=True)
def test_regularpolycollection_rotate():
    xx, yy = np.mgrid[:10, :10]
    xy_points = np.transpose([xx.flatten(), yy.flatten()])
    rotations = np.linspace(0, 2*np.pi, len(xy_points))

    fig, ax = plt.subplots()
    for xy, alpha in zip(xy_points, rotations):
        col = mcollections.RegularPolyCollection(
            4, sizes=(100,), rotation=alpha,
            offsets=[xy], transOffset=ax.transData)
        ax.add_collection(col, autolim=True)
    ax.autoscale_view()


@image_comparison(['regularpolycollection_scale.png'], remove_text=True)
def test_regularpolycollection_scale():
    # See issue #3860

    class SquareCollection(mcollections.RegularPolyCollection):
        def __init__(self, **kwargs):
            super().__init__(4, rotation=np.pi/4., **kwargs)

        def get_transform(self):
            """Return transform scaling circle areas to data space."""
            ax = self.axes

            pts2pixels = 72.0 / ax.figure.dpi

            scale_x = pts2pixels * ax.bbox.width / ax.viewLim.width
            scale_y = pts2pixels * ax.bbox.height / ax.viewLim.height
            return mtransforms.Affine2D().scale(scale_x, scale_y)

    fig, ax = plt.subplots()

    xy = [(0, 0)]
    # Unit square has a half-diagonal of `1/sqrt(2)`, so `pi * r**2` equals...
    circle_areas = [np.pi / 2]
    squares = SquareCollection(sizes=circle_areas, offsets=xy,
                               transOffset=ax.transData)
    ax.add_collection(squares, autolim=True)
    ax.axis([-1, 1, -1, 1])


def test_picking():
    fig, ax = plt.subplots()
    col = ax.scatter([0], [0], [1000], picker=True)
    fig.savefig(io.BytesIO(), dpi=fig.dpi)
    mouse_event = SimpleNamespace(x=325, y=240)
    found, indices = col.contains(mouse_event)
    assert found
    assert_array_equal(indices['ind'], [0])


def test_linestyle_single_dashes():
    plt.scatter([0, 1, 2], [0, 1, 2], linestyle=(0., [2., 2.]))
    plt.draw()


@image_comparison(['size_in_xy.png'], remove_text=True)
def test_size_in_xy():
    fig, ax = plt.subplots()

    widths, heights, angles = (10, 10), 10, 0
    widths = 10, 10
    coords = [(10, 10), (15, 15)]
    e = mcollections.EllipseCollection(
        widths, heights, angles,
        units='xy',
        offsets=coords,
        transOffset=ax.transData)

    ax.add_collection(e)

    ax.set_xlim(0, 30)
    ax.set_ylim(0, 30)


def test_pandas_indexing(pd):

    # Should not fail break when faced with a
    # non-zero indexed series
    index = [11, 12, 13]
    ec = fc = pd.Series(['red', 'blue', 'green'], index=index)
    lw = pd.Series([1, 2, 3], index=index)
    ls = pd.Series(['solid', 'dashed', 'dashdot'], index=index)
    aa = pd.Series([True, False, True], index=index)

    Collection(edgecolors=ec)
    Collection(facecolors=fc)
    Collection(linewidths=lw)
    Collection(linestyles=ls)
    Collection(antialiaseds=aa)


@pytest.mark.style('default')
def test_lslw_bcast():
    col = mcollections.PathCollection([])
    col.set_linestyles(['-', '-'])
    col.set_linewidths([1, 2, 3])

    assert col.get_linestyles() == [(0, None)] * 6
    assert col.get_linewidths() == [1, 2, 3] * 2

    col.set_linestyles(['-', '-', '-'])
    assert col.get_linestyles() == [(0, None)] * 3
    assert (col.get_linewidths() == [1, 2, 3]).all()


@pytest.mark.style('default')
def test_capstyle():
    col = mcollections.PathCollection([], capstyle='round')
    assert col.get_capstyle() == 'round'
    col.set_capstyle('butt')
    assert col.get_capstyle() == 'butt'


@pytest.mark.style('default')
def test_joinstyle():
    col = mcollections.PathCollection([], joinstyle='round')
    assert col.get_joinstyle() == 'round'
    col.set_joinstyle('miter')
    assert col.get_joinstyle() == 'miter'


@image_comparison(['cap_and_joinstyle.png'])
def test_cap_and_joinstyle_image():
    fig, ax = plt.subplots()
    ax.set_xlim([-0.5, 1.5])
    ax.set_ylim([-0.5, 2.5])

    x = np.array([0.0, 1.0, 0.5])
    ys = np.array([[0.0], [0.5], [1.0]]) + np.array([[0.0, 0.0, 1.0]])

    segs = np.zeros((3, 3, 2))
    segs[:, :, 0] = x
    segs[:, :, 1] = ys
    line_segments = LineCollection(segs, linewidth=[10, 15, 20])
    line_segments.set_capstyle("round")
    line_segments.set_joinstyle("miter")

    ax.add_collection(line_segments)
    ax.set_title('Line collection with customized caps and joinstyle')


@image_comparison(['scatter_post_alpha.png'],
                  remove_text=True, style='default')
def test_scatter_post_alpha():
    fig, ax = plt.subplots()
    sc = ax.scatter(range(5), range(5), c=range(5))
    sc.set_alpha(.1)


def test_scatter_alpha_array():
    x = np.arange(5)
    alpha = x / 5
    # With colormapping.
    fig, (ax0, ax1) = plt.subplots(2)
    sc0 = ax0.scatter(x, x, c=x, alpha=alpha)
    sc1 = ax1.scatter(x, x, c=x)
    sc1.set_alpha(alpha)
    plt.draw()
    assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
    assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
    # Without colormapping.
    fig, (ax0, ax1) = plt.subplots(2)
    sc0 = ax0.scatter(x, x, color=['r', 'g', 'b', 'c', 'm'], alpha=alpha)
    sc1 = ax1.scatter(x, x, color='r', alpha=alpha)
    plt.draw()
    assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
    assert_array_equal(sc1.get_facecolors()[:, -1], alpha)
    # Without colormapping, and set alpha afterward.
    fig, (ax0, ax1) = plt.subplots(2)
    sc0 = ax0.scatter(x, x, color=['r', 'g', 'b', 'c', 'm'])
    sc0.set_alpha(alpha)
    sc1 = ax1.scatter(x, x, color='r')
    sc1.set_alpha(alpha)
    plt.draw()
    assert_array_equal(sc0.get_facecolors()[:, -1], alpha)
    assert_array_equal(sc1.get_facecolors()[:, -1], alpha)


def test_pathcollection_legend_elements():
    np.random.seed(19680801)
    x, y = np.random.rand(2, 10)
    y = np.random.rand(10)
    c = np.random.randint(0, 5, size=10)
    s = np.random.randint(10, 300, size=10)

    fig, ax = plt.subplots()
    sc = ax.scatter(x, y, c=c, s=s, cmap="jet", marker="o", linewidths=0)

    h, l = sc.legend_elements(fmt="{x:g}")
    assert len(h) == 5
    assert_array_equal(np.array(l).astype(float), np.arange(5))
    colors = np.array([line.get_color() for line in h])
    colors2 = sc.cmap(np.arange(5)/4)
    assert_array_equal(colors, colors2)
    l1 = ax.legend(h, l, loc=1)

    h2, lab2 = sc.legend_elements(num=9)
    assert len(h2) == 9
    l2 = ax.legend(h2, lab2, loc=2)

    h, l = sc.legend_elements(prop="sizes", alpha=0.5, color="red")
    alpha = np.array([line.get_alpha() for line in h])
    assert_array_equal(alpha, 0.5)
    color = np.array([line.get_markerfacecolor() for line in h])
    assert_array_equal(color, "red")
    l3 = ax.legend(h, l, loc=4)

    h, l = sc.legend_elements(prop="sizes", num=4, fmt="{x:.2f}",
                              func=lambda x: 2*x)
    actsizes = [line.get_markersize() for line in h]
    labeledsizes = np.sqrt(np.array(l).astype(float)/2)
    assert_array_almost_equal(actsizes, labeledsizes)
    l4 = ax.legend(h, l, loc=3)

    loc = mpl.ticker.MaxNLocator(nbins=9, min_n_ticks=9-1,
                                 steps=[1, 2, 2.5, 3, 5, 6, 8, 10])
    h5, lab5 = sc.legend_elements(num=loc)
    assert len(h2) == len(h5)

    levels = [-1, 0, 55.4, 260]
    h6, lab6 = sc.legend_elements(num=levels, prop="sizes", fmt="{x:g}")
    assert_array_equal(np.array(lab6).astype(float), levels[2:])

    for l in [l1, l2, l3, l4]:
        ax.add_artist(l)

    fig.canvas.draw()


def test_EventCollection_nosort():
    # Check that EventCollection doesn't modify input in place
    arr = np.array([3, 2, 1, 10])
    coll = EventCollection(arr)
    np.testing.assert_array_equal(arr, np.array([3, 2, 1, 10]))


def test_collection_set_verts_array():
    verts = np.arange(80, dtype=np.double).reshape(10, 4, 2)
    col_arr = PolyCollection(verts)
    col_list = PolyCollection(list(verts))
    assert len(col_arr._paths) == len(col_list._paths)
    for ap, lp in zip(col_arr._paths, col_list._paths):
        assert np.array_equal(ap._vertices, lp._vertices)
        assert np.array_equal(ap._codes, lp._codes)

    verts_tuple = np.empty(10, dtype=object)
    verts_tuple[:] = [tuple(tuple(y) for y in x) for x in verts]
    col_arr_tuple = PolyCollection(verts_tuple)
    assert len(col_arr._paths) == len(col_arr_tuple._paths)
    for ap, atp in zip(col_arr._paths, col_arr_tuple._paths):
        assert np.array_equal(ap._vertices, atp._vertices)
        assert np.array_equal(ap._codes, atp._codes)


def test_blended_collection_autolim():
    a = [1, 2, 4]
    height = .2

    xy_pairs = np.column_stack([np.repeat(a, 2), np.tile([0, height], len(a))])
    line_segs = xy_pairs.reshape([len(a), 2, 2])

    f, ax = plt.subplots()
    trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
    ax.add_collection(LineCollection(line_segs, transform=trans))
    ax.autoscale_view(scalex=True, scaley=False)
    np.testing.assert_allclose(ax.get_xlim(), [1., 4.])


def test_singleton_autolim():
    fig, ax = plt.subplots()
    ax.scatter(0, 0)
    np.testing.assert_allclose(ax.get_ylim(), [-0.06, 0.06])
    np.testing.assert_allclose(ax.get_xlim(), [-0.06, 0.06])


def test_quadmesh_set_array():
    x = np.arange(4)
    y = np.arange(4)
    z = np.arange(9).reshape((3, 3))
    fig, ax = plt.subplots()
    coll = ax.pcolormesh(x, y, np.ones(z.shape))
    # Test that the collection is able to update with a 2d array
    coll.set_array(z)
    fig.canvas.draw()
    assert np.array_equal(coll.get_array(), z)

    # Check that pre-flattened arrays work too
    coll.set_array(np.ones(9))
    fig.canvas.draw()
    assert np.array_equal(coll.get_array(), np.ones(9))


def test_quadmesh_alpha_array():
    x = np.arange(4)
    y = np.arange(4)
    z = np.arange(9).reshape((3, 3))
    alpha = z / z.max()
    alpha_flat = alpha.ravel()
    # Provide 2-D alpha:
    fig, (ax0, ax1) = plt.subplots(2)
    coll1 = ax0.pcolormesh(x, y, z, alpha=alpha)
    coll2 = ax1.pcolormesh(x, y, z)
    coll2.set_alpha(alpha)
    plt.draw()
    assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
    assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)
    # Or provide 1-D alpha:
    fig, (ax0, ax1) = plt.subplots(2)
    coll1 = ax0.pcolormesh(x, y, z, alpha=alpha_flat)
    coll2 = ax1.pcolormesh(x, y, z)
    coll2.set_alpha(alpha_flat)
    plt.draw()
    assert_array_equal(coll1.get_facecolors()[:, -1], alpha_flat)
    assert_array_equal(coll2.get_facecolors()[:, -1], alpha_flat)


def test_alpha_validation():
    # Most of the relevant testing is in test_artist and test_colors.
    fig, ax = plt.subplots()
    pc = ax.pcolormesh(np.arange(12).reshape((3, 4)))
    with pytest.raises(ValueError, match="^Data array shape"):
        pc.set_alpha([0.5, 0.6])
        pc.update_scalarmappable()


def test_legend_inverse_size_label_relationship():
    """
    Ensure legend markers scale appropriately when label and size are
    inversely related.
    Here label = 5 / size
    """

    np.random.seed(19680801)
    X = np.random.random(50)
    Y = np.random.random(50)
    C = 1 - np.random.random(50)
    S = 5 / C

    legend_sizes = [0.2, 0.4, 0.6, 0.8]
    fig, ax = plt.subplots()
    sc = ax.scatter(X, Y, s=S)
    handles, labels = sc.legend_elements(
      prop='sizes', num=legend_sizes, func=lambda s: 5 / s
    )

    # Convert markersize scale to 's' scale
    handle_sizes = [x.get_markersize() for x in handles]
    handle_sizes = [5 / x**2 for x in handle_sizes]

    assert_array_almost_equal(handle_sizes, legend_sizes, decimal=1)


@pytest.mark.style('default')
@pytest.mark.parametrize('pcfunc', [plt.pcolor, plt.pcolormesh])
def test_color_logic(pcfunc):
    z = np.arange(12).reshape(3, 4)
    # Explicitly set an edgecolor.
    pc = pcfunc(z, edgecolors='red', facecolors='none')
    pc.update_scalarmappable()  # This is called in draw().
    # Define 2 reference "colors" here for multiple use.
    face_default = mcolors.to_rgba_array(pc._get_default_facecolor())
    mapped = pc.get_cmap()(pc.norm((z.ravel())))
    # Github issue #1302:
    assert mcolors.same_color(pc.get_edgecolor(), 'red')
    # Check setting attributes after initialization:
    pc = pcfunc(z)
    pc.set_facecolor('none')
    pc.set_edgecolor('red')
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), 'none')
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
    pc.set_alpha(0.5)
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 0.5]])
    pc.set_alpha(None)  # restore default alpha
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
    # Reset edgecolor to default.
    pc.set_edgecolor(None)
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_edgecolor(), mapped)
    pc.set_facecolor(None)  # restore default for facecolor
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), mapped)
    assert mcolors.same_color(pc.get_edgecolor(), 'none')
    # Turn off colormapping entirely:
    pc.set_array(None)
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_edgecolor(), 'none')
    assert mcolors.same_color(pc.get_facecolor(), face_default)  # not mapped
    # Turn it back on by restoring the array (must be 1D!):
    pc.set_array(z.ravel())
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), mapped)
    assert mcolors.same_color(pc.get_edgecolor(), 'none')
    # Give color via tuple rather than string.
    pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=(0, 1, 0))
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), mapped)
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
    # Provide an RGB array; mapping overrides it.
    pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 3)))
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), mapped)
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
    # Turn off the mapping.
    pc.set_array(None)
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), np.ones((12, 3)))
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
    # And an RGBA array.
    pc = pcfunc(z, edgecolors=(1, 0, 0), facecolors=np.ones((12, 4)))
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), mapped)
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])
    # Turn off the mapping.
    pc.set_array(None)
    pc.update_scalarmappable()
    assert mcolors.same_color(pc.get_facecolor(), np.ones((12, 4)))
    assert mcolors.same_color(pc.get_edgecolor(), [[1, 0, 0, 1]])


def test_LineCollection_args():
    with pytest.warns(MatplotlibDeprecationWarning):
        lc = LineCollection(None, 2.2, 'r', zorder=3, facecolors=[0, 1, 0, 1])
        assert lc.get_linewidth()[0] == 2.2
        assert mcolors.same_color(lc.get_edgecolor(), 'r')
        assert lc.get_zorder() == 3
        assert mcolors.same_color(lc.get_facecolor(), [[0, 1, 0, 1]])
    # To avoid breaking mplot3d, LineCollection internally sets the facecolor
    # kwarg if it has not been specified.  Hence we need the following test
    # for LineCollection._set_default().
    lc = LineCollection(None, facecolor=None)
    assert mcolors.same_color(lc.get_facecolor(), 'none')


def test_array_wrong_dimensions():
    z = np.arange(12).reshape(3, 4)
    pc = plt.pcolor(z)
    with pytest.raises(ValueError, match="^Collections can only map"):
        pc.set_array(z)
        pc.update_scalarmappable()
    pc = plt.pcolormesh(z)
    pc.set_array(z)  # 2D is OK for Quadmesh
    pc.update_scalarmappable()
