"""
Benchmark of jacobian related functions.
"""

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

from wejax.jacobian import autograd, findiff

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


class Single:
    """Single evaluation of a polynomial."""

    def _set_dim(self):
        n_freqs = 100_000
        df = 0.01
        self.freqs = np.arange(1, n_freqs + 1) * df

    def template(self, a, b):
        """A simple polynomial template"""
        channel = a**2 * self.freqs + 42j * jnp.cos(a) * b**2
        return jnp.stack([channel], axis=-1)

    def template_nojax(self, a, b):
        """A simple polynomial template, without jax"""
        channel = a**2 * self.freqs + 42j * np.cos(a) * b**2
        return np.stack([channel], axis=-1)

    def __init__(self):

        self._set_dim()
        a = np.ones((1,))
        b = np.ones((1,))

        self.apply_autograd = lambda: autograd(self.template, argnums=(0, 1))(a, b)
        self.apply_findiff = lambda: findiff(
            self.template, argnums=(0, 1), steps=[0.01, 0.1]
        )(a, b)
        self.apply_findiff_nojax = lambda: findiff(
            self.template_nojax, argnums=(0, 1), steps=[0.01, 0.1]
        )(a, b)

        # make a first evaluation
        self.apply_findiff()
        self.apply_autograd()

    def get_default_benchmark(self):
        """Return richbench inputs"""
        return [
            (
                self.apply_findiff,
                self.apply_autograd,
                "jac: using autograd instead of finite differences;",
            ),
            (
                self.apply_findiff_nojax,
                self.apply_findiff,
                "jac fin. diff.: using jax;",
            ),
        ]


class Vectorized(Single):
    """Vectorized evaluation of a polynomial using vmap or list comprehension."""

    def __init__(self):
        # pylint: disable=super-init-not-called
        Single._set_dim(self)
        n_col = 50
        va = jnp.ones((n_col,))
        vb = jnp.ones((n_col,))

        jacobian_findiff = findiff(self.template, argnums=(0, 1), steps=[0.01, 0.1])
        self.vect_findiff = lambda: jax.vmap(jacobian_findiff, in_axes=(0, 0))(
            va, vb
        ).block_until_ready()
        self.clist_findiff = lambda: [
            jacobian_findiff(_a, _b) for _a, _b in zip(va, vb)
        ]

        jacobian_autograd = autograd(self.template, argnums=(0, 1))
        self.vect_autograd = lambda: jax.vmap(jacobian_autograd, in_axes=(0, 0))(
            va, vb
        ).block_until_ready()
        self.clist_autograd = lambda: [
            jacobian_autograd(_a, _b) for _a, _b in zip(va, vb)
        ]
        # pylint: disable=not-callable
        self.jit_autograd = lambda: jax.jit(
            jax.vmap(jacobian_autograd, in_axes=(0, 0))
        )(va, vb).block_until_ready()

    def get_default_benchmark(self):
        """Return richbench inputs"""
        return [
            (
                self.clist_findiff,
                self.vect_findiff,
                "jac fin. diff.: vmap instead of list comp.;",
            ),
            (
                self.clist_autograd,
                self.vect_autograd,
                "jac autograd: vmap instead of list comp.;",
            ),
            (
                self.vect_autograd,
                self.jit_autograd,
                "jac autograd: jit on top of vmap;",
            ),
        ]


class SingleFFT:
    """A template with an fft"""

    def _set_dim(self):
        n_t = 100_000
        dt = 5
        self.times = np.arange(1, n_t + 1, dt)

    def wfm_template(self, ampl, phase):
        """A waveform template, with FFT"""
        phase = phase * self.times
        plus = -jnp.cos(phase) * ampl
        cross = -jnp.sin(phase) * 2.0 * ampl
        return jnp.stack([jnp.fft.fft(plus), jnp.fft.fft(cross)], axis=1)

    def wfm_template_nojax(self, ampl, phase):
        """A waveform template, with FFT, without JAX"""
        phase = phase * self.times
        plus = -np.cos(phase) * ampl
        cross = -np.sin(phase) * 2.0 * ampl
        return np.stack([np.fft.fft(plus), np.fft.fft(cross)], axis=1)

    def __init__(self):
        self._set_dim()

        ampl = np.ones((1,))
        phase = np.ones((1,)) * np.pi / 6.0

        self.fft_autograd = lambda: autograd(self.wfm_template, argnums=(0, 1))(
            ampl, phase
        )
        self.fft_findiff = lambda: findiff(
            self.wfm_template, argnums=(0, 1), steps=[0.1, 0.01]
        )(ampl, phase)
        self.fft_findiff_nojax = lambda: findiff(
            self.wfm_template_nojax, argnums=(0, 1), steps=[0.1, 0.01]
        )(ampl, phase)

        # make a first evaluation
        self.fft_findiff()
        self.fft_autograd()

    def get_default_benchmark(self):
        """Return richbench inputs"""
        return [
            (
                self.fft_findiff,
                self.fft_autograd,
                "jac fft: using autograd instead of finite differences; ",
            ),
            (
                self.fft_findiff_nojax,
                self.fft_findiff,
                "jac fft fin. diff.: using jax; ",
            ),
        ]


class VectorizedFFT(SingleFFT):
    """
    Vectorized usage of a template with an fft
    """

    def __init__(self):
        # pylint: disable=super-init-not-called
        self._set_dim()

        n_col = 50
        vampl = jnp.ones((n_col,))
        vphase = jnp.ones((n_col,)) * np.pi / 6.0

        fft_jacobian_findiff = findiff(
            self.wfm_template, argnums=(0, 1), steps=[0.1, 0.01]
        )
        self.fft_vect_findiff = lambda: jax.vmap(fft_jacobian_findiff, in_axes=(0, 0))(
            vampl, vphase
        ).block_until_ready()
        self.fft_clist_findiff = lambda: [
            fft_jacobian_findiff(_a, _b) for _a, _b in zip(vampl, vphase)
        ]

        fft_jacobian_autograd = autograd(self.wfm_template, argnums=(0, 1))
        self.fft_vect_autograd = lambda: jax.vmap(
            fft_jacobian_autograd, in_axes=(0, 0)
        )(vampl, vphase).block_until_ready()
        self.fft_clist_autograd = lambda: [
            fft_jacobian_autograd(_a, _b) for _a, _b in zip(vampl, vphase)
        ]
        # pylint: disable=not-callable
        self.fft_jit_autograd = lambda: jax.jit(
            jax.vmap(fft_jacobian_autograd, in_axes=(0, 0))
        )(vampl, vphase).block_until_ready()

    def get_default_benchmark(self):
        """Return richbench inputs"""
        return [
            (
                self.fft_clist_findiff,
                self.fft_vect_findiff,
                "jac fft fin. diff.: vmap instead of list comp.;",
            ),
            (
                self.fft_clist_autograd,
                self.fft_vect_autograd,
                "jac fft autograd: vmap instead of list comp.;",
            ),
            (
                self.fft_vect_autograd,
                self.fft_jit_autograd,
                "jac fft autograd: jit on top of vmap;",
            ),
        ]


__benchmarks__ = (
    Single().get_default_benchmark()
    + Vectorized().get_default_benchmark()
    + SingleFFT().get_default_benchmark()
    + VectorizedFFT().get_default_benchmark()
)
