#cython: boundscheck=False, wraparound=False, initializedcheck=False

cimport numpy as np
import numpy as np
import cython

from finesse.cymath cimport complex_t
from finesse.cymath.complex cimport conj, creal, cimag, cexp, carg, cabs
from finesse.cymath.complex cimport crotate, cnorm
from finesse.cymath.math cimport sqrt, sin, cos, float_eq
from finesse.cymath.homs cimport in_mask, field_index

from finesse.knm cimport KnmMatrix
from finesse.simulations.basematrix cimport CarrierSignalMatrixSimulation
from finesse.frequency cimport frequency_info_t
from finesse.detectors.workspace cimport (
    DetectorWorkspace,
    MaskedDetectorWorkspace,
    OutputFuncWrapper,
)
from finesse.element cimport BaseCValues

cimport finesse.constants

from libc.math cimport isnan
from libc.stdlib cimport calloc, realloc, free

ctypedef (double*,) ptr_tuple_1
ctypedef (double*, double*) ptr_tuple_2
ctypedef (double*, double*, double*) ptr_tuple_3
ctypedef (double*, double*, double*, double*) ptr_tuple_4
ctypedef (double*, double*, double*, double*, double*) ptr_tuple_5
ctypedef (double*, double*, double*, double*, double*, double*) ptr_tuple_6
ctypedef (double*, double*, double*, double*, double*, double*, double*) ptr_tuple_7
ctypedef (double*, double*, double*, double*, double*, double*, double*, double*) ptr_tuple_8
ctypedef (double*, double*, double*, double*, double*, double*, double*, double*, double*) ptr_tuple_9
ctypedef (double*, double*, double*, double*, double*, double*, double*, double*, double*, double*) ptr_tuple_10
ctypedef (double*, double*, double*, double*, double*, double*, double*, double*, double*, double*, double*) ptr_tuple_11

cdef extern from "constants.h":
    double complex COMPLEX_0

cdef:
    double NAN = 0.0/0.0


### PD0 workspace & output funcs ###

cdef class PD0Workspace(MaskedDetectorWorkspace):
    cdef public:
        int dc_node_id
        Py_ssize_t rhs_index
        Py_ssize_t size
        const complex_t[:,::1] K
        np.ndarray tmp


pd0_DC_output = OutputFuncWrapper.make_from_ptr(c_pd0_DC_output)
cdef c_pd0_DC_output(DetectorWorkspace dws):
    """This expects a `PD0Workspace` input. Fast computation of DC power at a node."""
    cdef:
        PD0Workspace ws = <PD0Workspace> dws
        Py_ssize_t i
        double tmp = 0

    for i in range(ws.rhs_index, ws.rhs_index + ws.size):
        tmp += cnorm(ws.sim.carrier.out_view[i])

    return tmp * ws.sim.model_data.EPSILON0_C * 0.5


pd0_DC_output_segmented = OutputFuncWrapper.make_from_ptr(c_pd0_DC_output_segmented)
cdef c_pd0_DC_output_segmented(DetectorWorkspace dws):
    """This expects a `PD0Workspace` input. Fast computation of DC power at a node
    with segmentation assumed in the photodiode. Where :math:`E` is the vector of
    complex fields, and :math:`K` is a Knm matrix describing the segementation of
    a detector, the output is:
    .. math::
        E^{\ast} \cdot (K \cdot E)
    """
    cdef:
        PD0Workspace ws = <PD0Workspace> dws
        Py_ssize_t i
        complex_t res = 0
        const complex_t[::1] E
        frequency_info_t *f1
        Py_ssize_t Nf = ws.sim.carrier.optical_frequencies.size
    for i in range(Nf):
        f1 = &ws.sim.carrier.optical_frequencies.frequency_info[i]
        rhs_idx = ws.sim.carrier.field_fast(ws.dc_node_id, f1.index, 0)
        E = ws.sim.carrier.out_view[rhs_idx:(rhs_idx + ws.sim.model_data.num_HOMs)]
        # segmented Knm multiply by the array of fields
        tmp = np.dot(ws.K, E)
        # inner product with conjugate
        res += np.vdot(tmp, E)
    # Only select the real part of this
    return res.real * ws.sim.model_data.EPSILON0_C * 0.5


