import numpy as np
import pytest
from pytest import approx

from pymt.component.grid import GridMixIn


class Port:
    def __init__(self, name, uses=None, provides=None):
        self._name = name
        self._uses = uses or []
        self._provides = provides or []

    def get_component_name(self):
        return self._name

    def get_input_item_count(self):
        return len(self._uses)

    def get_input_item_list(self):
        return self._uses

    def get_output_item_count(self):
        return len(self._provides)

    def get_output_item_list(self):
        return self._provides


def test_exchange_items():
    class Component(GridMixIn):
        def __init__(self):
            self._port = Port("test", uses=["invar"], provides=["outvar"])
            super().__init__()

    c = Component()
    assert c.input_items == ["invar"]
    assert c.output_items == ["outvar"]


def test_no_exchange_items():
    class Component(GridMixIn):
        def __init__(self):
            self._port = Port("test")
            super().__init__()

    c = Component()
    assert c.input_items == []
    assert c.output_items == []


def test_raster_1d():
    class RasterPort(Port):
        def get_grid_shape(self, grid_id):
            return (3,)

        def get_grid_spacing(self, grid_id):
            return (2.0,)

        def get_grid_origin(self, grid_id):
            return (3.0,)

    class Component(GridMixIn):
        def __init__(self):
            self._port = RasterPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_x("invar") == approx(np.array([3.0, 5.0, 7.0]))


