import os.path
import sys
import array
import time
import serial
import math

# Special characters to be escaped
LF   = 0x0A
CR   = 0x0D
ESC  = 0x1B
PLUS = 0x2B


#==============================================================================
def IsSpecial(data):
    return data in (LF, CR, ESC, PLUS)

    
#==============================================================================
def CheckError():
    ser.write("SYST:ERR?\n")
    ser.write("++read eoi\n")
    s = ser.read(100)
    print s

# load a waveform from a file
# note that the points have to be scaled from -2047 to + 2047
def load_file(fname):
    data = []
    with open(fname, 'r') as fd:
        for line in fd:
            pt = float(line)
            data.append(pt)

    max_pt = max(data)
    min_pt = min(data)
    scale_pt = 0
    if (math.fabs(max_pt)>math.fabs(min_pt)):
        scale_pt = math.fabs(max_pt)
    else:
        scale_pt = math.fabs(min_pt)


    # Waveform points are in short int (16-bit) array
    dev_data = array.array('h');

    # Create waveform data. Simple ramp up and down.    
    for pt in data:
        int_val = int(pt * (2047/scale_pt))
        print int_val
        dev_data.append(int_val)

    # Swap bytes so MSB is first. (Required on Windows)
    dev_data.byteswap()

    # Output data is in byte array
    outdata = array.array('B');

    # Build output data, escaping all special characters        
    for byte in dev_data.tostring():
        if IsSpecial(ord(byte)):
            outdata.append(ESC)
        outdata.append(ord(byte))                    

    return (len(data),outdata)



#==============================================================================
if __name__ == '__main__':

    if len( sys.argv ) != 4:
        print "Usage: ", os.path.basename( sys.argv[0] ), "<COM port> <HP33120A GPIB address> <filename>"
        sys.exit(1)

    # Prologix GPIB-USB Controller serial port
    comport = sys.argv[1];

    # HP33120A GPIB address
    addr = sys.argv[2]
    
    # Number of waveform data points
    fname = sys.argv[3]

    points, outdata = load_file(fname)
    
    # 8-16000 points required
    if points < 8:
        print "Too few points."
        exit(1)

    if points > 16000:
        print "Too many points."
        exit(1)

    try:

        # Open serial port
        ser = serial.Serial( comport, 9600, timeout=0.5 )

        # Set mode as CONTROLLER
        ser.write("++mode 1\n")
        
        # Set HP33120A address
        ser.write("++addr " + addr + "\n")

        # Turn off read-after-write to avoid "Query Unterminated" errors
        ser.write("++auto 0\n")

        # Do not append CR or LF to GPIB data
        ser.write("++eos 3\n")

        # Assert EOI with last byte to indicate end of data
        ser.write("++eoi 1\n")

        # Reset AWG 
        cmd = "*RST"
        print cmd
        ser.write(cmd + "\n")        

        time.sleep(1.0)
        CheckError()        

        # Format output data command. Use length of points array, NOT output array
        datalen = points * 2
        
        cmd = "DATA:DAC VOLATILE, #" + str(len(str(datalen))) + str(datalen)
        print cmd
        
        # Write binary block data
        ser.write(cmd)
        ser.write(outdata.tostring())

        # Terminate USB command string
        ser.write("\n")

        time.sleep(0.5)
        CheckError()        
        
        cmd = "DATA:COPY PULSE, VOLATILE"
        print cmd

        ser.write(cmd + "\n")
        time.sleep(2.0)     
        CheckError()        

        cmd = "FUNC:USER PULSE"
        print cmd

        ser.write(cmd + "\n")
        time.sleep(1)
        CheckError()        

        cmd = "FUNC:SHAP USER"
        print cmd

        ser.write(cmd + "\n")
        time.sleep(0.5)
        CheckError()        

        cmd = "OUTP:LOAD 50"
        print cmd

        ser.write(cmd + "\n")
        time.sleep(0.5)
        CheckError()        

        cmd = "FREQ 5000;VOLT 0.5"
        print cmd

        ser.write(cmd + "\n")
        time.sleep(0.5)
        CheckError()        

        
    except serial.SerialException, e:
        print e