pd0_DC_output_masked = OutputFuncWrapper.make_from_ptr(c_pd0_DC_output_masked)
cdef c_pd0_DC_output_masked(DetectorWorkspace dws):
    """This expects a `PD0Workspace` input. Fast computation of DC power at a node where
    a HOM mask has been applied to the PD."""
    cdef:
        PD0Workspace ws = <PD0Workspace> dws
        Py_ssize_t i, j, k
        Py_ssize_t freq_idx
        double tmp = 0

    for i in range(ws.sim.carrier.optical_frequencies.size):
        freq_idx = ws.rhs_index + i
        for j in range(ws.num_unmasked_HOMs):
            k = freq_idx + ws.unmasked_mode_indices[j]
            tmp += cnorm(ws.sim.carrier.out_view[k])

    return tmp * ws.sim.model_data.EPSILON0_C * 0.5


### PD1 workspace & output funcs ###

cdef class PD1Values(BaseCValues):
    def __init__(self):
        cdef ptr_tuple_2 ptr = (&self.f, &self.phase)
        cdef tuple params = ("f", "phase")
        self.setup(params, sizeof(ptr), <double**>&ptr)


cdef class PD1Workspace(MaskedDetectorWorkspace):
    def __cinit__(self, *args, **kwargs):
        self.num_pre_set_beats = 0
        self.pre_set_beats[0] = NULL
        self.pre_set_beats[1] = NULL

    def __init__(self, owner, sim, *, phase_offset=0, oinfo=None):
        super().__init__(owner, sim, PD1Values(), oinfo=oinfo)
        self.v = <PD1Values>self.values
        self.phase_offset = phase_offset

    cpdef update_beats(self):
        """
        Requires full initialisation and `self.update_parameter_values` calling
        before using
        """
        cdef:
            list beats = []
            frequency_info_t *f1
            frequency_info_t *f2
            Py_ssize_t i, j, n, prev_num_beats
            Py_ssize_t Nf = self.sim.carrier.optical_frequencies.size

        prev_num_beats = self.num_pre_set_beats

        for i in range(Nf):
            f1 = &self.sim.carrier.optical_frequencies.frequency_info[i]
            for j in range(Nf):
                f2 = &self.sim.carrier.optical_frequencies.frequency_info[j]
                if f1.f - f2.f == self.v.f:
                    beats.append((
                        self.sim.carrier.field_fast(self.dc_node_id, f1.index, 0),
                        self.sim.carrier.field_fast(self.dc_node_id, f2.index, 0),
                    ))

        self.num_pre_set_beats = len(beats)

        if self.num_pre_set_beats > 0:
            if prev_num_beats > 0:
                self.pre_set_beats[0] = <Py_ssize_t*>realloc(
                    self.pre_set_beats[0],
                    self.num_pre_set_beats
                    )
                self.pre_set_beats[1] = <Py_ssize_t*>realloc(
                    self.pre_set_beats[1],
                    self.num_pre_set_beats
                    )
            else:
                self.pre_set_beats[0] = <Py_ssize_t*>calloc(
                    self.num_pre_set_beats,
                    sizeof(Py_ssize_t)
                    )
                self.pre_set_beats[1] = <Py_ssize_t*>calloc(
                    self.num_pre_set_beats,
                    sizeof(Py_ssize_t)
                    )

            for n in range(self.num_pre_set_beats):
                self.pre_set_beats[0][n] = beats[n][0]
                self.pre_set_beats[1][n] = beats[n][1]

        elif prev_num_beats > 0:
            self.__dealloc__()

    def __dealloc__(self):
        if self.pre_set_beats[0] != NULL:
            free(self.pre_set_beats[0])
        if self.pre_set_beats[1] != NULL:
            free(self.pre_set_beats[1])


