import numpy as np
from bokeh.layouts import Column
from bokeh.models import Div, Toolbar

import holoviews as hv
from holoviews.operation import gridmatrix
from holoviews.streams import Stream
from holoviews.testing import assert_data_equal

from .test_plot import TestBokehPlot, bokeh_renderer


class TestGridPlot(TestBokehPlot):
    def test_grid_title(self):
        grid = hv.GridSpace(
            {
                (i, j): hv.HoloMap(
                    {a: hv.Image(np.random.rand(10, 10)) for a in range(3)}, kdims=["X"]
                )
                for i in range(2)
                for j in range(3)
            }
        )
        plot = bokeh_renderer.get_plot(grid)
        title = plot.handles["title"]
        assert isinstance(title, Div)
        text = (
            '<span style="color:black;font-family:Arial;'
            'font-style:bold;font-weight:bold;font-size:16pt">X: 0</span>'
        )
        assert title.text == text

    def test_grid_title_update(self):
        grid = hv.GridSpace(
            {
                (i, j): hv.HoloMap(
                    {a: hv.Image(np.random.rand(10, 10)) for a in range(3)}, kdims=["X"]
                )
                for i in range(2)
                for j in range(3)
            }
        )
        plot = bokeh_renderer.get_plot(grid)
        plot.update(1)
        title = plot.handles["title"]
        assert isinstance(title, Div)
        text = (
            '<span style="color:black;font-family:Arial;'
            'font-style:bold;font-weight:bold;font-size:16pt">X: 1</span>'
        )
        assert title.text == text

    def test_gridmatrix_overlaid_batched(self):
        ds = hv.Dataset(
            (["A"] * 5 + ["B"] * 5, np.random.rand(10), np.random.rand(10)), kdims=["a", "b", "c"]
        )
        gmatrix = gridmatrix(ds.groupby("a", container_type=hv.NdOverlay))
        plot = bokeh_renderer.get_plot(gmatrix)

        sp1 = plot.subplots[("b", "c")]
        assert sp1.state.xaxis[0].visible is False
        assert sp1.state.yaxis[0].visible is True
        sp2 = plot.subplots[("b", "b")]
        assert sp2.state.xaxis[0].visible is True
        assert sp2.state.yaxis[0].visible is True
        sp3 = plot.subplots[("c", "b")]
        assert sp3.state.xaxis[0].visible is True
        assert sp3.state.yaxis[0].visible is False
        sp4 = plot.subplots[("c", "c")]
        assert sp4.state.xaxis[0].visible is False
        assert sp4.state.yaxis[0].visible is False

    def test_gridspace_sparse(self):
        grid = hv.GridSpace(
            {
                (i, j): hv.Curve(range(i + j))
                for i in range(1, 3)
                for j in range(2, 4)
                if not (i == 1 and j == 2)
            }
        )
        plot = bokeh_renderer.get_plot(grid)
        size = bokeh_renderer.get_size(plot.state)
        assert size == (320, 311)

    def test_grid_shared_source_synced_update(self):
        hmap = hv.HoloMap(
            {
                i: hv.Dataset(
                    {chr(65 + j): np.random.rand(i + 2) for j in range(4)},
                    kdims=["A", "B", "C", "D"],
                )
                for i in range(3)
            }
        )

        # Create two holomaps of points sharing the same data source
        hmap1 = hmap.map(lambda x: hv.Points(x.clone(kdims=["A", "B"])), hv.Dataset)
        hmap2 = hmap.map(lambda x: hv.Points(x.clone(kdims=["D", "C"])), hv.Dataset)

        # Pop key (1,) for one of the HoloMaps and make GridSpace
        hmap2.pop(1)
        grid = hv.GridSpace({0: hmap1, 2: hmap2}, kdims=["X"]).opts(shared_datasource=True)

        # Get plot
        plot = bokeh_renderer.get_plot(grid)

        # Check plot created shared data source and recorded expected columns
        sources = plot.handles.get("shared_sources", [])
        source_cols = plot.handles.get("source_cols", {})
        assert len(sources) == 1
        source = sources[0]
        data = source.data
        cols = source_cols[id(source)]
        assert set(cols) == {"A", "B", "C", "D"}

        # Ensure the source contains the expected columns
        assert set(data.keys()) == {"A", "B", "C", "D"}

        # Update to key (1,) and check the source contains data
        # corresponding to hmap1 and filled in NaNs for hmap2,
        # which was popped above
        plot.update((1,))
        assert_data_equal(data["A"], hmap1[1].dimension_values(0))
        assert_data_equal(data["B"], hmap1[1].dimension_values(1))
        assert_data_equal(data["C"], np.full_like(hmap1[1].dimension_values(0), np.nan))
        assert_data_equal(data["D"], np.full_like(hmap1[1].dimension_values(0), np.nan))

    def test_grid_set_toolbar_location(self):
        grid = hv.GridSpace({0: hv.Curve([]), 1: hv.Points([])}, "X").opts(toolbar="left")
        plot = bokeh_renderer.get_plot(grid)
        assert isinstance(plot.state, Column)
        assert isinstance(plot.state.children[0].toolbar, Toolbar)

    def test_grid_disable_toolbar(self):
        grid = hv.GridSpace({0: hv.Curve([]), 1: hv.Points([])}, "X").opts(toolbar=None)
        plot = bokeh_renderer.get_plot(grid)
        assert isinstance(plot.state, Column)
        assert [p for p in plot.state.children if isinstance(p, Toolbar)] == []

    def test_grid_dimensioned_stream_title_update(self):
        stream = Stream.define("Test", test=0)()
        dmap = hv.DynamicMap(lambda test: hv.Curve([]), kdims=["test"], streams=[stream])
        grid = hv.GridMatrix({0: dmap, 1: hv.Curve([])}, "X")
        plot = bokeh_renderer.get_plot(grid)
        assert "test: 0" in plot.handles["title"].text
        stream.event(test=1)
        assert "test: 1" in plot.handles["title"].text
        plot.cleanup()
        assert stream._subscribers == []
