# Control functions for the HP8657B signal generators via GPIB

import time
import pyvisa as visa
import argparse

def connectHP(addrStr='7'):
  # Create the ResourceManager
  rm = visa.ResourceManager()

  # Connect to the signal generator
  my_instrument = rm.open_resource('GPIB0::'+addrStr+'::INSTR')

  # Assert GPIB-based control
  my_instrument.control_ren(visa.constants.VI_GPIB_REN_ASSERT)

  return my_instrument

def setFreqHP(my_instrument, freqMHz=1580.42):
  my_instrument.write("FR "+str(freqMHz)+" MZ")

def setRFOff(my_instrument):
  my_instrument.write("R2")

def setRFOn(my_instrument):
  my_instrument.write("R3")

def setAmpHP(my_instrument, ampdBm=-80.0):
  my_instrument.write("AP "+str(ampdBm)+"DM")

def defaultHP(my_instrument):
  # Set CW frequency with RF output off
  setRFOff(my_instrument)
  setFreqHP(my_instrument)
  setAmpHP(my_instrument)

def disconnectHP(my_instrument):
  # De-assert GPIB-based control
  my_instrument.control_ren(visa.constants.VI_GPIB_REN_DEASSERT)

  # Dis-connect from the signal generator
  my_instrument.close()

if __name__ == "__main__":
  parser = argparse.ArgumentParser()
  parser.add_argument('--addr', default='7', help='GPIB address # (def. 7)')
  parser.add_argument('--rfon', action='store_true', help='Enable RF')
  parser.add_argument('--rfoff', action='store_true', help='Disable RF')
  parser.add_argument('--rfdef', action='store_true', help='Default RF')
  parser.add_argument('--freq', default=None, help='Set frequency [MHz]')
  parser.add_argument('--amp', default=None, help='Set amplitude [dBm]')
  args = parser.parse_args()

  my_instrument = connectHP(args.addr)

  if args.rfdef:
    print("Go to default")
    defaultHP(my_instrument)
  elif args.rfon:
    print("Enable RF output")
    setRFOn(my_instrument)
  elif args.rfoff:
    print("Disable RF output")
    setRFOff(my_instrument)

  if args.freq is not None:
    print("Set frequency to {} MHz".format(args.freq))
    setFreqHP(my_instrument, float(args.freq))

  if args.amp is not None:
    print("Set amplitude to {} dBm".format(args.amp))
    setAmpHP(my_instrument, float(args.amp))

  disconnectHP(my_instrument)