# Copyright© 2017-2023 Quantum Interface (quantuminterface@ipe.kit.edu)
# Richard Gebauer, IPE, Karlsruhe Institute of Technology
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""This module contains the driver to control a signal generator on the QiController.

Signal generators are a part of the digital unit cell (see :class:`qiclib.hardware.unitcell`).
Each cell contains two signal generator, one for control and one for readout pulses.

The signal generator handles the generation of pulses in the complex-valued baseband
(as in-phase and quadrature components, I/Q) with a common frequency and well-defined,
stable phase relation. You can load up to 15 different pulse configurations into the
module, so-called trigger sets, each addressed by its index. The sequencer uses this
value to start the generation of the corresponding pulse in the signal generator.

.. todo:: Figure for structure of signal generator.

Pulses are defined by their envelope which is then multiplied within the module by a
reference oscillation (NCO) with the given internal frequency. It acts as global phase
reference and defines both pulse phase and frequency. This is necessary to obtain a
stable phase relation between consecutive pulses, which is crucial to define the
rotation axis for control pulses in the Bloch sphere. For this reason, directly storing
the full pulse including the oscillation is not possible.

For each trigger set, the following properties can be individually specified:

* The duration of the pulse in cycles. As four samples are handled each cycle (4ns),
    sample-accurate, sub-cycle precision (1ns) of the pulse can be achieved by padding the
    envelope values with zeros to fill up a full cycle.
* The phase offset of the pulse relative to the global phase reference (NCO) inside the
    signal generator module. Using this property, the rotation axis on the equator of the
    Bloch sphere can be selected.
* I and Q envelope sample values of the pulse.
* Option to hold the last envelope sample value until another trigger is received. This
    results in continuous wave (CW) operation and can also be used to create variable length
    pulse shapes like a trapezoid.
* Option to persist the phase offset in the global phase reference. This enables the
    user to perform a virtual Z rotation.

The envelope values are stored inside the 8kB large envelope memory. It can store 4096
real-valued samples (16bit each) corresponding to up to about 4µs of pulse data if I and
Q share the same envelope samples. Internally, trigger sets define the address offsets
to the envelope values in the memory. These can be the same address if the envelopes for
I and Q are identical, or different, if not (e.g. for DRAG pulses).

When a trigger arrives and selects one of the 15 trigger sets, the configuration of this
trigger set is activated and the execution of the sample players is started. After the
duration of the pulse has passed, the sample players either turn off the envelope output
or keep the last sample value, depending on the option given in the trigger set.

The I and Q envelope values are then forwarded to a complex multiplier. There, they will
be multiplied by the oscillating complex quadrature signal of the reference oscillator
(NCO). The oscillator acts as global phase reference for all pulses generated by this
signal generator. It has a configurable frequency and a global phase offset. The latter
can be temporarily shifted for a single pulse by the configuration in the trigger sets.
It can also be intentionally changed to perform arbitrary virtual Z rotations.
Synchronization of the frequency references between multiple signal generators is
performed automatically at the beginning of each experiment by simultaneously resetting
the internal phase of all oscillators to zero.

