import numpy as np

import holoviews as hv
from holoviews.core.ndmapping import MultiDimensionalMapping, UniformNdMapping
from holoviews.testing import assert_data_equal, assert_element_equal


class DimensionTest:
    def test_dimension_init(self):
        hv.Dimension("Test dimension")
        hv.Dimension("Test dimension", cyclic=True)
        hv.Dimension("Test dimension", cyclic=True, type=int)
        hv.Dimension("Test dimension", cyclic=True, type=int, unit="Twilight zones")

    def test_dimension_clone(self):
        dim1 = hv.Dimension("Test dimension")
        dim2 = dim1.clone(cyclic=True)
        assert dim2.cyclic is True

        dim3 = dim1.clone("New test dimension", unit="scovilles")
        assert dim3.name == "New test dimension"
        assert dim3.unit == "scovilles"

    def test_dimension_pprint(self):
        dim = hv.Dimension("Test dimension", cyclic=True, type=float, unit="Twilight zones")
        assert dim.pprint_value_string(3.23451) == "Test dimension: 3.2345 Twilight zones"
        assert dim.pprint_value_string(4.23441) == "Test dimension: 4.2344 Twilight zones"
        assert dim.pprint_value(3.23451, print_unit=True) == "3.2345 Twilight zones"
        assert dim.pprint_value(4.23441, print_unit=True) == "4.2344 Twilight zones"


