"""
Forward WDM (Wilson-Daubechies-Meyer) transform: time series → TF map.

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

Public API
----------
t2w_jax(M, m_H, signal, filter_taps, MM) -> (m_L, aligned_length, tf_map)
    JAX-JIT forward transform returning raw JAX arrays.  Safe to call inside
    ``@jax.jit`` / ``jax.vmap`` contexts.
t2w_numba(M, m_H, signal, filter_taps, MM) -> (m_L, aligned_length, tf_map)
    Numba-accelerated forward transform returning raw NumPy arrays.
"""
from functools import partial

import numpy as np

try:
    from numba import njit as _numba_njit

    HAS_NUMBA = True
except ImportError:
    def _numba_njit(*args, **kwargs):
        """No-op replacement for numba.njit when numba is unavailable."""

        def decorator(fn):
            return fn

        if len(args) == 1 and callable(args[0]):
            return args[0]
        return decorator

    HAS_NUMBA = False

try:
    import jax
    import jax.numpy as jnp

    jax.config.update("jax_enable_x64", True)
    HAS_JAX = True
except ImportError:
    jax = None
    jnp = None
    HAS_JAX = False


def _align_to_block_size(length, block_size):
    """Return the smallest multiple of `block_size` that is >= `length`."""
    remainder = length % block_size
    return length if remainder == 0 else length + (block_size - remainder)


def _mirror_extend_jax(signal, n_filter_taps, aligned_length):
    """
    Build mirror-extended time series used by the truncated Eq. (19)-(23) sums.
    """
    signal = jnp.asarray(signal, dtype=jnp.float64)
    n_input = int(signal.shape[0])
    if n_input <= 0:
        raise ValueError("WWS must contain at least one sample.")

    extended_len = aligned_length + 2 * n_filter_taps
    extended = jnp.zeros((extended_len,), dtype=jnp.float64)

    left_mirror_max = min(n_filter_taps, n_input - 1)
    left_idx = jnp.arange(left_mirror_max + 1, dtype=jnp.int32)
    extended = extended.at[n_filter_taps - left_idx].set(signal[left_idx])

    body_idx = jnp.arange(n_input, dtype=jnp.int32)
    extended = extended.at[n_filter_taps + body_idx].set(signal[body_idx])

    n_right = extended_len - n_filter_taps - n_input
    right_idx = jnp.arange(n_right, dtype=jnp.int32)
    extended = extended.at[n_filter_taps + n_input + right_idx].set(signal[n_input - right_idx - 1])

    return extended


def _mirror_extend_numpy(signal, n_filter_taps, aligned_length):
    """NumPy equivalent of :func:`_mirror_extend_jax`."""
    signal_arr = np.asarray(signal, dtype=np.float64).reshape(-1)
    n_input = int(signal_arr.shape[0])
    if n_input <= 0:
        raise ValueError("signal must contain at least one sample.")

    extended_len = aligned_length + 2 * n_filter_taps
    extended = np.zeros((extended_len,), dtype=np.float64)

    left_mirror_max = min(n_filter_taps, n_input - 1)
    for idx in range(left_mirror_max + 1):
        extended[n_filter_taps - idx] = signal_arr[idx]

    extended[n_filter_taps : n_filter_taps + n_input] = signal_arr

    n_right = extended_len - n_filter_taps - n_input
    for idx in range(n_right):
        extended[n_filter_taps + n_input + idx] = signal_arr[n_input - idx - 1]

    return extended


@_numba_njit(cache=True)
def _fill_eq23_all_numba(eq23_all, extended_signal, filter_taps, n_filter_taps, fft_size, mm_eff):
    """Fill Eq. (23) vectors for all time bins in-place."""
    n_time_bins = eq23_all.shape[0]
    n_blocks = (n_filter_taps - 1) // fft_size
    last_filter_idx = n_filter_taps - 1

    for t_idx in range(n_time_bins):
        center = n_filter_taps + t_idx * mm_eff
        for fft_bin in range(fft_size):
            value = 0.0
            for block_idx in range(n_blocks):
                block_start = block_idx * fft_size
                block_end = block_start + fft_size
                value += (
                    extended_signal[center + block_start + fft_bin] * filter_taps[block_start + fft_bin]
                )
                value += (
                    extended_signal[center - block_end + fft_bin] * filter_taps[block_end - fft_bin]
                )
            if fft_bin == 0:
                value += filter_taps[last_filter_idx] * extended_signal[center + last_filter_idx]
            eq23_all[t_idx, fft_bin] = value


