from qm.QuantumMachinesManager import QuantumMachinesManager
from qm import SimulationConfig, LoopbackInterface
from qm.qua import *
from typing import Union, Tuple, List, Dict, Any, Optional
import matplotlib.pyplot as plt
from scipy.signal import windows
from scipy.interpolate import splrep, splev, interp1d
import numpy as np
from functools import reduce

def zip_multi_streams(vars_streams,
                      buffer_args: Optional[Union[int, Tuple[int]]]=None, average=False):
    """
    Generates a chain of zipped streams to be used within the stream processing. The function
    :param vars_streams: a list of QUA streams to be zipped
    :param buffer_args: the size of the buffer. default is None
    :param average: If true will average each stream
    :return: a chain of streams e.g: a_stream.buffer(10).zip(b_stream.buffer(10)).zip(c_stream.buffer(10))....
    """
    last = None
    for i in range(len(vars_streams)):
        if buffer_args is not None:
            if isinstance(buffer_args, int):
                buffer_args_ = (buffer_args, )
            else:
                buffer_args_ = buffer_args

            current = vars_streams[i].buffer(*buffer_args_)
        else:
            current = vars_streams[i]

        if average:
            current = current.average()

        if last is None:
            last = current
        else:
            last = last.zip(current)
    return last

def wait_for_values(res_handles, var_list=None, count:int=1):
    if var_list is None:
        var_list = [var_name for var_name, _ in res_handles]
    
    for var in var_list:
        res_handles.get(var).wait_for_values(count)

def fetch_all(res_handles, var_list):
    """
    Method to simply fetch results of variables specified by `var_list` given a
    QM result handles `res_handles`.
    """
    return [res_handles.get(var).fetch_all() for var in var_list]

def update_config_mixer_dc_offset(qe, I_offset, Q_offset, config,
                                  print_details: bool=False):
    """
    Update the QM configuration `config` such that the DC offset of (I, Q)-pair
    of an element `qe` is updated to values specified by `I_offset`, `Q_offset`.
    """
    channel_info = config['elements'][qe]['mixInputs']
    conI, chI = channel_info['I']
    conQ, chQ = channel_info['Q']

    config['controllers'][conI]['analog_outputs'][chI]['offset'] = I_offset
    config['controllers'][conQ]['analog_outputs'][chQ]['offset'] = Q_offset

    if print_details:
        print("DC offset of", qe, "updated to (I, Q) =", (I_offset, Q_offset))


def find_closest_lo_frequency(qe, f_rf, mixer_dc_offset,
                              if_freq_min: float=5e6,
                              if_freq_max: float=460e6):
    """

    Find the optimal calibrated LO frequency given a (list of) RF frequency
    and the (list of) mixer dc offsets.

    Parameters
    ----------
    qe : a string or a list of strings
        DESCRIPTION.
    f_rf : a float or a list of float
        DESCRIPTION.
    mixer_dc_offset :  a list of dictionaries or a list of list of dictionaries
        DESCRIPTION.
    if_freq_min : float, optional
        DESCRIPTION. The default is 10e6.
    if_freq_max : float, optional
        DESCRIPTION. The default is 450e6.

    Raises
    ------
    Warning
        DESCRIPTION.
    ValueError
        DESCRIPTION.

    Returns
    -------
    TYPE
        DESCRIPTION.

    """                     
    if isinstance(qe, list):
        n_elem = len(qe)
        cal_lo_freq_list = []
        for n in range(n_elem):
            cal_lo_freq_list.append(
                np.array([cal_point["lo_frequency"] for cal_point in mixer_dc_offset[n]]))
        # list of common calibrated lo frequencies
        cal_lo_freqs = reduce(np.intersect1d, cal_lo_freq_list)
        if isinstance(f_rf, list):
            # lo frequency must be separated from all rf freqs by at least 10MHz
            try:
                available_lo_freq_mask = np.all(
                    [(if_freq_min < np.abs(f_rf_ - cal_lo_freqs)) for f_rf_ in f_rf],
                    axis=0)
                available_cal_lo_freqs = cal_lo_freqs[available_lo_freq_mask]
            except:
                raise Warning('WARNING:')
                available_cal_lo_freqs = cal_lo_freqs
            
            try:
                available_lo_freq_mask2 = np.all(
                    [(if_freq_max > np.abs(f_rf_ - available_cal_lo_freqs)) for f_rf_ in f_rf],
                    axis=0)
                available_cal_lo_freqs = available_cal_lo_freqs[available_lo_freq_mask2]
            except:
                raise Warning('WARNING:')
            
            if len(available_cal_lo_freqs) == 0:
                raise ValueError('Not possible to find calibrated LO frequency satisfying condition' + 
                                 '%.3f MHz < |f_rf - f_lo| < %.3f MHz' % (if_freq_min / 1e6, if_freq_max / 1e6))
            # find the LO frequency that minimizes the sum of detuning from all
            # specified RF frequencies.
            lo_idx = np.argmin(
                np.sum(np.abs([(f_rf_ - available_cal_lo_freqs)
                               for f_rf_ in f_rf]), axis=0))
            
            closest_lo_freq = available_cal_lo_freqs[lo_idx]
            
            return closest_lo_freq

        elif isinstance(f_rf, float):
            # if a float is given in `f_rf` it 
            return find_closest_lo_frequency(qe, [f_rf], mixer_dc_offset)
        
        else:
            raise ValueError('ERROR: Input arguments exception.')

    elif isinstance(qe, str):
        return find_closest_lo_frequency([qe], f_rf, [mixer_dc_offset])

