"""Test the fim module."""

import jax
import jax.numpy as jnp
import numpy as np
from jax import Array
from jax._src.typing import ArrayLike

from wejax.fim import (
    inner_product,
    matched_filter_gaussian_fim,
    univariate_inner_product,
)


def test_univariate_inner_product() -> None:
    """Test the univariate_inner_product function."""

    a = jnp.array([1, 2, 3, 4])
    b = jnp.array([4, 3, 2, 1])
    df = 0.1
    noise_psd = jnp.array([1, 1, 1, 1])

    result = univariate_inner_product(a, b, noise_psd, df=df)
    expected = 4 * jnp.real(jnp.trapezoid(a * jnp.conj(b) / noise_psd, dx=df))

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with a different df
    df = 0.2

    result = univariate_inner_product(a, b, noise_psd, df=df)
    expected = 4 * jnp.real(np.trapezoid(a * jnp.conj(b) / noise_psd, dx=df))

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with a different noise_psd
    noise_psd = jnp.array([1, 2, 3, 4])

    result = univariate_inner_product(a, b, df=df, noise_psd=noise_psd)
    expected = 4 * jnp.real(jnp.trapezoid(a * jnp.conj(b) / noise_psd, dx=df))

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with a different sum_func
    def jax_sum(a: ArrayLike, dx: float, axis: int = -1) -> Array:
        return jnp.sum(a, axis=axis) * dx

    result = univariate_inner_product(
        a, b, df=df, noise_psd=noise_psd, sum_func=jax_sum
    )
    expected = 4 * np.real(np.sum(a * np.conj(b) / noise_psd) * df)

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with additional dimensions
    a = jnp.array([[1, 2], [3, 4]])
    b = jnp.array([[4, 3], [2, 1]])
    noise_psd = jnp.array([[1, 2], [3, 4]])

    result = univariate_inner_product(a, b, df=df, noise_psd=noise_psd)
    expected = 4 * jnp.real(jnp.trapezoid(a * jnp.conj(b) / noise_psd, dx=df, axis=-1))

    assert result.shape == (2,)
    np.testing.assert_array_almost_equal(result, expected)


def test_univariate_inner_product_shape() -> None:
    """Test the univariate_inner_product function with different shapes."""

    df = 0.1
    a = jnp.ones((10,))
    b = jnp.ones((10,))
    noise_cov = jnp.ones((10,))
    assert univariate_inner_product(a, b, noise_cov, df=df).shape == ()

    a = jnp.ones((10,))
    b = jnp.ones((10,))
    noise_cov = jnp.ones((1,))
    assert univariate_inner_product(a, b, noise_cov, df=df).shape == ()

    a = jnp.ones((50, 10))
    b = jnp.ones((50, 10))
    noise_cov = jnp.ones((10,))
    assert univariate_inner_product(a, b, noise_cov, df=df).shape == (50,)

    a = jnp.ones((20, 10))
    b = jnp.ones((10))
    noise_cov = jnp.ones((10,))
    assert univariate_inner_product(a, b, noise_cov, df=df).shape == (20,)