@_numba_njit(cache=True)
def _fill_tf_map_numba(tf_map, fft_real, fft_imag, M, return_quadrature):
    """Build output TF map planes from FFT components in-place."""
    n_time_bins = fft_real.shape[0]
    n_layers = M + 1
    sqrt2 = np.sqrt(2.0)

    if return_quadrature:
        for t_idx in range(n_time_bins):
            parity_t = t_idx & 1
            for freq_idx in range(n_layers):
                parity_mask = ((freq_idx + parity_t) & 1) == 1
                if parity_mask:
                    tf_map[0, t_idx, freq_idx] = sqrt2 * fft_imag[t_idx, freq_idx]
                    tf_map[1, t_idx, freq_idx] = sqrt2 * fft_real[t_idx, freq_idx]
                else:
                    tf_map[0, t_idx, freq_idx] = sqrt2 * fft_real[t_idx, freq_idx]
                    tf_map[1, t_idx, freq_idx] = -sqrt2 * fft_imag[t_idx, freq_idx]
        return

    for t_idx in range(n_time_bins):
        for freq_idx in range(n_layers):
            real_part = fft_real[t_idx, freq_idx]
            imag_part = fft_imag[t_idx, freq_idx]
            tf_map[0, t_idx, freq_idx] = real_part * real_part + imag_part * imag_part
            tf_map[1, t_idx, freq_idx] = 0.0


if HAS_JAX:
    @partial(jax.jit, static_argnames=("M", "n_filter_taps", "mm_eff", "return_quadrature", "n_time_bins"))
    def _t2w_jax_impl(
        M,
        n_filter_taps,
        mm_eff,
        return_quadrature,
        n_time_bins,
        extended_signal,
        filter_taps,
    ):
        # M + 1 layers: DC (0), interior (1…M-1), Nyquist (M).
        n_layers = M + 1
        fft_size = 2 * M
        sqrt2 = jnp.sqrt(2.0)

        freq_idx = jnp.arange(n_layers, dtype=jnp.int32)
        fft_bins = jnp.arange(fft_size, dtype=jnp.int32)
        time_idx = jnp.arange(n_time_bins, dtype=jnp.int32)

        block_starts = jnp.arange(0, n_filter_taps - 1, fft_size, dtype=jnp.int32)
        block_ends = fft_size + block_starts
        center_idx = n_filter_taps + time_idx * mm_eff

        forward_offsets = block_starts[:, None] + fft_bins[None, :]
        mirror_offsets = -block_ends[:, None] + fft_bins[None, :]

        forward_filter = filter_taps[forward_offsets]
        mirror_filter = filter_taps[block_ends[:, None] - fft_bins[None, :]]

        last_filter_idx = block_ends[-1]
        center_filter = filter_taps[last_filter_idx]

        def build_eq23(center):
            forward_signal = extended_signal[center + forward_offsets]
            mirror_signal = extended_signal[center + mirror_offsets]

            eq23_vec = jnp.sum(
                forward_signal * forward_filter + mirror_signal * mirror_filter,
                axis=0,
            )
            return eq23_vec.at[0].add(center_filter * extended_signal[center + last_filter_idx])

        eq23_all = jax.vmap(build_eq23)(center_idx)

        spectrum = jnp.fft.rfft(eq23_all, n=fft_size, axis=-1)
        fft_real = spectrum.real
        fft_imag = spectrum.imag

        endpoint0 = fft_real[:, 0] / sqrt2
        endpointM = fft_real[:, M] / sqrt2
        fft_real = fft_real.at[:, 0].set(endpoint0).at[:, M].set(endpointM)
        fft_imag = fft_imag.at[:, 0].set(endpoint0).at[:, M].set(endpointM)

        if return_quadrature:
            parity_mask = ((freq_idx[None, :] + (time_idx[:, None] & jnp.int32(1))) & jnp.int32(1)) == 1
            primary = jnp.where(parity_mask, sqrt2 * fft_imag, sqrt2 * fft_real)
            quadrature = jnp.where(parity_mask, sqrt2 * fft_real, -sqrt2 * fft_imag)
            return jnp.stack((primary, quadrature), axis=0)

        power = fft_real * fft_real + fft_imag * fft_imag
        zero = jnp.zeros_like(power)
        return jnp.stack((power, zero), axis=0)


