import datetime as dt
from itertools import product

import numpy as np
import pandas as pd
import pytest

import holoviews as hv
from holoviews.testing import assert_element_equal


class DimensionedSelectionTest:
    def setup_method(self):
        self.img_fn = lambda: hv.Image(np.random.rand(10, 10))
        self.contour_fn = lambda: hv.Contours([np.random.rand(10, 2) for i in range(2)])
        self.datetime_fn = lambda: hv.Curve(
            (
                [dt.datetime(2000, 1, 1), dt.datetime(2000, 1, 2), dt.datetime(2000, 1, 3)],
                np.random.rand(3),
            ),
            "time",
            "x",
        )
        params = [list(range(3)) for i in range(2)]
        self.sanitized_map = hv.HoloMap(
            {i: hv.Image(i * np.random.rand(10, 10)) for i in range(1, 10)}, kdims=["A B"]
        )
        self.img_map = hv.HoloMap(
            {(i, j): self.img_fn() for i, j in product(*params)}, kdims=["a", "b"]
        )
        self.contour_map = hv.HoloMap(
            {(i, j): self.contour_fn() for i, j in product(*params)}, kdims=["a", "b"]
        )
        self.ndoverlay_map = self.img_map.overlay("b")
        self.overlay_map = self.img_map * self.contour_map
        self.layout_map = self.ndoverlay_map + self.contour_map
        self.duplicate_map = self.img_map.clone(kdims=["x", "y"])

        self.overlap1 = hv.HoloMap({i: self.img_fn() for i in range(5)})
        self.overlap2 = hv.HoloMap({i: self.img_fn() for i in range(10)})
        self.overlap_layout = self.overlap1 + self.overlap2

    def test_simple_holoselect(self):
        assert_element_equal(self.img_map.select(a=0, b=1), self.img_map[0, 1])

    def test_simple_holoslice(self):
        assert_element_equal(self.img_map.select(a=(1, 3), b=(1, 3)), self.img_map[1:3, 1:3])

    def test_function_holoslice(self):
        assert_element_equal(
            self.img_map.select(a=lambda x: 1 <= x < 3, b=lambda x: 1 <= x < 3),
            self.img_map[1:3, 1:3],
        )

    def test_sanitized_holoslice(self):
        assert_element_equal(self.sanitized_map.select(A_B=(1, 3)), self.sanitized_map[1:3])

    def test_simple_holo_ndoverlay_slice(self):
        assert_element_equal(
            self.ndoverlay_map.select(a=(1, 3), b=(1, 3)), self.ndoverlay_map[1:3, 1:3]
        )

    def test_deep_holoslice(self):
        selection = self.img_map.select(a=(1, 3), b=(1, 3), x=(None, 0), y=(None, 0))
        assert_element_equal(selection, self.img_map[1:3, 1:3, :0, :0])

    def test_deep_holooverlay_slice(self):
        map_slc = self.overlay_map[1:3, 1:3]
        img_slc = map_slc.map(lambda x: x[0:0.5, 0:0.5], [hv.Image, hv.Contours])
        selection = self.overlay_map.select(a=(1, 3), b=(1, 3), x=(0, 0.5), y=(0, 0.5))
        assert_element_equal(selection, img_slc)

    def test_deep_layout_nesting_slice(self):
        hmap1 = self.layout_map.HoloMap.I[1:3, 1:3, 0:0.5, 0:0.5]
        hmap2 = self.layout_map.HoloMap.II[1:3, 1:3, 0:0.5, 0:0.5]
        selection = self.layout_map.select(a=(1, 3), b=(1, 3), x=(0, 0.5), y=(0, 0.5))
        assert_element_equal(selection, hmap1 + hmap2)

    def test_spec_duplicate_dim_select(self):
        selection = self.duplicate_map.select(selection_specs=(hv.HoloMap,), x=(0, 1), y=(1, 3))
        assert_element_equal(selection, self.duplicate_map[0:1, 1:3])

    def test_duplicate_dim_select(self):
        selection = self.duplicate_map.select(x=(None, 0.5), y=(None, 0.5))
        assert_element_equal(selection, self.duplicate_map[:0.5, :0.5, :0.5, :0.5])

    def test_overlap_select(self):
        selection = self.overlap_layout.select(Default=(6, None))
        assert_element_equal(selection, self.overlap1.clone(shared_data=False) + self.overlap2[6:])

    def test_datetime_select(self):
        s, e = "1999-12-31", "2000-1-2"
        curve = self.datetime_fn()
        overlay = curve * self.datetime_fn()
        for el in [curve, overlay]:
            v = el.select(time=(s, e))
            assert_element_equal(v, el[s:e])
            assert_element_equal(
                el.select(time=(dt.datetime(1999, 12, 31), dt.datetime(2000, 1, 2))), el[s:e]
            )
            assert_element_equal(
                el.select(time=(pd.Timestamp(s), pd.Timestamp(e))),
                el[pd.Timestamp(s) : pd.Timestamp(e)],
            )

    def test_selection_spec_positional_error_message(self):
        s, e = "1999-12-31", "2000-1-2"
        curve = self.datetime_fn()
        with pytest.raises(ValueError, match="Use the selection_specs keyword"):
            curve.select((hv.Curve,), time=(s, e))