def test_inner_product() -> None:
    """Test the inner_product function."""

    a = jnp.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
    b = jnp.array([[13, 14, 15, 16], [17, 18, 19, 20], [21, 22, 23, 24]])
    df = 0.1
    noise_psd = jnp.array([1, 1, 1, 1])
    noise_cov = jnp.tile(jnp.diag(noise_psd), (3, 1, 1))

    result = inner_product(a, b, noise_cov, df=df)
    expected = 4 * np.real(
        np.trapezoid(
            np.einsum("...fi,...fij,...fj->...f", a, noise_cov, np.conj(b)),
            dx=df,
            axis=-1,
        )
    )

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with uncorrelated channels
    noise_psd = jnp.array([1, 2, 3, 4])
    noise_cov = jnp.tile(jnp.diag(noise_psd), (3, 1, 1))

    result = inner_product(a, b, noise_cov, df=df)
    inv_noise_psd = 1 / noise_psd
    inv_noise_cov = jnp.tile(jnp.diag(inv_noise_psd), (3, 1, 1))
    expected = 4 * np.real(
        np.trapezoid(
            np.einsum("...fi,...fij,...fj->...f", a, inv_noise_cov, np.conj(b)),
            dx=df,
            axis=-1,
        )
    )

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with full covariance matrix
    np_noise_cov = np.zeros((3, 4, 4), dtype=complex)
    np_noise_cov[:] += np.diag([1, 2, 3, 4])
    np_noise_cov[:, 0, 1] = 0.5 + 1j
    np_noise_cov[:, 1, 0] = 0.5 - 1j
    np_noise_cov[:, 1, 2] = 0.1 + 0.3j
    np_noise_cov[:, 2, 1] = 0.1 - 0.3j
    noise_cov = jnp.asarray(np_noise_cov)

    result = inner_product(a, b, noise_cov, df=df)
    inv_noise_cov = jnp.linalg.inv(noise_cov)
    expected = 4 * np.real(
        np.trapezoid(
            np.einsum("...fi,...fij,...fj->...f", a, inv_noise_cov, np.conj(b)),
            dx=df,
            axis=-1,
        )
    )

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with the inverse covariance matrix given

    result = inner_product(a, b, inv_noise_cov, df=df, should_invert_noise_cov=False)

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with a different sum_func
    def jax_sum(a: Array, dx: float, axis: int = -1) -> Array:
        return jnp.sum(a, axis=axis) * dx

    result = inner_product(a, b, df=df, noise_cov=noise_cov, sum_func=jax_sum)
    expected = 4 * np.real(
        np.sum(
            np.einsum("...fi,...fij,...fj->...f", a, inv_noise_cov, np.conj(b)),
            axis=-1,
        )
        * df
    )

    assert result.shape == ()
    np.testing.assert_approx_equal(result, expected)

    # Test with additional dimensions
    a = jnp.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
    b = jnp.array([[[8, 7], [6, 5]], [[4, 3], [2, 1]]])
    noise_psd = jnp.array([1, 2])
    noise_cov = jnp.tile(jnp.diag(noise_psd), (2, 2, 1, 1))

    result = inner_product(a, b, noise_cov, df=df)
    inv_noise_cov = jnp.linalg.inv(noise_cov)
    expected = 4 * np.real(
        np.trapezoid(
            np.einsum("...fi,...fij,...fj->...f", a, inv_noise_cov, np.conj(b)),
            dx=df,
            axis=-1,
        )
    )

    assert result.shape == (2,)
    np.testing.assert_array_almost_equal(result, expected)


def test_inner_product_shape() -> None:
    """Test the inner_product function with different shapes."""

    df = 0.1
    a = jnp.ones((10, 10))
    b = jnp.ones((10, 10))
    noise_cov = jnp.ones((10, 10))
    assert inner_product(a, b, noise_cov, df=df).shape == ()

    a = jnp.ones((10, 10))
    b = jnp.ones((10, 10))
    noise_cov = jnp.ones((1, 10, 10))
    assert inner_product(a, b, noise_cov, df=df).shape == ()

    a = jnp.ones((50, 10, 10))
    b = jnp.ones((50, 10, 10))
    noise_cov = jnp.ones((10, 10))
    assert inner_product(a, b, noise_cov, df=df).shape == (50,)

    a = jnp.ones((20, 10, 10))
    b = jnp.ones((10, 10))
    noise_cov = jnp.ones((10, 10))
    assert inner_product(a, b, noise_cov, df=df).shape == (20,)


def test_matched_filter_gaussian_fim_shape() -> None:
    """Test the matched_filter_gaussian_fim function with different shapes."""

    # Mock a Jacobian function
    n_freq = 10
    n_channels = 3
    n_params = 2

    def jac(a: float, b: float) -> Array:
        """Mock Jacobian function."""
        jac_shape = (n_freq, n_channels, n_params)
        return (a + b) * jnp.ones(jac_shape)

    df = 0.1
    a = 1.0
    b = 1.0
    noise_cov = jnp.ones((n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jac, noise_cov, df=df)
    assert fim(a, b).shape == (n_params, n_params)

    a = jnp.asarray(1.0)
    b = jnp.asarray(1.0)
    noise_cov = jnp.ones((n_freq, n_channels, n_channels))
    fim = matched_filter_gaussian_fim(jac, noise_cov, df=df)
    assert fim(a, b).shape == (n_params, n_params)

    a = jnp.ones((50,))
    b = jnp.ones((50,))
    noise_cov = jnp.ones((n_freq, n_channels, n_channels))
    vjac = jax.vmap(jac, in_axes=(0, 0))
    fim = matched_filter_gaussian_fim(vjac, noise_cov, df=df)
    assert fim(a, b).shape == (50, n_params, n_params)

    a = jnp.ones(
        (
            20,
            50,
        )
    )
    b = jnp.ones((50,))
    a, b = jnp.broadcast_arrays(a, b)
    noise_cov = jnp.ones((n_freq, n_channels, n_channels))
    vvjac = jax.vmap(vjac, in_axes=(0, 0))
    fim = matched_filter_gaussian_fim(vvjac, noise_cov, df=df)
    assert fim(a, b).shape == (20, 50, n_params, n_params)
