from typing import List, Union, Tuple
import numpy as np
import sys

# Entropy tools
from entropylab import EntropyContext
from entropylab_qpudb import QuaConfig

# home directory imports
sys.path.append(r'C:/Users/paint/QM/MMProcessor_v1')
from qpu_resolver import resolve
from qm_helper_functions import (find_closest_lo_frequency,
                                   update_config_lo_frequency,
                                   update_config_if_frequency)

def update_rf_lo_if_freqs(
        element: Union[int, List[int]], element_type: str,
        config: Union[dict, QuaConfig], context: EntropyContext,
        f_rf: Union[float, List[float]]=None, f_lo: float=None,
        f_if_min: float=5e6, f_if_max: float=480e6, print_details: bool=False,
        ) -> Tuple[Union[float, List[float]], float, Union[float, List[float]]]:
    """
    Update the LO and IF frequencies of a quantum element (or a list of
    quantum elements) given a RF frequency (or a list of RF frequencies).

    Parameters
    ----------
    element : Union[int, List[int]]
        An integer or a list of integers specifying the elements to update
        RF, LO, and IF frequencies.
    element_type : str
        The type of specified element(s). It must be either 'q' or 'res'.
    config : Union[dict, QuaConfig]
        A dictionary that represents QUA configuration.
    context : EntropyContext
        Entropy context.
    f_rf : Union[float, List[float]], optional
        A float or a list of floats specifying the RF frequencies of elements
        specified by `element`. If None, the RF frequency in QPU DB is used.
        The default is None.
    f_lo : float, optional
        A float specifying the LO frequency. If None, the LO frequency closest
        to the specified RF frequencies is found from the calibrated mixer
        DC offsets. The default is None.
    f_if_min : float, optional
        The minimum value of magnitude of IF frequencies.
        The default is 10e6.
    f_if_max : float, optional
        The maximum value of magnitude of IF frequencies.
        The default is 450e6.

    Returns
    -------
    f_rf_list : Union[float, List[float]]
        A float (a list of floats) specifying the RF frequency (frequencies).
    f_lo : float
        A float specifying the LO frequency.
    f_if_list : Union[float, List[float]]
        A float (a list of floats) specifying the IF frequency (frequencies).

    """


    qpu_db = context.get_resource('qpu_db')
    if element_type in ['q', 'qubit', 'q_xy', 'qubit_xy', 'xy']:
        elem_type = 'q'
        LO = context.get_resource('XY_LO')
    elif element_type in ['res', 'resonator', 'ro', 'readout', 'rr']:
        elem_type = 'res'
        LO = context.get_resource('RO_LO')

    if isinstance(element, list):
        if elem_type == 'q':
            qe_ch_list = [resolve.q(elem, channel='xy') for elem in element]
            qe_list = [resolve.q(elem) for elem in element]
            mixer_list = [resolve.q_mixer(elem) for elem in element]
        elif elem_type == 'res':
            qe_ch_list = [resolve.res(elem, channel='1') for elem in element]
            qe_list = [resolve.res(elem) for elem in element]
            mixer_list = [resolve.res_mixer(elem) for elem in element]

        if f_rf is None:
            # If f_rf is not specified, get the values from QPU DB
            f_rf_list = [qpu_db.get(qe, 'f_rf').value for qe in qe_list]
        else:
            if isinstance(f_rf, list):
                f_rf_list = f_rf
            elif isinstance(f_rf, float):
                # if f_rf is not a list, create a list of RF frequencies
                # whose length equals the number of elements
                f_rf_list = [f_rf] * len(element)
            else:
                raise ValueError(
                    'The specified f_rf must be either a float or a list of floats.'
                    )

        # get mixer calibration results from qpu_db
        mixer_dc_offset_list = [qpu_db.get("mixer_dc_offset", mixer).value
                                for mixer in mixer_list]
        mixer_iq_imbalance_list = [qpu_db.get("mixer_iq_imbalance", mixer).value
                                   for mixer in mixer_list]

        if f_lo is None:
            # If LO frequency is not specified, find the optimal calibrated
            # LO frequency that is closest to the list of RF frequencies specified
            f_lo = find_closest_lo_frequency(qe_list, f_rf_list,
                                             mixer_dc_offset_list)
        else:
            if isinstance(f_lo, float):
                pass
            else:
                raise ValueError('The specified f_lo must be a float.')

        # List of IF frequencies to update
        f_if_list = (np.array(f_rf_list) - f_lo).tolist()

        # update configuration
        for k, qe in enumerate(qe_ch_list):
            update_config_lo_frequency(qe, f_lo, config,
                                       mixer_dc_offset=mixer_dc_offset_list[k],
                                       print_details=print_details)
            update_config_if_frequency(qe, f_if_list[k], config,
                                       mixer_iq_imbalance=mixer_iq_imbalance_list[k],
                                       print_details=print_details)
            if elem_type =='res':
                update_config_lo_frequency(resolve.res(element[k], channel='2'),
                                           f_lo, config,
                                           print_details=print_details)
                update_config_if_frequency(resolve.res(element[k], channel='2'),
                                           f_if_list[k], config,
                                           print_details=print_details)

        # update RO_LO frequency
        # LO.set_freq(f_lo)

        if len(qe_list) == 1:
            # if only a single element is specified,
            # return f_rf: float, f_lo: float, f_if: float
            return f_rf_list[0], f_lo, f_if_list[0]
        elif len(qe_list) > 1:
            # if a list of elements is specified,
            # return f_rf_list: List[float], f_lo: float, f_if_list: List[float]
            return f_rf_list, f_lo, f_if_list

    elif isinstance(element, int):
        # If `element` is not a list but an index of an element,
        # promote it to a single-element list and call the function.
        return update_rf_lo_if_freqs([element], element_type,
                                     config, context, f_rf=f_rf, f_lo=f_lo,
                                     f_if_min=f_if_min, f_if_max=f_if_max)
    else:
        raise ValueError(
            'The specified element must be either an integer or a list of integers.'
            )
