import os
import sys
import time
import pyvisa
import numpy as np
import matplotlib.pyplot as plt


rm = pyvisa.ResourceManager()

address  = None
dev_list = []
devstr   = ""
inst     = None


def terminate():
    print("")
    exit()

def pint(s, defval = None):
    try:
        return int(s)
    except:
        return defval

def pfloat(s, defval = None):
    try:
        return float(s)
    except:
        return defval

def list(is_print = True):
    global dev_list

    dev_list = rm.list_resources()

    if is_print:
        print("  Device list:", dev_list)
    return dev_list

def get_devstr(addr):
    global dev_list

    dev_list = list(is_print = False)
#    print("dev_list=", dev_list)
    for s in dev_list:
        aa = s.split('::')
#        print("s=", s, aa)
        if len(aa) < 2:
            continue
        
        if pint(aa[1]) == addr:
            return s

    return None

def close_device():
    global inst
    
    if inst:
        inst.close()
        inst = None
        return True
    
    return False

def open_device(devname):
    global devstr, address, inst 

    close_device()

    inst = rm.open_resource(devname)
    if inst:
        print(f"  found device [{devname}]. assigned to current instrument handle")
        devstr = devname
        aa = devname.strip().split('::')
        address = pint(aa[1])
    else:
        print(f"  Error: Can not open device [{devname}]")

    return inst

def main():
    global rm, dev_list, devstr, address, inst

    dev_list = list()
    if len(dev_list) > 0:
        inst = open_device(dev_list[0])
        devname = dev_list[0]
    else:
        print("Error: No device found")
        terminate()

    if inst is None:
        return
        
    keithley = inst

# Keithley 2450の初期設定
    keithley.write("*RST")  # Keithleyをリセット
    keithley.write(":SOUR:FUNC VOLT")  # 電圧源モードに設定
    keithley.write(":SENS:FUNC 'CURR'")  # 電流測定モードに設定
    keithley.write(":SENS:CURR:RANG:AUTO ON")
    keithley.write(":SENS:VOLT:NPLC 1")  # 測定の積分時間を設定
#    keithley.write(":SENS:CURR:RANG 1.05e-6")  # 電流範囲を1μAに設定
    keithley.write(":OUTP ON")  # 出力をONにする

# トリガーモデルを設定: TIMMER1は受け付けないのでTIMER2にしている。とりあえずコメントアウト
#    keithley.write(":TRIG:TIM2:CLE")  # トリガーカウンタをクリア
#    keithley.write(":TRIG:TIM2:COUN 1")  # トリガーカウンタを1に設定
#    keithley.write(":TRIG:TIM2:SOUR IMM")  # イミディエイトトリガーを設定
#    keithley.write(":TRIG:TIM2:DIR SOUR")  # トリガーソースを指定。TIM2でもHeader errorになる

# 測定を開始
    keithley.write(":INIT")

# 測定データを取得
    Vlist = []
    Ilist = []
    for V in np.arange(-1.0, 1.001, 0.1):
        keithley.write(f":SOUR:VOLT {V:10.4g}")
        time.sleep(0.1)  # 安定化待ち時間

        ret = keithley.query(":READ? 'defbuffer1', SOUR, READ")  # 測定された電流値を取得
#       voltage = keithley.query(":READ? 'defbuffer1', READ, SOUR")  # 測定された電圧値を取得

        V, I = ret.split(',')
        V = pfloat(V)
        I = pfloat(I)
        Vlist.append(V)
        Ilist.append(I)
        print(f" set V={V:10.4g}  measured V={V:14.6g}  I={I:14.6g}")
#        print("Current:", current)

    keithley.write(":OUTP OFF")

# Keithley 2450をクローズ
    keithley.close()

    plt.plot(Vlist, Ilist)
#    plt.yscale('log')
    plt.show()
    

if __name__ == "__main__":
    main()