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

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

Public API
----------
w2t_jax(tf_map, m_layer, filter_taps, output_length)  -> jax.Array
    Inverse transform from the 00-phase coefficients.
w2tQ_jax(tf_map, m_layer, filter_taps, output_length) -> jax.Array
    Inverse transform from the 90-phase (quadrature) coefficients.
w2t_numba(tf_map, m_layer, filter_taps, output_length)  -> numpy.ndarray
    Numba-accelerated inverse transform from the 00-phase coefficients.
w2tQ_numba(tf_map, m_layer, filter_taps, output_length) -> numpy.ndarray
    Numba-accelerated inverse transform from the 90-phase coefficients.

Both functions return raw JAX arrays and are safe to use inside
``@jax.jit`` / ``jax.vmap`` contexts.
"""
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


if HAS_JAX:
    @partial(jax.jit, static_argnames=("M", "n_filter_taps", "n_time_bins"))
    def _w2t_jax_impl(M, n_filter_taps, n_time_bins, coeff_map, filter_taps):
        n_layers = M + 1
        fft_size = 2 * M
        sqrt2 = jnp.sqrt(2.0)
        endpoint_sign = -1.0 if (M & 1) else 1.0

        output_len = M * n_time_bins + 2 * n_filter_taps
        reconstructed = jnp.zeros((output_len,), dtype=jnp.float64)

        time_idx = jnp.arange(n_time_bins, dtype=jnp.int32)
        tap_idx = jnp.arange(n_filter_taps, dtype=jnp.int32)

        real_bins = jnp.zeros((n_time_bins, n_layers), dtype=jnp.float64)
        imag_bins = jnp.zeros((n_time_bins, n_layers), dtype=jnp.float64)

        if M > 1:
            interior_freq = jnp.arange(1, M, dtype=jnp.int32)
            interior_sign = jnp.where((interior_freq & jnp.int32(1)) == 1, -1.0, 1.0)
            interior_coeff = coeff_map[:, 1:M]
            odd_mask = ((time_idx[:, None] + interior_freq[None, :]) & jnp.int32(1)) == 1

            imag_vals = jnp.where(odd_mask, interior_coeff / sqrt2, 0.0)
            real_vals = jnp.where(odd_mask, 0.0, interior_coeff / sqrt2 * interior_sign[None, :])
            real_bins = real_bins.at[:, 1:M].set(real_vals)
            imag_bins = imag_bins.at[:, 1:M].set(imag_vals)

        real0 = jnp.where((time_idx & jnp.int32(1)) == 0, coeff_map[:, 0], 0.0)
        realM = jnp.where(((time_idx + M) & jnp.int32(1)) == 0, endpoint_sign * coeff_map[:, M], 0.0)
        real_bins = real_bins.at[:, 0].set(real0)
        real_bins = real_bins.at[:, M].set(realM)

        spectrum = real_bins + 1j * imag_bins
        eq23_all = jnp.fft.irfft(spectrum, n=fft_size, axis=-1).real * fft_size

        read_start = M * (time_idx & jnp.int32(1))
        write_center = n_filter_taps + time_idx * M

        forward_idx = (read_start[:, None] + tap_idx[None, :]) % fft_size
        forward_write = write_center[:, None] + tap_idx[None, :]
        forward_vals = filter_taps[None, :] * eq23_all[time_idx[:, None], forward_idx]
        reconstructed = reconstructed.at[forward_write.reshape(-1)].add(forward_vals.reshape(-1))

        if n_filter_taps > 1:
            tap_tail = tap_idx[1:]
            backward_idx = (read_start[:, None] - tap_tail[None, :]) % fft_size
            backward_write = write_center[:, None] - tap_tail[None, :]
            backward_vals = filter_taps[None, 1:] * eq23_all[time_idx[:, None], backward_idx]
            reconstructed = reconstructed.at[backward_write.reshape(-1)].add(backward_vals.reshape(-1))

        return reconstructed[n_filter_taps:-n_filter_taps]


@_numba_njit(cache=True)
def _fill_w2t_spectrum_numba(real_bins, imag_bins, coeff_map, M):
    """Fill real/imag irfft bins for 00-phase inverse transform."""
    n_time_bins = coeff_map.shape[0]
    n_layers = M + 1
    sqrt2 = np.sqrt(2.0)
    endpoint_sign = -1.0 if (M & 1) else 1.0

    for t_idx in range(n_time_bins):
        parity_t = t_idx & 1

        real_bins[t_idx, 0] = coeff_map[t_idx, 0] if parity_t == 0 else 0.0
        real_bins[t_idx, M] = endpoint_sign * coeff_map[t_idx, M] if ((t_idx + M) & 1) == 0 else 0.0

        for freq_idx in range(1, n_layers - 1):
            coeff_value = coeff_map[t_idx, freq_idx]
            odd_mask = ((t_idx + freq_idx) & 1) == 1
            if odd_mask:
                imag_bins[t_idx, freq_idx] = coeff_value / sqrt2
                real_bins[t_idx, freq_idx] = 0.0
            else:
                interior_sign = -1.0 if (freq_idx & 1) else 1.0
                real_bins[t_idx, freq_idx] = (coeff_value / sqrt2) * interior_sign
                imag_bins[t_idx, freq_idx] = 0.0


@_numba_njit(cache=True)
def _fill_w2tQ_spectrum_numba(real_bins, imag_bins, coeff_map, M):
    """Fill real/imag irfft bins for 90-phase inverse transform."""
    n_time_bins = coeff_map.shape[0]
    n_layers = M + 1
    sqrt2 = np.sqrt(2.0)

    for t_idx in range(n_time_bins):
        parity_t = t_idx & 1

        real_bins[t_idx, 0] = coeff_map[t_idx, 0] if parity_t == 1 else 0.0
        real_bins[t_idx, M] = coeff_map[t_idx, M] if ((t_idx + M) & 1) == 1 else 0.0

        for freq_idx in range(1, n_layers - 1):
            coeff_value = coeff_map[t_idx, freq_idx]
            odd_mask = ((t_idx + freq_idx) & 1) == 1
            if odd_mask:
                real_bins[t_idx, freq_idx] = coeff_value / sqrt2
                imag_bins[t_idx, freq_idx] = 0.0
            else:
                imag_sign = 1.0 if (freq_idx & 1) else -1.0
                real_bins[t_idx, freq_idx] = 0.0
                imag_bins[t_idx, freq_idx] = (coeff_value / sqrt2) * imag_sign


@_numba_njit(cache=True)
def _accumulate_inverse_numba(reconstructed, eq23_all, filter_taps, M):
    """Accumulate symmetric filter contributions for inverse synthesis."""
    n_time_bins = eq23_all.shape[0]
    n_filter_taps = filter_taps.shape[0]
    fft_size = 2 * M

    for t_idx in range(n_time_bins):
        read_start = M * (t_idx & 1)
        write_center = n_filter_taps + t_idx * M

        for tap_idx in range(n_filter_taps):
            read_idx = (read_start + tap_idx) % fft_size
            reconstructed[write_center + tap_idx] += filter_taps[tap_idx] * eq23_all[t_idx, read_idx]

        for tap_idx in range(1, n_filter_taps):
            read_idx = (read_start - tap_idx) % fft_size
            reconstructed[write_center - tap_idx] += filter_taps[tap_idx] * eq23_all[t_idx, read_idx]


if HAS_JAX:
    @partial(jax.jit, static_argnames=("M", "n_filter_taps", "n_time_bins"))
    def _w2tQ_jax_impl(M, n_filter_taps, n_time_bins, coeff_map, filter_taps):
        n_layers = M + 1
        fft_size = 2 * M
        sqrt2 = jnp.sqrt(2.0)

        output_len = M * n_time_bins + 2 * n_filter_taps
        reconstructed = jnp.zeros((output_len,), dtype=jnp.float64)

        time_idx = jnp.arange(n_time_bins, dtype=jnp.int32)
        tap_idx = jnp.arange(n_filter_taps, dtype=jnp.int32)

        real_bins = jnp.zeros((n_time_bins, n_layers), dtype=jnp.float64)
        imag_bins = jnp.zeros((n_time_bins, n_layers), dtype=jnp.float64)

        if M > 1:
            interior_freq = jnp.arange(1, M, dtype=jnp.int32)
            imag_sign = jnp.where((interior_freq & jnp.int32(1)) == 1, 1.0, -1.0)
            interior_coeff = coeff_map[:, 1:M]
            odd_mask = ((time_idx[:, None] + interior_freq[None, :]) & jnp.int32(1)) == 1

            real_vals = jnp.where(odd_mask, interior_coeff / sqrt2, 0.0)
            imag_vals = jnp.where(odd_mask, 0.0, interior_coeff / sqrt2 * imag_sign[None, :])
            real_bins = real_bins.at[:, 1:M].set(real_vals)
            imag_bins = imag_bins.at[:, 1:M].set(imag_vals)

        real0 = jnp.where((time_idx & jnp.int32(1)) == 1, coeff_map[:, 0], 0.0)
        realM = jnp.where(((time_idx + M) & jnp.int32(1)) == 1, coeff_map[:, M], 0.0)
        real_bins = real_bins.at[:, 0].set(real0)
        real_bins = real_bins.at[:, M].set(realM)

        spectrum = real_bins + 1j * imag_bins
        eq23_all = jnp.fft.irfft(spectrum, n=fft_size, axis=-1).real * fft_size

        read_start = M * (time_idx & jnp.int32(1))
        write_center = n_filter_taps + time_idx * M

        forward_idx = (read_start[:, None] + tap_idx[None, :]) % fft_size
        forward_write = write_center[:, None] + tap_idx[None, :]
        forward_vals = filter_taps[None, :] * eq23_all[time_idx[:, None], forward_idx]
        reconstructed = reconstructed.at[forward_write.reshape(-1)].add(forward_vals.reshape(-1))

        if n_filter_taps > 1:
            tap_tail = tap_idx[1:]
            backward_idx = (read_start[:, None] - tap_tail[None, :]) % fft_size
            backward_write = write_center[:, None] - tap_tail[None, :]
            backward_vals = filter_taps[None, 1:] * eq23_all[time_idx[:, None], backward_idx]
            reconstructed = reconstructed.at[backward_write.reshape(-1)].add(backward_vals.reshape(-1))

        return reconstructed[n_filter_taps:-n_filter_taps]


def w2t_jax(tf_map, m_layer, filter_taps, output_length=None):
    """
    Inverse WDM transform from 00-phase coefficients (JAX backend).

    Reconstructs the original time series from the 00-phase column of the
    WDM TF map.  The result is a raw JAX array; for a gwpy
    :class:`~gwpy.timeseries.TimeSeries` wrapper use :meth:`WDM.w2t`.

    Parameters
    ----------
    tf_map : array-like
        Flattened TF coefficient array produced by :func:`t2w_jax`,
        shape ``(2 * n_time * (M+1),)``.
    m_layer : int
        Number of frequency layers, i.e. ``M + 1``.
    filter_taps : array-like, shape (m_H,)
        WDM prototype filter coefficients.
    output_length : int or None
        If given, truncate the output to this many samples.

    Returns
    -------
    jax.Array, shape ``(output_length,)`` or ``(M * n_time,)``
    """
    if not HAS_JAX:
        raise ImportError("JAX is not installed; w2t_jax is unavailable.")

    M = m_layer - 1
    n_layers = M + 1
    fft_size = 2 * M

    coeff_arr = jnp.asarray(tf_map, dtype=jnp.float64).reshape(-1)
    n_time_bins = int(coeff_arr.size) // (fft_size + 2)

    coeff_map = coeff_arr[:n_time_bins * n_layers].reshape(n_time_bins, n_layers)
    filt = jnp.asarray(filter_taps, dtype=jnp.float64)

    result = _w2t_jax_impl(
        M=M,
        n_filter_taps=len(filt),
        n_time_bins=n_time_bins,
        coeff_map=coeff_map,
        filter_taps=filt,
    )
    if output_length is not None:
        result = result[: int(output_length)]
    return result


def w2tQ_jax(tf_map, m_layer, filter_taps, output_length=None):
    """
    Inverse WDM transform from 90-phase (quadrature) coefficients (JAX backend).

    Reconstructs the original time series using the 90-phase column of the WDM
    TF map.  The input ``tf_map`` must contain both 00 and 90 quadratures
    (i.e. produced with ``MM < 0`` in :func:`t2w_jax`).

    Parameters
    ----------
    tf_map : array-like
        Flattened TF coefficient array produced by :func:`t2w_jax` with
        ``MM < 0``, shape ``(2 * n_time * (M+1),)``.
    m_layer : int
        Number of frequency layers, i.e. ``M + 1``.
    filter_taps : array-like, shape (m_H,)
        WDM prototype filter coefficients.
    output_length : int or None
        If given, truncate the output to this many samples.

    Returns
    -------
    jax.Array, shape ``(output_length,)`` or ``(M * n_time,)``
    """
    if not HAS_JAX:
        raise ImportError("JAX is not installed; w2tQ_jax is unavailable.")

    M = m_layer - 1
    n_layers = M + 1

    coeff_arr = jnp.asarray(tf_map, dtype=jnp.float64).reshape(-1)
    n_time_bins = int(coeff_arr.size) // (2 * n_layers)
    quadrature_offset = n_time_bins * n_layers

    if int(coeff_arr.size) < 2 * quadrature_offset:
        raise ValueError("w2tQ requires both 00 and 90 quadrature coefficient maps.")

    coeff_map = coeff_arr[quadrature_offset:2 * quadrature_offset].reshape(n_time_bins, n_layers)
    filt = jnp.asarray(filter_taps, dtype=jnp.float64)

    result = _w2tQ_jax_impl(
        M=M,
        n_filter_taps=len(filt),
        n_time_bins=n_time_bins,
        coeff_map=coeff_map,
        filter_taps=filt,
    )
    if output_length is not None:
        result = result[: int(output_length)]
    return result


def w2t_numba(tf_map, m_layer, filter_taps, output_length=None):
    """
    Inverse WDM transform from 00-phase coefficients (Numba backend).

    Parameters
    ----------
    tf_map : array-like
        Flattened TF coefficient array produced by :func:`t2w_numba` or
        :func:`t2w_jax`, shape ``(2 * n_time * (M+1),)``.
    m_layer : int
        Number of frequency layers, i.e. ``M + 1``.
    filter_taps : array-like, shape (m_H,)
        WDM prototype filter coefficients.
    output_length : int or None
        If given, truncate the output to this many samples.

    Returns
    -------
    numpy.ndarray, shape ``(output_length,)`` or ``(M * n_time,)``
    """
    M = int(m_layer) - 1
    if M <= 0:
        raise ValueError("w2t_numba requires m_layer >= 2.")

    n_layers = M + 1
    fft_size = 2 * M

    coeff_arr = np.asarray(tf_map, dtype=np.float64).reshape(-1)
    n_time_bins = int(coeff_arr.size) // (fft_size + 2)
    coeff_map = coeff_arr[: n_time_bins * n_layers].reshape(n_time_bins, n_layers)

    filt = np.asarray(filter_taps, dtype=np.float64).reshape(-1)
    n_filter_taps = int(filt.shape[0])

    real_bins = np.zeros((n_time_bins, n_layers), dtype=np.float64)
    imag_bins = np.zeros((n_time_bins, n_layers), dtype=np.float64)
    _fill_w2t_spectrum_numba(real_bins, imag_bins, coeff_map, M)

    spectrum = real_bins + 1j * imag_bins
    eq23_all = np.fft.irfft(spectrum, n=fft_size, axis=-1).real * fft_size

    output_len = M * n_time_bins + 2 * n_filter_taps
    reconstructed = np.zeros((output_len,), dtype=np.float64)
    _accumulate_inverse_numba(reconstructed, eq23_all, filt, M)

    result = reconstructed[n_filter_taps:-n_filter_taps]
    if output_length is not None:
        result = result[: int(output_length)]
    return result


def w2tQ_numba(tf_map, m_layer, filter_taps, output_length=None):
    """
    Inverse WDM transform from 90-phase coefficients (Numba backend).

    Parameters
    ----------
    tf_map : array-like
        Flattened TF coefficient array produced by :func:`t2w_numba` or
        :func:`t2w_jax` with ``MM < 0``, shape ``(2 * n_time * (M+1),)``.
    m_layer : int
        Number of frequency layers, i.e. ``M + 1``.
    filter_taps : array-like, shape (m_H,)
        WDM prototype filter coefficients.
    output_length : int or None
        If given, truncate the output to this many samples.

    Returns
    -------
    numpy.ndarray, shape ``(output_length,)`` or ``(M * n_time,)``
    """
    M = int(m_layer) - 1
    if M <= 0:
        raise ValueError("w2tQ_numba requires m_layer >= 2.")

    n_layers = M + 1

    coeff_arr = np.asarray(tf_map, dtype=np.float64).reshape(-1)
    n_time_bins = int(coeff_arr.size) // (2 * n_layers)
    quadrature_offset = n_time_bins * n_layers

    if int(coeff_arr.size) < 2 * quadrature_offset:
        raise ValueError("w2tQ_numba requires both 00 and 90 quadrature coefficient maps.")

    coeff_map = coeff_arr[quadrature_offset : 2 * quadrature_offset].reshape(n_time_bins, n_layers)
    filt = np.asarray(filter_taps, dtype=np.float64).reshape(-1)
    n_filter_taps = int(filt.shape[0])

    real_bins = np.zeros((n_time_bins, n_layers), dtype=np.float64)
    imag_bins = np.zeros((n_time_bins, n_layers), dtype=np.float64)
    _fill_w2tQ_spectrum_numba(real_bins, imag_bins, coeff_map, M)

    fft_size = 2 * M
    spectrum = real_bins + 1j * imag_bins
    eq23_all = np.fft.irfft(spectrum, n=fft_size, axis=-1).real * fft_size

    output_len = M * n_time_bins + 2 * n_filter_taps
    reconstructed = np.zeros((output_len,), dtype=np.float64)
    _accumulate_inverse_numba(reconstructed, eq23_all, filt, M)

    result = reconstructed[n_filter_taps:-n_filter_taps]
    if output_length is not None:
        result = result[: int(output_length)]
    return result
