"""
Tests the transformers in :mod:`pyrate.common.raster_datasets.transformer_base` and in
:mod:`pyrate.common.raster_datasets.transformers_concrete`.
"""

# Testing
from unittest import TestCase

# Scientific
from numpy import array
from numpy import empty
from numpy import float32
from numpy import int16
from numpy.testing import assert_array_equal
from numpy import uint16
from numpy import uint32
from pandas import DataFrame
from pandas import Series
from pandas.testing import assert_frame_equal
from pandas.testing import assert_series_equal

# Module under test
from pyrate.common.raster_datasets.transformers_concrete import BathymetricTransformer
from pyrate.common.raster_datasets.transformers_concrete import ConstantTransformer

# Graph generation
from pyrate.plan.graph import create_earth_graph
from pyrate.plan.graph import GeoNavigationGraph
from pyrate.plan.graph import min_required_frequency

# CI/Testing helpers
from ... import _open_test_geo_dataset


class TestGetNodePropertiesWithConstantTransformer(TestCase):
    """Ensure that the :meth:`pyrate.plan.graph.GeoNavigationGraph.append_properties` works correctly."""

    def test_get_node_properties_empty_coordinates(self) -> None:
        """Tests getting properties for a graph without nodes."""
        graph = GeoNavigationGraph.from_coordinates_degrees(
            latitudes=empty((0,)), longitudes=empty((0,)), edges=empty((0, 2)), node_radius=111.111
        )
        transformers = [ConstantTransformer(42, uint32, "prop_1"), ConstantTransformer(43, uint16, "prop_2")]
        graph.append_properties(transformers)
        self.assertEqual(len(graph.node_properties), 0)
        assert_array_equal(graph.node_properties.columns, ["prop_1", "prop_2"])

    def test_get_node_properties_no_transformers(self) -> None:
        """Tests getting properties without a transformer."""
        graph = GeoNavigationGraph.from_coordinates_degrees(
            latitudes=array([0, 1]), longitudes=array([0, 0]), edges=array([[0, 1]]), node_radius=111.111
        )
        graph.append_properties([])  # empty!
        self.assertEqual(len(graph.node_properties), 2)
        assert_array_equal(graph.node_properties.columns, [])

    def test_get_node_properties_single_transformer(self) -> None:
        """Tests getting properties using only a single transformer."""
        graph = GeoNavigationGraph.from_coordinates_degrees(
            latitudes=array([0, 1]),
            longitudes=array([0, 0]),
            edges=array([[0, 1]]),
            node_radius=0.0,  # some weird radius
        )
        # now we use `append_property` to append a single one
        graph.append_property(ConstantTransformer(33, uint32, "prop_1"))
        self.assertEqual(len(graph.node_properties), 2)
        assert_frame_equal(graph.node_properties, DataFrame(data={"prop_1": [33, 33]}, dtype=uint32))

    def test_get_node_properties_single_transformer_str_datatype(self) -> None:
        """Tests getting properties using only a single transformer and a string datatype."""
        graph = GeoNavigationGraph.from_coordinates_degrees(
            latitudes=array([0]),
            longitudes=array([0]),
            edges=array([[0, 0]]),  # edge to itself
            node_radius=111.111,
        )
        # now we use `append_property` to append a single one
        data_type = "U10"  # must give string data type explicitly and not with np.str or "U"
        graph.append_property(ConstantTransformer("content", data_type, "prop_1"))
        self.assertEqual(len(graph.node_properties), 1)
        assert_frame_equal(graph.node_properties, DataFrame(data={"prop_1": ["content"]}, dtype=data_type))

    def test_get_node_properties_multiple_transformers(self) -> None:
        """Tests getting properties using multiple transformers."""
        graph = GeoNavigationGraph.from_coordinates_degrees(
            latitudes=array([0, 1]), longitudes=array([0, 0]), edges=array([[0, 1]]), node_radius=111.111
        )
        # now we use `append_property` to append a single one
        graph.append_properties(
            [ConstantTransformer(33, uint32, "prop_1"), ConstantTransformer(99, int16, "prop_2")]
        )
        self.assertEqual(len(graph.node_properties), 2)
        assert_array_equal(graph.node_properties.columns, ["prop_1", "prop_2"])
        assert_series_equal(
            graph.node_properties["prop_1"], Series(data=[33, 33], dtype=uint32, name="prop_1")
        )
        assert_series_equal(
            graph.node_properties["prop_2"], Series(data=[99, 99], dtype=int16, name="prop_2")
        )


class TestBathymetricTransformer(TestCase):
    """Tests :class:`pyrate.common.raster_datasets.transformers_concrete.BathymetricTransformer`."""

    def test_all_modes(self) -> None:
        """Tests all modes at once."""

        # create a coarse grid
        distance_meters = 1000_000
        graph = create_earth_graph(min_required_frequency(distance_meters, in_meters=True))

        # fetch properties
        modes = list(BathymetricTransformer.Modes)
        graph.append_property(BathymetricTransformer(_open_test_geo_dataset(), modes))
        properties = graph.node_properties

        # check that the returned properties are all floats
        self.assertTrue((properties.dtypes == float32).all())

    def test_no_data(self) -> None:
        """Tests that querying for data where there are no data points in the result range raises an error."""
        for mode in list(BathymetricTransformer.Modes):
            with self.subTest(mode.name), self.assertRaises(ValueError):
                with BathymetricTransformer(_open_test_geo_dataset(), [mode]) as transformer:
                    # This works by querying for a point (at 1e-3°N 1e-3°E), where there is no data point
                    # within 1e-9 meters in the underlying dataset
                    # This should trigger an exception (e.g. because the average depth over zero data
                    # points is not clearly)
                    transformer.get_transformed_at_nodes(
                        latitudes=array([1e-3]), longitudes=array([1e-3]), radius=1e-9
                    )