#This script generates an increasing ramp of SampleNum point on analog channels and 
#a 16 bit counter on digital lines, then send the mixed waveform to the AWG using
#the Binary Block Data.
#The Digital Pod A/C represents the lower part, the Pod B/D the higher one.

#You can set the Number of points, the Amplitude, the Offset and the Sampling Frequency

#####################################################################################

#Parameters
InstrumentAddr = "TCPIP0::ARB-RIDER::inst0::INSTR"
Amplitude = 5
Offset = -2.5
SampleNum = 65536*4
SampleFreq = 250000000  #250MHz

#####################################################################################

import visa
import struct

#Create a Analog Waveform
AnalogSamples = []
i = 0
while (i<SampleNum) :
    AnalogSamples.append((Amplitude * i / SampleNum) + Offset)
    i=i+1

#Create Digital Waveform (16-bit counter)
DigitalSamples = []
i = 0
d = 0
while (i < SampleNum) :
    if ((i % 4) == 0) :
        DigitalSamples.append(d)
        d = d + 1
        if (d == 65536) :
            d = 0
    i=i+1    

#Create an array of byte to contain the mixed waveform
byteSignal = bytearray(6 * SampleNum)

#Loop on Samples
for i in range(0, SampleNum) :
    #Loop on byte of each samples
    byteAnSample = bytearray(struct.pack('f', AnalogSamples[i]))
    for j in range (0,4) :
        byteSignal[6 * i + j] = byteAnSample[j]
    #Add the digital data
    byteDigSample = DigitalSamples[(int)(i // 4)].to_bytes(2,'little')
    byteSignal[6 * i + 4] = byteDigSample[0]
    byteSignal[6 * i + 5] = byteDigSample[1]

    
try:
    #Search the VISA instrument
    rm = visa.ResourceManager()

    #Connect with th AWG (insert the AWG name)
    AWG = rm.open_resource(InstrumentAddr)

    #Set the VISA timeout
    AWG.timeout = 60000 #ms, it is 60s

    #Identify the instrument
    print("You are connected to the instrument:\n",AWG.query("*IDN?"))

    #Reset the instrument
    AWG.write("*RST")

    #Set the functionality
    AWG.write("AWGControl:CONFigure:SEQuencer:MODE SIN")
    AWG.write("AWGControl:CHANnel:FUNCTionality ARB, LOWSP")

    #Set the Continuous run mode
    AWG.write("AWGControl:RMODe CONT")
    
    #Create a new Waveform
    strCmd = 'WLISt:WAVeform:NEW "TestWfm",' + str(SampleNum)
    AWG.write(strCmd)


    #Send analog waveform
    strCmd = 'WLISt:WAVeform:DATA "TestWfm",0,' + str(SampleNum) + ','
    AWG.write_binary_values(strCmd, byteSignal, datatype='B', is_big_endian=False)

    #Assign the waveform to channels
    AWG.write('SOURCE1:WAVeform "TestWfm"')
    AWG.write('SOURCE2:WAVeform "TestWfm"')

    #Set the sampling frequency
    strCmd = "TIMING:SAMPLEFReq " + str(SampleFreq)
    AWG.write(strCmd)
    
    #Start the AWG
    AWG.write("AWGC:UPDATE")
    AWG.write("AWGC:RUN")

    #Query for errors until receive: "0,No error"
    print("Errors:")
    err = AWG.query("SYST:ERR?")
    print (err)
    while (err.find("No error") < 0) :
        err = AWG.query('SYST:ERR?')
        print (err)                   
  
    #Close the VISA resources
    AWG.close()
    rm.close()

except Exception as e:
    print("Error: instrument communication: ",e)


input("Press enter to exit")