class NdIndexableMappingTest:
    def setup_method(self):
        self.init_items_1D_list = [(1, "a"), (5, "b")]
        self.init_item_list = [((1, 2.0), "a"), ((5, 3.0), "b")]
        self.init_item_odict = dict([((1, 2.0), "a"), ((5, 3.0), "b")])
        self.dimension_labels = ["intdim", "floatdim"]
        self.dim1 = hv.Dimension("intdim", type=int)
        self.dim2 = hv.Dimension("floatdim", type=float)
        self.time_dimension = hv.Dimension

    def test_idxmapping_init(self):
        MultiDimensionalMapping()

    def test_idxmapping_init_item_odict(self):
        MultiDimensionalMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])

    def test_idxmapping_init_item_list(self):
        MultiDimensionalMapping(self.init_item_list, kdims=[self.dim1, self.dim2])

    def test_idxmapping_init_dimstr(self):
        MultiDimensionalMapping(self.init_item_odict, kdims=self.dimension_labels)

    def test_idxmapping_init_dimdict(self):
        MultiDimensionalMapping(
            self.init_item_odict,
            kdims=[dict(name=label) for label in self.dimension_labels],
        )

    def test_idxmapping_init_dimensions(self):
        MultiDimensionalMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])

    def test_idxmapping_dimension_labels(self):
        idxmap = MultiDimensionalMapping(self.init_item_odict, kdims=[self.dim1, "floatdim"])
        assert [d.name for d in idxmap.kdims] == self.dimension_labels

    def test_idxmapping_ndims(self):
        dims = [self.dim1, self.dim2, "strdim"]
        idxmap = MultiDimensionalMapping(kdims=dims)
        assert idxmap.ndims == len(dims)

    def test_idxmapping_key_len_check(self):
        try:
            MultiDimensionalMapping(initial_items=self.init_item_odict)
            raise AssertionError("Invalid key length check failed.")
        except KeyError:
            pass

    def test_idxmapping_nested_update(self):
        data1 = [(0, "a"), (1, "b")]
        data2 = [(2, "c"), (3, "d")]
        data3 = [(2, "e"), (3, "f")]

        ndmap1 = MultiDimensionalMapping(data1, kdims=[self.dim1])
        ndmap2 = MultiDimensionalMapping(data2, kdims=[self.dim1])
        ndmap3 = MultiDimensionalMapping(data3, kdims=[self.dim1])

        ndmap_list = [(0.5, ndmap1), (1.5, ndmap2)]
        nested_ndmap = MultiDimensionalMapping(ndmap_list, kdims=[self.dim2])
        nested_ndmap[(0.5,)].update(dict([(0, "c"), (1, "d")]))
        assert list(nested_ndmap[0.5].values()) == ["c", "d"]

        nested_ndmap[1.5] = ndmap3
        assert list(nested_ndmap[1.5].values()) == ["e", "f"]

    def test_ndmapping_slice_lower_bound_inclusive_int(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[1:].keys() == [(1, 2.0), (5, 3.0)]

    def test_ndmapping_slice_lower_bound_inclusive2_int(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[1:6].keys() == [(1, 2.0), (5, 3.0)]

    def test_ndmapping_slice_upper_bound_exclusive_int(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[1:5].keys() == [(1, 2.0)]

    def test_ndmapping_slice_upper_bound_exclusive2_int(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[:5].keys() == [(1, 2.0)]

    def test_ndmapping_slice_lower_bound_inclusive_float(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[:, 2.0:].keys() == [(1, 2.0), (5, 3.0)]

    def test_ndmapping_slice_lower_bound_inclusive2_float(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[:, 2.0:5.0].keys() == [(1, 2.0), (5, 3.0)]

    def test_ndmapping_slice_upper_bound_exclusive_float(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[:, :3.0].keys() == [(1, 2.0)]

    def test_ndmapping_slice_upper_bound_exclusive2_float(self):
        ndmap = hv.NdMapping(self.init_item_odict, kdims=[self.dim1, self.dim2])
        assert ndmap[:, 0.0:3.0].keys() == [(1, 2.0)]

    def test_idxmapping_unsorted(self):
        data = [("B", 1), ("C", 2), ("A", 3)]
        ndmap = MultiDimensionalMapping(data, sort=False)
        assert ndmap.keys() == ["B", "C", "A"]

    def test_idxmapping_unsorted_clone(self):
        data = [("B", 1), ("C", 2), ("A", 3)]
        ndmap = MultiDimensionalMapping(data, sort=False).clone()
        assert ndmap.keys() == ["B", "C", "A"]

    def test_idxmapping_groupby_unsorted(self):
        data = [(("B", 2), 1), (("C", 2), 2), (("A", 1), 3)]
        grouped = hv.NdMapping(data, sort=False, kdims=["X", "Y"]).groupby("Y")
        assert grouped.keys() == [2, 1]
        assert grouped.values()[0].keys() == ["B", "C"]
        assert grouped.last.keys() == ["A"]

    def test_idxmapping_reindex(self):
        data = [((0, 0.5), "a"), ((1, 0.5), "b")]
        ndmap = MultiDimensionalMapping(data, kdims=[self.dim1, self.dim2])

        reduced_dims = ["intdim"]
        reduced_ndmap = ndmap.reindex(reduced_dims)

        assert [d.name for d in reduced_ndmap.kdims] == reduced_dims

    def test_idxmapping_redim(self):
        data = [((0, 0.5), "a"), ((1, 0.5), "b")]
        ndmap = MultiDimensionalMapping(data, kdims=[self.dim1, self.dim2])
        redimmed = ndmap.redim(intdim="Integer")
        assert redimmed.kdims == [
            hv.Dimension("Integer", type=int),
            hv.Dimension("floatdim", type=float),
        ]

    def test_idxmapping_redim_range_aux(self):
        data = [((0, 0.5), "a"), ((1, 0.5), "b")]
        ndmap = MultiDimensionalMapping(data, kdims=[self.dim1, self.dim2])
        redimmed = ndmap.redim.range(intdim=(-9, 9))
        assert redimmed.kdims == [
            hv.Dimension(
                "intdim",
                type=int,
                range=(
                    -9,
                    9,
                ),
            ),
            hv.Dimension("floatdim", type=float),
        ]

    def test_idxmapping_redim_type_aux(self):
        data = [((0, 0.5), "a"), ((1, 0.5), "b")]
        ndmap = MultiDimensionalMapping(data, kdims=[self.dim1, self.dim2])
        redimmed = ndmap.redim.type(intdim=str)
        assert redimmed.kdims == [
            hv.Dimension("intdim", type=str),
            hv.Dimension("floatdim", type=float),
        ]

    def test_idxmapping_add_dimension(self):
        ndmap = MultiDimensionalMapping(self.init_items_1D_list, kdims=[self.dim1])
        ndmap2d = ndmap.add_dimension(self.dim2, 0, 0.5)

        assert list(ndmap2d.keys()) == [(0.5, 1), (0.5, 5)]
        assert ndmap2d.kdims == [self.dim2, self.dim1]

    def test_idxmapping_apply_key_type(self):
        data = dict([(0.5, "a"), (1.5, "b")])
        ndmap = MultiDimensionalMapping(data, kdims=[self.dim1])

        assert list(ndmap.keys()) == [0, 1]

    def test_setitem_nested_1(self):
        nested1 = MultiDimensionalMapping([("B", 1)])
        ndmap = MultiDimensionalMapping([("A", nested1)])
        nested2 = MultiDimensionalMapping([("B", 2)])
        ndmap["A"] = nested2
        assert ndmap["A"] == nested2

    def test_setitem_nested_2(self):
        nested1 = MultiDimensionalMapping([("B", 1)])
        ndmap = MultiDimensionalMapping([("A", nested1)])
        nested2 = MultiDimensionalMapping([("C", 2)])
        nested_clone = nested1.clone()
        nested_clone.update(nested2)
        ndmap.update({"A": nested2})
        assert ndmap["A"].data == nested_clone.data


class UniformNdMappingTest:
    def test_collapse_nested(self):
        inner1 = UniformNdMapping({1: hv.Dataset([(1, 2)], ["x", "y"])}, "Y")
        inner2 = UniformNdMapping({1: hv.Dataset([(3, 4)], ["x", "y"])}, "Y")
        outer = UniformNdMapping({1: inner1, 2: inner2}, "X")
        collapsed = outer.collapse()
        expected = hv.Dataset([(1, 1, 1, 2), (2, 1, 3, 4)], ["X", "Y", "x", "y"])
        assert_element_equal(collapsed, expected)


class HoloMapTest:
    def setup_method(self):
        self.xs = range(11)
        self.y_ints = [i * 2 for i in range(11)]
        self.ys = np.linspace(0, 1, 11)
        self.columns = hv.Dataset(
            np.column_stack([self.xs, self.y_ints]), kdims=["x"], vdims=["y"]
        )

    def test_holomap_redim(self):
        hmap = hv.HoloMap(
            {
                i: hv.Dataset({"x": self.xs, "y": self.ys * i}, kdims=["x"], vdims=["y"])
                for i in range(10)
            },
            kdims=["z"],
        )
        redimmed = hmap.redim(x="Time")
        assert redimmed.dimensions("all", True) == ["z", "Time", "y"]

    def test_holomap_redim_nested(self):
        hmap = hv.HoloMap(
            {
                i: hv.Dataset({"x": self.xs, "y": self.ys * i}, kdims=["x"], vdims=["y"])
                for i in range(10)
            },
            kdims=["z"],
        )
        redimmed = hmap.redim(x="Time", z="Magnitude")
        assert redimmed.dimensions("all", True) == ["Magnitude", "Time", "y"]

    def test_columns_collapse_heterogeneous(self):
        collapsed = hv.HoloMap(
            {
                i: hv.Dataset({"x": self.xs, "y": self.ys * i}, kdims=["x"], vdims=["y"])
                for i in range(10)
            },
            kdims=["z"],
        ).collapse("z", np.mean)
        expected = hv.Dataset({"x": self.xs, "y": self.ys * 4.5}, kdims=["x"], vdims=["y"])
        assert_element_equal(collapsed, expected)

    def test_columns_sample_homogeneous(self):
        samples = self.columns.sample([0, 5, 10]).dimension_values("y")
        assert_data_equal(samples, np.array([0, 10, 20]))

    def test_holomap_map_with_none(self):
        hmap = hv.HoloMap(
            {
                i: hv.Dataset({"x": self.xs, "y": self.ys * i}, kdims=["x"], vdims=["y"])
                for i in range(10)
            },
            kdims=["z"],
        )
        mapped = hmap.map(lambda x: x if x.range(1)[1] > 0 else None, hv.Dataset)
        assert_element_equal(hmap[1:10], mapped)

    def test_holomap_hist_two_dims(self):
        hmap = hv.HoloMap(
            {
                i: hv.Dataset({"x": self.xs, "y": self.ys * i}, kdims=["x"], vdims=["y"])
                for i in range(10)
            },
            kdims=["z"],
        )
        hists = hmap.hist(dimension=["x", "y"])
        assert hists["right"].last.kdims == ["y"]
        assert hists["top"].last.kdims == ["x"]

    def test_holomap_collapse_overlay_no_function(self):
        hmap = hv.HoloMap(
            {
                (1, 0): hv.Curve(np.arange(8)) * hv.Curve(-np.arange(8)),
                (2, 0): hv.Curve(np.arange(8) ** 2) * hv.Curve(-(np.arange(8) ** 3)),
            },
            kdims=["A", "B"],
        )
        assert_element_equal(
            hmap.collapse(),
            hv.Overlay(
                [
                    (
                        ("Curve", "I"),
                        hv.Dataset(
                            {
                                "A": np.concatenate([np.ones(8), np.ones(8) * 2]),
                                "B": np.zeros(16),
                                "x": np.tile(np.arange(8), 2),
                                "y": np.concatenate([np.arange(8), np.arange(8) ** 2]),
                            },
                            kdims=["A", "B", "x"],
                            vdims=["y"],
                        ),
                    ),
                    (
                        ("Curve", "II"),
                        hv.Dataset(
                            {
                                "A": np.concatenate([np.ones(8), np.ones(8) * 2]),
                                "B": np.zeros(16),
                                "x": np.tile(np.arange(8), 2),
                                "y": np.concatenate([-np.arange(8), -(np.arange(8) ** 3)]),
                            },
                            kdims=["A", "B", "x"],
                            vdims=["y"],
                        ),
                    ),
                ]
            ),
        )

    def test_holomap_collapse_overlay_max(self):
        hmap = hv.HoloMap(
            {
                (1, 0): hv.Curve(np.arange(8)) * hv.Curve(-np.arange(8)),
                (2, 0): hv.Curve(np.arange(8) ** 2) * hv.Curve(-(np.arange(8) ** 3)),
            },
            kdims=["A", "B"],
        )
        assert_element_equal(
            hmap.collapse(function=np.max),
            hv.Overlay(
                [
                    (
                        ("Curve", "I"),
                        hv.Curve(
                            {"x": np.arange(8), "y": np.arange(8) ** 2}, kdims=["x"], vdims=["y"]
                        ),
                    ),
                    (
                        ("Curve", "II"),
                        hv.Curve(
                            {"x": np.arange(8), "y": -np.arange(8)}, kdims=["x"], vdims=["y"]
                        ),
                    ),
                ]
            ),
        )
