#!/usr/bin/env python

# Richard Piersall
#
# VNA
#   Date Modified: 03-01-2016
#   Purpose:
#       Interface with HP 8753A or 8753D
#       Connect to VNA through NI USB-GPIB-HS adapter
#       Pull trace information
#


#TODO: Make it easier to switch averaging on and off
#TODO: Add options to the in-program functions
#TODO: Add ability to run, save, and recall calibrations
#TODO: Remove timeout on IDN? step

# Standard modules Utilized
import sys
import time

# Esoteric Modules Utilized
try:
    import visa
except ImportError as e:
    sys.exit("\nERROR: Requires '{0}' module to run\n".format(e.name))

# Setup defaults for each instrument
defaults =  {   "8753A":    {   "POIN":   "1601",
                                "STAR":   "0.3",
                                "STOP":   "3000"
                            },
                "8753D":    {   "POIN":   "1601",
                                "STAR":   "0.03",
                                "STOP":   "6000"
                            }
            }



def main():
    # Attempt to initialize an instrument
    try:
        VNA = Initialize()
        print("Instrument open. Manual control disabled")
    except:
        return

    # Grab the identity of the VNA
    model = Identify(VNA)

    # Ask if user wishes to "preset", bringing machine back to default settings
    choice = raw_input("Preset device? (0 for no, 1 for yes, default is no): ")
    if choice == '1': Preset(VNA)

    # Begin main operation loop
    while True:
        # Prompt the user
        choice = raw_input("Select function ('h' for help): ")

        if choice == '':
            break
        elif choice == 'h':
            Help()
        elif choice == 'p':
            Preset(VNA)
        elif choice == 'c':
            Config(VNA, model)
        elif choice == 'C':
            ChangeMeasurement(VNA)
        elif choice == 'r':
            Reset(VNA)
        elif choice == 'm':
            Manual(VNA)
        elif choice == 's':
            Save(VNA)
        # elif choice == 'S':
        #     Sparams(VNA)
        elif choice == 't':
            Terminal(VNA)

    VNA.close()
    print("\nThanks for using the script!\n")
    return



def Help():
    print("h - Print list of available commands")
    print("p - Preset instrument to default settings")
    print("c - Configure the instrument output")
    print("C - Change measurement type")
    print("r - Reset the sampling")
    print("m - Enter manual control mode")
    print("s - Save the current on-screen trace")
    # print("S - Record full s-parameter set")
    print("t - Open a terminal to the instrument's GPIB interface")
    print("\n")

def Preset(VNA):
    print("Presetting Instrument to defaults...")
    VNA.write('PRES')
    print("Done.")

def Config(VNA, model):
    print("Clearing the instrument")
    VNA.clear()

    default = defaults[model]

    points = raw_input("Enter number of data "+\
                       "points (default = {}): ".format(default["POIN"]))
    if points == '': points = '1601'
    VNA.write('POIN{0}'.format(points))

    minfreq = raw_input("Enter minimum frequency "+\
                        "in MHz (default = {}): ".format(default["STAR"]))
    if minfreq == '': minfreq = '0.03'
    maxfreq = raw_input("Enter maximum frequency "+\
                        "in MHz (default = {}): ".format(default["STOP"]))
    if maxfreq == '': maxfreq = '6000'
    VNA.write('STAR{0}MHZ'.format(minfreq))
    VNA.write('STOP{0}MHZ'.format(maxfreq))

    VNA.write('SCAL10')
    VNA.write('REFV0')
    VNA.write('REFP10')

    avgfac = raw_input("Enter averaging factor, 0 or 1 for no averaging "+\
    "(default = 0): ")
    if avgfac == '': avgfac = '0'
    if int(avgfac) < 2:
        VNA.write('AVEROOFF')
    else:
        VNA.write('AVEROON')
        VNA.write('AVERFACT{0}'.format(avgfac))

    mode = raw_input("Select (s)ingle or (c)ontinuous capture (default = c): ")
    if mode == 's':
        VNA.write('SING')
    else:
        VNA.write('CONT')

    form = raw_input("Enter data format (1-4, default = 4): ")
    if form == '': form = '4'
    VNA.write('FORM{0}'.format(form))

