"""Test the jacobian module."""

import jax
import jax.numpy as jnp
import numpy as np
from jax import Array

from wejax.jacobian import autograd, findiff

jax.config.update("jax_enable_x64", True)


def _template_1(a: float) -> Array:
    """Univariate, single-channel, 10-frequency-point template."""
    return jnp.tile(a**2, (10, 1))


def _jacobian_1(a: float) -> Array:
    """Jacobian for :func:`_template_1`."""
    return jnp.tile(2 * a, (10, 1, 1))


def _template_2(a: float) -> Array:
    """Univariate, 3-channel, 10-frequency-point template."""
    channels = jnp.stack([2 * a**2 - 5 * a + 10, jnp.sin(a), jnp.exp(1j * a)])
    return jnp.tile(channels, (10, 1))


def _jacobian_2(a: float) -> Array:
    """Jacobian for :func:`_template_2`."""
    jac1 = 4 * a - 5
    jac2 = jnp.cos(a)
    jac3 = 1j * jnp.exp(1j * a)
    jac = jnp.stack([jac1, jac2, jac3])[:, None]
    return jnp.tile(jac, (10, 1, 1))


def _template_3(a: float, b: float) -> Array:
    """Bivariate, 2-channel, 10-frequency-point template."""
    channels = jnp.stack([a**2, 3 * a * b**2])
    return jnp.tile(channels, (10, 1))


def _jacobian_3(a: float, b: float) -> Array:
    """Jacobian for :func:`_template_3`."""
    jac1 = [2 * a, 0]
    jac2 = [3 * b**2, 6 * a * b]
    jac = jnp.array([jac1, jac2])
    return jnp.tile(jac, (10, 1, 1))


def test_autograd_jacobian() -> None:
    """Test the autograd_jacobian function."""

    jac_1 = autograd(_template_1, argnums=(0,))

    result = jac_1(42.0)
    expected = _jacobian_1(42.0)
    np.testing.assert_array_equal(result, expected)

    result = jax.vmap(jac_1)(jnp.array([1.0, 2.0]))
    expected = jnp.stack([_jacobian_1(1.0), _jacobian_1(2.0)])
    np.testing.assert_array_equal(result, expected)

    jac_2 = autograd(_template_2, argnums=(0,))

    result = jac_2(42.0)
    expected = _jacobian_2(42.0)
    np.testing.assert_array_equal(result, expected)

    result = jax.vmap(jac_2)(jnp.array([1.0, 2.0]))
    expected = jnp.stack([_jacobian_2(1.0), _jacobian_2(2.0)])
    np.testing.assert_array_equal(result, expected)

    jac_3 = autograd(_template_3, argnums=(0, 1))

    result = jac_3(42.0, 43.0)
    expected = _jacobian_3(42.0, 43.0)
    np.testing.assert_array_equal(result, expected)


def test_findiff_jacobian() -> None:
    """Test the findiff_jacobian function."""

    jac_1 = findiff(_template_1, argnums=(0,), steps=1e-6)

    result = jac_1(42.0)
    expected = _jacobian_1(42.0)
    np.testing.assert_array_almost_equal(result, expected, decimal=5)

    result = jax.vmap(jac_1)(jnp.array([1.0, 2.0]))
    expected = jnp.stack([_jacobian_1(1.0), _jacobian_1(2.0)])
    np.testing.assert_array_almost_equal(result, expected, decimal=5)

    jac_2 = findiff(_template_2, argnums=(0,), steps=1e-6)

    result = jac_2(42.0)
    expected = _jacobian_2(42.0)
    np.testing.assert_array_almost_equal(result, expected)

    result = jax.vmap(jac_2)(jnp.array([1.0, 2.0]))
    expected = jnp.stack([_jacobian_2(1.0), _jacobian_2(2.0)])
    np.testing.assert_array_almost_equal(result, expected, decimal=5)

    jac_3 = findiff(_template_3, argnums=(0, 1), steps=(1e-3, 1e-5))

    result = jac_3(42.0, 43.0)
    expected = _jacobian_3(42.0, 43.0)
    np.testing.assert_array_almost_equal(result, expected, decimal=5)