def update_config_lo_frequency(qe, lo_freq, config, mixer_dc_offset=None,
                               print_details: bool=False):
    """
    Update the QM configuration `config` such that the dc offset of analog output
    channels corresponding to a (I, Q)-pair of an element `qe` is updated.
    The offsets corresponding to LO frequency specified by `lo_freq` is found
    by searching the lookup table `dc_offset_list`. The `dc_offset_list` must
    be a list of dictionaries containing the LO frequencies and DC offsets in
    the following form:
    [{'lo_frequency': 4.7e9, 'I': -0.005625, 'Q': 0.0084375}},
     {'lo_frequency': 3.9e9, 'I': -0.0117187, 'Q': -0.005390625},
     ... ]
    """
    if print_details:
        print("Update LO frequency of " + qe + " to %.4f GHz" % (lo_freq / 1e9))
    if mixer_dc_offset is not None:
        offset = next((x for x in mixer_dc_offset if x['lo_frequency'] == lo_freq),
                      None)
        I_offset, Q_offset = offset['I'], offset['Q']
        update_config_mixer_dc_offset(qe, I_offset, Q_offset, config)
    
    config['elements'][qe]['mixInputs']['lo_frequency'] = lo_freq

def find_qe_by_mixer(mixer, config):
    
    elements = config['elements']  
    qe = next((elem for elem, props in elements.items() if ('mixInputs' in props.keys()) and (props['mixInputs']['mixer'] == mixer)),None)
    return qe

def find_closest_mixer_correction(lo_freq, if_freq, mixer_iq_imbalance):
    cals = [x for x in mixer_iq_imbalance if float(x['lo_frequency']) == float(lo_freq)]
    cal_if_freqs = np.array([x['if_frequency'] for x in cals])
    amps = np.array([x['amp'] for x in cals])
    phases = np.array([x['phase'] for x in cals])

    abs_detuning = np.abs(cal_if_freqs - if_freq)
    if np.any(abs_detuning < 40e6):

        closest_if_freq_idx = np.argmin(abs_detuning)
        amp = amps[closest_if_freq_idx]
        phase = phases[closest_if_freq_idx]

        correction = iq_imbalance_corr(amp, phase)
    else:
        correction = [1.0, 0.0, 0.0, 1.0]
    return correction

