import numpy as np

from holoviews.element import QuadMesh, Image, Dataset

from .testplot import TestMPLPlot, mpl_renderer


class TestQuadMeshPlot(TestMPLPlot):

    def test_quadmesh_invert_axes(self):
        arr = np.array([[0, 1, 2], [3, 4,  5]])
        qmesh = QuadMesh(Image(arr)).opts(plot=dict(invert_axes=True))
        plot = mpl_renderer.get_plot(qmesh)
        artist = plot.handles['artist']
        self.assertEqual(artist.get_array().data, arr.T[:, ::-1].flatten())

    def test_quadmesh_update_cbar(self):
        xs = ys = np.linspace(0, 6, 10)
        zs = np.linspace(1, 2, 5)
        XS, YS, ZS = np.meshgrid(xs, ys, zs)
        values = np.sin(XS) * ZS
        ds = Dataset((xs, ys, zs, values.T), ['x', 'y', 'z'], 'values')
        hmap = ds.to(QuadMesh).options(colorbar=True, framewise=True)
        plot = mpl_renderer.get_plot(hmap)
        self.assertEqual(plot.handles['cbar'].get_clim(), (-0.9989549170979283, 0.9719379013633127))
        plot.update(3)
        self.assertEqual(plot.handles['cbar'].get_clim(), (-1.7481711049213744, 1.7008913273857973))
