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)

[((255, 9, 255, 34), 'pytest.raises', 'pytest.raises', ({(255, 23, 255, 33): 'IndexError'}, {}), '(IndexError)', False, 'import pytest\n'), ((69, 38, 69, 63), 'numpy.array', 'np.array', ({(69, 47, 69, 62): '[3.0, 5.0, 7.0]'}, {}), '([3.0, 5.0, 7.0])', True, 'import numpy as np\n'), ((91, 32, 91, 76), 'numpy.array', 'np.array', ({(91, 41, 91, 75): '[[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]]'}, {}), '([[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]])', True, 'import numpy as np\n'), ((92, 32, 92, 76), 'numpy.array', 'np.array', ({(92, 41, 92, 75): '[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]'}, {}), '([[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]])', True, 'import numpy as np\n'), ((93, 43, 93, 77), 'numpy.array', 'np.array', ({(93, 52, 93, 76): '[0, 1, 4, 3, 1, 2, 5, 4]'}, {}), '([0, 1, 4, 3, 1, 2, 5, 4])', True, 'import numpy as np\n'), ((94, 37, 94, 53), 'numpy.array', 'np.array', ({(94, 46, 94, 52): '[4, 8]'}, {}), '([4, 8])', True, 'import numpy as np\n'), ((115, 8, 117, 9), 'numpy.array', 'np.array', ({(116, 12, 116, 84): '[[[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]]])', True, 'import numpy as np\n'), ((120, 8, 122, 9), 'numpy.array', 'np.array', ({(121, 12, 121, 84): '[[[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]]])', True, 'import numpy as np\n'), ((125, 8, 127, 9), 'numpy.array', 'np.array', ({(126, 12, 126, 84): '[[[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]]])', True, 'import numpy as np\n'), ((149, 32, 149, 76), 'numpy.array', 'np.array', ({(149, 41, 149, 75): '[[0.0, 3.0, 4.0], [0.0, 3.0, 4.0]]'}, {}), '([[0.0, 3.0, 4.0], [0.0, 3.0, 4.0]])', True, 'import numpy as np\n'), ((150, 32, 150, 76), 'numpy.array', 'np.array', ({(150, 41, 150, 75): '[[2.0, 2.0, 2.0], [7.0, 7.0, 7.0]]'}, {}), '([[2.0, 2.0, 2.0], [7.0, 7.0, 7.0]])', True, 'import numpy as np\n'), ((159, 19, 159, 59), 'numpy.array', 'np.array', ({(159, 28, 159, 58): '[0.0, 1.0, 2.0, 0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 2.0, 0.0, 1.0, 2.0])', True, 'import numpy as np\n'), ((162, 19, 162, 59), 'numpy.array', 'np.array', ({(162, 28, 162, 58): '[0.0, 1.0, 2.0, 1.0, 2.0, 3.0]'}, {}), '([0.0, 1.0, 2.0, 1.0, 2.0, 3.0])', True, 'import numpy as np\n'), ((171, 32, 171, 72), 'numpy.array', 'np.array', ({(171, 41, 171, 71): '[0.0, 1.0, 2.0, 0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 2.0, 0.0, 1.0, 2.0])', True, 'import numpy as np\n'), ((172, 32, 172, 72), 'numpy.array', 'np.array', ({(172, 41, 172, 71): '[0.0, 1.0, 2.0, 1.0, 2.0, 3.0]'}, {}), '([0.0, 1.0, 2.0, 1.0, 2.0, 3.0])', True, 'import numpy as np\n'), ((178, 19, 178, 54), 'numpy.array', 'np.array', ({(178, 28, 178, 53): '[0.0, 1.0, 0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 0.0, 1.0, 2.0])', True, 'import numpy as np\n'), ((181, 19, 181, 54), 'numpy.array', 'np.array', ({(181, 28, 181, 53): '[0.0, 0.0, 1.0, 1.0, 0.0]'}, {}), '([0.0, 0.0, 1.0, 1.0, 0.0])', True, 'import numpy as np\n'), ((184, 19, 184, 50), 'numpy.array', 'np.array', ({(184, 28, 184, 49): '[0, 1, 3, 2, 4, 3, 1]'}, {}), '([0, 1, 3, 2, 4, 3, 1])', True, 'import numpy as np\n'), ((187, 19, 187, 35), 'numpy.array', 'np.array', ({(187, 28, 187, 34): '[4, 7]'}, {}), '([4, 7])', True, 'import numpy as np\n'), ((196, 32, 196, 67), 'numpy.array', 'np.array', ({(196, 41, 196, 66): '[0.0, 1.0, 0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 0.0, 1.0, 2.0])', True, 'import numpy as np\n'), ((197, 32, 197, 67), 'numpy.array', 'np.array', ({(197, 41, 197, 66): '[0.0, 0.0, 1.0, 1.0, 0.0]'}, {}), '([0.0, 0.0, 1.0, 1.0, 0.0])', True, 'import numpy as np\n'), ((206, 19, 206, 44), 'numpy.array', 'np.array', ({(206, 28, 206, 43): '[0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 2.0])', True, 'import numpy as np\n'), ((223, 19, 223, 44), 'numpy.array', 'np.array', ({(223, 28, 223, 43): '[0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 2.0])', True, 'import numpy as np\n'), ((240, 19, 240, 44), 'numpy.array', 'np.array', ({(240, 28, 240, 43): '[0.0, 1.0, 2.0]'}, {}), '([0.0, 1.0, 2.0])', True, 'import numpy as np\n')]