# Needed packages
import visa
import numpy as np
import matplotlib.pyplot as plt
import time
import math

# Identify and open connection to oscilloscope
rm = visa.ResourceManager()
# Check the position of USB address using the command rm.list_resources(). Below it is at 0.
keysight=rm.list_resources()[0]
inst = rm.open_resource(keysight)

# Setup oscilloscope
scale=50E-6
timewindow=10*scale
inst.write(':TIMebase:SCALe '+str(scale)) # 10xscale=time range
#inst.write(':TIMebase:RANGe 0.001') # Instead of time scale, one can use time range as shown here
inst.write(':TIMebase:DELay 0') # Center oscilloscope window around 0
inst.write(':WAVeform:POINts:MODE RAW') # Aquires raw data - not sure if this is needed

N=20000
SR=2.50e9
dt=1/SR
step=math.ceil(40e-9/dt)

timer=time.time()
file1=open('DATA/N10000_0.00005_trace1_Htest.txt','w') #BS1
file2=open('DATA/N10000_0.00005_trace2_Htest.txt','w') #BS2
file3=open('DATA/N10000_0.00005_trace3_Htest.txt','w') #Trigger

for i in range(N):
    inst.write(':SINGle') # Runs a single trace
    inst.write(':STOP') # Stop is needed, else it will keep running after single run
    inst.write(':WAVeform:SOURce CHANnel1') # Focus on channel 1
    trace=inst.query_binary_values(':WAVeform:DATA?', container=np.array, datatype='B')
    tmp=np.where(trace>(max(trace)+min(trace))/2)[0]
    file1.write(str((tmp[1]*dt+i*timewindow)*1E6)+'\n')
    for j in range(len(tmp)-1):
        if tmp[j+1]>(tmp[j]+step):
            file1.write(str((tmp[j+1]*dt+i*timewindow)*1E6)+'\n')

    inst.write(':WAVeform:SOURce CHANnel2') # Focus on channel 2
    trace=inst.query_binary_values(':WAVeform:DATA?', container=np.array, datatype='B')
    tmp=np.where(trace>(max(trace)+min(trace))/2)[0]
    file2.write(str((tmp[1]*dt+i*timewindow)*1E6)+'\n')
    for j in range(len(tmp)-1):
        if tmp[j+1]>(tmp[j]+step):
            file2.write(str((tmp[j+1]*dt+i*timewindow)*1E6)+'\n')

    inst.write(':WAVeform:SOURce CHANnel3') # Focus on channel 3
    trace=inst.query_binary_values(':WAVeform:DATA?', container=np.array, datatype='B')
    tmp=np.where(trace>(max(trace)+min(trace))/2)[0]
    file3.write(str((tmp[1]*dt+i*timewindow)*1E6)+'\n')
    for j in range(len(tmp)-1):
        if tmp[j+1]>(tmp[j]+step):
            file3.write(str((tmp[j+1]*dt+i*timewindow)*1E6)+'\n')

    print(str(i+1))

file1.close()
file2.close()
file3.close()
elapse1=time.time()-timer
print(str(elapse1))