pd1_AC_output = OutputFuncWrapper.make_from_ptr(c_pd1_AC_output)
cdef object c_pd1_AC_output(DetectorWorkspace dws):
    """This expects a `PD1Workspace` input"""
    cdef:
        PD1Workspace ws = <PD1Workspace> dws
        frequency_info_t *f1,
        frequency_info_t *f2
        Py_ssize_t i, j, k
        Py_ssize_t m # mode index
        complex_t tmp
        Py_ssize_t Nf = ws.sim.carrier.optical_frequencies.size
        double demod_phase = (ws.phase_offset + ws.v.phase) * finesse.constants._DEG2RAD
    ws.Aij = 0

    for i in range(Nf):
        f1 = &ws.sim.carrier.optical_frequencies.frequency_info[i]
        for k in range(ws.num_unmasked_HOMs):
            m = ws.unmasked_mode_indices[k]
            tmp = ws.sim.carrier.get_out_fast(ws.dc_node_id, f1.index, m)
            ws.Aij += conj(tmp) * ws.sim.signal.get_out_fast(ws.ac_node_id, f1.audio_upper_index, m)
            # lower sideband conjugate is modelled
            ws.Aij += tmp * ws.sim.signal.get_out_fast(ws.ac_node_id, f1.audio_lower_index, m)

    # Factor of two because of signal scaling
    if ws.output_real:
        tmp = ws.Aij*cexp(-1j * demod_phase)
        return 2 * tmp.real * ws.sim.model_data.EPSILON0_C * 0.5
    else:
        return 2 * ws.Aij * ws.sim.model_data.EPSILON0_C * 0.5


pd1_DC_output = OutputFuncWrapper.make_from_ptr(c_pd1_DC_output)
cdef object c_pd1_DC_output(DetectorWorkspace dws):
    """This expects a `PD1Workspace` input"""
    cdef:
        PD1Workspace ws = <PD1Workspace> dws
        frequency_info_t *f1
        frequency_info_t *f2
        Py_ssize_t i, j, k, n
        Py_ssize_t m # mode index
        complex_t tmp
        Py_ssize_t Nf = ws.sim.carrier.optical_frequencies.size
        double demod_phase = (ws.phase_offset + ws.v.phase) * finesse.constants._DEG2RAD

    ws.Aij = 0

    if ws.num_pre_set_beats > 0:
        for n in range(ws.num_pre_set_beats):
            # Get the RHS view indices
            i = ws.pre_set_beats[0][n]
            j = ws.pre_set_beats[1][n]
            for k in range(ws.num_unmasked_HOMs):
                m = ws.unmasked_mode_indices[k]
                ws.Aij += ws.sim.carrier.out_view[i+m] * conj(ws.sim.carrier.out_view[j+m])
    else:
        # Do dumb thing and check everything if we have/can not optimise
        for i in range(ws.sim.carrier.optical_frequencies.size):
            f1 = &ws.sim.carrier.optical_frequencies.frequency_info[i]
            for j in range(ws.sim.carrier.optical_frequencies.size):
                f2 = &ws.sim.carrier.optical_frequencies.frequency_info[j]
                if float_eq(f1.f - f2.f, ws.v.f):
                    for k in range(ws.num_unmasked_HOMs):
                        m = ws.unmasked_mode_indices[k]
                        ws.Aij += (ws.sim.carrier.get_out_fast(ws.dc_node_id, f1.index, m) *
                            conj(ws.sim.carrier.get_out_fast(ws.dc_node_id, f2.index, m)))

    if demod_phase:
        tmp = ws.Aij*cexp(-1j * demod_phase)
    else:
        tmp = ws.Aij

    if ws.output_real:
        return tmp.real * ws.sim.model_data.EPSILON0_C * 0.5
    else:
        return tmp * ws.sim.model_data.EPSILON0_C * 0.5


### PD2 workspace & output funcs ###

cdef class PD2Values(BaseCValues):
    def __init__(self):
        cdef ptr_tuple_4 ptr = (&self.f1, &self.phase1, &self.f2, &self.phase2)
        cdef tuple params = ("f1", "phase1", "f2", "phase2")
        self.setup(params, sizeof(ptr), <double**>&ptr)


cdef class PD2Workspace(MaskedDetectorWorkspace):
    def __init__(self, owner, sim):
        super().__init__(owner, sim, PD2Values())
        self.v = <PD2Values>self.values

