"""
Time-delay filters and delayed-pixel amplitudes for WDM.

References
----------
- V. Necula et al., J. Phys.: Conf. Ser. 363 (2012) 012032, Sec. 4 Eq. (27)-(32).
- `docs/ref/WDM.cc` C++ methods:
  `getTDFilter1`, `getTDFilter2`, `setTDFilter`, `getPixelAmplitude`,
  `getTDamp`, `getTDvec`.

Paper equations used here
-------------------------
Eq. (27): delayed TF coefficient
    w_nm(delta_t) = sum_{l,k} w_lk * < g_nm(t+delta_t), g_lk(t) >

Eq. (29): same-band overlap (k = m)
    <g_nm(t+delta_t), g_{n+l,m}(t)> -> term with exp(i * m * DeltaOmega * delta_t) * T_l

Eq. (30): same-band integral family
    T_l = integral_{-Omega}^{Omega} exp(i*w*(M*l*tau + delta_t)) * |phi(w)|^2 dw

Eq. (31): neighbor-band overlap (k = m +/- 1)
    <g_nm(t+delta_t), g_{n+l,m+/-1}(t)> -> term with exp(i*(m+/-1/2)*DeltaOmega*delta_t) * T'_l
    and extra sqrt(2) endpoint factor when sub-band index touches {0, M}.

Eq. (32): cross-band integral family
    T'_l = integral_{-Omega}^{Omega} exp(i*w*(M*l*tau + delta_t))
           * phi*(w-DeltaOmega/2) * phi(w+DeltaOmega/2) dw
"""

import math
from dataclasses import dataclass

import numpy as np

try:
    import numba
    HAS_NUMBA = True
except ImportError:
    numba = None
    HAS_NUMBA = False

