# -*- coding: iso-8859-1 -*-

TXEnableDelay = 2e-6
TXEnableValue = 0b0001  # TTL line blanking RF amplifier (bit 0)
TXPulseValue  = 0b0010  # TTL line triggering RF pulses  (bit 1)
ADCSensitivity = 2	# voltage span for ADC

def experiment(): # ZZ-exchange with T2-selection  
    
  # set up acqusition parameters:      
    pars = {}
    pars['P90'] = 1.7e-6	# 90-degree pulse length (s)
    pars['SF'] = 338.7e6	# spectrometer frequency (Hz)
    pars['O1'] = -60e3  	# offset from SF (Hz)
    pars['SW'] = 200e3 		# spectral width (Hz)
    pars['SI'] = 128	 	# acquisition points
    pars['D1'] = 10e-6   	# echo delay (s)
    pars['D2'] = 100e-6         # z-storage duration (s)
    pars['NS'] = 16      	# number of scans
    pars['DS'] = 0      	# number of dummy scans
    pars['RD'] = 3      	# delay between scans (s)
    pars['DEAD1'] = 5e-6	# receiver dead time (s)  
    pars['PHA'] = 0     	# receiver phase (degree)
    pars['DATADIR'] = '/home/fprak/Students/'	# data directory
    pars['OUTFILE'] = None                            	# output file name
    
  # specify a variable parameter:
    pars['VAR_PAR'] = 'D2'      	# variable parameter name
    start = 20e-6          	# first value
    stop  = 1e-3      	 	# last value
    steps = 20  	      	# number of values
    log_scale = True		# log scale flag 
    stag_range = False		# staggered range flag
    
  # check parameters for safety:     
    if pars['PHA'] < 0:
        pars['PHA'] = 360 + pars['PHA']
        
    if pars['P90'] > 20e-6:
        raise Exception("Pulse too long!!!")                    

    if pars['NS']%16 != 0:
        pars['NS'] = (pars['NS'] / 16 + 1) * 16 
        print 'Number of scans changed to', pars['NS'], 'due to phase cycling'

    # check whether a variable parameter is named: 
    var_key = pars.get('VAR_PAR')
    if var_key == 'P90' and (start > 20e-6 or stop > 20e-6):
        raise Exception("Pulse too long!!!")

    if pars['NS']%8 != 0:
        pars['NS'] = int(round(pars['NS'] / 8) + 1) * 8 
        print 'Number of scans changed to ',pars['NS'],' due to phase cycling'

    # start the experiment:
    if var_key:
        # this is an arrayed experiment:
        if log_scale:
            array = log_range(start,stop,steps)
        else:
            array = lin_range(start,stop,steps)

        if stag_range:
            array = staggered_range(array, size = 2)

        # estimate the experiment time:
        if var_key == 'D1':
            seconds = (sum(array)*2 + (pars['D2'] + pars['RD']) * steps) * (pars['NS'] + pars['DS'])
        elif var_key == 'D2':
            seconds = (sum(array) + (pars['D1']*2 + pars['RD']) * steps) * (pars['NS'] + pars['DS'])
        elif var_key == 'RD':
            seconds = (sum(array) + (pars['D1']*2 + pars['D2']) * steps) * (pars['NS'] + pars['DS'])
        else:
            seconds = (pars['D1']*2 + pars['D2'] + pars['RD']) * steps * (pars['NS']+ pars['DS'])
        m, s = divmod(seconds, 60)
        h, m = divmod(m, 60)
        print '%s%02d:%02d:%02d' % ('Experiment time estimated: ', h, m, s)

        # loop for a variable parameter:
        for index, pars[var_key] in enumerate(array):
            print 'Arrayed experiment for '+var_key+': run = '+str(index+1)+\
                             ' out of '+str(array.size)+': value = '+str(pars[var_key])
            # loop for accumulation:
            for run in xrange(pars['NS']+pars['DS']):
                yield t2zz_experiment(pars, run) 
            synchronize()
    else:
       # estimate the experiment time:
        seconds = (pars['D1']*2 + pars['D2'] + pars['RD']) * (pars['NS']+ pars['DS'])
        m, s = divmod(seconds, 60)
        h, m = divmod(m, 60)
        print '%s%02d:%02d:%02d' % ('Experiment time estimated: ', h, m, s)

        # loop for accumulation:
        for run in xrange(pars['NS']+pars['DS']):
            yield t2zz_experiment(pars, run)

