#!/usr/bin/env python
#===============================================================================
# CERN (BE-CO-HT)
# Agilent 33250A pulse generator controller
#===============================================================================
# author: Maciej Suminski (maciej.suminski@cern.ch)
# date of creation: 02.02.2017
# version: 1.0
#
# description:
#   This script is used to configure an Agilent 33250A signal generator
#   to generate an arbitrary burst of pulses via serial port interface.
#
#   Be sure that both the generator and the serial port have the same settings:
#   - **RTS/CTS handshake enabled**
#   - the same baud rate (recommended 115200)
#   - parity bit disabled
#   - 1 bit start, 1 bit stop,
#
#===============================================================================
# GNU LESSER GENERAL PUBLIC LICENSE
#===============================================================================
# This source file is free software; you can redistribute it and/or modify it
# under the terms of the GNU Lesser General Public License as published by the
# Free Software Foundation; either version 2.1 of the License, or (at your
# option) any later version. This source is distributed in the hope that it
# will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty
# of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Lesser General Public License for more details. You should have
# received a copy of the GNU Lesser General Public License along with this
# source; if not, download it from http://www.gnu.org/licenses/lgpl-2.1.html
#===============================================================================
# last changes:
#===============================================================================
#  TODO: -
#===============================================================================

import sys
import serial
import getopt
import math

# default parameters
device = '/dev/ttyUSB0'
baud = 115200
count = 2000000
freq = 2e6 # Hz
width = 250e-9 # s
samples_count = 500 # number of samples for arbitrary waveform

print('Agilent 33250A pulse burst generator\r\n')

def usage():
    print('(c) CERN 2017')
    print('Maciej Suminski <maciej.suminski@cern.ch>')
    print('')
    print('usage: %s [opts]' % sys.argv[0])
    print('Options:')
    print('-b|--baud=       serial port baud rate (default: %d)' % baud)
    print('-d|--device=     serial port device path (default: %s' % device)
    print('')
    print('-c|--count=      number of requested pulses (default: %d)' % count)
    print('-f|--freq=       frequency of the pulses [Hz] (default: %G)' % freq)
    print('-w|--width=      pulse width [s] (default: %G)' % width)
    print('-s|--samples=    arbitrary waveform samples count (default: %d)' % samples_count)
    print('')
    print('-h|--help        shows this information')

# generator limits
# max number of pulses in a single burst, higher than that is done using arbitrary waveform
MAX_COUNT=1000000

try:
    opts, args = getopt.getopt(sys.argv[1:], 'b:c:d:f:hw:',
            ['baud=', 'count=', 'device=', 'freq=', 'help', 'width='])
except getopt.GetoptError as err:
    print(str(err))
    usage()
    sys.exit(2)

for opt, arg in opts:
    if opt in ('-b', '--baud'):
        baud = int(arg)
    elif opt in ('-c', '--count'):
        count = int(arg)
    elif opt in ('-d', '--device'):
        device = arg
    elif opt in ('-f', '--freq'):
        freq = float(arg)
    elif opt in ('-s', '--samples'):
        samples = int(arg)

        if samples_count > 65536:
            print('ERROR: Samples number limit is 65536')
            sys.exit(2)

    elif opt in ('-w', '--width'):
        width = float(arg)
    elif opt in ('-h', '--help'):
        usage()
        sys.exit()
    else:
        assert False, 'unhandled option'

###############################################################################
# Connect to the generator and reset it

try:
    print('Connecting to %s @ %s' % (device, baud))
    ser = serial.Serial(device, baud, timeout=1, rtscts=1)
except serial.serialutil.SerialException as err:
    print(str(err))
    sys.exit(2)

# flush buffer that may contain an invalid command
ser.write(b'\r\n')

# verify the device is there
ser.write(b'*IDN?\r\n')
ident = ser.readline().decode()

if(ident == ''):
    print('ERROR: Could not identify the generator')
    sys.exit(2)

print('Read device ID: %s' % ident)

if not '33250A' in ident:
    print('ERROR: Invalid generator type')
    sys.exit(2)

# reset to a known state
ser.write(b'OUTP OFF\r\n')
ser.write(b'*RST\r\n')

