from pulse_templates.coherent_control.single_qubit_gates.phase_offsets_charac import phase_offset_charac
from core_tools.sweeps.pulse_lib_wrappers.PSB_exp import run_qubit_exp
from good_morning.fittings.fit_resonance import fit_resonance
from core_tools.sweeps.sweeps import scan_generic, do0D, do1D
from dev_V2.six_qubit_QC.system import six_dot_sample
from core_tools.utility.variable_mgr.var_mgr import variable_mgr
import qcodes as qc
import numpy as np
from pulse_lib.segments.utility.looping import linspace
from pulse_templates.coherent_control.single_qubit_gates.single_qubit_gates import single_qubit_gate_spec

def get_target(obj, name):
	name = name.split('.')

	for operator in name:
		obj = getattr(obj, operator)
	
	return obj

def CROT_cali_meas(pair, target, ancilla, old_freq, time, flip):

	s = six_dot_sample(qc.Station.default.pulse)
	var_mgr = variable_mgr()

	pair = str(int(pair))
	s.init(pair[0], pair[1])

	s.wait(100).add(s.manip)
	s.pre_pulse.add(s.manip)
	s.wait(100).add(s.manip)

	gate_set = getattr(s, f'q{pair}')
	target_gate_set = getattr(s, f'q{pair[target-1]}')
	ancilla_gate_set = getattr(s, f'q{pair[ancilla-1]}')

	if flip:
		ancilla_gate_set.X180.add(s.manip)
		# target_gate_set.X180.add(s.manip)

	scan_range = single_qubit_gate_spec(f'qubit{target}_MW', 
                                    linspace(old_freq*1e9-10e6, old_freq*1e9+10e6, 100, axis= 0, name='freq', unit='Hz'),
                                    time, getattr(var_mgr, f'q{pair[target-1]}_MW_power'), padding=2)

	gate_to_test = getattr(gate_set, f'CROT{target}{ancilla}')

	gate_to_test.add(s.manip, gate_spec = scan_range)
	s.wait(100).add(s.manip)
	s.read(pair[target-1])
	sequence, minstr, name = run_qubit_exp(f'crot-z_crot_cal_q{pair}_target{pair[target-1]}', s.segments(), s.measurement_manager)

	return sequence, minstr, name

def CROT_calib(pair, target, ancilla, plot=False):
	'''
	calibrates a single qubit phase for a target qubit

	Args:
		target (int) : qubit number to target (1-6)
		ancillary (list) : qubit that need to be initialized to make this work
	'''
	var_mgr = variable_mgr()
	pair = str(int(pair))

	old_crot = getattr(var_mgr, f'crot{pair[target-1]}{pair[ancilla-1]}')
	old_z_crot = getattr(var_mgr, f'z_crot{pair[target-1]}{pair[ancilla-1]}')
	crot_time = getattr(var_mgr, f'pi_crot{pair[target-1]}_m3dbm')
	z_crot_time = getattr(var_mgr, f'pi_z_crot{pair[target-1]}_m3dbm')

	gate_time= min(crot_time,z_crot_time)	

	sequence, minstr, name = CROT_cali_meas(pair, target, ancilla, old_crot, gate_time, 0)
	ds_one = scan_generic(sequence, minstr, name=name).run()
	sequence, minstr, name = CROT_cali_meas(pair, target, ancilla, old_z_crot, gate_time, 1)
	ds_two = scan_generic(sequence, minstr, name=name).run()

	frequency_one = ds_one(f'read{pair[target-1]}').x()
	probabilities_one = ds_one(f'read{pair[target-1]}').y()
	fit_freq_one = fit_resonance(frequency_one, probabilities_one, plot=plot)
	
	frequency_two = ds_two(f'read{pair[target-1]}').x()
	probabilities_two = ds_two(f'read{pair[target-1]}').y()
	fit_freq_two = fit_resonance(frequency_two, probabilities_two, plot=plot)

	z_crot_res = min(fit_freq_one, fit_freq_two)
	crot_res = max(fit_freq_one, fit_freq_two)

	old_z_crot = getattr(var_mgr, f'z_crot{pair[target-1]}{pair[ancilla-1]}')
	setattr(var_mgr, f'z_crot{pair[target-1]}{pair[ancilla-1]}', round(z_crot_res*1e-9,6))
	old_crot = getattr(var_mgr, f'crot{pair[target-1]}{pair[ancilla-1]}')
	setattr(var_mgr, f'crot{pair[target-1]}{pair[ancilla-1]}', round(crot_res*1e-9,6))
	print(f'calibrated z_crot_res for qubit pair {pair}, target {pair[target-1]} '+
		f'Old z_crot_res : {round(old_z_crot,6)} \n New z_crot_res : {round(z_crot_res*1e-9,6)} \n')
	print(f'calibrated crot_res for qubit pair {pair}, target {pair[target-1]} '+
		f'Old crot_res : {round(old_crot,6)} \n New z_crot_res : {round(crot_res*1e-9,6)} \n')