import numpy as np
import pytest
from numpy.testing import assert_array_almost_equal

from findiff import Gradient, Divergence, Curl, Laplacian


class TestGradient:

    def test_3d_gradient_on_scalar_func(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.sin(X) * np.sin(Y) * np.sin(Z)
        grad_f_ex = np.array(
            [
                np.cos(X) * np.sin(Y) * np.sin(Z),
                np.sin(X) * np.cos(Y) * np.sin(Z),
                np.sin(X) * np.sin(Y) * np.cos(Z),
            ]
        )
        grad = Gradient(spac=h, acc=4)
        grad_f = grad(f)
        assert_array_almost_equal(grad_f, grad_f_ex)

    def test_spacing_with_h(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.sin(X) * np.sin(Y) * np.sin(Z)
        grad_f_ex = np.array(
            [
                np.cos(X) * np.sin(Y) * np.sin(Z),
                np.sin(X) * np.cos(Y) * np.sin(Z),
                np.sin(X) * np.sin(Y) * np.cos(Z),
            ]
        )
        grad = Gradient(h=h, acc=4)
        grad_f = grad(f)
        assert_array_almost_equal(grad_f, grad_f_ex)

    def test_3d_gradient_on_scalar_func_non_uni(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.sin(X) * np.sin(Y) * np.sin(Z)
        grad_f_ex = np.array(
            [
                np.cos(X) * np.sin(Y) * np.sin(Z),
                np.sin(X) * np.cos(Y) * np.sin(Z),
                np.sin(X) * np.sin(Y) * np.cos(Z),
            ]
        )
        grad = Gradient(coords=axes, acc=4)
        grad_f = grad(f)
        assert_array_almost_equal(grad_f, grad_f_ex)

    def test_3d_gradient_on_vector_func_should_fail(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.array(
            [np.sin(X) * np.sin(Y) * np.sin(Z), np.sin(X) * np.sin(Y) * np.sin(Z)]
        )
        grad = Gradient(spac=h, acc=4)
        with pytest.raises(ValueError):
            grad(f)


class TestDivergence:

    def test_3d_divergence_on_vector_func(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.array([np.sin(X) * np.sin(Y) * np.sin(Z)] * 3)
        assert f.shape == (3, 50, 50, 50)
        div_f_ex = (
            np.cos(X) * np.sin(Y) * np.sin(Z)
            + np.sin(X) * np.cos(Y) * np.sin(Z)
            + np.sin(X) * np.sin(Y) * np.cos(Z)
        )

        div = Divergence(spac=h, acc=4)
        div_f = div(f)
        assert_array_almost_equal(div_f, div_f_ex)

    def test_3d_divergence_on_vector_func_of_wrong_dim(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.array([np.sin(X) * np.sin(Y) * np.sin(Z)] * 3)
        assert f.shape == (3, 50, 50, 50)
        div = Divergence(spac=[1, 1], acc=4)
        with pytest.raises(ValueError):
            div(f)


class TestCurl:

    def test_curl_on_3d_vector_func(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.array([np.sin(X) * np.sin(Y) * np.sin(Z)] * 3)
        assert f.shape == (3, 50, 50, 50)
        curl_f_ex = np.array(
            [
                np.sin(X) * np.cos(Y) * np.sin(Z) - np.sin(X) * np.sin(Y) * np.cos(Z),
                np.sin(X) * np.sin(Y) * np.cos(Z) - np.cos(X) * np.sin(Y) * np.sin(Z),
                np.cos(X) * np.sin(Y) * np.sin(Z) - np.sin(X) * np.cos(Y) * np.sin(Z),
            ]
        )
        curl = Curl(spac=h, acc=4)
        curl_f = curl(f)
        assert_array_almost_equal(curl_f, curl_f_ex)

    def test_curl_for_2d(self):
        axes, h, [X, Y, Z] = init_mesh(3, (50, 50, 50))
        f = np.array([np.sin(X) * np.sin(Y) * np.sin(Z)] * 3)
        assert f.shape == (3, 50, 50, 50)
        with pytest.raises(ValueError):
            Curl(spac=[1, 1], acc=4)

    def test_curl_for_2d_function(self):
        axes, h, [X, Y] = init_mesh(2, (50, 50))
        f = np.array([np.sin(X) * np.sin(Y)] * 2)
        curl = Curl(spac=[1, 1, 1], acc=4)
        with pytest.raises(ValueError):
            curl(f)

    def test_laplacian(self):
        axes, h, [X, Y] = init_mesh(2, (50, 50))
        f = X**2 + Y**2
        laplace = Laplacian(h)
        np.testing.assert_array_almost_equal(laplace(f), 4 * np.ones_like(X))


def init_mesh(ndims, npoints):
    axes = [np.linspace(-1, 1, npoints[k]) for k in range(ndims)]
    h = [x[1] - x[0] for x in axes]
    mesh = np.meshgrid(*axes, indexing="ij")
    return axes, h, mesh