def update_config_if_frequency(qe, if_freq, config, correction=None,
                               mixer_iq_imbalance=None, print_details: bool=False):
    """
    Update the QM configuration `config` such that the IF frequency of an 
    lement `qe` is updated. This updates the 'intermediate_frequency' under
    'element' and adds a mixer calibration point at this IF frequency and
    LO frequency specified by 'lo_frequency' under 'mixInputs' with correction
    specified by `correction` or a list `mixer_iq_imbalance`.
    
    """
    if print_details:
        print("Update IF frequency of " + qe + " to %.4f MHz" % (if_freq / 1e6))
    elem = config['elements'][qe]
    elem['intermediate_frequency'] = if_freq

    mix_inputs = elem['mixInputs']
    lo_freq = mix_inputs['lo_frequency']
    mixer = mix_inputs['mixer']
    mixer_config = config['mixers'][mixer]

    if correction is None:
        if mixer_iq_imbalance is None:
            correction = [1.0, 0.0, 0.0, 1.0]
        else:
            correction = find_closest_mixer_correction(lo_freq, if_freq,
                                                       mixer_iq_imbalance)

    # existing mixer calibration inside config at the specified LO and IF frequency
    existing_cal = next(
        (cal for cal in mixer_config if
         (float(cal['lo_frequency']) == float(lo_freq) and
          float(cal['intermediate_frequency']) == float(if_freq))
         ), None)
    try:
        mixer_config.remove(existing_cal)
    except:
        pass
    # correction = [1, 0, 0, 1]
    if print_details:
        print("Mixer correction of " + qe + " updated to:", correction)
    mixer_config.append({'intermediate_frequency': int(if_freq),
                         'lo_frequency': int(lo_freq),
                         'correction': correction})
    config['mixers'][mixer] = mixer_config

WaveformSample = Union[List[float], float]

