"""
WDM prototype filter construction.

Reference
---------
V. Necula et al., J. Phys.: Conf. Ser. 363, 012032 (2012), Sec. 2
  https://doi.org/10.1088/1742-6596/363/1/012032

The WDM prototype filter phi[k] is the discretely sampled inverse Fourier
transform of the generalized Meyer window defined by Eq. (13)-(16) and (33)
of the paper.  Two backends are provided:

  * ``get_filter_jax``  — JAX ``lax.fori_loop`` implementation (preferred).
  * ``get_filter_numba`` — Numba ``@njit`` implementation (fallback).

The public :func:`get_filter` dispatcher picks JAX when available.
"""
import math
from functools import partial

import numpy as np

try:
    from numba import njit as _numba_njit
    HAS_NUMBA = True
except ImportError:
    # Provide a transparent no-op so the @_numba_njit decorator below is valid
    # even when numba is not installed.  The function will run as plain Python.
    def _numba_njit(*args, **kwargs):
        """No-op replacement for numba.njit when numba is not available."""
        def decorator(fn):
            return fn
        # Handle both @_numba_njit and @_numba_njit(cache=True) call styles.
        if len(args) == 1 and callable(args[0]):
            return args[0]
        return decorator
    HAS_NUMBA = False

try:
    import jax
    import jax.numpy as jnp
    from jax import lax

    # Match the float64 behavior of the original NumPy/Numba implementation.
    jax.config.update("jax_enable_x64", True)
    HAS_JAX = True
except ImportError:
    jax = None
    jnp = None
    lax = None
    HAS_JAX = False


@_numba_njit(cache=True)
def get_filter_numba(M, K, BetaOrder, n, Cos, CosSize):
    # Reference:
    #   V. Necula et al., J. Phys.: Conf. Ser. 363 (2012) 012032
    #   doi:10.1088/1742-6596/363/1/012032
    #
    # Parameters follow Eq. (33) with normalized Nyquist angular frequency Omega = pi:
    #   B = Omega/K = pi/K
    #   2A + B = Omega/M  =>  A = (K-M)pi/(2KM)
    # These constants define the generalized Meyer window in Eq. (13)-(15).
    transition_width_B = math.pi / K
    flat_half_width_A = (K - M) * math.pi / (2.0 * K * M)
    k_squared = K ** 2
    # Normalization for the discrete phi[k] samples used in Eq. (19)-(23).
    global_norm = math.sqrt(2 * M) / math.pi

    # Initialize the filter array
    phi_taps = np.zeros(n)

    # Precomputed cosine-series coefficients of nu_n(x) from Eq. (16).
    beta_cos_coeffs = Cos[BetaOrder]
    n_beta_coeffs = CosSize[BetaOrder]
    coeffs = np.copy(beta_cos_coeffs)

    # i = 0 term: analytic limit of inverse Fourier transform of Eq. (13)-(14).
    phi_taps[0] = (flat_half_width_A + coeffs[0] * transition_width_B) * global_norm
    coeffs[0] /= math.sqrt(2)

    for sample_idx in range(1, n):
        sample_float = float(sample_idx)
        sample_sq = sample_float * sample_float
        even_sum = 0.0
        odd_sum = 0.0

        # Closed-form discrete inverse Fourier evaluation of Eq. (13)-(14),
        # split into even/odd cosine-series parts of nu_n from Eq. (16).
        is_multiple_of_k = (sample_idx % K) == 0
        resonant_harmonic = sample_idx // K if is_multiple_of_k else -1
        for harmonic_idx in range(n_beta_coeffs):
            if is_multiple_of_k and harmonic_idx == resonant_harmonic:
                # Skip singular term; handled by the explicit resonant contribution below.
                continue

            coeff = sample_float / (sample_sq - harmonic_idx ** 2 * k_squared)
            if harmonic_idx & 1:
                odd_sum += coeff * coeffs[harmonic_idx]
            else:
                even_sum += coeff * coeffs[harmonic_idx]

        # Transition-band contribution (A,B from Eq. (15), parameterization Eq. (33)).
        transition_integral = 0.0
        if is_multiple_of_k and resonant_harmonic < n_beta_coeffs:
            transition_integral = (
                coeffs[resonant_harmonic]
                * transition_width_B
                / 2
                * math.cos(sample_idx * flat_half_width_A)
            )

        transition_integral += 2 * (
            even_sum * math.sin(sample_idx * transition_width_B / 2) * math.cos(sample_idx * math.pi / (2 * M))
            - odd_sum * math.sin(sample_idx * math.pi / (2 * M)) * math.cos(sample_idx * transition_width_B / 2)
        )

        phi_taps[sample_idx] = global_norm * (
            math.sin(sample_idx * flat_half_width_A) / sample_float
            + math.sqrt(2) * transition_integral
        )

    return phi_taps