At the output of the multiplier, one obtains the digital pulse in the baseband. Before
it is forwarded towards the DAC, the I and Q quadratures can be independently calibrated
by a scaling factor. Factors from zero to two (excluding) are possible with high
precision. If the multiplication product exceeds the valid range of the 16bit data
width, it will be clipped and a saturation flag will indicate this. A corresponding
message will be generated within the status report until it is manually cleared.
"""

from __future__ import annotations

import warnings
from typing import Any

import numpy as np
import numpy.typing as npt

import qiclib.packages.grpc.datatypes_pb2 as dt
import qiclib.packages.grpc.pulsegen_pb2 as proto
import qiclib.packages.grpc.pulsegen_pb2_grpc as grpc_stub
import qiclib.packages.utility as util
from qiclib.hardware.platform_component import (
    PlatformComponent,
    platform_attribute,
    platform_attribute_collector,
)
from qiclib.packages.servicehub import ServiceHubCall


@platform_attribute_collector
class PulseGen(PlatformComponent):
    """A signal generator within a digital unit cell of the QiController.

    Attributes
    ----------
    local_oscillator: QkitInstrument, optional
        Reference to the driver of an external local oscillator. If this is set, the
        `frequency` property will also consider the external local oscillator and adapt
        it accordingly.

        The instrument class instance must provide the following three methods:

        * get_frequency()
        * set_frequency(frequency_in_hertz: float)
        * set_status(enable: bool)
    invert_sideband: bool, optional
        If the sidebands of the IQ mixer connected to the signal generator are
        inverted (by default False).

        This value needs to be adapted by the user if a local oscillator is specified in
        the signal generator using `PulseGen.local_oscillator`. It is only necessary to
        correctly calculate the resulting frequency when using the `frequency` property.
    """

    def __init__(
        self, name: str, connection, controller, qkit_instrument=True, index=0
    ):
        super().__init__(name, connection, controller, qkit_instrument)
        self._stub = grpc_stub.PulseGenServiceStub(self._conn.channel)
        self._index = index
        self._component = dt.EndpointIndex(value=self._index)

        # TODO Numbers to appropriate constants
        self._triggerset: list[TriggerSet | None] = [None]
        self._triggerset.extend(
            TriggerSet(
                f"{self.name} Trigger {i:02d}",
                self._conn,
                self._qip,
                qkit_instrument=False,
                triggerset=i,
                pulsegen=self,
            )
            for i in range(1, 15)
        )
        # Triggerset 15 is designated for NCO Sync and will by default
        # be configured as that. If access is nevertheless really necessary:
        self._triggerset_nco = TriggerSet(
            f"{self.name} Trigger 15 NCO",
            self._conn,
            self._qip,
            qkit_instrument=False,
            triggerset=15,
            pulsegen=self,
        )

        self.local_oscillator = None
        """
        Reference to the driver of an external local oscillator. If this is set, the
        `frequency` property will also consider the external local oscillator and adapt
        it accordingly.

        The instrument class instance must provide the following three methods:

        * get_frequency()
        * set_frequency(frequency_in_hertz: float)
        * set_status(enable: bool)

        """

        self.invert_sideband = False  # TODO better name
        """
        If the sidebands of the IQ mixer connected to the signal generator are
        inverted (by default False).

        This value needs to be adapted by the user if a local oscillator is specified in
        the signal generator using `PulseGen.local_oscillator`. It is only necessary to
        correctly calculate the resulting frequency when using the `frequency` property.

        """

        self._pulses: dict[str, dict[str, Any]] = {}

    @property
    def triggerset(self) -> list[TriggerSet | None]:
        """A list containing the 14 triggersets (indices 1-14) of the signal generator.

        The index corresponds to the trigger value as specified within the sequencer.
        Therefore, index 0 is `None` as this value is reserved for performing no
        operation.

        Triggerset 15 is reserved as NCO sync and cannot be accessed using this list.

        For information about the configuration possibilities of the trigger sets, please
        refer to the `TriggerSet` class.
        """
        return self._triggerset

    ##################################################################
    # High level commands that allow to load a pulse dictionary      #
    ##################################################################

    @property
    @platform_attribute
    def pulses(self) -> dict[str, dict[str, Any]]:
        """A dictionary of pulses loaded into the signal generator.

        Each pulse description (dictionary value) has a name (dictionary key).
        The name can later be used to identify the parameters of the loaded pulse.

        Each pulse description must contain at least the following entries:

         * `'trigger'`: The number of the trigger set where to write the pulse.
         * `'length'`: The duration of the pulse.

        It may further contain the following optional entries:

         * `'amplitude'`: The amplitude factor of the envelope (defaults to 1).
         * `'phase'`: The phase of the pulse (defaults to 0).
         * `'hold'`: If the last value of the pulse should be hold until the next trigger
           (defaults to False).
         * `'envelope'`: An array containing the envelope values of the pulse
           (defaults to a rectangular pulse if unset).
           If this is set, the `'length'` entry will be ignored. It is nevertheless
           advised to set it correctly, as other routines might depend on the length
           being correctly specified.

        .. note::
            This dictionary is only meaningful if `load_pulses()` or direct assignment to
            this property is used to configure the signal generator.
            There is no automatic update if parameters are set or changed directly.

        Example
        -------

        .. code-block:: python

            # Load pulses into the manipulation signal generator of the first unit cell:
            qic.cell[0].manipulation.pulses = {
                "Rx pi": {
                    "trigger": 1,
                    "length": 200e-9,
                },
                "Rx pi/2": {
                    "trigger": 2,
                    "length": 200e-9,
                    "amplitude": 0.5,
                },
                "Ry pi": {
                    "trigger": 3,
                    "length": 200e-9,
                    "phase": np.pi / 2,
                },
                "CW on": {
                    "trigger": 14,
                    "length": 4e-9,
                    "hold": True,
                },
                "CW off": {
                    "trigger": 15,
                    "length": 4e-9,
                    "amplitude": 0,
                },
            }
        """
        return self._pulses

    @pulses.setter
    def pulses(self, pulses: dict[str, dict[str, Any]]):
        """loads the given pulses

        :param pulses:
            Dictionary containing the name as key and pulse as values
        """
        self.load_pulses(pulses)

    def load_pulses(
        self, pulse_dict: dict[str, dict[str, Any]], drag_amplitude: float | None = None
    ):
        """Loads a set of pulses, given as dictionary into the signal generator.

        For information about the required structure of the pulse dictionary argument,
        please refer to the `PulseGen.pulses` property that holds the last loaded
        dictionary. Direct assignment to this property is equivalent to calling this
        method.

        .. note::
            This will erase/overwrite all existing pulses on the signal generator!

        :param pulse_dict:
            A dictionary holding key value pairs of name and pulse descriptions.
            See `PulseGen.pulses` for information about the object structure.
        :param drag_amplitude:
            If given, square shaped pulses will be replaced by DRAG pulses with given
            amplitude factor.
        """
        # Reset signal generator envelope memory
        self.reset_env_mem()

        # Loading pulse shapes into the signal generator
        for pulse in pulse_dict.values():
            self.triggerset[pulse["trigger"]].load_pulse(
                pulseform=pulse.get("amplitude", 1)
                * (
                    pulse.get("envelope", None)
                    or util.generate_pulseform(
                        pulse["length"], drag_amplitude=drag_amplitude
                    )
                ),
                phase=pulse.get("phase", 0),
                hold=pulse.get("hold", False),
            )

        # Persist the pulse dictionary for later use
        self._pulses = pulse_dict

    ##################################################################
    # High level commands that take into account the analog frontend #
    ##################################################################

    @property
    @platform_attribute
    def frequency(self) -> float:
        r"""The resulting RF frequency of the output after IQ mixing with an external
        local oscillator in Hz. This only works when the local oscillator is given as
        the `local_oscillator` attribute.

        Then, depending on the :python:`invert_sideband` attribute, this will either be

         - :math:`f = f_\mathrm{LO} + f_\mathrm{IF}` when :python:`False` (default), or
         - :math:`f = f_\mathrm{LO} - f_\mathrm{IF}` when :python:`invert_sideband` is :python:`True`.

        If no :python:`local_oscillator` is set, :math:`f_\mathrm{LO}` defaults to 0 and the
        (internal) intermediate frequency will be returned. Setting the intermediate
        frequency is only possible using `internal_frequency` or `if_frequency`.

        :raises AttributeError:
            when setting the frequency but no `local_oscillator` is specified.
        """
        try:
            lo_frequency = self.local_oscillator.get_frequency()
        except AttributeError:
            lo_frequency = 0

        if self.invert_sideband:
            return lo_frequency - self.internal_frequency
        return lo_frequency + self.internal_frequency

    @frequency.setter
    def frequency(self, frequency: float):
        """sets the frequency of the local oscillator (LO)

        :param frequency: float
            the frequency

        :raises AttributeError:
            if no LO is specified
        """
        if self.invert_sideband:
            lo_frequency = frequency + self.internal_frequency
        else:
            lo_frequency = frequency - self.internal_frequency

        try:
            self.local_oscillator.set_frequency(lo_frequency)
        except AttributeError as e:
            raise AttributeError(
                "Cannot set frequency if no LO is specified. "
                "To set the IF frequency use if_frequency instead."
            ) from e

    @property
    @platform_attribute
    def if_frequency(self) -> float:
        """The IF frequency of the QiController output before IQ mixing in Hz.

        In contrast to `PulseGen.internal_frequency`, the RF frequency is kept constant
        when this value is changed. This means that the local oscillator frequency will
        be adapted accordingly (if one is specified).
        """
        return self.internal_frequency

    @if_frequency.setter
    def if_frequency(self, if_frequency: float):
        """Sets the internal frequency of the output

        :param if_frequency:
            the internal frequency
        """
        # Remember RF frequency that should stay constant
        rf_frequency = self.frequency
        # Set the new IF frequency in the module
        self.internal_frequency = if_frequency
        # Try to adapt RF frequency
        try:
            self.frequency = rf_frequency
        except AttributeError:
            # Probably no LO is specified, in that case we cannot adapt the RF frequency
            pass

    def output(self, enable: bool):
        """Enables or disables the output of the PulseGen and the local oscillator (if present)."""
        self.nco_enable(enable)
        try:
            self.local_oscillator.set_status(enable)
        except AttributeError:
            pass

    ##################################################################
    # Internal commands modifying and accessing the hardware module  #
    ##################################################################

    @ServiceHubCall(errormsg="Error resetting the signal generator")
    def reset(self):
        """Resets the module including all configuration values and trigger sets."""
        self._stub.Reset(self._component)

    @ServiceHubCall(errormsg="Error resetting the signal generator status flags")
    def reset_status_flags(self):
        """Resets the status flags of the signal generator."""
        self._stub.ResetStatusFlags(self._component)

    @ServiceHubCall(errormsg="Error obtaining the signal generator status flags")
    def get_status_flags(self) -> proto.StatusFlags:
        """Obtains the status flags of the signal generator.

        The following flags will be returned as object properties:

        - ``saturation``: If the amplitude calibration resulted in signal saturation at
            signal generator output (output was compressed/clipped).

        """
        return self._stub.GetStatusFlags(self._component)

    def check_status(self, raise_exceptions: bool = True) -> bool:
        """Check if any status flags are set and report them, if any.

        :param raise_exceptions: bool, optional
            if assigned status flags will result in an exception or just in a warning
            output, by default True.

        :return:
            if everything is fine and no status flags have been raised.

        :raises Warning:
            if a status flag was set and `raise_exceptions` was set to True.
        """
        status = self.get_status_flags()

        if status.saturation:
            self.reset_status_flags()
            if raise_exceptions:
                raise Warning(
                    "Amplitude calibration in signal generator lead to saturation!"
                )

            warnings.warn(
                "Amplitude calibration in signal generator lead to saturation!"
            )

        return not status.saturation

    @ServiceHubCall(errormsg="Error changing NCO status of Pulse Generator")
    def nco_enable(self, enable=True):
        """Enables or disables the internal oscillator and thus the module output.

        .. note::
            By default, the NCO is disabled meaning that no output can be observed when
            triggering the module without turning it on.

        :param enable:
            If the oscillator should be enabled or disabled.
        """
        # self._stub.EnableNCO(proto.NCOEnable(cindex=self._component, value=enable))

    @ServiceHubCall(
        errormsg="Failed to execute a manual trigger for the signal generator"
    )
    def trigger_manually(self, triggerset: int):
        """Manually triggers the Pulse Generator module with the given triggerset.

        .. note::
            Be aware that the internal oscillator (NCO) needs to be enabled for the output
            to be visible, see :meth:`PulseGen.nco_enable`.

        :param triggerset: number
            The trigger set to play (can be a number between 1 and 15).
        """
        self._stub.TriggerManually(
            proto.IndexSet(
                cindex=self._component, tindex=proto.TriggerSetIndex(value=triggerset)
            )
        )

    @property
    @ServiceHubCall
    def internal_frequency(self) -> float:
        """The internal frequency in Hz used for the generated microwave pulses."""
        return self._stub.GetNCOFrequency(self._component).value

    @internal_frequency.setter
    @ServiceHubCall
    def internal_frequency(self, internal_frequency: float):
        self._stub.SetNCOFrequency(
            proto.Frequency(cindex=self._component, value=internal_frequency)
        )

    @property
    @platform_attribute
    @ServiceHubCall(
        errormsg="Could not read the amplitude calibration of signal generator"
    )
    def amplitude_calibration(self) -> tuple[float, float]:
        """A tuple containing calibration values of the I and Q output amplitudes.

        Valid values are in the range of [0, 2) representing multiplication factors.

        Setting this value to high (typically >1) can lead to saturation and clipping
        effects of the output. This will also be indicated by the status flags, see
        :meth:`PulseGen.get_status_flags`.
        """
        response = self._stub.GetAmplitudeCalibration(self._component)
        return response.i_factor, response.q_factor

    @amplitude_calibration.setter
    @ServiceHubCall(errormsg="Failed to set amplitude calibration of signal generator")
    def amplitude_calibration(self, amplitude_calibration: tuple[float, float]):
        if len(amplitude_calibration) != 2:
            raise ValueError("(I, Q) tuple for calibration factors expected.")

        self._stub.SetAmplitudeCalibration(
            proto.AmplitudeCalibration(
                cindex=self._component,
                i_factor=amplitude_calibration[0],
                q_factor=amplitude_calibration[1],
            )
        )

    @property
    @platform_attribute
    @ServiceHubCall
    def global_phase(self) -> float:
        """The global reference phase of the generator.

        It can be shifted by trigger sets, e.g., to realize virtual Z gates.
        """
        return self._stub.GetGlobalPhase(self._component).value

    @ServiceHubCall(errormsg="Could not reset envelope memory of the signal generator")
    def reset_env_mem(self):
        """Resets the envelope memory erasing all stored pulse forms."""
        self._stub.ResetEnvelopeMemory(self._component)

    def get_configuration_dict(self) -> dict[str, Any]:
        """Returns a dictionary containing all configuration values which will not be
        automatically overwritten for all experiments (esp. calibration values).

        :return:
            containing configuration and calibration values of this signal generator.
        """
        return {
            "internal_frequency": self.internal_frequency,
            "amplitude_calibration": self.amplitude_calibration,
        }

    ##################################
    ### DEPRECATED FUNCTION REGION ###
    ##################################

    @property
    def nco_frequency(self) -> float:
        """The internal frequency in Hz used for the generated microwave pulses.

        .. deprecated::
            For more intuitive handling, this has been replaced by `PulseGen.internal_frequency`.
            It will be removed completely in a future version.
        """
        warnings.warn(
            "Deprecation Warning: nco_frequency is deprecated, please use internal_frequency instead."
        )
        return self.internal_frequency

    @nco_frequency.setter
    def nco_frequency(self, nco_frequency):
        warnings.warn(
            "Deprecation Warning: nco_frequency is deprecated, please use internal_frequency instead."
        )
        self.internal_frequency = nco_frequency


@platform_attribute_collector
class TriggerSet(PlatformComponent):
    """Single trigger set within a signal generator of the QiController.

    Trigger sets are typically accessed only via the `PulseGen.triggerset` list.
    """

    def __init__(
        self,
        name: str,
        connection,
        controller,
        qkit_instrument,
        pulsegen: PulseGen,
        triggerset=0,
    ):
        super().__init__(name, connection, controller, qkit_instrument)
        self._pulsegen: PulseGen = pulsegen
        self._stub = self._pulsegen._stub

        self._index = triggerset
        self._indexset = proto.IndexSet(
            cindex=self._pulsegen._component,
            tindex=proto.TriggerSetIndex(value=self._index),
        )

    @property
    @platform_attribute
    @ServiceHubCall(errormsg="Could not read NCO phase offset of Triggerset.")
    def phase_offset(self) -> float:
        """Phase offset of the internal reference oscillation for this pulse.
        It is defined with respect to the global phase of the oscillator.

        This value is given in radian and in the interval [0, 2pi).
        If a different value is set, it will be wrapped to end up in this range.
        """
        return self._stub.GetPhaseOffset(self._indexset).value

    @phase_offset.setter
    @ServiceHubCall(errormsg="Could not set NCO phase offset of Triggerset.")
    def phase_offset(self, value: float):
        self._stub.SetPhaseOffset(proto.PhaseOffset(index=self._indexset, value=value))

    @property
    @platform_attribute
    @ServiceHubCall(errormsg="Could not read duration of the Triggerset.")
    def duration(self) -> float:
        """Duration of the loaded pulse in seconds.

        This time determines how many samples are read from the envelope memory before
        stopping the output again (or halting the last value if this option is set).

        The accuracy is in whole clock cycles of the QiController and deviating values
        will be rounded up to next higher cycle count during assignment."""
        return self._stub.GetDuration(self._indexset).value

    @duration.setter
    @ServiceHubCall(errormsg="Could not set duration of the Triggerset.")
    def duration(self, duration: float):
        self._stub.SetDuration(proto.Duration(index=self._indexset, value=duration))

    @property
    @platform_attribute
    @ServiceHubCall(errormsg="Could not read trigger offset of Triggerset.")
    def trigger_offset(self) -> int:
        """Time offset between trigger signal and generation of the pulse in seconds.

        This is helpful if you want to trigger two modules at the same time but delay
        one of the modules with respect to another.
        """
        return 0  # self._stub.GetTriggerOffset(self._indexset).offset

    @trigger_offset.setter
    @ServiceHubCall(errormsg="Could not set trigger offset of Triggerset.")
    def trigger_offset(self, trigger_offset: int):
        # self._stub.SetTriggerOffset(
        #    proto.TOffset(index=self._indexset, offset=trigger_offset)
        # )
        pass

    @ServiceHubCall(errormsg="Could not load the pulse into the Triggerset.")
    def load_pulse(
        self,
        pulseform: npt.ArrayLike,
        phase: float = 0.0,
        offset: float = 0.0,
        hold: bool = False,
        shift_phase: bool = False,
    ):
        """Loads a given pulse shape into the trigger set.

        The shape will be stored in the envelope memory.
        An exception will be raised once this memory is full.
        You can reset the memory using `PulseGen.reset_env_mem()`.
        This will erase the pulse shapes of all previously loaded trigger sets.

        :param pulseform:
            The pulse shape envelope to be loaded, as numpy array.
            This can be complex to perform intra-pulse phase modulations.
        :param phase:
            Phase offset of the internal oscillator for this pulse with respect to the
            global phase of the oscillator in radian. It defaults to 0.

            For more information, see `TriggerSet.phase_offset`.
        :param offset:
            Time offset between trigger signal and generation of the pulse in seconds.
            It defaults to 0.

            For more information, see `TriggerSet.trigger_offset`.
        :param hold:
            If the last value of the envelope should be kept after playing the pulse,
            by default False (i.e. output is muted again after the pulse).

            This can be used to output a continuous wave (CW). The output will only be
            stopped once a new trigger command is received at the signal generator.
        :param shift_phase:
            If the global phase should be shifted by the given phase offset (True)
            or just adapted for this one pulse (False).
            This can be used to realize virtual Z gates.
        """

        pulseform_i = np.real(pulseform)
        pulseform_q = np.imag(pulseform)
        if not np.any(pulseform_q):
            # No imaginary part present, so only real envelope
            pulseform_q = []

        self._stub.LoadPulse(
            proto.Pulse(
                index=self._indexset,
                i=pulseform_i,
                q=pulseform_q,
                phase=phase,
                offset=offset,
                hold=hold,
                shift_phase=shift_phase,
            )
        )

    def trigger_manually(self):
        """Manually triggers the signal generator with this trigger set, resulting in
        the output of the configured pulse.

        .. note ::
            Be aware that the internal oscillator (NCO) needs to be enabled for the output
            to be visible, see :meth:`PulseGen.nco_enable()`.
        """
        self._pulsegen.trigger_manually(self._index)