if HAS_NUMBA:
    @numba.njit(cache=True)
    def _fill_td_filter1_loop(out, fourier, n_fourier, K, K2, L, size, A, B, M_int, g_norm):
        """
        JIT-compiled inner loop for get_td_filter1.

        Fills out[1:size] in-place.  The caller is responsible for setting
        out[0] and for dividing fourier[0] by sqrt(2) before calling this
        function (matching the C++ getTDFilter1 setup).
        """
        for i in range(1, size):
            di = float(i) / float(L)
            i2 = di * di
            sum_even = 0.0
            sum_odd = 0.0
            for j in range(n_fourier):
                if j * K * L == i:
                    continue
                coeff = di / (i2 - float(j) * float(j) * K2)
                if j & 1:
                    sum_odd += coeff * fourier[j]
                else:
                    sum_even += coeff * fourier[j]

            int_ab = 0.0
            if i % (K * L) == 0:
                harmonic = i // (K * L)
                if harmonic < n_fourier:
                    int_ab = fourier[harmonic] * B / 2.0 * math.cos(di * A)

            int_ab += 2.0 * (
                sum_even * math.sin(di * B / 2.0) * math.cos(di * math.pi / (2.0 * float(M_int)))
                - sum_odd * math.sin(di * math.pi / (2.0 * float(M_int))) * math.cos(di * B / 2.0)
            )

            out[i] = g_norm * (math.sin(di * A) / di + math.sqrt(2.0) * int_ab)

    @numba.njit(cache=True)
    def _fill_td_filter2_loop(out, fourier, n_fourier, K, K2, L, size, B, g_norm):
        """
        JIT-compiled inner loop for get_td_filter2.

        Fills out[1:size] in-place.  The caller is responsible for setting
        out[0] and for dividing fourier[0] by sqrt(2) before calling this
        function (matching the C++ getTDFilter2 setup).
        """
        for i in range(1, size):
            di = float(i) / float(L)
            i2 = di * di
            summation = 0.0
            for j in range(0, n_fourier + 1, 2):
                if j * K * L == i:
                    continue
                summation += di / (i2 - float(j) * float(j) * K2) * fourier[j]

            summation *= 2.0 * math.sin(di * B / 2.0)

            if i % (2 * K * L) == 0:
                harmonic = i // (K * L)
                if harmonic <= n_fourier:
                    if (i // (2 * K * L)) & 1:
                        summation -= fourier[harmonic] * B / 2.0
                    else:
                        summation += fourier[harmonic] * B / 2.0

            out[i] = g_norm * math.sqrt(2.0) * summation


@dataclass
class TDFilterBank:
    """
    Time-delay filter bank used for delayed WDM amplitudes.

    Parameters
    ----------
    M : int
        WDM max frequency-layer index.
    L : int
        Upsampling factor for delay step dt = tau / L.
    n_coeffs : int
        Half-width of symmetric TD filters.
    T0 : np.ndarray
        First TD filter family from Eq. (30), shape (2*M*L+1, 2*n_coeffs+1).
    Tx : np.ndarray
        Cross-term TD filter family from Eq. (32), shape (2*M*L+1, 2*n_coeffs+1).
    sinTD : np.ndarray
        sin(k*pi/(L*M)), used by Eq. (29)-style phase mixing.
    cosTD : np.ndarray
        cos(k*pi/(L*M)), used by Eq. (29)-style phase mixing.
    sinTDx : np.ndarray
        sin(k*pi/(2*L*M)), used by neighbor-band terms.
    """

    M: int
    L: int
    n_coeffs: int
    T0: np.ndarray
    Tx: np.ndarray
    sinTD: np.ndarray
    cosTD: np.ndarray
    sinTDx: np.ndarray

    @property
    def max_delay(self):
        # Number of delay indices in one pixel (C++: J = M * LWDM).
        return self.M * self.L


def _trunc_division(numerator, denominator):
    """Integer division that truncates toward zero (matches C++ int division)."""
    return int(math.trunc(numerator / denominator))


def get_td_filter1(M, K, beta_order, n, L, Cos2, Cos2Size):
    """
    Compute the first TD integral family used in Eq. (29)-(31) reduction.

    This follows the C++ reference `WDM::getTDFilter1` and computes T_l
    in paper Sec. 4 Eq. (30):
      T_l = integral_{-Omega}^{Omega} exp(i*w*(M*l*tau + delta_t)) * |phi(w)|^2 dw

    In the discrete code:
      delta_t = i / L  (in units of input sampling interval tau)
      B = pi/K, A = (K-M) * pi / (2*K*M)
    and the Fourier-series expansion of |phi(w)|^2 is used exactly as in `WDM.cc`.
    """
    n_fourier = int(Cos2Size[beta_order])
    fourier = np.array(Cos2[beta_order], copy=True, dtype=np.float64)

    B = math.pi / K
    A = (K - M) * math.pi / (2.0 * K * M)
    K2 = float(K * K)
    g_norm = 2.0 * M / math.pi

    size = n * L
    out = np.zeros(size, dtype=np.float64)

    out[0] = (A + fourier[0] * B) * g_norm
    fourier[0] /= math.sqrt(2.0)

    if HAS_NUMBA:
        _fill_td_filter1_loop(out, fourier, n_fourier, K, K2, L, size, A, B, int(M), g_norm)
    else:
        for i in range(1, size):
            di = i / float(L)
            i2 = di * di
            sum_even = 0.0
            sum_odd = 0.0

            for j in range(n_fourier):
                if j * K * L == i:
                    continue
                coeff = di / (i2 - j * j * K2)
                if j & 1:
                    sum_odd += coeff * fourier[j]
                else:
                    sum_even += coeff * fourier[j]

            int_ab = 0.0
            if i % (K * L) == 0:
                harmonic = i // (K * L)
                if harmonic < n_fourier:
                    int_ab = fourier[harmonic] * B / 2.0 * math.cos(di * A)

            int_ab += 2.0 * (
                sum_even * math.sin(di * B / 2.0) * math.cos(di * math.pi / (2.0 * M))
                - sum_odd * math.sin(di * math.pi / (2.0 * M)) * math.cos(di * B / 2.0)
            )

            out[i] = g_norm * (math.sin(di * A) / di + math.sqrt(2.0) * int_ab)

    return out


def get_td_filter2(M, K, beta_order, n, L, SinCos, SinCosSize):
    """
    Compute the second TD integral family used in Eq. (31) cross-band terms.

    This follows the C++ reference `WDM::getTDFilter2` and computes T'_l
    in paper Sec. 4 Eq. (32):
      T'_l = integral_{-Omega}^{Omega} exp(i*w*(M*l*tau + delta_t))
             * phi*(w-DeltaOmega/2) * phi(w+DeltaOmega/2) dw

    In the discrete code:
      delta_t = i / L, B = pi/K, and only even Fourier harmonics enter this
      cross-term kernel, matching the C++ implementation.
    """
    n_fourier = int(SinCosSize[beta_order])
    fourier = np.array(SinCos[beta_order], copy=True, dtype=np.float64)

    B = math.pi / K
    K2 = float(K * K)
    g_norm = M / math.pi

    size = n * L
    out = np.zeros(size, dtype=np.float64)

    out[0] = fourier[0] * B * g_norm
    fourier[0] /= math.sqrt(2.0)

    if HAS_NUMBA:
        _fill_td_filter2_loop(out, fourier, n_fourier, K, K2, L, size, B, g_norm)
    else:
        # C++ uses j <= n_fourier and step 2.
        for i in range(1, size):
            di = i / float(L)
            i2 = di * di
            summation = 0.0

            for j in range(0, n_fourier + 1, 2):
                if j * K * L == i:
                    continue
                summation += di / (i2 - j * j * K2) * fourier[j]

            summation *= 2.0 * math.sin(di * B / 2.0)

            if i % (2 * K * L) == 0:
                harmonic = i // (K * L)
                if harmonic <= n_fourier:
                    if (i // (2 * K * L)) & 1:
                        summation -= fourier[harmonic] * B / 2.0
                    else:
                        summation += fourier[harmonic] * B / 2.0

            out[i] = g_norm * math.sqrt(2.0) * summation

    return out


def set_td_filter(M, K, beta_order, n_coeffs, L, Cos2, SinCos, Cos2Size, SinCosSize):
    """
    Build TD filters used for delayed WDM amplitudes.

    Reference:
      V. Necula et al., J. Phys.: Conf. Ser. 363 (2012) 012032
      Eq. (27)-(32) (Sec. 4, WDM time-delay filters)

    This mirrors C++ `WDM::setTDFilter` (non-SSE data path).

    With J = M*L and delay index dT in [-J, J], filter tables are stored as:
      T0[dT, 0]   = T_{|dT|}
      T0[dT,+j]   = T_{dT + J*j}
      T0[dT,-j]   = T_{J*j - dT}

      Tx[dT, 0]   = (1/2) * T'_{|dT|}
      Tx[dT,+/-j] = (1/2) * T'_{dT +/- J*j} with C++ sign convention for j mod 4
                    that absorbs i^l into C_l factors used in Eq. (29)/(31).
    """
    if M % 2 != 0:
        # C++ marks endpoint handling as unreliable for odd M in getPixelAmplitude.
        raise ValueError("TD filters for odd M are not supported (C++ implementation notes this limitation).")
    if L <= 0:
        raise ValueError("L must be a positive integer.")
    if n_coeffs <= 0:
        raise ValueError("n_coeffs must be a positive integer.")

    J = M * L
    center_col = n_coeffs

    filt1 = get_td_filter1(M, K, beta_order, M * (n_coeffs + 1) + 1, L, Cos2, Cos2Size)
    filt2 = get_td_filter2(M, K, beta_order, M * (n_coeffs + 1) + 1, L, SinCos, SinCosSize)

    T0 = np.zeros((2 * J + 1, 2 * n_coeffs + 1), dtype=np.float64)
    Tx = np.zeros((2 * J + 1, 2 * n_coeffs + 1), dtype=np.float64)

    for dT in range(-J, J + 1):
        dT_idx = dT + J
        T0[dT_idx, center_col] = filt1[abs(dT)]
        # Eq. (31) uses sin(...)cos(...): the precomputed T'_l table is divided by 2
        # because sin(2x) = 2 sin(x) cos(x).
        Tx[dT_idx, center_col] = filt2[abs(dT)] / 2.0

        for j in range(1, n_coeffs + 1):
            pos_col = n_coeffs + j
            neg_col = n_coeffs - j
            T0[dT_idx, pos_col] = filt1[dT + J * j]
            T0[dT_idx, neg_col] = filt1[J * j - dT]

            tx_pos = filt2[dT + J * j] / 2.0
            tx_neg = filt2[J * j - dT] / 2.0

            # Sign convention copied from C++ to move i^l factors into C_l form.
            mod = j % 4
            if mod == 1:
                tx_neg *= -1.0
            elif mod == 2:
                tx_pos *= -1.0
                tx_neg *= -1.0
            elif mod == 3:
                tx_pos *= -1.0

            Tx[dT_idx, pos_col] = tx_pos
            Tx[dT_idx, neg_col] = tx_neg

    sinTD = np.sin(np.arange(2 * L * M, dtype=np.float64) * math.pi / (L * M))
    cosTD = np.cos(np.arange(2 * L * M, dtype=np.float64) * math.pi / (L * M))
    sinTDx = np.sin(np.arange(4 * L * M, dtype=np.float64) * math.pi / (2.0 * L * M))

    return TDFilterBank(
        M=M,
        L=L,
        n_coeffs=n_coeffs,
        T0=T0,
        Tx=Tx,
        sinTD=sinTD,
        cosTD=cosTD,
        sinTDx=sinTDx,
    )


def _extract_plane_samples(tf_plane, n, m, n_coeffs):
    """
    Extract TF samples [n-k, ..., n, ..., n+k] for one frequency layer with zero padding.
    """
    n_time = tf_plane.shape[0]
    offsets = np.arange(-n_coeffs, n_coeffs + 1, dtype=np.int32)
    time_idx = n + offsets
    valid = (time_idx >= 0) & (time_idx < n_time)
    out = np.zeros_like(offsets, dtype=np.float64)
    out[valid] = tf_plane[time_idx[valid], m]
    return out


def _sum_even_odd(samples, kernel, even_mask):
    """
    Sum symmetric l-contributions split by l parity (even/odd).

    Used by Eq. (29) and Eq. (31) terms:
      sum_even = sum_{l even} x_{n+l,m} * T_l
      sum_odd  = sum_{l odd}  x_{n+l,m} * T_l
    """
    sum_even = float(np.sum(samples[even_mask] * kernel[even_mask]))
    sum_odd = float(np.sum(samples[~even_mask] * kernel[~even_mask]))
    return sum_even, sum_odd


def _sum_even_odd_for_layer(tf_plane, n, layer, n_coeffs, kernel, even_mask):
    """Convenience wrapper: sample one layer then apply `_sum_even_odd`."""
    samples = _extract_plane_samples(tf_plane, n, layer, n_coeffs)
    return _sum_even_odd(samples, kernel, even_mask)


def _apply_edge_parity_rule(sum_even, sum_odd, is_edge_layer, cancel_even):
    """
    Apply C++ endpoint parity rule for layers m in {0, M}.

    For edge layers, either even-l or odd-l contribution is suppressed depending on
    (n parity) XOR (quadrature flag), matching `WDM.cc::getPixelAmplitude`.
    """
    if not is_edge_layer:
        return sum_even, sum_odd

    if cancel_even:
        return 0.0, sum_odd
    return sum_even, 0.0


def _normalize_pixel_index(pixel_index, n_pixels):
    """
    Map pixel index from either 00/90 half into primary [0, n_pixels) interval.
    """
    j = int(pixel_index)
    if j >= n_pixels:
        j -= n_pixels
    if j < 0 or j >= n_pixels:
        raise ValueError(f"pixel_index must be in [0, {2*n_pixels-1}], got {pixel_index}.")
    return j


def _split_delay_index(delay_index, J):
    """
    Decompose global delay index k into:
      k = wdm_shift * J + local_delay, with local_delay in (-J, J)
    using C++ truncation-toward-zero integer division semantics.
    """
    wdm_shift = _trunc_division(int(delay_index), J)
    local_delay = int(delay_index) - wdm_shift * J
    return wdm_shift, local_delay


def get_pixel_amplitude(tf00, tf90, td_filters, m, n, dT, quad=False):
    """
    Compute one delayed pixel amplitude for layer m and time index n.

    This is the Python equivalent of C++ `WDM::getPixelAmplitude` and follows
    paper Sec. 4:
    - same-band term from Eq. (29) using T_l (filter family `T0`)
    - neighbor-band terms from Eq. (31) using T'_l (filter family `Tx`)

    Discrete structure used by C++ and mirrored here:
      1) Build even/odd-l sums over symmetric l in [-n_coeffs, n_coeffs]
         using T0 (same-band) and Tx (neighbor-band).
      2) Apply phase factors:
           cos(m*pi*dt/M), sin(m*pi*dt/M), sin((2m+/-1)*pi*dt/(2M)),
         where dt = dT/L.
      3) Apply parity/sign rules from Eq. (29)/(31) via n parity and (m+n) parity.
      4) Apply Eq. (31) endpoint multiplier sqrt(2) for bands touching 0 or M.
    """
    M = td_filters.M
    L = td_filters.L
    J = M * L
    n_coeffs = td_filters.n_coeffs
    offsets = np.arange(-n_coeffs, n_coeffs + 1, dtype=np.int32)
    even_mask = (offsets % 2) == 0

    is_odd_n = bool(n & 1)
    is_odd_mn = bool((m + n) & 1)
    cancel_even = is_odd_n ^ bool(quad)
    is_edge_layer = (m == 0) or (m == M)

    if m < 0 or m > M:
        raise ValueError(f"m must be in [0, {M}], got {m}.")
    if dT < -J or dT > J:
        raise ValueError(f"dT must be in [-{J}, {J}], got {dT}.")
    if n < 0 or n >= tf00.shape[0]:
        return 0.0

    if is_edge_layer and cancel_even:
        return 0.0

    plane = tf90 if quad else tf00
    t0_row = td_filters.T0[dT + J]
    tx_row = td_filters.Tx[dT + J]
    dt = dT / float(L)

    # Eq. (29): same-band contribution <g_nm(t+dt), g_{n+l,m}(t)> via T_l.
    # Implemented as even/odd-l sums followed by:
    #   res0 = sum_even*cos(m*pi*dt/M) + (-1)^n * sum_odd*sin(m*pi*dt/M).
    same_even, same_odd = _sum_even_odd_for_layer(plane, n, m, n_coeffs, t0_row, even_mask)
    same_even, same_odd = _apply_edge_parity_rule(same_even, same_odd, is_edge_layer, cancel_even)

    same_phase = m * math.pi * dt / M
    if is_odd_n:
        result = same_even * math.cos(same_phase) - same_odd * math.sin(same_phase)
    else:
        result = same_even * math.cos(same_phase) + same_odd * math.sin(same_phase)

    # Eq. (31): cross-band term with (m-1), via T'_l and phase
    # sin((2m-1)*pi*dt/(2M)), with parity signs exactly as in C++.
    if m > 0:
        low_even, low_odd = _sum_even_odd_for_layer(plane, n, m - 1, n_coeffs, tx_row, even_mask)
        low_even, low_odd = _apply_edge_parity_rule(low_even, low_odd, m == 1, cancel_even)

        low_phase = (2 * m - 1) * dt * math.pi / (2.0 * M)
        if is_odd_n:
            low_term = (low_even - low_odd) * math.sin(low_phase)
        else:
            low_term = (low_even + low_odd) * math.sin(low_phase)

        # Eq. (31) note: multiply by sqrt(2) when m or (m-1) touches {0, M}.
        if m == 1 or m == M:
            low_term *= math.sqrt(2.0)
        if is_odd_mn:
            result -= low_term
        else:
            result += low_term

    # Eq. (31): cross-band term with (m+1), via T'_l and phase
    # sin((2m+1)*pi*dt/(2M)), with parity signs exactly as in C++.
    if m < M:
        high_even, high_odd = _sum_even_odd_for_layer(plane, n, m + 1, n_coeffs, tx_row, even_mask)
        high_even, high_odd = _apply_edge_parity_rule(high_even, high_odd, m == M - 1, cancel_even)

        high_phase = (2 * m + 1) * dt * math.pi / (2.0 * M)
        if is_odd_n:
            high_term = (high_even + high_odd) * math.sin(high_phase)
        else:
            high_term = (high_even - high_odd) * math.sin(high_phase)

        # Eq. (31) note: multiply by sqrt(2) when m or (m+1) touches {0, M}.
        if m == 0 or m == M - 1:
            high_term *= math.sqrt(2.0)
        if is_odd_mn:
            result -= high_term
        else:
            result += high_term

    return float(result)


def get_td_amp(tf00, tf90, td_filters, pixel_index, delay_index, mode="p"):
    """
    Return delayed amplitude or power for one TF pixel and one delay index.

    This follows C++ `WDM::getTDamp`.

    Relation to paper:
    - delayed TF coefficients are defined by Eq. (27)
    - inner products used internally are Eq. (29)/(31) through `get_pixel_amplitude`

    Discrete delay decomposition:
      J = M*L, delay_index = wdm_shift*J + local_delay, with local_delay in (-J, J).
      This maps a global delay to local filter index + time-bin shift exactly as in C++.

    Mode equations:
      mode='a'  -> delayed 00 amplitude
      mode='A'  -> delayed 90 amplitude
      mode='p'  -> p = (a00^2 + a90^2)/2
      and odd `wdm_shift` parity triggers the same sign/quad swap used in `WDM.cc`.
    """
    if mode not in ("a", "A", "p", "P"):
        raise ValueError("mode must be one of 'a', 'A', 'p', 'P'.")

    M = td_filters.M
    M1 = M + 1
    J = td_filters.max_delay
    n_time = tf00.shape[0]
    n_pixels = n_time * M1

    j = _normalize_pixel_index(pixel_index, n_pixels)
    n = j // M1
    m = j % M1

    wdm_shift, local_delay = _split_delay_index(delay_index, J)
    shifted_n = n - wdm_shift

    if mode == "a":
        if wdm_shift & 1:
            val = get_pixel_amplitude(tf00, tf90, td_filters, m, shifted_n, local_delay, quad=True)
            return -val if ((n + m) & 1) else val
        return get_pixel_amplitude(tf00, tf90, td_filters, m, shifted_n, local_delay, quad=False)

    if mode == "A":
        if wdm_shift & 1:
            val = get_pixel_amplitude(tf00, tf90, td_filters, m, shifted_n, local_delay, quad=False)
            return val if ((n + m) & 1) else -val
        return get_pixel_amplitude(tf00, tf90, td_filters, m, shifted_n, local_delay, quad=True)

    a00 = get_pixel_amplitude(tf00, tf90, td_filters, m, shifted_n, local_delay, quad=False)
    a90 = get_pixel_amplitude(tf00, tf90, td_filters, m, shifted_n, local_delay, quad=True)
    return 0.5 * (a00 * a00 + a90 * a90)


def get_td_vec(tf00, tf90, td_filters, pixel_index, K, mode="p"):
    """
    Return delayed amplitudes/power over delay range [-K, ..., K].

    For mode in {'a','A'} this mirrors C++ layout:
      [a00(-K..K), a90(-K..K)] with total length 2*(2K+1).
    For mode in {'p','P'}:
      [p(-K..K)] with total length (2K+1).

    This is a direct vectorized wrapper over Eq. (27) evaluated via `get_td_amp`
    for each integer delay index k in [-K, K].
    """
    if K < 0:
        raise ValueError("K must be >= 0.")

    delays = range(-K, K + 1)
    if mode in ("a", "A"):
        a00 = np.array([get_td_amp(tf00, tf90, td_filters, pixel_index, k, "a") for k in delays], dtype=np.float64)
        a90 = np.array([get_td_amp(tf00, tf90, td_filters, pixel_index, k, "A") for k in delays], dtype=np.float64)
        return np.concatenate([a00, a90])

    return np.array([get_td_amp(tf00, tf90, td_filters, pixel_index, k, "p") for k in delays], dtype=np.float64)


def split_tf_map(time_frequency_map, M):
    """
    Convert TF map input to (tf00, tf90) with shape (N_time, M+1).

    WDM convention in this package:
      TF complex map stores w00 + i*w90.
    This helper returns:
      tf00 = Re(TF), tf90 = Im(TF),
    with orientation normalized to (N_time, M+1).
    """
    if isinstance(time_frequency_map, tuple) and len(time_frequency_map) == 2:
        tf00 = np.asarray(time_frequency_map[0], dtype=np.float64)
        tf90 = np.asarray(time_frequency_map[1], dtype=np.float64)
        if tf00.shape != tf90.shape:
            raise ValueError("Tuple TF planes must have identical shapes.")
        if tf00.ndim != 2:
            raise ValueError("Tuple TF planes must be 2D arrays.")
        if tf00.shape[1] == M + 1:
            return tf00, tf90
        if tf00.shape[0] == M + 1:
            return tf00.T, tf90.T
        raise ValueError("Could not infer TF plane orientation from tuple input.")

    data = time_frequency_map.data if hasattr(time_frequency_map, "data") else time_frequency_map
    arr = np.asarray(data)
    if arr.ndim != 2:
        raise ValueError("time_frequency_map must be a 2D complex array or TimeFrequencyMap.")

    if not np.iscomplexobj(arr):
        raise ValueError("time_frequency_map data must be complex (00 + i*90).")

    if arr.shape[0] == M + 1:
        return arr.real.T.astype(np.float64), arr.imag.T.astype(np.float64)
    if arr.shape[1] == M + 1:
        return arr.real.astype(np.float64), arr.imag.astype(np.float64)

    raise ValueError(
        f"Could not infer TF orientation for M={M}; expected one axis to be {M+1}, got {arr.shape}."
    )