def create_IQ_wf_samples(
        pulse_dict: Dict[str, Any]
        ) -> Tuple[WaveformSample, WaveformSample]:
    """
    Create samples of I and Q waveforms of operations whose type and
    properties are specified by `pulse_dict`.

    Parameters
    ----------
    pulse_dict : dict
        A dictionary containing attributes necessary for creating pulse 
        IQ waveforms. The keywords of pulse_dict are 'type' and 'properties'
        specifying the type and the properties of the pulse.

        Example:
            
        pulse_dict = {
            'type': 'drag_cosine',
            'properties': {'amp': 0.4, 'len': 20, 'beta': 1.0}
            }
        
    Raises
    ------
    ValueError
        The type of the pulse must be either 'square', 'cosine' or 'gaussian'.
        Otherwise ValueError is raised.

    Returns
    -------
    I_sample : float or list of float
        A list or float that specifies the samples or sample of the I waveform.
    Q_sample : float or list of float
        A list or float that specifies the samples or sample of the Q waveform.

    """

    type_ = pulse_dict['type']
    props_ = pulse_dict['properties']

    if type_ not in ['square', 'cosine', 'gaussian', 'drag_cosine',
                     'drag_gaussian', 'square_with_padding', 'clear',
                     'gauss_convolved_square', 'gauss_convolved_square_with_kick']:
        raise ValueError('The specified waveform type ' + str(type_) +
                         ' is not supported.')

    if type_ == 'square':
        A, N = props_['amp'], props_['len']
        I_sample = A
        Q_sample = 0.0

    if type_ == 'cosine':
        A, N = props_['amp'], props_['len']

        n = np.arange(N, dtype=int)

        I_sample = (A / 2 * (1 - np.cos(2 * np.pi * n / (N - 1)))).tolist()
        Q_sample = 0.0

    if type_ == 'gaussian':
        A, N, sigma_ = props_['amp'], props_['len'], props_['sigma']

        n = np.arange(N, dtype=int)
        n0 = 0.5 * (n[0] + n[-1])

        I_sample = (A * np.exp(-(n - n0) ** 2 / (2 * sigma_ ** 2))).tolist()
        Q_sample = 0.0

    if type_ == 'drag_cosine':
        A, N = props_['amp'], props_['len']
        beta_ = props_['beta']
        df = props_['detuning']

        n = np.arange(N, dtype=int)

        z = (A / 2 * (1 - np.cos(2 * np.pi * n / (N - 1))) +
             1j * beta_ * 2 * np.pi / (N - 1) * A / 2 * np.sin(2 * np.pi * n / (N - 1)))

        z *= np.exp(1j * 2 * np.pi * df * n * 1e-9)
        
        I_sample = z.real.tolist()
        Q_sample = z.imag.tolist()

    if type_ == 'drag_gaussian':
        A, N, sigma_ = props_['amp'], props_['len'], props_['sigma']
        beta_ = props_['beta']
        df = props_['detuning']

        n = np.arange(N, dtype=int)
        n0 = 0.5 * (n[0] + n[-1])

        z = (A * np.exp(-(n - n0) ** 2 / (2 * sigma_ ** 2)) +
             1j * beta_ * - (n - n0) / sigma_ ** 2 * A * np.exp(-(n - n0) ** 2 / (2 * sigma_ ** 2)))
        z *= np.exp(1j * 2 * np.pi * df * n * 1e-9)

        I_sample = z.real.tolist()
        Q_sample = z.imag.tolist()

    if type_ == 'square_with_padding':
        A, N, M = props_['amp'], props_['len'], props_['padding']
        
        I_sample = [A] * (N - M) + [0.0] * M
        Q_sample = 0.0

    if type_ == 'gauss_convolved_square':
        A, N, M1, M2 = props_['amp'], props_['len'], props_['delay'], props_['padding']
        sigma = props_['sigma']     
        
        gauss_env = windows.gaussian(N, sigma)
        I_sample = np.convolve(np.array([0.0] * M1 + [A] * (N - M1 - M2) + [0.0] * M2),
                               gauss_env / np.sum(gauss_env), mode='same').tolist()
        Q_sample = 0.0

    if type_ == 'gauss_convolved_square_with_kick':
        A, N, M1, M2 = props_['amp'], props_['len'], props_['delay'], props_['padding']
        K, L = props_['kick_factor'], props_['kick_len']
        sigma = props_['sigma']     
        
        gauss_env = windows.gaussian(N, sigma)
        I_sample = np.convolve(np.array([0.0] * M1 + [K * A] * L + [A] * (N - L - M1 - M2) + [0.0] * M2),
                               gauss_env / np.sum(gauss_env), mode='same').tolist()
        Q_sample = 0.0

    if type_ == 'clear':
        A, N, M = props_['amp'], props_['len'], props_['kick_len']
        
        Ai1 = props_['initial_kick1_amp']
        Ai2 = props_['initial_kick2_amp']
        Af1 = props_['final_kick1_amp']
        Af2 = props_['final_kick2_amp']
    
        I_sample = ([Ai1] * M + [Ai2] * M +
                    [A] * (N - 4 * M) +
                    [Af1] * M + [Af2] * M)
        Q_sample = 0.0

    return I_sample, Q_sample

def update_op_from_pulse_dict(pulse_dict: Dict[str, Any],
                              op: str, qe: str, config, debug: bool=False):
    '''
    

    Parameters
    ----------
    pulse_dict : Dict[str, Any]
        DESCRIPTION.
    op : str
        DESCRIPTION.
    qe : str
        DESCRIPTION.
    config : TYPE
        DESCRIPTION.

    Returns
    -------
    None.

    '''
    pulse = config['elements'][qe]['operations'][op]
    wf_ = config['pulses'][pulse]['waveforms']

    # update pulse length in config
    pulse_len = pulse_dict['properties']['len']
    config['pulses'][pulse]['length'] = pulse_len

    I_sample, Q_sample = create_IQ_wf_samples(pulse_dict)
    if debug:
        plt.figure()
        plt.plot(I_sample, label='I')
        plt.plot(Q_sample, label='Q')
        plt.title(', '.join([op, qe]))
        plt.xlabel('Time (ns)')
        plt.ylabel('Waveform (V)')
        plt.legend()
    
    wfI_name, wfQ_name = wf_['I'], wf_['Q']

    if wfI_name not in ['zero_wf']:
        update_config_waveform(wfI_name, I_sample, config)
    if wfQ_name not in ['zero_wf']:
        update_config_waveform(wfQ_name, Q_sample, config)