# the pulse program: 

def t2zz_experiment(pars, run):
    e=Experiment()

    dummy_scans = pars.get('DS')
    if dummy_scans:
        run -= dummy_scans
        
    pars['PROG'] = 't2zz_experiment'
  
    # eliminates T1 recovery and contaminating echo signals:
    pars['PH1'] = [0, 180, 90, 270, 180, 0, 270, 90,       180, 0, 270, 90, 0, 180, 90, 270]
    pars['PH3'] = [90, 90, 180, 180, 270, 270, 0, 0,       270, 270, 0, 0, 90, 90, 180, 180]	
    pars['PH4'] = [0, 0, 90, 90, 180, 180, 270, 270,       180, 180, 270, 270, 0, 0, 90, 90]
    pars['PH5'] = [0, 0, 90, 90, 180, 180, 270, 270,       0, 0, 90, 90, 180, 180, 270, 270]
    pars['PH2'] = [0, 180, 90, 270, 180, 0, 270, 90,       0, 180, 90, 270, 180, 0, 270, 90]    # data routing

    # read in variables:
    P90 = pars['P90'] 
    SF = pars['SF']
    O1 = pars['O1']   
    RD = pars['RD']
    DEAD1 = pars['DEAD1'] 
    D1 = pars['D1']
    D2 = pars['D2']
    PH1 = pars['PH1'][run%len(pars['PH1'])]
    PH3 = pars['PH3'][run%len(pars['PH3'])]
    PH4 = pars['PH4'][run%len(pars['PH4'])]
    PH5 = pars['PH5'][run%len(pars['PH5'])]
    PH2 = pars['PH2'][run%len(pars['PH2'])]
    PHA = pars['PHA']

    # sampling parameters:
    SI = pars['SI']
    SW = pars['SW']
    while SW <= 10e6 and SI < 256*1024:
        SI *= 2
        SW *= 2
    
    # run the pulse sequence: 
    e.wait(RD)                                  	# delay between scans
    e.set_frequency(SF+O1, phase=PH1)
    e.ttl_pulse(TXEnableDelay, value=TXEnableValue)
    e.ttl_pulse(P90, value=TXEnableValue|TXPulseValue)	# 1st pulse

    e.wait(D1-P90-TXEnableDelay)     	              	# echo delay
    e.set_phase(PH3)

    e.ttl_pulse(TXEnableDelay, value=TXEnableValue)		
    e.ttl_pulse(P90, value=TXEnableValue|TXPulseValue)	# 2nd pulse (90- or 180-degree)
    
    e.wait(D1-P90-TXEnableDelay)     	              	# echo delay
    e.set_phase(PH4)
    
    e.ttl_pulse(TXEnableDelay, value=TXEnableValue)		
    e.ttl_pulse(P90, value=TXEnableValue|TXPulseValue)  # 3rd pulse (z-storage pulse)

    e.wait(D2-P90-TXEnableDelay)    	              	# mixing time
    e.set_phase(PH5)

    e.ttl_pulse(TXEnableDelay, value=TXEnableValue)		
    e.ttl_pulse(P90, value=TXEnableValue|TXPulseValue)	# 4th pulse (readout pulse)

    e.set_phase(PHA)                                    # set phase for receiver 
    e.wait(DEAD1)                                       # wait for coil ringdown
    e.record(SI, SW, sensitivity=ADCSensitivity)       	# acquire signal

    # write experiment attributes:  
    for key in pars.keys():
        e.set_description(key, pars[key])		# acqusition parameters
    e.set_description('run', run)               	# current scan
    e.set_description('rec_phase', -PH2)		# current receiver phase

    return e