"""
Unit test of the (non-rendering) exporters and importers.
"""

import numpy as np

import holoviews as hv
from holoviews.core.io import Deserializer, Pickler, Serializer, Unpickler
from holoviews.testing import assert_element_equal


class TestSerialization:
    """
    Test the basic serializer and deserializer (i.e. using pickle),
    including metadata access.
    """

    def setup_method(self):
        self.image1 = hv.Image(np.array([[1, 2], [4, 5]]))
        self.image2 = hv.Image(np.array([[5, 4], [3, 2]]))

    def test_serializer_save(self, tmp_path):
        Serializer.save(
            self.image1,
            tmp_path / "test_serializer_save.pkl",
            info={"info": "example"},
            key={1: 2},
        )

    def test_serializer_save_no_file_extension(self, tmp_path):
        Serializer.save(
            self.image1,
            tmp_path / "test_serializer_save_no_ext",
            info={"info": "example"},
            key={1: 2},
        )
        assert tmp_path / "test_serializer_save_no_ext.pkl" in tmp_path.iterdir()

    def test_serializer_save_and_load_data_1(self, tmp_path):
        Serializer.save(self.image1, tmp_path / "test_serializer_save_and_load_data_1.pkl")
        loaded = Deserializer.load(tmp_path / "test_serializer_save_and_load_data_1.pkl")
        assert_element_equal(loaded, self.image1)

    def test_serializer_save_and_load_data_2(self, tmp_path):
        Serializer.save(self.image2, tmp_path / "test_serializer_save_and_load_data_2.pkl")
        loaded = Deserializer.load(tmp_path / "test_serializer_save_and_load_data_2.pkl")
        assert_element_equal(loaded, self.image2)

    def test_serializer_save_and_load_key(self, tmp_path):
        input_key = {"test_key": "key_val"}
        Serializer.save(
            self.image1, tmp_path / "test_serializer_save_and_load_data.pkl", key=input_key
        )
        key = Deserializer.key(tmp_path / "test_serializer_save_and_load_data.pkl")
        assert key == input_key

    def test_serializer_save_and_load_info(self, tmp_path):
        input_info = {"info": "example"}
        Serializer.save(
            self.image1, tmp_path / "test_serializer_save_and_load_data.pkl", info=input_info
        )
        info = Deserializer.info(tmp_path / "test_serializer_save_and_load_data.pkl")
        assert info["info"] == input_info["info"]

    def test_serialize_deserialize_1(self, tmp_path):
        data, _ = Serializer(self.image1)
        obj = Deserializer(data)
        assert_element_equal(obj, self.image1)

    def test_serialize_deserialize_2(self, tmp_path):
        data, _ = Serializer(self.image2)
        obj = Deserializer(data)
        assert_element_equal(obj, self.image2)


class TestBasicPickler:
    """
    Test pickler and unpickler using the .hvz format, including
    metadata access.
    """

    def setup_method(self):
        self.image1 = hv.Image(np.array([[1, 2], [4, 5]]))
        self.image2 = hv.Image(np.array([[5, 4], [3, 2]]))

    def test_pickler_save(self, tmp_path) -> None:
        Pickler.save(
            self.image1, tmp_path / "test_pickler_save.hvz", info={"info": "example"}, key={1: 2}
        )

    def test_pickler_save_no_file_extension(self, tmp_path):
        Pickler.save(
            self.image1,
            tmp_path / "test_pickler_save_no_ext",
            info={"info": "example"},
            key={1: 2},
        )
        assert tmp_path / "test_pickler_save_no_ext.hvz" in tmp_path.iterdir()

    def test_pickler_save_and_load_data_1(self, tmp_path):
        Pickler.save(self.image1, tmp_path / "test_pickler_save_and_load_data_1.hvz")
        loaded = Unpickler.load(tmp_path / "test_pickler_save_and_load_data_1.hvz")
        assert_element_equal(loaded, self.image1)

    def test_pickler_save_and_load_data_2(self, tmp_path):
        Pickler.save(self.image2, tmp_path / "test_pickler_save_and_load_data_2.hvz")
        loaded = Unpickler.load(tmp_path / "test_pickler_save_and_load_data_2.hvz")
        assert_element_equal(loaded, self.image2)

    def test_pickler_save_and_load_key(self, tmp_path):
        input_key = {"test_key": "key_val"}
        Pickler.save(self.image1, tmp_path / "test_pickler_save_and_load_data.hvz", key=input_key)
        key = Unpickler.key(tmp_path / "test_pickler_save_and_load_data.hvz")
        assert key == input_key

    def test_pickler_save_and_load_info(self, tmp_path):
        input_info = {"info": "example"}
        Pickler.save(
            self.image1, tmp_path / "test_pickler_save_and_load_data.hvz", info=input_info
        )
        info = Unpickler.info(tmp_path / "test_pickler_save_and_load_data.hvz")
        assert info["info"] == input_info["info"]

    def test_serialize_deserialize_1(self):
        data, _ = Pickler(self.image1)
        obj = Unpickler(data)
        assert_element_equal(obj, self.image1)

    def test_serialize_deserialize_2(self):
        data, _ = Pickler(self.image2)
        obj = Unpickler(data)
        assert_element_equal(obj, self.image2)


