import pytest
import numpy as np
import rna.plotting

try:
    import tfields
except ImportError:
    tfields = None


def test_add_subplot():
    # Test 2D axes
    ax_2 = rna.plotting.add_subplot(2)
    assert rna.plotting.axes_dim(ax_2) == 2
    ax_2_implicit = rna.plotting.add_subplot()
    assert rna.plotting.axes_dim(ax_2_implicit) == 2

    # Test 3D axes
    ax_3 = rna.plotting.add_subplot(3)
    assert rna.plotting.axes_dim(ax_3) == 3

    # Test invalid dimension
    with pytest.raises(NotImplementedError):
        rna.plotting.add_subplot(4)


@rna.plotting.plot_signature
def my_plot_function(axes, *data, **kwargs):
    return axes, data, kwargs


class Foo:
    @rna.plotting.plot_signature
    def my_plot_method(self, axes, *data, **kwargs):
        return axes, data, kwargs


@pytest.fixture(params=[Foo().my_plot_method, my_plot_function])
def plotter(request):
    return request.param


def test_plot_signature_ax_equals(plotter):
    axes = rna.plotting.add_subplot(2)
    assert rna.plotting.axes_dim(axes) == 2

    data = (1, 2)
    kwargs = {"a": 1}
    ret_ax, ret_data, ret_kwargs = plotter(*data, ax=axes, **kwargs)
    assert ret_ax is axes
    assert ret_data == data
    assert ret_kwargs == kwargs


def test_plot_signature_ax_first(plotter):
    axes = rna.plotting.add_subplot(2)
    assert rna.plotting.axes_dim(axes) == 2

    data = (1, 2)
    kwargs = {"a": 1}
    ret_ax, ret_data, ret_kwargs = plotter(axes, *data, **kwargs)
    assert ret_ax is axes
    assert ret_data == data
    assert ret_kwargs == kwargs


@pytest.mark.skipif(tfields is None, reason="Requires tfields which is not installed.")
def test_plot_mesh_2d_vector_rotated(request):
    """
    Test, if rotation argument works in plot_mesh
    """
    # build a cube Cube with edge length of 2 units
    vertices = np.array(
        [
            [-1.0, -1.0, -5.0],
            [-1.0, -1.0, -3.0],
            [-1.0, 1.0, -5.0],
            [-1.0, 1.0, -3.0],
            [-1.0, -1.0, -5.0],
            [-1.0, -1.0, -3.0],
            [1.0, -1.0, -5.0],
            [1.0, -1.0, -3.0],
            [-1.0, -1.0, -5.0],
            [-1.0, 1.0, -5.0],
            [1.0, -1.0, -5.0],
            [1.0, 1.0, -5.0],
            [1.0, -1.0, -5.0],
            [1.0, -1.0, -3.0],
            [1.0, 1.0, -5.0],
            [1.0, 1.0, -3.0],
            [-1.0, 1.0, -5.0],
            [-1.0, 1.0, -3.0],
            [1.0, 1.0, -5.0],
            [1.0, 1.0, -3.0],
            [-1.0, -1.0, -3.0],
            [-1.0, 1.0, -3.0],
            [1.0, -1.0, -3.0],
            [1.0, 1.0, -3.0],
        ]
    )
    faces = np.array(
        [
            [0, 1, 2],
            [1, 2, 3],
            [4, 5, 6],
            [5, 6, 7],
            [8, 9, 10],
            [9, 10, 11],
            [12, 13, 14],
            [13, 14, 15],
            [16, 17, 18],
            [17, 18, 19],
            [20, 21, 22],
            [21, 22, 23],
        ]
    )
    rotation = np.array(
        [
            [0.85, -0.15, 0.5],
            [0.5, 0.5, -0.71],
            [-0.15, 0.85, 0.5],
        ]
    )
    rna.plotting.plot_mesh(vertices, faces, dim=2, edgecolor="k", rotation=rotation)
    rna.plotting.show()