def find_iq_waveform_from_op(qe, op, config):
    pulse = config['elements'][qe]['operations'][op]
    wf = config['pulses'][pulse]['waveforms']
    wfI, wfQ = wf['I'], wf['Q']
    return wfI, wfQ

def update_config_waveform(wf_name: str, samples, config):
    """
    Update a waveform specified by `wf_name` in the configuration `config`
    according to the sample or samples specified by `samples`. If the
    waveform with name `wf_name` is not defined in the configuration, a new
    waveform with the specified name is initialized in `config` and updated.

    Parameters
    ----------
    wf_name : string
        A string that specifies the name of the waveform.
    samples : list or float
        A list or float that specifies the samples or sample of the waveform.
        The waveform is treated as an arbitrary waveform if a list is given;
        it is treated as a constant waveform if a float is given..
    config : dictionary
        A dictionary containing the QM configuration.

    Raises
    ------
    ValueError
        Error is raised when `samples` is not an instance of float or list.

    Returns
    -------
    None.

    """
    if isinstance(samples, list): 
        update_config_arbitrary_waveform(wf_name, samples, config)

    elif isinstance(samples, float):
        update_config_constant_waveform(wf_name, samples, config)

    else:
        raise ValueError("`samples` must be either a list or a float.")

def update_config_arbitrary_waveform(wf_name, samples, config):
    """
    Update a waveform specified by `wf_name` in the configuration `config`
    to an arbitrary waveform whose samples are specified by `samples`. If the
    waveform with name `wf_name` is not defined in the configuration, a new
    waveform with the specified name is initialized in `config` and updated.

    Parameters
    ----------
    wf_name : string
        A string that specifies the name of the waveform.
    samples : list
        A list that specifies the samples of the arbitrary waveform.
    config : dictionary
        A dictionary containing the QM configuration.

    Returns
    -------
    None.

    """
    if wf_name not in config['waveforms'].keys():
        config['waveforms'][wf_name] = {}

    waveform = config['waveforms'][wf_name]   
    waveform.pop('sample', None) # remove key 'sample'
    waveform['type'] = 'arbitrary'
    waveform['samples'] = samples

def update_config_constant_waveform(wf_name, sample, config):
    """
    Update a waveform specified by `wf_name` in the configuration `config`
    to a constant waveform whose sample is specified by `sample`. If the
    waveform with name `wf_name` is not defined in the configuration, a new
    waveform with the specified name is initialized in `config` and updated.    

    Parameters
    ----------
    wf_name : string
        A string that specifies the name of the waveform..
    sample : float
        A float that specifies the sample of the constant waveform.
    config : dictionary
        A dictionary containing the QM configuration.

    Returns
    -------
    None.

    """
    if wf_name not in config['waveforms'].keys():
        config['waveforms'][wf_name] = {}
        
    waveform = config['waveforms'][wf_name] 
    waveform.pop('samples', None) # remove key 'sample'
    waveform['type'] = 'constant'
    waveform['sample'] = sample


def iq_imbalance_corr(amp_, phase_):
    """
    Create a correction matrix for an IQ mixer with amplitude imbalance
    specified by `amp` and phase imbalance specified by `phase`. See
    https://en.wikipedia.org/wiki/IQ_imbalance for details.

    Parameters
    ----------
    amp : float
        amplitude imbalance.
    phase : float
        phase imbalance.

    Returns
    -------
    list
        List of matrix elements that corrects the IQ imbalance of the IQ mixer.

    """
    c = np.cos(phase_ / 2)
    s = np.sin(phase_ / 2)
    
    # normalization factor
    N = 1 / ((1 - amp_ ** 2) * (2 * c ** 2 - 1))
    return [
        float(x) if np.abs(x) < 2 else 1.9999 * np.sign(x)
        for x in [N * (1 - amp_) * c, N * (1 + amp_) * s,
                  N * (1 - amp_) * s, N * (1 + amp_) * c]
        ]

