#!/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
#   (SCPI commands).
#
#   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,
#
#   SCPI commands for the generator are described in the Agilent 33250A manual
#   (http://cp.literature.agilent.com/litweb/pdf/33250-90002.pdf, Chapter 4)
#
#===============================================================================
# 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

class PulseGen:
    """
    Class to configure a pulse generator (Agilent 33250A).
    """

    def __init__(self, device="/dev/ttyUSB0", baud=115200):
        self.ser = serial.Serial(device, baud, timeout=1, rtscts=1)

        # generator limits
        # max number of pulses in a single burst, higher than that has to be
        # done using arbitrary waveform
        self.MAX_PULSE_COUNT = 1000000

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

        # verify the device is there
        ident = self.ident()

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

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

        self.reset()

    def __del__(self):
        if self.ser is not None:
            # disable generator output
            self.ser.write(b'OUTP OFF\r\n')
            self.ser.close()

    def ident(self):
        self.ser.write(b'*IDN?\r\n')
        return self.ser.readline().decode()

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

        # clear errors
        self.ser.write(b'*CLS\r\n')

    def pulse(self, count, freq, width, samples):
        # 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 <= self.MAX_PULSE_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')
            self.ser.write(b'PULS:TRAN %G\r\n' % edge)
            self.ser.write(b'PULS:WIDT %G\r\n' % width)
            self.ser.write(b'PULS:PER %G\r\n' % (1.0 / freq))
            self.ser.write(b'FUNC:SHAP PULS\r\n')
        else:
            # store the original values to compare them against the requested ones
            orig_count = count
            orig_width = width

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

            if width > period * 0.9:
                raise RuntimeError('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 * duty_cycle)
            zero_count = samples - 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

            # TODO these warnings should be reported in some other way
            if width != orig_width:
                print('WARNING: Pulse width rounded to %G (requested %G)' % (width, orig_width))

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

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


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

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

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

        # check errors queue and printout error messages if any
        error = False
        error_msg = ''

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

            if status.startswith('+0'): # no more errors
                break
            else:
                error = True
                error_msg = error_msg + status

        if error:
            raise RuntimeError('Generator error(s) occurred: %s' % error_msg)

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

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

        # trigger the output
        self.ser.write(b'*TRG\r\n')

    def wait(self):
        # wait for the operation to complete
        self.ser.write(b'*OPC?\r\n')

        while True:
            status = self.ser.readline().decode()

            if len(status) > 0:
                if not status.startswith('1'):
                    raise RuntimeError('Operation interrupted (status: %s)' % status)
                break

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

if __name__ == '__main__':
    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)
        print('')
        print('-h|--help        shows this information')

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

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

    # parse options
    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 > 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'


    print('Connecting to %s @ %s' % (device, baud))
    gen = PulseGen(device, baud)
    gen.reset()
    gen.pulse(count, freq, width, samples)
    gen.wait()
    print('Finished successfully')
