# Python code to sweep current from a Keithly 2400.  
# Records data in a text file
# See 'Variables' section to change filenames and sweep parameters

# Import functions
import os
from pylab import *
import win32api, win32con
from time import  strftime
from visa import instrument

# Set keithley current and read voltage
def keithleySetCurrent(start,stop,points,compliance_V,GPIB):
    num_points = points+1
    step = float(stop-start)/(num_points-1)

    voltages = [0]*num_points
    currents = [0]*num_points

    try:
        k2400 = instrument('GPIB1::' + str(GPIB),timeout = 20)
    except:
        try:
            k2400 = instrument('GPIB::' + str(GPIB),timeout = 20)
        except:
            print ' no GPIB'
            return
    k2400.write('*rst; status:preset; *cls')
    k2400.write('system:azero:state on')
    k2400.write('source:delay:auto on')
    k2400.write('display:digits 7')
    k2400.write('sense:voltage:NPLCycles 0.01')    
    k2400.write('sour:func:mode curr')
    k2400.write('sens:volt:prot:lev ' + str(compliance_V))
    k2400.write('sens:volt:rang:auto on')
    k2400.write('sens:func "volt"')
    k2400.write('form:elem volt')
#    k2400.write('outp on')
    for i in range(0,num_points):
        currents[i] = start + step*i
    
#Set sweep parameters
    k2400.write('sour:swe:poin ' + str(float(num_points)))
    k2400.write('sour:curr:star ' +str(float(start)/1000))
    k2400.write('sour:curr:stop ' +str(float(stop)/1000))
    k2400.write('sour:curr:mode swe')
    k2400.write('sour:swe:rang auto')
    k2400.write('sour:swe:spac lin')
    k2400.write('trig:coun ' + str(float(num_points)))
    k2400.write('sour:del .01')
    k2400.write('outp on')
    k2400.write('read?')
    voltages = k2400.read_values()
    k2400.write('outp off')
    
    return currents,voltages

#################
### VARIABLES ###
#################
path = 'C:/Users/femto/Desktop/Jared Hulme/E2/Wide Laser Bar/'
filename = 'SGDBR_10.txt'
userName = 'Jared Hulme'
fileHeader = 'Laser IV curves for Sweeper 2.2 Chip E2 Wide Laser Bar'
start = 0           # (mA)
stop = 250          # (mA)
points = 500        # Number of data points (actually it adds one extra)
compliance_V = 4    # (V)
GPIB = 12           # GPIB Address
#############################

# If you have the same filename this asks if you want to overwrite it.
reply = win32con.IDYES
if(os.path.exists(path + filename)):
    reply = win32api.MessageBox(0, "Do you want to overwrite " + filename + '?', "File Already Exists", win32con.MB_YESNO)
if reply == win32con.IDNO:
    print 'Filename already exists. Process Aborted'
else:
    currents, voltages = keithleySetCurrent(start,stop,points,compliance_V,GPIB)
    
    try:
        # Record data in a text file
        
        f = open(path + filename,'w')
        f.write('######################################################################################\n' +\
        '# ' + fileHeader + '\n' + \
        '# Data Taken: ' + strftime("%Y-%m-%d %H:%M:%Ss") + '\n' + \
        '# Taken by:' + userName + '\n' +\
        '# \n'
        '####################################################################################### \n\n')
        f.write('Voltage (V)\tCurrent (mA)\n')
        for i in range(0,len(voltages)-1):
            f.write('%.6f' % voltages[i] + '\t' + '%.6f' % currents[i] + '\n')
        f.write('%.6f' % voltages[-1] + '\t' + '%.6f' % currents[-1])
        f.close()
        print(filename + ' Complete')
    except IOError:
        print("File Save Error")
    
    plot(currents,voltages)
    show()