# Control functions for the HP8648C signal generators via GPIB

import time
import pyvisa as visa
import argparse

def connectHP(addrStr='18', verbose=False):
  # Create the ResourceManager
  rm = visa.ResourceManager()

  # Connect to the signal generator
  my_instrument = rm.open_resource('GPIB0::'+addrStr+'::INSTR')
  if verbose:
    print(my_instrument.query("*IDN?"))

  # 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("FREQ:CW "+str(freqMHz)+" MHZ")

def setRFOff(my_instrument):
  my_instrument.write("OUTP:STAT OFF")

def setRFOn(my_instrument):
  my_instrument.write("OUTP:STAT ON")

def setAmpHP(my_instrument, ampdBm=-80.0):
  if ampdBm > 5.0:
    print("CAUTION:")
    print("CW power limited to 5dBm due to the 1dBm input power limit on downstream amplifier")
    ampdBm = 5.0
  my_instrument.write("POW:AMPL "+str(ampdBm)+"DBM")

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='18', help='GPIB address # (def. 18)')
  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)