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, 'pytest.raises', 'pytest.raises', (['IndexError'], {}), False, 'import pytest\n'), (69, 'numpy.array', 'np.array', (['[3.0, 5.0, 7.0]'], {}), True, 'import numpy as np\n'), (91, 'numpy.array', 'np.array', (['[[0.0, 1.0, 2.0], [0.0, 1.0, 2.0]]'], {}), True, 'import numpy as np\n'), (92, 'numpy.array', 'np.array', (['[[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]]'], {}), True, 'import numpy as np\n'), (93, 'numpy.array', 'np.array', (['[0, 1, 4, 3, 1, 2, 5, 4]'], {}), True, 'import numpy as np\n'), (94, 'numpy.array', 'np.array', (['[4, 8]'], {}), True, 'import numpy as np\n'), (115, '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]]]'], {}), True, 'import numpy as np\n'), (120, '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]]]'], {}), True, 'import numpy as np\n'), (125, '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]]]'], {}), True, 'import numpy as np\n'), (149, 'numpy.array', 'np.array', (['[[0.0, 3.0, 4.0], [0.0, 3.0, 4.0]]'], {}), True, 'import numpy as np\n'), (150, 'numpy.array', 'np.array', (['[[2.0, 2.0, 2.0], [7.0, 7.0, 7.0]]'], {}), True, 'import numpy as np\n'), (159, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n'), (162, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 1.0, 2.0, 3.0]'], {}), True, 'import numpy as np\n'), (171, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n'), (172, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0, 1.0, 2.0, 3.0]'], {}), True, 'import numpy as np\n'), (178, 'numpy.array', 'np.array', (['[0.0, 1.0, 0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n'), (181, 'numpy.array', 'np.array', (['[0.0, 0.0, 1.0, 1.0, 0.0]'], {}), True, 'import numpy as np\n'), (184, 'numpy.array', 'np.array', (['[0, 1, 3, 2, 4, 3, 1]'], {}), True, 'import numpy as np\n'), (187, 'numpy.array', 'np.array', (['[4, 7]'], {}), True, 'import numpy as np\n'), (196, 'numpy.array', 'np.array', (['[0.0, 1.0, 0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n'), (197, 'numpy.array', 'np.array', (['[0.0, 0.0, 1.0, 1.0, 0.0]'], {}), True, 'import numpy as np\n'), (206, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n'), (223, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n'), (240, 'numpy.array', 'np.array', (['[0.0, 1.0, 2.0]'], {}), True, 'import numpy as np\n')]