

# =============================================================================
# 1. Import classes and modules
# =============================================================================

# =============================================================================
# #If local installation of dopes instead of using PyPI (https://pypi.org/project/dopes/)
# import sys
# dopes_path = 'path/to/dopes'        
# if dopes_path not in sys.path:
#     sys.path.insert(0, dopes_path)
# =============================================================================


import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colormaps

import dopes.equipment_control.equipment as eq
import dopes.equipment_control.k2400 as k2400
import dopes.equipment_control.meerstetterTEC as TEC

import time

# =============================================================================
# 2. List  available connections (dopes use pyvisa package for communicate with most equipments)
# =============================================================================
rm=eq.resource_manager()
list_connections= eq.available_connections()
print("Available connections: %s"%str(list_connections))

# =============================================================================
# 3. Connection to the equipment
# =============================================================================
myk2400=k2400.k2400("GPIB0::27::INSTR",timeout=10e3)
myTEC=TEC.meerstetterTEC()
myTEC.initialize()

# =============================================================================
# 4. Measurement parameters
# =============================================================================
units={"voltage":"V", "current":"A", "resistance":"Ohms"}        

source_mode = "voltage"
measurement_mode = "current"
compliance = 5e-3
autozero = True
nplc = 0.2
digits = 6
continuous_trigger = False
disp_enable = True

n_cycle=1
vmin=0
vmax=7
n_sample=301
double=True
bias_list=np.linspace(vmin,vmax,n_sample)
if double:
    bias_list=np.append(bias_list,bias_list[-2::-1])

temperature_list=np.linspace(10,30,5)
wait_for_stabilization=True

prefix='path_to_file_with_prefix'


# =============================================================================
# 5. Initialization of the equipment
# =============================================================================
myk2400.initialize(source_mode=source_mode, measurement_mode=measurement_mode,
                   compliance=compliance, autozero=autozero, digits=digits,
                   continuous_trigger=continuous_trigger,disp_enable=disp_enable, nplc=nplc)
myk2400.set_source(vmin)
myk2400.set_output("ON")
time.sleep(1)
myTEC=TEC.meerstetterTEC()
myTEC.initialize()

# =============================================================================
# 6. Measurement script
# =============================================================================
for temperature in temperature_list:
    file_path = prefix+'%ddegC_rep%d.txt'%(temperature,n_cycle)

    t_init=time.time()

    myTEC.set_temp(temperature,wait_for_stabilization=True)
    time.sleep(5)

    data=np.zeros((len(bias_list),n_cycle))

    for cycle in range(n_cycle):
        print("- Cycle %d ..."%cycle,end="")
        tcycle=time.time()
        i=0
        for bias_source in bias_list:
            myk2400.set_source(bias_source)
            data[i,cycle]=float(myk2400.read_single())
            i+=1
        print(" done in %d s"%(time.time()-tcycle))

    myk2400.set_output("OFF")
    print("Measurement done in %d s"%(time.time()-t_init))

    # =============================================================================
    # 7. Save data
    # =============================================================================
    header="%s (%s)"%(source_mode,units[source_mode])

    for cycle in range(n_cycle):
        header+=", %s (%s)"%(measurement_mode,units[measurement_mode])
        
    data_to_write=np.append(np.transpose([bias_list]),data,axis=1)
    eq.write_in_file(file_path,data_to_write,header=header,date=True,overwrite=True)

    # =============================================================================
    # 8. Plot data
    # =============================================================================
    cmap=colormaps['coolwarm']
    color_list=cmap(np.linspace(0,1,n_cycle))
    figure_path=prefix+'%ddegC_rep%d.png'%(temperature,n_cycle)

    fig,ax=plt.subplots()
    ax.set_xlabel("Voltage (V)")
    ax.set_ylabel("Current (A)")
    ax.set_title("temperature: %.1f"%temperature)
    for cycle in range(n_cycle):
        ax.plot(bias_list,abs(data[:,cycle]),color=color_list[cycle])
    plt.savefig(figure_path,dpi=200)
    plt.show(block=False)


# =============================================================================
# 9. Close connection
# =============================================================================
myk2400.close_connection()
myTEC.close_connection()