pd2_DC_output = OutputFuncWrapper.make_from_ptr(c_pd2_DC_output)
cdef object c_pd2_DC_output(DetectorWorkspace dws):
    """This expects a `PD2Workspace` input"""
    cdef:
        PD2Workspace ws = <PD2Workspace> dws
        frequency_info_t *f1
        frequency_info_t *f2
        Py_ssize_t i, j, k, n
        Py_ssize_t m # mode index
        complex_t tmp, phs
        double f_plus = ws.v.f1 + ws.v.f2
        double f_minus = ws.v.f2 - ws.v.f1
        Py_ssize_t Nf = ws.sim.carrier.optical_frequencies.size

    ws.z1 = 0
    ws.z2 = 0

    for i in range(Nf):
        f1 = &ws.sim.carrier.optical_frequencies.frequency_info[i]
        for j in range(Nf):
            f2 = &ws.sim.carrier.optical_frequencies.frequency_info[j]
            if f1.f - f2.f == f_plus:
                for k in range(ws.num_unmasked_HOMs):
                    m = ws.unmasked_mode_indices[k]
                    ws.z1 += (ws.sim.carrier.get_out_fast(ws.dc_node_id, f1.index, m) *
                        conj(ws.sim.carrier.get_out_fast(ws.dc_node_id, f2.index, m)))

            if f1.f - f2.f == f_minus:
                for k in range(ws.num_unmasked_HOMs):
                    m = ws.unmasked_mode_indices[k]
                    ws.z2 += (ws.sim.carrier.get_out_fast(ws.dc_node_id, f1.index, m) *
                        conj(ws.sim.carrier.get_out_fast(ws.dc_node_id, f2.index, m)))

    if ws.v.phase1:
        phs = cexp(-1j * ws.v.phase1 * finesse.constants._DEG2RAD)
        tmp = ws.z1 * phs + ws.z2 * conj(phs)
    else:
        tmp = ws.z1 + ws.z2

    if ws.v.phase2:
        tmp *= cexp(-1j * ws.v.phase2 * finesse.constants._DEG2RAD)

    # Factor 1/2 from demod gain
    # And a mystery factor of two to match up with v2
    if ws.output_real:
        return tmp.real * ws.sim.model_data.EPSILON0_C * 0.5 * 0.5
    else:
        return tmp * ws.sim.model_data.EPSILON0_C * 0.5 * 0.5


pd2_AC_output = OutputFuncWrapper.make_from_ptr(c_pd2_AC_output)
cdef object c_pd2_AC_output(DetectorWorkspace dws):
    """This expects a `PD2Workspace` input"""
    cdef:
        PD2Workspace ws = <PD2Workspace> dws
        frequency_info_t *f1
        frequency_info_t *f2
        Py_ssize_t i, j, k
        Py_ssize_t m # mode index
        complex_t tmp, phs, car1, car2
        double df
        double f_plus = ws.v.f1 + ws.v.f2
        double f_minus = ws.v.f2 - ws.v.f1
        Py_ssize_t Nf = ws.sim.carrier.optical_frequencies.size

    ws.z1 = 0
    ws.z2 = 0

    # Here we loop over each of the carrier beats and apply the relevant signal sideband
    # beats. This should be the equivalent code in readout elements but they fill the maxtrix
    # instead.
    for i in range(Nf):
        f1 = &ws.sim.carrier.optical_frequencies.frequency_info[i]
        for j in range(Nf):
            f2 = &ws.sim.carrier.optical_frequencies.frequency_info[j]
            df = f1.f - f2.f
            if float_eq(df, -ws.v.f1):
                for k in range(ws.num_unmasked_HOMs):
                    m = ws.unmasked_mode_indices[k]
                    car1 = ws.sim.carrier.get_out_fast(ws.dc_node_id, f1.index, m)
                    ws.z2 += car1 * ws.sim.signal.get_out_fast(ws.ac_node_id, f2.audio_lower_index, m)

                    car2 = ws.sim.carrier.get_out_fast(ws.dc_node_id, f2.index, m)
                    ws.z2 += ws.sim.signal.get_out_fast(ws.ac_node_id, f1.audio_upper_index, m) * conj(car2)

            elif float_eq(df, ws.v.f1):
                for k in range(ws.num_unmasked_HOMs):
                    m = ws.unmasked_mode_indices[k]
                    car1 = ws.sim.carrier.get_out_fast(ws.dc_node_id, f1.index, m)
                    ws.z1 += car1 * ws.sim.signal.get_out_fast(ws.ac_node_id, f2.audio_lower_index, m)

                    car2 = ws.sim.carrier.get_out_fast(ws.dc_node_id, f2.index, m)
                    ws.z1 += ws.sim.signal.get_out_fast(ws.ac_node_id, f1.audio_upper_index, m) * conj(car2)

    phs = cexp(-1j * ws.v.phase1 * finesse.constants._DEG2RAD)
    # 1/2 from cos(X)*cos(Y) = 1/2(cos(X+Y)+cos(X-Y)) expansion
    tmp = (ws.z1 * phs + ws.z2 * conj(phs)) * 0.5

    if ws.v.phase2:
        tmp *= cexp(-1j * ws.v.phase2 * finesse.constants._DEG2RAD)

    # Factor of two because of signal scaling
    if ws.output_real:
        return 2 * tmp.real * ws.sim.model_data.EPSILON0_C * 0.5
    else:
        return 2 * tmp * ws.sim.model_data.EPSILON0_C * 0.5