def plot_1D_sweep_with_res_handles(res_handles, x_axis, x_label):
    """
    Simple plotting of the results of 1D sweep QM program given the result
    handles `res_handles`. The magitude and phase of the data is plotted against
    the sweeping parameter specified by `x_axis` with label `x_label`. Also,
    the data is plotted on the complex (I, Q) plane.
    """
    fig = plt.figure()
    I, Q, n = fetch_all(res_handles, ["I", "Q", "n"])
    sig = I + 1j * Q

    plt.subplot(2, 2, 1)
    plt.plot(x_axis, np.abs(sig) / 1e-3)
    plt.xlabel(x_label)
    plt.ylabel("Voltage (mV)")

    plt.subplot(2, 2, 3)
    plt.plot(x_axis, np.unwrap(np.angle(sig)))
    plt.xlabel(x_label)
    plt.ylabel("Phase (rad)")

    plt.subplot(1, 2, 2)
    plt.plot(I / 1e-3, Q / 1e-3)
    plt.xlabel("I (mV)")
    plt.ylabel("Q (mV)")
    plt.axis('equal')
    fig.tight_layout()

    print("Number of Averages: ", n)
    return fig

def execute_wait_for_all_and_fetch(prog, config):
    """
    Execute a QM program specified by `prog` with configuration `config`,
    wait for all values, and fetch the result in the form of complex signal
    `sig = I + 1j * Q`.
    """
    qmm = QuantumMachinesManager()
    qm = qmm.open_qm(config)
    job = qm.execute(prog, data_limit=0, duration_limit=0)

    res_handles = job.result_handles
    res_handles.wait_for_all_values()
    I, Q, n = fetch_all(res_handles)
    sig = I + 1j * Q
    return sig

def update_op_waveform_from_gaussian(qe, q_op, amp, config):
    pulse = config['elements'][qe]['operations'][q_op]
    pi_wf = config['pulses'][pulse]['waveforms']
    gauss_wf = config['pulses']['gaussian_pulse']['waveforms']

    for iq in ['I', 'Q']:
        wf_type = config['waveforms'][gauss_wf[iq]]['type']
        config['waveforms'][pi_wf[iq]]['type'] = wf_type
        if wf_type == 'constant':
            config['waveforms'][pi_wf[iq]]['sample'] = config['waveforms'][gauss_wf[iq]]['sample']
        elif wf_type == 'arbitrary':
            gauss_wf_samples = np.array(config['waveforms'][gauss_wf[iq]]['samples'])
            config['waveforms'][pi_wf[iq]]['samples'] = (amp * gauss_wf_samples).tolist()

def update_op_waveform_from_general_op(qe, q_op0, q_op1, amp, config):
    """
    Update waveform of operation `q_op1` from that of `q_op0` scaled by a
    factor given by `amp`
    """
    pulse0 = config['elements'][qe]['operations'][q_op0]
    pulse_wf0 = config['pulses'][pulse0]['waveforms']

    pulse1 = config['elements'][qe]['operations'][q_op1]
    pulse_wf1 = config['pulses'][pulse1]['waveforms']

    for iq in ['I', 'Q']:
        wf_type = config['waveforms'][pulse_wf0[iq]]['type']
        config['waveforms'][pulse_wf1[iq]]['type'] = wf_type
        if wf_type == 'constant':
            config['waveforms'][pulse_wf1[iq]].pop('samples', None) # remove key 'samples'
            config['waveforms'][pulse_wf1[iq]]['sample'] = amp * config['waveforms'][pulse_wf0[iq]]['sample']
        elif wf_type == 'arbitrary':
            config['waveforms'][pulse_wf1[iq]].pop('sample', None) # remove key 'sample'
            pulse_wf0_samples = np.array(config['waveforms'][pulse_wf0[iq]]['samples'])
            config['waveforms'][pulse_wf1[iq]]['samples'] = (amp * pulse_wf0_samples).tolist()