def ChangeMeasurement(VNA):
    parameter = raw_input("Select S-parameter to measure: ")
    validChoices = ['s11','S11','s21','S21','s12','S12','s22','S22']
    if parameter in validChoices:
        VNA.write(parameter)
    else:
        print("Invalid selection")

def Reset(VNA):
    VNA.write('REST')

def Manual(VNA):
    VNA.control_ren(0)
    raw_input("Manual control enabled. Press enter to reassert PC control...")
    VNA.control_ren(1)
    Reset(VNA)

def Save(VNA):
    trace = QueryTrace(VNA)
    SaveTrace(trace)

def Sparams(VNA):
    VNA.write('SING?')
    singflag = int(VNA.read())
    VNA.write('CONT?')
    contflag = int(VNA.read())
    VNA.write('AVERO?')
    avgflag  = int(VNA.read())
    VNA.write('AVERFACT?')
    avgcount = int(float(VNA.read()))

    timeout  = VNA.timeout
    del VNA.timeout

    if singflag:
        tracerequest = 'SING'
    else:
        tracerequest = 'NUMG{}'.format(str(avgcount))
        VNA.write('AVEROON')

    sparams = ['S11','S21','S12','S22']

    for parameter in sparams:
        # See page ~268 in programming guide
        VNA.write(parameter)
        VNA.write('AVERREST')
        VNA.write('OPC')
        VNA.write(tracerequest)
        VNA.write('OPC?')
        OPCresult = VNA.read()
        print OPCresult
        if OPCresult != 1:
            print "Operation not completed"
            return
        trace = QueryTrace(VNA)
        SaveTrace(trace)

    if singflag:
        VNA.write('SING')
    else:
        VNA.write('CONT')
    VNA.timeout = timeout

def Terminal(VNA):
    while True:
        command = raw_input()
        if command == '': break
        VNA.write(command)
        if command[-1:] == '?':
            print("VNA response: {0}".format(VNA.read()))



def Initialize():
    # Instantiate resource manager
    resMgr      = visa.ResourceManager()
    # Get list of resources which have a GPIB bus
    resources   = resMgr.list_resources(query=u'GPIB?*INSTR')

    # Check for length of instrument list
    n = len(resources)
    if n < 1:
        print("No Instruments available. Terminating script.\n")
        raise IOError

    # If nonzero, display the list of instruments
    else:
        print("Instruments:")
        for index,instrument in enumerate(resources):
            print("\t{0}: {1}".format(index,instrument))

    # Prompt the user to pick one
    choice = raw_input("Choice ({}-{}, default = 0): ".format(0, n-1))

    if choice == '': choice = '0'

    try:
        # Attempt to cast to an int
        choice = int(choice)
        # Call a script based on input
    except (KeyError, ValueError) as exception:
        print("Invalid Selection. Terminating script.\n")
        raise exception

    # Open the instrument
    VNA = resMgr.open_resource(resources[choice])

    return VNA

def Identify(VNA):
    VNA.write('IDN?')
    identity = VNA.read()
    model = identity.split(',')[1]
    model = model.encode('ascii')
    return model

def QueryTrace(VNA):
    VNA.write('OUTPLIML')
    limits = VNA.read()
    VNA.write('OUTPFORM')
    values = VNA.read()
    trace = (limits, values)
    return trace

def ParseTrace(trace):
    limits = trace[0]
    values = trace[1]

    frequencies = []
    magnitudes = []

    limits = limits.encode('ascii')
    values = values.encode('ascii')

    limitsSplit = limits.splitlines()
    valuesSplit = values.splitlines()

    for line in limitsSplit:
        entries   = line.split(',')
        frequency = entries[0]
        frequencies.append(frequency)

    for line in valuesSplit:
        entries   = line.split(',')
        magnitude = entries[0]
        magnitudes.append(magnitude)

    return frequencies, magnitudes

def SaveTrace(trace):
    frequencies,magnitudes = ParseTrace(trace)
    if len(frequencies) != len(magnitudes):
        print("Invalid trace data. Unable to record trace")
        return

    filename = raw_input("Enter filename sans extension (default=<datetime>): ")
    if filename == "": filename = time.strftime("%Y-%m-%d_%I-%M-%S")
    filename = filename +".csv"
    file = open(filename, 'w+')

    for point in range(len(frequencies)):
        string = frequencies[point] + ',' + magnitudes[point] + '\n'
        file.write(string)




if __name__ == "__main__":
    main()