# ### Generic (slower) PD workspace & output funcs ###
# TODO ddb this all needs changing over to use the CarrierSignalMatrixSimulation
# class, which as none of this is used anywhere really I can't be bothered doing
# right now, as I'm not sure if any of this have any valid test cases
# cdef class PDWorkspace(DetectorWorkspace):
#     cdef:
#         public BaseSimulation DC
#         public BaseSimulation AC
#         public Py_ssize_t dc_node_id
#         public Py_ssize_t ac_node_id
#         public int[:, ::1] homs
#         const int[:, ::1] mask
#         bint mode_is_network
#         double[:, :, :, ::1] beat_factors
#         Py_ssize_t num_demod
#         Py_ssize_t num_mixes
#         double[:, ::1] fmix_view
#         double EPSILON0_C
#         complex_t[::1] zview

#     def __init__(self,
#                 owner,
#                 sims,
#                 int N_demods,
#                 double EPSILON0_C,
#                 const int[:, ::1] mask,
#                 bint mode_is_network,
#                 beat_factors=None):

#         super().__init__(owner, sims)
#         self.EPSILON0_C = EPSILON0_C
#         self.mask = mask
#         self.mode_is_network = mode_is_network
#         self.beat_factors = beat_factors
#         # Computing values and allocating arrays that will be used throughout
#         # the simulation output calculations
#         self.num_demod = N_demods
#         self.num_mixes = 1 + self.num_demod * (self.num_demod - 1) / 2
#         self.fmix_view = np.ones((self.num_mixes, self.num_demod))
#         self.zview = np.zeros(self.num_mixes, dtype=np.complex128)
#         # TODO ddb - I think the mix table only needs to be made once
#         # Might need to update if the demod frequencies change
#         _create_mix_table(self.num_demod, self.fmix_view)

# def mixer_pd_output(PDWorkspace ws,
#                     double[::1] freqs,
#                     double[::1] phases,
#                     long long[::1] max_phase_indices
# ):
#     cdef:
#         Py_ssize_t mix_index
#         double f_ref
#         complex_t dm_amp
#         complex_t power
#         # Flag for whether this pd is computing purely a transfer function
#         # TODO: Is this really needed, or just an artifact from the past?
#         int transfer = 0

#     for mix_index in range(ws.num_mixes):
#         f_ref = _get_reference_freq(ws.num_demod, mix_index, ws.fmix_view, freqs)
#         if ws.beat_factors is not None:
#             ws.zview[mix_index] = compute_amplitude_beats(ws, &transfer, f_ref)
#         else:
#             ws.zview[mix_index] = compute_amplitude(ws, &transfer, f_ref)

#     dm_amp = _set_demodulation_phase(
#         ws.num_demod, ws.zview, ws.fmix_view, ws.num_mixes, phases, max_phase_indices
#     )
#     power = _get_demodulation_signal(
#         ws.num_demod, dm_amp, phases, max_phase_indices, ws.EPSILON0_C, ws.mode_is_network, transfer
#     )

#     return power


# cdef _create_mix_table(Py_ssize_t num_demod, double[:, ::1] fmix_view):
#     cdef:
#         int num_mixes = 1

#         Py_ssize_t demod_index_outer, demod_index_inner, mix_permutation_index

#         np.ndarray[double, ndim=1] fmix_tmp = np.ones(num_demod)
#         double[::1] fmix_tmp_view = fmix_tmp

