"""
JaxGB module
============

Fast waveform and response generator for galactic binaries, as seen by the
LISA instrument. Galactic binaries are described by the following 8
parameters: frequency (f0), frequency evolution (fdot), amplitude
(ampl), sky location (ra, dec), polarisation (psi), inclination
(incl) and initial phase (phi0).

Usage
-----

.. code-block:: python

    # Catalogue parameter definition and definition of epoch `t_init`.
    catalogue_pd = pd.DataFrame(
        {
        "Frequency": #...,
        # other parameters...
        "InitialPhase": #...,
        }
    )
    t_init = 1.0e3  # catalogue epoch

    # Definition of `JaxGB` simulator, with simulation `t0` (simulation epoch).
    t0 = 1.0e4  # simulation t0
    myjaxgb = JaxGB(..., t0=t0, ...)

    # Use GBObject to create the jaxgb input parameter array, which also
    # accounts for the evalution of `f0` and `phi0` for the particular epoch
    # of simulation `t0`.

    # Step 1: define GBObject, holding all source parameters
    catalogue_as_gbo = GBObject.from_pandas_dataframe(catalogue_pd, t_init=t_init)

    # Step 2: get parameters in JaxGB input format, evaluated at `t0`
    params = catalogue_as_gbo.to_jaxgb_array(t0=t0)

    # Step 3: calculate TDI variables in frequency domain
    myjaxgb.get_tdi(jnp.array(params), tdi_generation=1.5, tdi_combination="XYZ")


.. autoclass:: JaxGB
    :members:

.. autoclass:: JaxGBaccurate
    :members:
"""

import jax
import jax.numpy as jnp
import numpy as np
from lisaconstants import c  # pylint: disable=no-name-in-module
from lisaorbits import Orbits
from scipy.signal.windows import tukey

from .tdi import to_tdi_combination, to_tdi_generation


