import sys, time, threading
import numpy as np
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from typing import Optional, List
from tqdm import tqdm

# Entropy Tools
from entropylab import EntropyContext
from entropylab_qpudb import QuaConfig, QuaCalNode
from qm.simulate import SimulationConfig

# QUA Tools
from qm.qua import *

from qpu_resolver import resolve
from qm_helper_functions import (fetch_all, 
                                 update_config_lo_frequency,
                                 update_config_if_frequency,
                                 wait_for_values, get_approx_snr,
                                 update_config_waveform,
                                 find_iq_waveform_from_op)
from helper_functions import simple_resonance_idx_finder
from node_utils import update_rf_lo_if_freqs

sys.path.append(r"Z:\Discord\Qanalysis")
from Qanalysis import *
from labber_helper_functions import create_uniqueLabberLogFile, Labber

class QubitSpecNode(QuaCalNode):
    def __init__(self, element: int, xy_f_rf: Optional[float]=None,
                 span: Optional[float]=800e6, df: Optional[float]=1e6,
                 wait_time: Optional[int]=100000, n_avg: Optional[int]=20,
                 plot: Optional[bool]=False, LabberLog: Optional[bool]=True,
                 xy_op: Optional[str]='spectroscopy',
                 ro_op: Optional[str]='spectroscopy', overlap: bool=True,
                 debug: Optional[bool]=False, subtract_bg: Optional[bool]=True,
                 stop_early: Optional[int]=0, callback=None, **kwargs):
        super().__init__(**kwargs)
        self.element = element
        self.q_xy = resolve.q(element, channel='xy')
        self.res = resolve.res(element, channel='1')
        self.xy_f_rf0 = xy_f_rf
        self.span = 2 * int(span / 2)
        self.df = int(df)
        self.wait_time = wait_time
        self.n_avg = n_avg
        self.plot = plot
        self.xy_op = xy_op
        self.ro_op = ro_op
        self.callback = callback
        self.subtract_bg = subtract_bg
        self.overlap = overlap
        self.stop_early = stop_early
        self.LabberLog = LabberLog
        self.debug = debug

        self.node_type = type(self).__name__
        self.node_label = '[' + self.node_type + ']' + self._label
    def prepare_config(self, config: QuaConfig, context: EntropyContext) -> QuaConfig:
        # PREPARE CONFIG for running program, do not set qpu_db here!
        if self.debug:
            print('NODE INFO: ' + self._label + ' (' + self.node_type + ')' + 
                  ', PREPARE CONFIG')

        return config 

    def _fetch_sig(self, res_handles):
        try:
            I, Q = fetch_all(res_handles, ['I', 'Q'])
            self.sig = I + 1j * Q
            if self.subtract_bg:
                I_bg, Q_bg = fetch_all(res_handles, ['I_bg', 'Q_bg'])
                sig_bg = I_bg + 1j * Q_bg
                self.sig = self.sig - sig_bg
        except:
            time.sleep(1)
            self._fetch_sig(res_handles)
    
    def _qua_program(self):
        f_if_min, f_if_max, df = \
            int(self.xy_f_if - self.span / 2), int(self.xy_f_if + self.span / 2), int(self.df)
        f_if_vec = np.arange(f_if_min, f_if_max, df)

        with program() as q_spec_prog:
            i = declare(int)
            f = declare(int) # frequency

            I, Q = declare(fixed), declare(fixed)
            I_st, Q_st = declare_stream(), declare_stream()

            if self.subtract_bg:
                I_bg, Q_bg = declare(fixed), declare(fixed)
                I_bg_st, Q_bg_st = declare_stream(), declare_stream()

            count = declare(int, value=0) # averaging counter
            count_st = declare_stream()

            with for_(i, 0, i < self.n_avg, i + 1):
                align()
                wait(self.wait_time, self.q_xy, self.res)
                if self.subtract_bg:
                    # background data
                    measure(self.ro_op, self.res, None,
                            dual_demod.full('integW_cos', 'out1',
                                            'integW_sine', 'out2', I_bg),
                            dual_demod.full('integW_minus_sine', 'out1',
                                            'integW_cos', 'out2', Q_bg))
                    save(I_bg, I_bg_st)
                    save(Q_bg, Q_bg_st)                
                
                with for_(f, f_if_min, f < f_if_max, f + df):
                    update_frequency(self.q_xy, f)
                    align()
                    wait(self.wait_time, self.q_xy, self.res)
                    play(self.xy_op, self.q_xy)
                    if not self.overlap:
                        align()
                    measure(self.ro_op, self.res, None,
                            dual_demod.full('integW_cos', 'out1',
                                            'integW_sine', 'out2', I),
                            dual_demod.full('integW_minus_sine', 'out1',
                                            'integW_cos', 'out2', Q))
                    save(I, I_st)
                    save(Q, Q_st)
                
                assign(count, count + 1)
                save(count, count_st)
            with stream_processing():
                I_st.buffer(len(f_if_vec)).average().save('I')
                Q_st.buffer(len(f_if_vec)).average().save('Q')
                if self.subtract_bg:
                    I_bg_st.average().save('I_bg')
                    Q_bg_st.average().save('Q_bg')
                count_st.save('data_count')
            
        return q_spec_prog

    def _initialize_params(self, config: QuaConfig, context: EntropyContext):
        qpu_db = context.get_resource('qpu_db')

        # update RF, LO, and IF frequencies
        self.xy_f_rf, self.xy_f_lo, self.xy_f_if = update_rf_lo_if_freqs(
            self.element, 'q', config, context, f_rf=self.xy_f_rf0)
        qpu_db.set(resolve.q(self.element), 'f_rf', self.xy_f_rf)
        qpu_db.set(resolve.q(self.element), 'f_lo', self.xy_f_lo)
        qpu_db.set(resolve.q(self.element), 'f_if', self.xy_f_if)
        # qpu_db.set('XY_LO', 'frequency', self.xy_f_lo)

        self.ro_f_rf, self.ro_f_lo, self.ro_f_if = update_rf_lo_if_freqs(
            self.element, 'res', config, context)
        qpu_db.set(resolve.res(self.element), 'f_rf', self.ro_f_rf)
        qpu_db.set(resolve.res(self.element), 'f_lo', self.ro_f_lo)
        qpu_db.set(resolve.res(self.element), 'f_if', self.ro_f_if)
        # qpu_db.set('RO_LO', 'frequency', self.ro_f_lo)
        
        qpu_db.commit('Updated ' + ', '.join([resolve.q(self.element),
                                              resolve.res(self.element)]) +
                      ' RF LO IF frequencies.')

    def run_program(self, config: QuaConfig, context: EntropyContext):
        if self.debug:
            print('NODE INFO: ' + self._label + ' (' + self.node_type + ')' + 
                  ', RUN PROGRAM')
            
        self._initialize_params(config, context)
        qmm = context.get_resource("qmm")

        f_if_min, f_if_max, df = \
            int(self.xy_f_if - self.span / 2), int(self.xy_f_if + self.span / 2), int(self.df)
        f_if_vec = np.arange(f_if_min, f_if_max, df)
        f_rf_vec = self.xy_f_lo + f_if_vec
        self.f_rf_vec = f_rf_vec

        qm = qmm.open_qm(config)
        # self.job = qm.execute(self._qua_program())
        self.job = qm.simulate(self._qua_program(), SimulationConfig(1000))
        res_handles = self.job.result_handles
        res_handles.wait_for_all_values()
        
        #     self._fetch_sig(res_handles)
        # self.proj_sig = EasyReadout(self.sig, readout_type='magnitude').projected_data

        # res_idx = simple_resonance_idx_finder(self.proj_sig)
        # f_q = f_rf_vec[res_idx]
        # if self.plot:
        #     plt.figure()
        #     plt.title("Qubit Spectroscopy " + resolve.q(self.element) + ": %.4f GHz" % (f_q / 1e9))
        #     plt.plot(f_rf_vec / 1e9, self.proj_sig)
        #     plt.axvline(f_q / 1e9, color='black', ls='--', lw=1)
        #     plt.xlabel("Frequency (GHz)")
        #     plt.ylabel("Magnitude")
        #     plt.pause(0.1)
        #     plt.draw()

        # if self.LabberLog:
        #     filenm = f'q{self.element}_spectroscopy'
        #     log = create_uniqueLabberLogFile(
        #         filenm, [dict(name='Signal', x_name='XY Frequency',
        #                       x_unit='Hz', complex=True, vector=True)])
        #     log.addEntry({'Signal': Labber.getTraceDict(self.sig, 
        #                                                 x0=self.f_rf_vec[0],
        #                                                 x1=self.f_rf_vec[-1])})

        # # qpu db updating
        qpu_db = context.get_resource('qpu_db')
        # qpu_db.set(resolve.q(self.element), 'f_rf', f_q, 'FINE')
        # qpu_db.set(resolve.q(self.element), 'f_lo', self.xy_f_lo, 'FINE')
        # qpu_db.set(resolve.q(self.element), 'f_if', f_q - self.xy_f_lo, 'FINE')
        # # qpu_db.set('XY_LO', 'frequency', self.xy_f_lo, 'FINE')
        # qpu_db.commit()

        if self.callback is not None:
            self.callback(self)

        print("Calibrated Qubit Frequency:")
        print(
            "\t" +
            ', '.join(
                ["RF: %.4f GHz" % (qpu_db.q(self.element).f_rf.value / 1e9),
                 "LO: %.4f GHz" % (qpu_db.q(self.element).f_lo.value / 1e9),
                 "IF: %.4f MHz" % (qpu_db.q(self.element).f_if.value / 1e6)]
                )
            )

    def update_config(self, config: QuaConfig, context: EntropyContext) -> QuaConfig:
        if self.debug:
            # UPDATE CONFIG from qpu_db
            print('NODE INFO: ' + self._label + ' (' + self.node_type + ')' + 
                  ', UPDATE CONFIG')

        return config