if HAS_JAX:
    @partial(jax.jit, static_argnames=("M", "K", "n", "n_beta_coeffs"))
    def _get_filter_jax_impl(M, K, n, coeffs, n_beta_coeffs):
        # Reference:
        #   V. Necula et al., J. Phys.: Conf. Ser. 363 (2012) 012032
        #   doi:10.1088/1742-6596/363/1/012032
        transition_width_B = jnp.pi / K
        flat_half_width_A = (K - M) * jnp.pi / (2.0 * K * M)
        k_squared = K ** 2
        global_norm = jnp.sqrt(2.0 * M) / jnp.pi

        coeffs = coeffs.astype(jnp.float64)
        phi0 = (flat_half_width_A + coeffs[0] * transition_width_B) * global_norm
        coeffs = coeffs.at[0].set(coeffs[0] / jnp.sqrt(2.0))

        active_coeffs = coeffs[:n_beta_coeffs]
        harmonic_idx = jnp.arange(n_beta_coeffs, dtype=jnp.int32)
        odd_mask = (harmonic_idx & jnp.int32(1)) == 1
        harmonic_sq_k2 = (harmonic_idx.astype(jnp.float64) ** 2) * k_squared

        phi_taps = jnp.zeros((n,), dtype=jnp.float64)
        phi_taps = phi_taps.at[0].set(phi0)

        def sample_body(sample_idx, taps):
            sample_float = sample_idx.astype(jnp.float64)
            sample_sq = sample_float * sample_float

            is_multiple_of_k = (sample_idx % K) == 0
            resonant_harmonic = sample_idx // K

            coeff = sample_float / (sample_sq - harmonic_sq_k2)
            skip_resonant = jnp.logical_and(is_multiple_of_k, harmonic_idx == resonant_harmonic)
            contributions = jnp.where(skip_resonant, 0.0, coeff * active_coeffs)

            odd_sum = jnp.sum(jnp.where(odd_mask, contributions, 0.0))
            even_sum = jnp.sum(jnp.where(odd_mask, 0.0, contributions))

            def _resonant_term(idx):
                return active_coeffs[idx] * transition_width_B / 2.0 * jnp.cos(sample_float * flat_half_width_A)

            resonant_term = lax.cond(
                jnp.logical_and(is_multiple_of_k, resonant_harmonic < n_beta_coeffs),
                _resonant_term,
                lambda _: jnp.float64(0.0),
                resonant_harmonic,
            )

            transition_integral = resonant_term + 2.0 * (
                even_sum * jnp.sin(sample_float * transition_width_B / 2.0) * jnp.cos(sample_float * jnp.pi / (2.0 * M))
                - odd_sum * jnp.sin(sample_float * jnp.pi / (2.0 * M)) * jnp.cos(sample_float * transition_width_B / 2.0)
            )

            tap_value = global_norm * (
                jnp.sin(sample_float * flat_half_width_A) / sample_float
                + jnp.sqrt(2.0) * transition_integral
            )
            return taps.at[sample_idx].set(tap_value)

        return lax.fori_loop(1, n, sample_body, phi_taps)


def get_filter_jax(M, K, beta_order, n, Cos, CosSize):
    """
    WDM prototype filter construction (JAX backend).

    Implements the same analytic formula as :func:`get_filter_numba` using
    ``jax.lax.fori_loop``.  Returns a NumPy array for downstream compatibility.

    Parameters
    ----------
    M : int
        Number of frequency bands.
    K : int
        Meyer window transition-band parameter (slope sharpness).
    beta_order : int
        Index selecting the beta polynomial order from ``Cos`` / ``CosSize``.
    n : int
        Maximum number of filter taps to compute.
    Cos, CosSize :
        Pre-loaded Fourier coefficient tables from
        :func:`~wdm_wavelet.fourier_coeff.load_fourier_coeff`.

    Returns
    -------
    numpy.ndarray, shape (n,)
        WDM prototype filter taps phi[0], phi[1], ..., phi[n-1].
    """
    if not HAS_JAX:
        raise ImportError("JAX is not installed; get_filter_jax is unavailable.")

    beta_cos_coeffs = jnp.asarray(Cos[beta_order], dtype=jnp.float64)
    n_beta_coeffs = int(CosSize[beta_order])
    return np.asarray(_get_filter_jax_impl(M, K, n, beta_cos_coeffs, n_beta_coeffs))


def get_filter(M, K, beta_order, n, Cos, CosSize):
    """
    WDM prototype filter construction (auto-dispatch).

    Uses the JAX backend when available; falls back to Numba (or plain Python
    if Numba is also absent).

    Parameters
    ----------
    M, K, beta_order, n, Cos, CosSize :
        See :func:`get_filter_jax` for full documentation.

    Returns
    -------
    numpy.ndarray, shape (n,)
    """
    if HAS_JAX:
        return get_filter_jax(M, K, beta_order, n, Cos, CosSize)
    return get_filter_numba(M, K, beta_order, n, Cos, CosSize)
