"""
Benchmark of inner products and other functions used by fim.
"""

import jax
import jax.numpy as jnp
import numpy as np

from wejax.fim import inner_product

N = 10000
M = 10
D = 2
a = np.random.randn(M, N, D)
b = np.random.randn(M, N, D)[..., [0], :, :]
DF = 0.1
noise_psd = np.ones((N, D, D))


def inner_product_nojax(a, b, df, noise_cov, should_invert_noise_cov, sum_func):
    if should_invert_noise_cov:
        inv_noise_cov = np.linalg.inv(noise_cov)  # (..., Nf, Nc, Nc)
    else:
        inv_noise_cov = noise_cov  # (..., Nf, Nc, Nc)
    a, b = np.broadcast_arrays(a, b)  # Contract over the channel axis
    subscripts = "...i, ...ij, ...j -> ..."
    contracted = np.einsum(subscripts, a, inv_noise_cov, np.conj(b))
    summed = 4 * np.real(
        sum_func(contracted, dx=df, axis=-1)
    )  # Sum over the frequencies
    return summed


def jax_sum(integrand, dx, axis):
    return jnp.sum(integrand, axis=axis) * dx


sum_ip = lambda: inner_product(
    a, b, df=DF, noise_cov=noise_psd, should_invert_noise_cov=False, sum_func=jax_sum
).block_until_ready()
trapz_ip = lambda: inner_product(
    a,
    b,
    df=DF,
    noise_cov=noise_psd,
    should_invert_noise_cov=False,
    sum_func=jnp.trapezoid,
).block_until_ready()
nojax_ip = lambda: inner_product_nojax(
    a,
    b,
    df=DF,
    noise_cov=noise_psd,
    should_invert_noise_cov=False,
    sum_func=np.trapezoid,
)

# make a first evaluation
sum_ip()
trapz_ip()

size_str = f"size {M},{D},{N}"
__benchmarks__ = [
    (
        sum_ip,
        trapz_ip,
        f"inner product with jax: using trapz instead of sum; {size_str}",
    ),
    (nojax_ip, trapz_ip, f"inner product with trapz: using jax; {size_str}"),
]