#     for demod_index_outer in range(1, num_demod):
#         for mix_permutation_index in range(demod_index_outer, 0, -1):
#             num_mixes += 1
#             fmix_tmp_view[mix_permutation_index - 1] *= -1

#             for demod_index_inner in range(num_demod):
#                 fmix_view[num_mixes - 1][demod_index_inner] = fmix_tmp_view[demod_index_inner]


# def dc_pd_output(PDWorkspace ws):
#     cdef:
#         complex_t power = COMPLEX_0
#         int transfer = 0

#     if ws.beat_factors is not None:
#         power = compute_amplitude_beats(ws, &transfer)
#     else:
#         power = compute_amplitude(ws, &transfer)

#     return power.real * ws.EPSILON0_C * 0.5


# cdef double _get_reference_freq(
#     Py_ssize_t num_demod,
#     Py_ssize_t mix_index,
#     double[:, ::1] fmix_view,
#     double[::1] freqs
# ) nogil:
#     cdef:
#         double f_ref = 0.0
#         Py_ssize_t demod_index

#     for demod_index in range(num_demod):
#         f_ref += fmix_view[mix_index][demod_index] * freqs[demod_index]

#     return f_ref



# cdef complex_t _set_demodulation_phase(
#     Py_ssize_t num_demod,
#     complex_t[::1] zview,
#     double[:, ::1] fmix_view,
#     Py_ssize_t num_mixes,
#     double[::1] phases,
#     long long[::1] max_phase_indices
# ) nogil:
#     cdef:
#         Py_ssize_t demod_index, mix_index
#         complex_t demod_amp = COMPLEX_0

#         complex_t z1, z2
#         double phi_max

#         double* phi = <double*> calloc(num_mixes, sizeof(double))

#     for demod_index in range(num_demod - 1):
#         z1 = COMPLEX_0
#         z2 = COMPLEX_0
#         # max phase
#         if in_index_array(max_phase_indices, demod_index):
#             for mix_index in range(num_mixes):
#                 if fmix_view[mix_index][demod_index] < 0.0:
#                     z1 += zview[mix_index]
#                 else:
#                     z2 += zview[mix_index]

#             phi_max = 0.5 * (carg(z2) - carg(z1))
#             phases[demod_index] = phi_max

#             for mix_index in range(num_mixes):
#                 phi[mix_index] -= phi_max * fmix_view[mix_index][demod_index]
#         # user defined phase
#         else:
#             for mix_index in range(num_mixes):
#                 phi[mix_index] -= (
#                     phases[demod_index] * fmix_view[mix_index][demod_index]
#                 )

#     for mix_index in range(num_mixes):
#         demod_amp += crotate(zview[mix_index], phi[mix_index])

#     free(phi)

#     return demod_amp


# cdef complex_t _get_demodulation_signal(
#     Py_ssize_t num_demod,
#     complex_t demod_amp,
#     double[::1] phases,
#     long long[::1] max_phase_indices,
#     double epsilon0_c,
#     bint mode_is_network,
#     int transfer
# ) nogil:
#     cdef:
#         complex_t z = demod_amp

#     if not mode_is_network:
#         if in_index_array(max_phase_indices, num_demod - 1):
#             z = cabs(demod_amp)
#         else:
#             z = (
#                 creal(demod_amp) * cos(phases[num_demod - 1])
#                 + cimag(demod_amp) * sin(phases[num_demod - 1])
#             )

#     z *= 0.5**(num_demod - 1)

#     if transfer == 2:
#         z *= 2

#     return z * epsilon0_c * 0.5

# cdef complex_t compute_amplitude(
#     PDWorkspace ws,
#     int *transfer,
#     double f_ref=NAN
# ):
#     cdef:
#         complex_t amplitude = COMPLEX_0

#         Py_ssize_t field_idx
#         Py_ssize_t N = ws.homs.shape[0]
#         int n, m
#         bint mixed = False

#     if not isnan(f_ref):
#         mixed = True

#     for field_idx in range(N):
#         n = ws.homs[field_idx][0]
#         m = ws.homs[field_idx][1]

#         if in_mask(n, m, ws.mask):
#             continue