class JaxGB:
    """Galactic binary fast response generation.

    Parameters
    ----------
    orbits
        Orbits instance, based on lisaorbits.
    t_obs
        Duration (s).
    t0
        Start time (s).
    n
        Number of points for slow response evaluation (even-valued integer!).
    """

    def __init__(
        self,
        orbits: Orbits,
        *,
        t_obs: float = 6.2914560e7,
        t0: float = 0.0,
        n: int = 128,
    ) -> None:

        # duration (s)
        self.t_obs = t_obs

        # t0 (s)
        self.t0 = t0

        # number of points for slow response evaluation
        self.n = n

        # time grid
        self.tm = np.linspace(0.0, self.t_obs, num=self.n, endpoint=False)

        # time step slow response
        self.dtm = self.t_obs / self.n

        # initialize orbits and calculate positions of M = 3 spacecraft at time
        # of reception
        self.orbits = orbits
        ltt = self._get_ltt_simple()
        self.arm_length = ltt * c
        self.position = (
            self.orbits.compute_position(self.t0 + self.tm, [1, 2, 3])
            .swapaxes(0, 1)
            .swapaxes(1, 2)
        )  # (3, 3, n)
        self.fstar = c / (self.arm_length * 2 * np.pi)

    def _get_ltt_simple(self) -> jax.Array:
        """Get light-travel times: either pick "L" from orbits file or calculate the an
        ltt average using ``compute_ltt``.

        Returns
        -------
        jax.Array of shape ()
            LTT from orbits file, or averaged calculated LTT (both are scalar
            quantities).
        """

        if hasattr(self.orbits, "L"):
            return jnp.array(self.orbits.L / c)
        if hasattr(self.orbits, "original_attrs"):
            if "L" in self.orbits.original_attrs:
                return jnp.array(self.orbits.original_attrs["L"] / c)
        ltt = self.orbits.compute_ltt(self.t0 + self.tm).T  # (6, n)
        return jnp.mean(ltt)

    def get_kmin(self, f0: jax.Array) -> jax.Array:
        """Calculates kmin, the index of the first heterodyned frequency bin in the
        full frequency grid. Given the observation time baseline `t_obs`, the
        corresponding Fourier frequencies are integer multiples of `1 / t_obs`. The
        nearest frequency bin number to `f0` is therefore `rint(f0 * t_obs)`. This
        bin is centred on the grid of heterodyned Fourier frequencies, of length `n`,
        hence kmin (i.e., the first bin) is shifted by `n/2`.

        Note: `f0` needs to be evaluated at simulation epoch `self.t0`, which in general
        differs from the catalogue epoch `t_init`. Use the :mod:`params.GBObject` to
        calculate source frequencies at simulation epoch `t0` from catalogue frequencies
        (specified at `t_init`).

        Assuming `params` is the `JAXGB`-compatible array of parameters, and `myjaxgb`
        is a `JaxGB` instance, we have:
        .. code-block:: python

            myjaxgb.get_kmin(f0=params[:, 0])
            # `params[:, 0]` needs to contain the catalogue frequencies at the
            # simulation reference time `t0`, which is in general different from the
            # catalogue epoch `t_init`.

        Parameters
        ----------
        f0
            Frequency f0 for source(s), at epoch self.t0 (Hz).

        Returns
        -------
        jax.Array of shape (nsrc,)
            Array with kmin values (integer).
        """

        q = jnp.rint(f0 * self.t_obs)
        return jnp.round(q - self.n / 2).astype(np.int32)

    def get_frequency_grid(self, kmin: jax.Array) -> jax.Array:
        """Return the frequency grid associated to the returned signal.

        Parameters
        ----------
        kmin
            Index of the first frequency bin.

        Returns
        -------
        jax.Array of shape (nsrc, n)
            Array with frequency values (in Hz).
        """
        return (1 / self.t_obs) * (jnp.arange(self.n) + kmin[:, None])

    def _construct_slow_part(
        self,
        f0_fdot_phi0: jax.Array,
        k: jax.Array,
        dpdc: tuple[jax.Array, jax.Array],
        pol_tensors: tuple[jax.Array, jax.Array],
    ) -> jax.Array:
        """This method implements the slow, heterodyned response in time domain. It
        accepts the f0, fdot, and phi0 of all the sources, their direction (k),
        the hplus/hcross strain amplitudes (Dp, Dc) and polarization tensors
        (eplus, ecross). The output is an array of shape (nsrc, 6, n) [other equations
        need to be included here].

        Parameters
        ----------
        f0_fdot_phi0
            array (nsrc, 3) listing f0, fdot, and phi0 values of all sources.
        k
            array (nsrc, 3) with source k-vectors.
        dpdc
            tuple with hplus and hcross strain amplitudes (complex-valued).
        pol_tensors
            tuple with Eplus, Ecross (polarization tensors).

        Returns
        -------
        jax.Array of shape (nsrc, 6, n)
            Time-domain slow heterodyned response.
        """

        nsrc = len(f0_fdot_phi0)

        r = jnp.zeros((4, 3, self.n))  # r convention: 12,13,23,31
        r = r.at[0].set(self.position[1] - self.position[0])  # (3, n)
        r = r.at[1].set(self.position[2] - self.position[0])  # (3, n)
        r = r.at[2].set(self.position[2] - self.position[1])  # (3, n)
        r = r.at[3].set(-r[1])  # (3, n)
        r /= self.arm_length

        # kdotr convention: 12, 21, 13, 31, 23, 32
        kdotr = jnp.zeros((nsrc, 6, self.n))

        kdotr = kdotr.at[:, ::2].set(jnp.dot(k.T, r[:3]))
        kdotr = kdotr.at[:, 1::2].set(-kdotr[:, ::2])

        kdotp = jnp.dot(k.T, self.position)
        kdotp /= c

        f0, fdot, phi0 = f0_fdot_phi0[:, 0], f0_fdot_phi0[:, 1], f0_fdot_phi0[:, 2]
        xi = self.tm - kdotp
        fi = f0[:, None, None] + fdot[:, None, None] * xi
        fonfs = fi / self.fstar  # Ratio of true frequency to transfer frequency

        ### compute transfer f-n
        q = jnp.rint(f0 * self.t_obs)  # index of nearest Fourier bin
        df = 2.0 * np.pi * (q / self.t_obs)
        om = 2.0 * np.pi * f0

        # Reduce to (3, 3, n).
        # A = (eplus @ r * r * dp + ecross @ r * r * dc).sum(axis=2)
        idx = jnp.array([0, 2, 3])
        eplus, ecross = pol_tensors
        aij = (
            jnp.swapaxes(jnp.dot(eplus, r[idx]), 1, 2)
            * r[idx]
            * dpdc[0][:, None, None, None]
            + jnp.swapaxes(jnp.dot(ecross, r[idx]), 1, 2)
            * r[idx]
            * dpdc[1][:, None, None, None]
        )
        asum = aij.sum(axis=2)  # A convention: 12, 23, 31
        arg_s = (phi0[:, None] + (om - df)[:, None] * self.tm[None, :])[:, None, :] + (
            np.pi * fdot[:, None, None] * (xi**2)
        )
        kdotp = om[:, None, None] * kdotp - arg_s

        # gs convention: 12, 23, 31, 21, 32, 13
        idx = jnp.array([0, 4, 3, 1, 5, 2])
        arg_ij = (
            0.5 * fonfs[:, [0, 1, 2, 1, 2, 0]] * (1 + kdotr[:, idx])
        )  # note the plus-sign!
        gs = (
            0.5j
            * (f0[:, None, None] + fdot[:, None, None] * self.tm[None, None, :])
            / self.fstar
            * jnp.sinc(arg_ij / np.pi)
            * jnp.exp(-1j * (arg_ij + kdotp[:, [0, 1, 2, 1, 2, 0]]))
            * asum[:, [0, 1, 2, 0, 1, 2]]
        )  # last term returns a frequency-dependent sign, related to the windowing

        # shift convention: 12, 23, 31, 21, 32, 13 to 12, 23, 31, 13, 32, 21
        gs = gs[:, [0, 1, 2, 5, 4, 3]]

        return gs  # (nsrc, 6, n)

    def _get_gs(self, params: jax.Array) -> jax.Array:
        """Calculate slow response in time domain. Generic for single-link and
        TDI response calculation.

        Parameters
        ----------
        params
            Parameter array of source parameters, as defined in :mod:`jaxgb.params`.

        Returns
        -------
        jax.Array of shape (nsrc, 6, n)
            Slow response in time domain.
        """

        # sign flip phase
        params = params.at[:, 7].set(params[:, 7] * -1)

        # equatorial latitude to theta
        params = params.at[:, 4].set(np.pi / 2 - params[:, 4])

        cosiota = jnp.cos(params[:, 6])
        cosps = jnp.cos(2 * params[:, 5])
        sinps = jnp.sin(2 * params[:, 5])
        aplus = params[:, 2] * (1.0 + cosiota * cosiota)
        across = -2.0 * params[:, 2] * cosiota
        dpdc = (
            aplus * cosps - 1.0j * across * sinps,
            -aplus * sinps - 1.0j * across * cosps,
        )

        sinph = jnp.sin(params[:, 3])
        cosph = jnp.cos(params[:, 3])
        sinth = jnp.sin(params[:, 4])
        costh = jnp.cos(params[:, 4])

        u = jnp.array([costh * cosph, costh * sinph, -sinth], ndmin=2)
        v = jnp.array([sinph, -cosph, jnp.zeros_like(sinph)], ndmin=2)
        k = jnp.array([-sinth * cosph, -sinth * sinph, -costh], ndmin=2)

        eplus = (v[None, :, :] * v[:, None, :] - u[None, :, :] * u[:, None, :]).T
        ecross = (u[None, :, :] * v[:, None, :] + v[None, :, :] * u[:, None, :]).T
        pol_tensors = (eplus, ecross)

        gs = self._construct_slow_part(params[:, [0, 1, 7]], k, dpdc, pol_tensors)

        return gs  # (nsrc, 6, n)

    def _compute_response(self, gs: jax.Array) -> jax.Array:
        """Compute single-link response from gs.

        Parameters
        ----------
        gs
            Slow (heterodyned) response in time domain.

        Returns
        -------
        jax.Array of shape (nsrc, 6, n)
            Frequency-domain response vector.
        """
        yf_slow = jnp.fft.fft(gs, axis=2)

        # yf convention: 12, 23, 31, 13, 32, 21
        yf = jnp.fft.fftshift(yf_slow, axes=2)  # (nsrc, 6, n)
        fctr = 0.5 * self.dtm  # normalization of (half-sided) fourier transform
        return yf * fctr

    def _compute_tdi_xyz(
        self, gs: jax.Array, f0: jax.Array, fdot: jax.Array
    ) -> jax.Array:
        """Compute TDI X, Y, Z from gs.

        Parameters
        ----------
        gs
            Slow (heterodyned) response in time domain.

        Returns
        -------
        jax.Array of shape (nsrc, 3, n)
            Frequency-domain TDI response vector.
        """

        f = f0[:, None] + fdot[:, None] * self.tm[None, :]
        oml = f / np.mean(self.fstar)  # constant arm length assumed
        soml = jnp.sin(oml)
        fctr = jnp.exp(-1.0j * oml)
        fctr2 = -2.0j * soml * fctr

        # gs convention: 12, 23, 31, 13, 32, 21
        xsl = gs[:, 5] - gs[:, 2] + (gs[:, 0] - gs[:, 3]) * fctr
        ysl = gs[:, 4] - gs[:, 0] + (gs[:, 1] - gs[:, 5]) * fctr
        zsl = gs[:, 3] - gs[:, 1] + (gs[:, 2] - gs[:, 4]) * fctr

        xyzsl = fctr2[:, None, :] * jnp.concatenate(
            [xsl[:, None, :], ysl[:, None, :], zsl[:, None, :]], axis=1
        )
        xyzf_slow = jnp.fft.fft(xyzsl, axis=2)
        xyzf = jnp.fft.fftshift(xyzf_slow, axes=2)

        fctr3 = 0.5 * self.dtm
        xyzf *= fctr3
        return xyzf

    def get_link_responses(self, params: jax.Array) -> jax.Array:
        """Calculate the 6 single-link responses in frequency domain.

        Parameters
        ----------
        params
            Parameter array of shape (nsrc, 8) of source parameters, as
            defined in :mod:`jaxgb.params`.

        Returns
        -------
        jax.Array of shape (nsrc, 6, n)
            Single-link response vector.
        """
        params = params.copy()
        gs = self._get_gs(params)
        yf = self._compute_response(gs)
        return yf

    def get_tdi(
        self,
        params: jax.Array,
        tdi_generation: float = 1.5,
        tdi_combination: str = "XYZ",
    ) -> tuple[jax.Array, jax.Array, jax.Array]:
        """Calculate TDI XYZ in frequency domain.

        Parameters
        ----------
        params
            Parameter array of shape (nsrc, 8) of source parameters, as
            defined in :mod:`jaxgb.params`.
        tdi_generation
            TDI generation (either 1.5 or 2).
        tdi_combination
            TDI combinations (either "XYZ" or "AET").

        Returns
        -------
        tuple of arrays, each of shape (nsrc, n)
            TDI frequency-domain response channels (X, Y, Z) or
            (A, E, T).
        """
        _params = params.copy()
        if len(_params.shape) == 1:
            _params = _params.reshape(1, -1)
        gs = self._get_gs(_params)
        tdi = self._compute_tdi_xyz(gs, _params[:, 0], _params[:, 1])

        kmin = self.get_kmin(f0=_params[:, 0])
        freqs = self.get_frequency_grid(kmin)

        tdi = to_tdi_generation(
            tdi, freqs, to=tdi_generation, arm_length=self.arm_length
        )
        tdi = to_tdi_combination(tdi, tdi_combination)
        x = tdi[:, 0, :]
        y = tdi[:, 1, :]
        z = tdi[:, 2, :]
        if len(params.shape) == 1:
            return x[0], y[0], z[0]
        return x, y, z

    def sum_tdi(
        self,
        params: jax.Array,
        kmin: int | None = None,
        kmax: int | None = None,
        **kwargs,
    ) -> tuple[jax.Array, jax.Array, jax.Array]:
        r"""Sum TDI response channels for a set of source parameters.

        Each source is simulated independently, and the TDI responses
        are then summed (after aligning the frequency grids) and
        returned.

        Parameters
        ----------
        params
            Parameter array with shape (nsrc, 8) of source parameters,
            as defined in :mod:`jaxgb.params`.
        kmin
            First index of the frequency grid. If None, takes the
            first index of the lowest frequency source.
        kmax
            Last index of the frequency grid. If None, takes the
            last index of the highest frequency source.
        kwargs
             Extra keywords arguments passed to :meth:`get_tdi`.

        Returns
        -------
        tuple (jax.Array, jax.Array, jax.Array, ), each of shape (nfreq)
            TDI frequency-domain response channels (X, Y, Z) or
            (A, E, T) [strain].

        """
        kmins = self.get_kmin(f0=params[:, 0])
        kmin = jnp.min(kmins).item() if kmin is None else kmin
        kmax = jnp.max(kmins).item() + self.n if kmax is None else kmax

        def _add_to(tdi_one, tdi_sum, _kmin):
            def _update_fpoint(idx, val):
                _mask = ((idx + _kmin) >= 0) & ((idx + _kmin) < kmax)
                return jax.lax.dynamic_update_index_in_dim(
                    val,
                    val[:, idx + _kmin] + _mask * tdi_sum[:, idx],
                    idx + _kmin,
                    axis=1,
                )

            return jax.lax.fori_loop(0, self.n, _update_fpoint, tdi_one)

        def _add_one(idx, tdi):
            _x, _y, _z = self.get_tdi(
                jax.lax.dynamic_slice_in_dim(params, idx, 1, axis=0),
                **kwargs,
            )
            return _add_to(tdi, jnp.vstack([_x, _y, _z]), kmins[idx] - kmin)

        _tdi = jnp.zeros((3, kmax - kmin), dtype=np.complex128)
        _tdi = jax.lax.fori_loop(0, params.shape[0], _add_one, _tdi)
        return tuple(_tdi)