class TestPicklerAdvanced:
    """
    Test advanced pickler and unpickler functionality supported by the
    .hvz format.
    """

    def setup_method(self):
        self.image1 = hv.Image(np.array([[1, 2], [4, 5]]))
        self.image2 = hv.Image(np.array([[5, 4], [3, 2]]))

    def test_pickler_save_layout(self, tmp_path):
        Pickler.save(
            self.image1 + self.image2,
            tmp_path / "test_pickler_save_layout",
            info={"info": "example"},
            key={1: 2},
        )

    def test_pickler_save_load_layout(self, tmp_path):
        Pickler.save(
            self.image1 + self.image2,
            tmp_path / "test_pickler_save_load_layout",
            info={"info": "example"},
            key={1: 2},
        )
        loaded = Unpickler.load(tmp_path / "test_pickler_save_load_layout.hvz")
        assert_element_equal(loaded, self.image1 + self.image2)

    def test_pickler_save_load_layout_entries(self, tmp_path):
        Pickler.save(
            self.image1 + self.image2,
            tmp_path / "test_pickler_save_load_layout_entries",
            info={"info": "example"},
            key={1: 2},
        )
        entries = Unpickler.entries(tmp_path / "test_pickler_save_load_layout_entries.hvz")
        assert entries == ["Image.I", "Image.II"]

    def test_pickler_save_load_layout_entry1(self, tmp_path):
        Pickler.save(
            self.image1 + self.image2,
            tmp_path / "test_pickler_save_load_layout_entry1",
            info={"info": "example"},
            key={1: 2},
        )
        entries = Unpickler.entries(tmp_path / "test_pickler_save_load_layout_entry1.hvz")
        assert "Image.I" in entries, "Entry 'Image.I' missing"
        loaded = Unpickler.load(
            tmp_path / "test_pickler_save_load_layout_entry1.hvz", entries=["Image.I"]
        )
        assert_element_equal(loaded, self.image1)

    def test_pickler_save_load_layout_entry2(self, tmp_path):
        Pickler.save(
            self.image1 + self.image2,
            tmp_path / "test_pickler_save_load_layout_entry2",
            info={"info": "example"},
            key={1: 2},
        )
        entries = Unpickler.entries(tmp_path / "test_pickler_save_load_layout_entry2.hvz")
        assert "Image.II" in entries, "Entry 'Image.II' missing"
        loaded = Unpickler.load(
            tmp_path / "test_pickler_save_load_layout_entry2.hvz", entries=["Image.II"]
        )
        assert_element_equal(loaded, self.image2)

    def test_pickler_save_load_single_layout(self, tmp_path):
        single_layout = hv.Layout([self.image1])
        Pickler.save(
            single_layout,
            tmp_path / "test_pickler_save_load_single_layout",
            info={"info": "example"},
            key={1: 2},
        )

        entries = Unpickler.entries(tmp_path / "test_pickler_save_load_single_layout.hvz")
        assert entries == ["Image.I(L)"]
        loaded = Unpickler.load(
            tmp_path / "test_pickler_save_load_single_layout.hvz", entries=["Image.I(L)"]
        )
        assert_element_equal(single_layout, loaded)