def get_approx_snr(samples, baseline_portion=0.3):
    """
    Extract the signal-to-noise ratio from real-valued samples specified by
    `samples`. The noise level is found from the standard deviation of the
    baseline samples (according to the portion specified by `baseline_portion`)
    and the signal level is found from the distance of max or min sample from
    the mean of baseline. This works the best when there is a well defined
    peak in a noisy background (which improves under averaging).

    Parameters
    ----------
    samples : numpy.ndarray
        Real-valued samples.
    baseline_portion : float, optional
        The portion of samples to be taken as the baseline. The default is 0.3.

    Returns
    -------
    float
        Estimate of signal-to-noise ratio (in amplitude) of samples.

    """
    N = len(samples)
    low_samples = np.partition(samples, int(baseline_portion * N))
    high_samples = np.partition(samples, -int(baseline_portion * N))
    mean = np.mean(samples)
    
    if np.abs(mean - np.mean(low_samples)) < np.abs(mean - np.mean(high_samples)):
        baseline = low_samples
    else:
        baseline = high_samples
    
    signal = np.max([np.max(samples) - np.mean(baseline),
                     np.mean(baseline) - np.min(samples)])
    noise = np.std(baseline)
    
    return signal / noise

def update_config_time_of_flight(rr, tof, config):
    config['elements'][rr]['time_of_flight'] = tof
    
def plot_multiple_1D_sweep_with_res_handles(res_handles, x_data, x_label, data_labels=None):
    I, Q, n = fetch_all(res_handles)
    sig = I + 1j * Q
    nrow, ncol = sig.shape

    fig = plt.figure()
    plt.subplot(2, 2, 1)
    for idx in range(ncol):
        plt.plot(x_data, np.abs(sig[:, idx]) / 1e-3)
    plt.xlabel(x_label)
    plt.ylabel("Voltage (mV)")

    plt.subplot(2, 2, 3)
    for idx in range(ncol):
        plt.plot(x_data, np.unwrap(np.angle(sig[:, idx])))
    plt.xlabel(x_label)
    plt.ylabel("Phase (rad)")

    plt.subplot(1, 2, 2)
    for idx in range(ncol):
        plt.plot(I[:, 0] / 1e-3, Q[:, 0] / 1e-3)
        plt.xlabel("I (mV)")
        plt.ylabel("Q (mV)")
    plt.axis('equal')
    
    print("Number of averages:", n)

    fig.tight_layout()
    return fig

def add_integW_to_operation(iw_dict: Dict[str, List[float]],
                            ro_op: str, res: str, config):
    """
    Add a collection of integration weights in `iw_dict` to operatio

    Parameters
    ----------
    iw_dict : Dict[str, List[float]]
        DESCRIPTION.
    ro_op : str
        DESCRIPTION.
    res : str
        DESCRIPTION.
    config : TYPE
        DESCRIPTION.

    Returns
    -------
    None.

    """
    pulse = config['elements'][res]['operations'][ro_op]
    pulse_config = config['pulses'][pulse]

    if 'integration_weights' not in pulse_config:
        pulse_config['integration_weights'] = {}
    for iw_name in iw_dict.keys():
        iw_label = '_'.join([ro_op, iw_name, res])
        pulse_config['integration_weights'][iw_name] = iw_label
        config['integration_weights'][iw_label] = iw_dict[iw_name]