class JaxGBaccurate(JaxGB):
    """Galactic binary less fast but more accurate response generation. Considers
    time-varying arm lengths and reduces edge artefacts (ringing).

    Parameters
    ----------

    window
        Dimensionless extension factor for slow response calculation, for
        windowing purposes. The slow response will be calculated on the domain
        [t0 - window * t_obs, t0 + t_obs + window * t_obs], i.e., the extension is
        symmetric around the time interval of interest [t0, t0 + t_obs].

        window needs to have a strictly positive integer or half-integer value. Note
        that using window also modifies the kmin values:
        df = 1/(t_obs + 2 * window * t_obs).
    """

    def __init__(self, window: float = 0.5, **kwargs) -> None:
        super().__init__(**kwargs)

        # window value (strictly positive integer or half-integer)
        self.window = window

        # length of extended time grid (with windowing)
        self.t_ext = self.t_obs + 2.0 * self.window * self.t_obs

        # time grid with windowing [-window*t_obs, t_obs + 2*window*t_obs]
        self.tm = (
            np.linspace(
                0,
                self.t_ext,
                num=self.n,
                endpoint=False,
            )
            - self.window * self.t_obs
        )

        # time step slow response
        self.dtm = self.t_ext / self.n

        # calculate positions of M = 3 spacecraft at time of emission
        ltt = self.orbits.compute_ltt(self.t0 + self.tm).T  # (6, n), recalculate
        self.position = (
            self.orbits.compute_position(self.t0 + self.tm, [1, 2, 3])
            .swapaxes(0, 1)
            .swapaxes(1, 2)
        )  # (M, 3, n)
        self.positionemi = jnp.array(
            [
                self.orbits.compute_position(self.t0 + self.tm - ltt[0], [2])
                .squeeze()
                .T,
                self.orbits.compute_position(self.t0 + self.tm - ltt[1], [3])
                .squeeze()
                .T,
                self.orbits.compute_position(self.t0 + self.tm - ltt[2], [1])
                .squeeze()
                .T,
                self.orbits.compute_position(self.t0 + self.tm - ltt[3], [3])
                .squeeze()
                .T,
                self.orbits.compute_position(self.t0 + self.tm - ltt[4], [2])
                .squeeze()
                .T,
                self.orbits.compute_position(self.t0 + self.tm - ltt[5], [1])
                .squeeze()
                .T,
            ]
        )  # (6, 3, n)
        self.arm_length = ltt * c
        fstar = c / (self.arm_length * 2 * np.pi)
        self.fstar = fstar[None, :, :]
        # convention: 12, 23, 31, 13, 32, 21

    def get_kmin(self, f0: jax.Array) -> jax.Array:
        """Calculates kmin, the index of the first heterodyned frequency bin in the
        full frequency grid. For windowed response calculations, the observation time
        baseline is extended from `t_obs` to `t_ext = t_obs + 2 * window * t_obs`, so
        the corresponding Fourier frequencies are integer multiples of `1 / t_ext`.
        The nearest frequency bin number to `f0` is therefore `rint(f0 * t_ext)`. This
        bin is centred on the grid of heterodyned Fourier frequencies, of length `n`,
        hence kmin (i.e., the first bin) is shifted by `n/2`.

        If `t_init` is different from `t0`, `f0` needs to be corrected for the linear
        drift proportional to `fdot`.

        Parameters
        ----------
        f0
            Frequency f0 for source(s) (Hz)

        Returns
        -------
        jax.Array of shape (nsrc,)
            Array with kmin values (integer)
        """
        q = jnp.rint(f0 * self.t_ext)
        return jnp.round(q - self.n / 2).astype(np.int32)

    def get_frequency_grid(self, kmin: jax.Array) -> jax.Array:
        """Return the frequency grid associated to the returned signal.

        Parameters
        ----------
        kmin
            Index of the first frequency bin.

        Returns
        -------
        jax.Array of shape (nsrc, n)
            Array with frequency values (in Hz).
        """
        return (1 / self.t_ext) * (jnp.arange(self.n) + kmin[:, None])

    def _construct_slow_part(
        self,
        f0_fdot_phi0: jax.Array,
        k: jax.Array,
        dpdc: tuple[jax.Array, jax.Array],
        pol_tensors: tuple[jax.Array, jax.Array],
    ) -> jax.Array:
        """This method implements the slow, heterodyned response in time domain. It
        accepts the f0, fdot, and phi0 of all the sources, their direction (k),
        the hplus/hcross strain amplitudes (Dp, Dc) and polarization tensors
        (eplus, ecross). The output is an array of shape  (nsrc, 6, n) [other equations
        need to be included here].

        For this implementation, we include tukey windowing immediately here, which
        tapers off the slow response at the edges of the extended domain.

        Parameters
        ----------
        f0_fdot_phi0
            array (nrsc, 3) listing f0, fdot, and phi0 values of all sources.
        k
            array (nsrc, 3) with source k-vectors.
        dpdc
            tuple with hplus and hcross strain amplitudes (complex-valued).
        pol_tensors
            tuple with Eplus, Ecross (polarization tensors).

        Returns
        -------
        jax.Array of shape (nsrc, 6, n)
            Time-domain slow heterodyned response.
        """

        r = jnp.zeros((6, 3, self.n))  # r convention: 12, 23, 31, 13, 32, 21
        r = r.at[0].set(self.position[0] - self.positionemi[0])  ## [3xn]
        r = r.at[1].set(self.position[1] - self.positionemi[1])  ## [3xn]
        r = r.at[2].set(self.position[2] - self.positionemi[2])  ## [3xn]
        r = r.at[3].set(self.position[0] - self.positionemi[3])  ## [3xn]
        r = r.at[4].set(self.position[2] - self.positionemi[4])  ## [3xn]
        r = r.at[5].set(self.position[1] - self.positionemi[5])  ## [3xn]
        r /= self.arm_length[:, None, :]

        # kdotr convention: 12, 23, 31, 13, 32, 21
        kdotr = jnp.dot(k.T, r)
        kdotp = jnp.dot(k.T, self.position)
        kdotp /= c

        f0, fdot, phi0 = f0_fdot_phi0[:, 0], f0_fdot_phi0[:, 1], f0_fdot_phi0[:, 2]
        xi = self.tm - kdotp
        fi = f0[:, None, None] + fdot[:, None, None] * xi
        # Ratio of true frequency to transfer frequency
        fonfs = fi[:, [0, 1, 2, 0, 2, 1]] / self.fstar

        ### compute transfer f-n
        q = jnp.rint(f0 * self.t_ext)  # index of nearest Fourier bin
        df = 2.0 * np.pi * (q / self.t_ext)
        om = 2.0 * np.pi * f0

        # Reduce to (3, 3, n).
        # A = (eplus @ r * r * dp + ecross @ r * r * dc).sum(axis=2)
        eplus, ecross = pol_tensors
        aij = (
            jnp.swapaxes(jnp.dot(eplus, r), 1, 2) * r * dpdc[0][:, None, None, None]
            + jnp.swapaxes(jnp.dot(ecross, r), 1, 2) * r * dpdc[1][:, None, None, None]
        )
        asum = aij.sum(axis=2)  # A convention: 12, 23, 31, 13, 32, 21
        arg_s = (phi0[:, None] + (om - df)[:, None] * self.tm[None, :])[:, None, :] + (
            np.pi * fdot[:, None, None] * (xi**2)
        )
        kdotp = om[:, None, None] * kdotp - arg_s

        # gs convention: 12, 23, 31, 13, 32, 21
        idx = jnp.array([0, 1, 2, 0, 2, 1])
        arg_ij = 0.5 * fonfs * (1 - kdotr)  # note the minus-sign!
        gs = (
            0.5j
            * (f0[:, None, None] + fdot[:, None, None] * self.tm[None, None, :])
            / self.fstar
            * jnp.sinc(arg_ij / np.pi)
            * jnp.exp(-1j * (arg_ij + kdotp[:, idx]))
            * asum
            * jnp.exp(
                -2.0j * np.pi * q[:, None, None] * self.window * self.t_obs / self.t_ext
            )
        )  # last term returns a frequency-dependent sign, related to the windowing

        # finally, apply windowing
        w = tukey(len(self.tm), alpha=0.45) ** 4
        gs *= w[None, None, :]

        return gs  # (nsrc, 6, n)