def test_raster_2d():
    class RasterPort(Port):
        def get_grid_shape(self, grid_id):
            return (2, 3)

        def get_grid_spacing(self, grid_id):
            return (2.0, 1.0)

        def get_grid_origin(self, grid_id):
            return (0.0, 0.0)

    class Component(GridMixIn):
        def __init__(self):
            self._port = RasterPort("test-2d", uses=["invar"], provides=["outvar"])
            super().__init__()

    c = Component()
    assert c.name == "test-2d"
    assert c.get_grid_type(0) == "RASTER"
    assert c.get_x(0) == approx(np.array([[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]]))
    assert c.get_y(0) == approx(np.array([[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]))
    assert np.all(c.get_connectivity(0) == np.array([0, 1, 4, 3, 1, 2, 5, 4]))
    assert np.all(c.get_offset(0) == np.array([4, 8]))


def test_raster_3d():
    class RasterPort(Port):
        def get_grid_shape(self, grid_id):
            return (2, 2, 3)

        def get_grid_spacing(self, grid_id):
            return (1.0, 2.0, 1.0)

        def get_grid_origin(self, grid_id):
            return (0.0, 0.0, 0.0)

    class Component(GridMixIn):
        def __init__(self):
            self._port = RasterPort("test-3d", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_x(0) == approx(
        np.array(
            [[[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]], [[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]]]
        )
    )
    assert c.get_y(0) == approx(
        np.array(
            [[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]], [[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]]
        )
    )
    assert c.get_z(0) == approx(
        np.array(
            [[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]]
        )
    )


def test_rectilinear():
    class RectilinearPort(Port):
        def get_grid_shape(self, grid_id):
            return (2, 3)

        def get_grid_x(self, grid_id):
            return (0.0, 3.0, 4)

        def get_grid_y(self, grid_id):
            return (2.0, 7.0)

    class Component(GridMixIn):
        def __init__(self):
            self._port = RectilinearPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_grid_type(0) == "RECTILINEAR"
    assert c.get_x(0) == approx(np.array([[0.0, 3.0, 4.0], [0.0, 3.0, 4.0]]))
    assert c.get_y(0) == approx(np.array([[2.0, 2.0, 2.0], [7.0, 7.0, 7.0]]))


def test_structured():
    class StructuredPort(Port):
        def get_grid_shape(self, grid_id):
            return (2, 3)

        def get_grid_x(self, grid_id):
            return np.array([0.0, 1.0, 2.0, 0.0, 1.0, 2.0])

        def get_grid_y(self, grid_id):
            return np.array([0.0, 1.0, 2.0, 1.0, 2.0, 3.0])

    class Component(GridMixIn):
        def __init__(self):
            self._port = StructuredPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_grid_type(0) == "STRUCTURED"
    assert c.get_x(0) == approx(np.array([0.0, 1.0, 2.0, 0.0, 1.0, 2.0]))
    assert c.get_y(0) == approx(np.array([0.0, 1.0, 2.0, 1.0, 2.0, 3.0]))


def test_unstructured():
    class UnstructuredPort(Port):
        def get_grid_x(self, grid_id):
            return np.array([0.0, 1.0, 0.0, 1.0, 2.0])

        def get_grid_y(self, grid_id):
            return np.array([0.0, 0.0, 1.0, 1.0, 0.0])

        def get_grid_connectivity(self, grid_id):
            return np.array([0, 1, 3, 2, 4, 3, 1])

        def get_grid_offset(self, grid_id):
            return np.array([4, 7])

    class Component(GridMixIn):
        def __init__(self):
            self._port = UnstructuredPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_grid_type(0) == "UNSTRUCTURED"
    assert c.get_x(0) == approx(np.array([0.0, 1.0, 0.0, 1.0, 2.0]))
    assert c.get_y(0) == approx(np.array([0.0, 0.0, 1.0, 1.0, 0.0]))


def test_get_grid_shape_is_none():
    class UnstructuredPort(Port):
        def get_grid_shape(self, grid_id):
            return None

        def get_grid_x(self, grid_id):
            return np.array([0.0, 1.0, 2.0])

    class Component(GridMixIn):
        def __init__(self):
            self._port = UnstructuredPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_grid_type(0) == "UNSTRUCTURED"


def test_get_grid_shape_raises():
    class UnstructuredPort(Port):
        def get_grid_shape(self, grid_id):
            raise NotImplementedError("get_grid_shape")

        def get_grid_x(self, grid_id):
            return np.array([0.0, 1.0, 2.0])

    class Component(GridMixIn):
        def __init__(self):
            self._port = UnstructuredPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_grid_type(0) == "UNSTRUCTURED"


def test_structured_1d():
    class RectilinearPort(Port):
        def get_grid_shape(self, grid_id):
            return (2, 3)

        def get_grid_x(self, grid_id):
            return np.array([0.0, 1.0, 2.0])

        def get_grid_y(self, grid_id):
            raise NotImplementedError("get_grid_y")

        def get_grid_z(self, grid_id):
            raise NotImplementedError("get_grid_z")

    class Component(GridMixIn):
        def __init__(self):
            self._port = RectilinearPort("test", uses=["invar"])
            super().__init__()

    c = Component()
    assert c.get_grid_type(0) == "RECTILINEAR"
    with pytest.raises(IndexError):
        c.get_z(0)

[((7100, 7125), 'pytest.raises', 'pytest.raises', (['IndexError'], {}), '(IndexError)\n', (7113, 7125), False, 'import pytest\n'), ((1613, 1638), 'numpy.array', 'np.array', (['[3.0, 5.0, 7.0]'], {}), '([3.0, 5.0, 7.0])\n', (1621, 1638), True, 'import numpy as np\n'), ((2214, 2258), 'numpy.array', 'np.array', (['[[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]]'], {}), '([[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]])\n', (2222, 2258), True, 'import numpy as np\n'), ((2292, 2336), 'numpy.array', 'np.array', (['[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]'], {}), '([[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]])\n', (2300, 2336), True, 'import numpy as np\n'), ((2381, 2415), 'numpy.array', 'np.array', (['[0, 1, 4, 3, 1, 2, 5, 4]'], {}), '([0, 1, 4, 3, 1, 2, 5, 4])\n', (2389, 2415), True, 'import numpy as np\n'), ((2454, 2470), 'numpy.array', 'np.array', (['[4, 8]'], {}), '([4, 8])\n', (2462, 2470), True, 'import numpy as np\n'), ((2974, 3061), 'numpy.array', 'np.array', (['[[[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]], [[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]]]'], {}), '([[[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]], [[0.0, 1.0, 2.0], [0.0, 1.0, \n    2.0]]])\n', (2982, 3061), True, 'import numpy as np\n'), ((3126, 3213), 'numpy.array', 'np.array', (['[[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]], [[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]]'], {}), '([[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]], [[0.0, 0.0, 0.0], [2.0, 2.0, \n    2.0]]])\n', (3134, 3213), True, 'import numpy as np\n'), ((3278, 3365), 'numpy.array', 'np.array', (['[[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]]'], {}), '([[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]], [[1.0, 1.0, 1.0], [1.0, 1.0, \n    1.0]]])\n', (3286, 3365), True, 'import numpy as np\n'), ((3917, 3961), 'numpy.array', 'np.array', (['[[0.0, 3.0, 4.0], [0.0, 3.0, 4.0]]'], {}), '([[0.0, 3.0, 4.0], [0.0, 3.0, 4.0]])\n', (3925, 3961), True, 'import numpy as np\n'), ((3995, 4039), 'numpy.array', 'np.array', (['[[2.0, 2.0, 2.0], [7.0, 7.0, 7.0]]'], {}), '([[2.0, 2.0, 2.0], [7.0, 7.0, 7.0]])\n', (4003, 4039), True, 'import numpy as np\n'), ((4226, 4266), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 2.0, 0.0, 1.0, 2.0])\n', (4234, 4266), True, 'import numpy as np\n'), ((4326, 4366), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 1.0, 2.0, 3.0]'], {}), '([0.0, 1.0, 2.0, 1.0, 2.0, 3.0])\n', (4334, 4366), True, 'import numpy as np\n'), ((4622, 4662), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 2.0, 0.0, 1.0, 2.0])\n', (4630, 4662), True, 'import numpy as np\n'), ((4696, 4736), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 1.0, 2.0, 3.0]'], {}), '([0.0, 1.0, 2.0, 1.0, 2.0, 3.0])\n', (4704, 4736), True, 'import numpy as np\n'), ((4857, 4892), 'numpy.array', 'np.array', (['[0.0, 1.0, 0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 0.0, 1.0, 2.0])\n', (4865, 4892), True, 'import numpy as np\n'), ((4952, 4987), 'numpy.array', 'np.array', (['[0.0, 0.0, 1.0, 1.0, 0.0]'], {}), '([0.0, 0.0, 1.0, 1.0, 0.0])\n', (4960, 4987), True, 'import numpy as np\n'), ((5058, 5089), 'numpy.array', 'np.array', (['[0, 1, 3, 2, 4, 3, 1]'], {}), '([0, 1, 3, 2, 4, 3, 1])\n', (5066, 5089), True, 'import numpy as np\n'), ((5154, 5170), 'numpy.array', 'np.array', (['[4, 7]'], {}), '([4, 7])\n', (5162, 5170), True, 'import numpy as np\n'), ((5430, 5465), 'numpy.array', 'np.array', (['[0.0, 1.0, 0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 0.0, 1.0, 2.0])\n', (5438, 5465), True, 'import numpy as np\n'), ((5499, 5534), 'numpy.array', 'np.array', (['[0.0, 0.0, 1.0, 1.0, 0.0]'], {}), '([0.0, 0.0, 1.0, 1.0, 0.0])\n', (5507, 5534), True, 'import numpy as np\n'), ((5733, 5758), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 2.0])\n', (5741, 5758), True, 'import numpy as np\n'), ((6214, 6239), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 2.0])\n', (6222, 6239), True, 'import numpy as np\n'), ((6656, 6681), 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0]'], {}), '([0.0, 1.0, 2.0])\n', (6664, 6681), True, 'import numpy as np\n')]