class TimeDiffCalibrator:

    @staticmethod
    def _default_config(freq, con_name):
        return {

            'version': 1,

            'controllers': {
                con_name: {
                    'type': 'opx1',
                    'analog_outputs': {
                        1: {'offset': 0.0},
                        2: {'offset': 0.0}
                    },
                    'digital_outputs': {},
                    'analog_inputs': {
                        1: {'offset': 0.0},
                        2: {'offset': 0.0}
                    }
                }
            },

            'elements': {

                'rr': {
                    'mixInputs': {
                        'I': (con_name, 1),
                        'Q': (con_name, 2),
                        'lo_frequency': 0,
                        'mixer': 'mixer'
                    },
                    'intermediate_frequency': freq,
                    'operations': {
                        'readout': 'readout',
                    },
                    "outputs": {
                        'out1': (con_name, 1),
                        'out2': (con_name, 2)
                    },
                    'time_of_flight': 32,
                    'smearing': 0
                },

            },

            "pulses": {

                'readout': {
                    'operation': 'measurement',
                    'length': 1000,
                    'waveforms': {
                        'I': 'const_wf',
                        'Q': 'zero_wf',

                    },
                    'integration_weights': {
                        'integW1': 'integW1',
                        'integW2': 'integW2',
                    },
                    'digital_marker': 'ON'
                },
            },

            'waveforms': {

                'const_wf': {
                    'type': 'constant',
                    'sample': 0.4
                },

                'zero_wf': {
                    'type': 'constant',
                    'sample': 0.0
                },

            },

            'digital_waveforms': {
                'ON': {
                    'samples': [(1, 0)]
                }
            },

            'integration_weights': {

                'integW1': {
                    'cosine': [1.0] * int(1000 / 4),
                    'sine': [0.0] * int(1000 / 4)
                },

                'integW2': {
                    'cosine': [0.0] * int(1000 / 4),
                    'sine': [1.0] * int(1000 / 4)
                },
            },

            'mixers': {

                'mixer': [
                    {'intermediate_frequency': freq, 'lo_frequency': 0, 'correction': [1, 0, 0, 1]}
                ]
            }
        }

    @staticmethod
    def calibrate(qmm, con_name):
        with program() as cal_phase:
            I1 = declare(fixed)
            Q1 = declare(fixed)
            I2 = declare(fixed)
            Q2 = declare(fixed)
            I = declare(fixed)
            Q = declare(fixed)

            measure('readout', 'rr', 'adc', demod.full('integW1', I1, 'out1'),
                    demod.full('integW2', Q1, 'out1'),
                    demod.full('integW1', I2, 'out2'),
                    demod.full('integW2', Q2, 'out2'))
            assign(I, I1 + Q2)
            assign(Q, -Q1 + I2)
            save(I, 'I')
            save(Q, 'Q')

        freq = 1.4567e6
        qm = qmm.open_qm(TimeDiffCalibrator._default_config(freq, con_name))
        job = qm.simulate(cal_phase,
                          SimulationConfig(500, simulation_interface=LoopbackInterface([('con1', 1, 'con1', 1),
                                                                                        ('con2', 2, 'con2', 2)])))

        adc1 = job.result_handles.adc_input1.fetch_all()['value']
        adc2 = job.result_handles.adc_input2.fetch_all()['value']
        adc_ts = job.result_handles.adc_input1.fetch_all()['timestamp']

        I = job.result_handles.get("I").fetch_all()['value'][0]
        Q = job.result_handles.get("Q").fetch_all()['value'][0]

        sig = (adc1 + 1j * adc2) * np.exp(-1j * 2 * np.pi * freq * 1e-9 * adc_ts)
        d = np.sum(sig)

        I_ = np.real(d)
        Q_ = np.imag(d)

        time_diff_ns = np.angle((I + 1j * Q) / (I_ + 1j * Q_)) / 1e-9 / 2 / np.pi / freq
        time_diff = np.round(time_diff_ns / 4) * 4
        return time_diff


# def python_demodulation(IF_freq: float, Wc: List[float], Ws: List[float], S: List[int], ts0: int):
#     '''
#     demodulation
#     :param IF_freq: IF frequency in Hz
#     :param Wc: list of cosine weights.
#     :param Ws: list of sine weights
#     :param S: The signal
#     :param ts0: The time stamp in ns of the first sample
#     :return: demodulation value
#     '''

#     phi0 = 2 * np.pi * IF_freq * ts0 * 1e-9
#     Sum = 0
#     delta_phi = 2 * np.pi * IF_freq * 1e-9
    
#     # 
#     Phi = [i * delta_phi + phi0 for i in range(4)]

#     for i in range(int(len(S) / 4)):
#         # A vector of four samples (4ns)
#         Svec = np.array([S[4 * i + s] for s in range(4)])
#         Sum = Sum + np.sum(np.multiply(np.multiply(Wc[i], np.cos(Phi)) +
#                                        np.multiply(Ws[i], np.sin(Phi)), Svec))
#         for j in range(4):
#             Phi[j] = Phi[j] + 4 * delta_phi

#     return 16 * Sum * 2 ** -28