#!/usr/bin/env python

# Copyright (c) 2021
#      Daniel O'Connor <darius@dons.net.au>.  All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY AUTHOR AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL AUTHOR OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
# OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
# SUCH DAMAGE.
#

import sys
import usb
import usb488

#
# Tested with a Rigol DSG815 connected via USB
#

def main():
    u = usb488.USB488Device()
    res = u.ask('*IDN?')
    print('Found device ID: ' + res)

    setup(u)

def freqsetup(u, swstart, swstop, swpoints, swoffset, swlevel):
    '''Setup a frequency sweep'''

    # Output off
    wrcheck(u, ':OUTP:STATE', 'OFF')
    # Single sweep
    wrcheck(u, ':SWE:MODE', 'SING')
    # Frequency sweep
    wrcheck(u, ':SWE:STATE', 'FREQ')
    # Step (not list)
    wrcheck(u, ':SWE:TYPE', 'STEP')
    # In single sweep so it starts when SOU:SWE:EXE is sent
    wrcheck(u, ':SWE:SWE:TRIG:TYPE', 'AUTO')
    # Sweep a point every external trigger
    wrcheck(u, ':SWE:POIN:TRIG:TYPE', 'EXT')
    # Sweep up
    wrcheck(u, ':SWE:DIR', 'FWD')
    # Ramp (vs triangle)
    wrcheck(u, ':SWE:STEP:SHAP', 'RAMP')
    # Linear sweep
    wrcheck(u, ':SWE:STEP:SPAC', 'LIN')
    # Start/stop/points in sweep
    wrcheck(u, ':SWE:STEP:STAR:FREQ', '%fHz' % (swstart + swoffset,))
    wrcheck(u, ':SWE:STEP:STOP:FREQ', '%fHz' % (swstop + swoffset,))
    wrcheck(u, ':SWE:STEP:POIN', '%d' % (swpoints,))
    # Output level
    wrcheck(u, ':SWE:STEP:STAR:LEV', '%fdBm' % (swlevel,))
    wrcheck(u, ':SWE:STEP:STOP:LEV', '%fdBm' % (swlevel,))
    wrcheck(u, ':LEV', '%fdBm' % (swlevel,))

def powersetup(u, swstart, swstop, swpoints, swfreq, swoffset):
    '''Setup a power sweep'''

    # Output off
    wrcheck(u, ':OUTP:STATE', 'OFF')
    # Single sweep
    wrcheck(u, ':SWE:MODE', 'SING')
    # Level sweep
    wrcheck(u, ':SWE:STATE', ' LEV')
    # Step (not list)
    wrcheck(u, ':SWE:TYPE', 'STEP')
    # In single sweep so it starts when SOU:SWE:EXE is sent
    wrcheck(u, ':SWE:SWE:TRIG:TYPE', 'AUTO')
    # Sweep a point every external trigger
    wrcheck(u, ':SWE:POIN:TRIG:TYPE', 'EXT')
    # Sweep up
    wrcheck(u, ':SWE:DIR', 'FWD')
    # Ramp (vs triangle)
    wrcheck(u, ':SWE:STEP:SHAP', 'RAMP')
    # Linear sweep
    wrcheck(u, ':SWE:STEP:SPAC', 'LIN')
    # Start/stop/points in sweep
    wrcheck(u, ':LEV', '%fdBm' % (swstart,))
    wrcheck(u, ':SWE:STEP:STAR:LEV', '%fdBm' % (swstart,))
    wrcheck(u, ':SWE:STEP:STOP:LEV', '%fdBm' % (swstop,))
    wrcheck(u, ':SWE:STEP:POIN', '%d' % (swpoints,))
    # Output frequency
    wrcheck(u, ':SWE:STEP:STAR:FREQ', '%fHz' % (swfreq + swoffset,))
    wrcheck(u, ':SWE:STEP:STOP:FREQ', '%fHz' % (swfreq + swoffset,))
    wrcheck(u, ':FREQ', '%fHz' % (swfreq,))
    # Level sweep

    # XXX: need to send this twice for some reason
    # isn't needed for frequency sweep
    wrcheck(u, ':SWE:STATE', ' LEV')

def arm(u):
    # Output on
    wrcheck(u, ':OUTP:STATE', 'ON')
    # Arm it
    u.chkcmd(':SWE:EXEC')

def wrcheck(u, name, value):
    '''Helper function to set a value then read it back and print it out'''
    u.chkcmd(name + ' ' + value)
    readval = u.ask(name + '?')
    print(name + ' -> ' + readval + ' (' + value + ')')
    #sys.stdin.readline()

if __name__ == '__main__':
    main()