# clear errors
while True:
    ser.write(b'SYST:ERR?\r\n')
    status = ser.readline().decode()

    if status == '':
        print('ERROR: Could not read the error queue')
        sys.exit(2)

    if status.startswith('+0'):
        break

###############################################################################
# Select the way of generating pulses
# 33250A has a limit of 1e6 pulses in a single burst. To overcome this
# limitation, one can create an arbitrary waveform containing n pulses
# and then request n times pulses less.
if(count <= MAX_COUNT):
    # normal case, just use the pulse generator, no tricks needed
    div = 1
    real_count = count

    # rising/falling edge time
    edge = 5e-9

    print('Configuring pulse waveform')
    ser.write(b'PULS:TRAN %G\r\n' % edge)
    ser.write(b'PULS:WIDT %G\r\n' % width)
    ser.write(b'PULS:PER %G\r\n' % (1.0 / freq))
    ser.write(b'FUNC:SHAP PULS\r\n')
else:
    orig_count = count
    orig_width = width

    # deal with the limitation described above
    # number of pulses in the arbitrary waveform
    div = int(math.floor(count / MAX_COUNT)) + 1
    period = 1.0 / freq

    if width > period * 0.9:
        print('ERROR: Pulse width is greater than 90% of the pulse period')
        sys.exit(2)

    # generate the string that represents the requested waveform
    duty_cycle = width / period
    one_count = int(samples_count * duty_cycle)
    zero_count = samples_count - one_count

    # generator DAC values mapping
    # see VOLT and VOLT:OFFS commands below to see the details
    one_value = b', 2047'
    zero_value = b', -2047'

    # single pulse (take into account the duty cycle)
    chunk = b'%s%s' % (zero_count * zero_value, one_count * one_value )
    # full waveform (repeat the pulse 'div' times)
    pattern = chunk * div

    # update all variables basing on the divider
    freq = freq / div
    count = int(count / div)
    width = duty_cycle * period
    # the actual number of pulses is rounded to the number of pulses in
    # the arbitrary waveform
    real_count = count * div

    if width != orig_width:
        print('WARNING: Pulse width had to be rounded to %G (requested %G)' % (width, orig_width))

    if real_count != orig_count:
        print('WARNING: Number of pulses had to be rounded to %d (requested %d)' % (real_count, orig_count))

    # set the arbitrary waveform (a number of consecutive pulses)
    print('Configuring arbitrary waveform')
    ser.write(b'FREQ %G\r\n' % freq)
    ser.write(b'DATA:DAC VOLATILE %s\r\n' % pattern)
    # set the arbitrary waveform to the one we have just uploaded
    ser.write(b'FUNC:USER VOLATILE\r\n')
    # switch the generator to arbitrary waveform
    ser.write(b'FUNC USER\r\n')


###############################################################################
# Common part

# peak-to-peak voltage 5V, offset = 2.5V
# makes arbitrary waveform values mapping:
# -2047 => 0V
# 2047  => +5V
ser.write(b'VOLT 5\r\n')
ser.write(b'VOLT:OFFS 2.5\r\n')

# configure software trigger (executed by *TRG command)
ser.write(b'TRIG:SOUR BUS\r\n')

# configure burst mode (number of cycles, software triggered, 0 phase)
ser.write(b'BURS:MODE TRIG\r\n')
ser.write(b'BURS:NCYC %d\r\n' % count)
ser.write(b'BURS:PHAS 0\r\n')
ser.write(b'BURS:STAT ON\r\n')

# check errors queue
error = False
while True:
    ser.write(b'SYST:ERR?\r\n')
    status = ser.readline().decode()

    if status.startswith('+0'): # no more errors
        break

    print('ERROR: Generator error occurred: %s' % status)
    error = True

if error:
    sys.exit(2) # we should not proceed, the results are unknown

print('Sending %d pulses (%g s) at %g Hz' % (real_count, width, freq * div))

# enable output for load = 50 ohms
ser.write(b'OUTP:LOAD 50\r\n')
ser.write(b'OUTP ON\r\n')

# trigger the output
ser.write(b'*TRG\r\n')
ser.write(b'*WAI\r\n')
ser.write(b'OUTP OFF\r\n')
ser.close()

print('Finished successfully')