def t2w_jax(M, m_H, signal, filter_taps, MM):
    """
    Forward WDM transform (JAX backend).

    Converts a 1-D time series into a 2-D (quadrature, time, frequency) TF map
    using the WDM prototype filter defined by ``filter_taps``.

    This function is safe to call inside ``@jax.jit`` / ``jax.vmap`` contexts;
    all outputs are JAX arrays.  For a user-friendly wrapper that returns a
    :class:`~wdm_wavelet.types.TimeFrequencyMap`, use :meth:`WDM.t2w` instead.

    Parameters
    ----------
    M : int
        Number of frequency bands (Nyquist layer index).
    m_H : int
        Filter half-length; must satisfy ``(m_H - 1) % (2*M) == 0``.
    signal : array-like, shape (N,)
        Input time series.
    filter_taps : array-like, shape (>= m_H,)
        WDM prototype filter coefficients (only the first ``m_H`` are used).
    MM : int
        Decimation stride.  ``MM <= 0`` selects ``MM = M`` (full-rate quadrature
        output); ``MM > 0`` selects power-only (no 90° quadrature) with stride MM.

    Returns
    -------
    m_L : int
        Filter half-length used in this call (``m_H`` for quadrature, 0 otherwise).
    aligned_length : int
        Signal length after block alignment.
    tf_map : jax.Array, shape ``(2, n_time, M+1)``
        ``tf_map[0]`` — 00-phase (or power) coefficients.
        ``tf_map[1]`` — 90-phase coefficients (zeros when ``MM >= 0``).
    """
    if not HAS_JAX:
        raise ImportError("JAX is not installed; t2w_jax is unavailable.")

    n_filter_taps = m_H
    fft_size = 2 * M
    if n_filter_taps < 2 or (n_filter_taps - 1) % fft_size != 0:
        raise ValueError(
            "t2w_jax requires (m_H - 1) to be divisible by (2 * M) and m_H >= 2."
        )

    original_mm = MM
    if MM <= 0:
        MM = M

    aligned_length = _align_to_block_size(len(signal), MM)
    n_time_bins = aligned_length // MM
    return_quadrature = original_mm < 0
    m_L = m_H if return_quadrature else 0

    extended_signal = _mirror_extend_jax(signal, n_filter_taps, aligned_length)
    filt = jnp.asarray(filter_taps[:n_filter_taps], dtype=jnp.float64)

    tf_map = _t2w_jax_impl(
        M=M,
        n_filter_taps=n_filter_taps,
        mm_eff=MM,
        return_quadrature=return_quadrature,
        n_time_bins=n_time_bins,
        extended_signal=extended_signal,
        filter_taps=filt,
    )

    return m_L, aligned_length, tf_map


def t2w_numba(M, m_H, signal, filter_taps, MM):
    """
    Forward WDM transform (Numba backend).

    Parameters
    ----------
    M : int
        Number of frequency bands (Nyquist layer index).
    m_H : int
        Filter half-length; must satisfy ``(m_H - 1) % (2*M) == 0``.
    signal : array-like, shape (N,)
        Input time series.
    filter_taps : array-like, shape (>= m_H,)
        WDM prototype filter coefficients (only first ``m_H`` are used).
    MM : int
        Decimation stride. ``MM <= 0`` selects ``MM = M``.

    Returns
    -------
    m_L : int
        Filter half-length used in this call (``m_H`` for quadrature, 0 otherwise).
    aligned_length : int
        Signal length after block alignment.
    tf_map : numpy.ndarray, shape ``(2, n_time, M+1)``
        ``tf_map[0]`` — 00-phase (or power) coefficients.
        ``tf_map[1]`` — 90-phase coefficients (zeros when ``MM >= 0``).
    """
    n_filter_taps = int(m_H)
    fft_size = 2 * M
    if n_filter_taps < 2 or (n_filter_taps - 1) % fft_size != 0:
        raise ValueError(
            "t2w_numba requires (m_H - 1) to be divisible by (2 * M) and m_H >= 2."
        )

    original_mm = MM
    if MM <= 0:
        MM = M

    signal_arr = np.asarray(signal, dtype=np.float64).reshape(-1)
    aligned_length = _align_to_block_size(int(signal_arr.shape[0]), int(MM))
    n_time_bins = aligned_length // MM
    return_quadrature = original_mm < 0
    m_L = m_H if return_quadrature else 0

    extended_signal = _mirror_extend_numpy(signal_arr, n_filter_taps, aligned_length)
    filt = np.asarray(filter_taps, dtype=np.float64).reshape(-1)[:n_filter_taps]

    eq23_all = np.zeros((n_time_bins, fft_size), dtype=np.float64)
    _fill_eq23_all_numba(eq23_all, extended_signal, filt, n_filter_taps, fft_size, MM)

    spectrum = np.fft.rfft(eq23_all, n=fft_size, axis=-1)
    fft_real = np.asarray(spectrum.real, dtype=np.float64)
    fft_imag = np.asarray(spectrum.imag, dtype=np.float64)

    sqrt2 = np.sqrt(2.0)
    fft_real[:, 0] /= sqrt2
    fft_real[:, M] /= sqrt2
    fft_imag[:, 0] = fft_real[:, 0]
    fft_imag[:, M] = fft_real[:, M]

    tf_map = np.zeros((2, n_time_bins, M + 1), dtype=np.float64)
    _fill_tf_map_numba(tf_map, fft_real, fft_imag, M, return_quadrature)

    return m_L, aligned_length, tf_map