#         amplitude += compute_mode_amplitude(
#                 ws.sim.carrier,
#                 ws.sim.signal,
#                 ws.dc_node_id,
#                 ws.ac_node_id,
#                 field_idx,
#                 field_idx,
#                 transfer,
#                 f_ref,
#                 mixed
#         )

#     return amplitude


# cdef complex_t compute_amplitude_beats(
#     PDWorkspace ws,
#     int *transfer,
#     double f_ref=NAN
# ):
#     cdef:
#         complex_t amplitude = COMPLEX_0

#         Py_ssize_t N = ws.homs.shape[0]
#         Py_ssize_t field_idx_outer, field_idx_inner
#         int n1, m1, n2, m2
#         double pd_factor
#         bint mixed = False

#     if ~isnan(f_ref):
#         mixed = True

#     for field_idx_outer in range(N):
#         n1 = ws.homs[field_idx_outer][0]
#         m1 = ws.homs[field_idx_outer][1]

#         # TODO ddb : probably worth changing ws.mask to a list of included
#         # mode indicies rather than checking all the time
#         if in_mask(n1, m1, ws.mask):
#             continue

#         for field_idx_inner in range(N):
#             n2 = ws.homs[field_idx_inner][0]
#             m2 = ws.homs[field_idx_inner][1]

#             pd_factor = ws.beat_factors [n1][m1][n2][m2]
#             if pd_factor == 0.0:
#                 continue

#             amplitude += compute_mode_amplitude(
#                     ws.sim.carrier,
#                     ws.sim.signal,
#                     ws.dc_node_id,
#                     ws.ac_node_id,
#                     field_idx_outer,
#                     field_idx_inner,
#                     transfer,
#                     f_ref,
#                     mixed
#             ) * pd_factor

#     return amplitude


# cdef complex_t compute_mode_amplitude(
#     BaseSimulation DC,
#     BaseSimulation AC,
#     Py_ssize_t dc_node_id,
#     Py_ssize_t ac_node_id,
#     Py_ssize_t field_idx_outer,
#     Py_ssize_t field_idx_inner,
#     int *transfer,
#     double f_ref,
#     bint mixed
# ) nogil:
#     cdef:
#         Py_ssize_t i, j

#         complex_t inner, outer
#         complex_t amplitude = COMPLEX_0

#         frequency_info_t f1, f2

#     for i in range(DC.num_frequencies):
#         f1 = DC.frequency_info[i]
#         outer = DC.get_out_fast(dc_node_id, f1.index, field_idx_outer)
#         for j in range(DC.num_frequencies):
#             f2 = DC.frequency_info[j]
#             if mixed:
#                 if not float_eq(f1.f - f2.f, f_ref):
#                     continue
#             else:
#                 if not float_eq(f1.f, f2.f):
#                     continue

#             transfer[0] |= 1
#             inner = DC.get_out_fast(dc_node_id, f2.index, field_idx_inner)
#             amplitude += outer * conj(inner)

#         if not mixed or AC is None:
#             continue

#         for j in range(AC.num_frequencies):
#             f2 = AC.frequency_info[j]

#             if not float_eq(f1.f - f2.f, f_ref):
#                 continue

#             transfer[0] |= 2

#             inner = AC.get_out_fast(ac_node_id, f2.index, field_idx_inner)

#             if f2.audio_order < 0:
#                 amplitude += outer * inner
#             else:
#                 amplitude += outer * conj(inner)

#     if AC is not None and mixed:
#         for i in range(AC.num_frequencies):
#             f1 = AC.frequency_info[i]
#             outer = AC.get_out_fast(ac_node_id, f1.index, field_idx_outer)
#             if f1.audio_order < 0:
#                 outer = conj(outer)
#             for j in range(DC.num_frequencies):
#                 f2 = DC.frequency_info[j]

#                 if not float_eq(f1.f - f2.f, f_ref):
#                     continue

#                 transfer[0] |= 2

#                 inner = DC.get_out_fast(dc_node_id, f2.index, field_idx_inner)

#                 amplitude += outer * conj(inner)

#     return amplitude


# cdef bint in_index_array(long long[::] indices, Py_ssize_t idx) nogil:
#     cdef:
#         Py_ssize_t i
#         Py_ssize_t N = indices.shape[0]

#     for i in range(N):
#         if indices[i] == idx:
#             return True

#     return False
