from qolab.hardware.multimeter.hp3457a import HP3457A
from qolab.hardware.multimeter.bk_5491 import BK_5491
from qolab.hardware.multimeter.hp_34401 import HP_34401
from qolab.hardware.vacuum_gauge.mks390 import MKS390
from qolab.hardware.power_supply.psw25045 import PSW25045
from qolab.hardware.power_supply.gpp3610h import GPP3610H


from qolab.hardware.scope.sds800xhd import SDS800XHD

import pyvisa
import numpy as np

# Example usage
if __name__ == "__main__":
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    import time

    mpl.style.use('custom-style')

    rm = pyvisa.ResourceManager()
    address_dict = {
        'scope': 'TCPIP0::192.168.110.191::inst0::INSTR',
        'dmm_hpa': 'visa://192.168.194.15/GPIB1::22::INSTR',
        'dmm_hp': 'visa://192.168.194.15/ASRL9::INSTR',
        'dmm_bk': 'visa://192.168.194.15/ASRL12::INSTR',
        'vacgauge': 'visa://192.168.194.15/ASRL13::INSTR',
        'psu_anode': 'visa://192.168.194.15/ASRL10::INSTR',
        'psu_disp': 'TCPIP0::192.168.2.213::1026::SOCKET'       
    }
    #----------------------------------------------
    # Open visa resources and initialize intruments
    dmm_hpa = HP3457A(
        rm.open_resource(address_dict['dmm_hpa'])
    )
    dmm_bk = BK_5491(
        rm.open_resource(address_dict['dmm_bk'])
    )

    dmm_hp = HP_34401(
        rm.open_resource(address_dict['dmm_hp'])
    )

    scope = SDS800XHD(
        rm.open_resource(address_dict['scope'])
    )

    gauge = MKS390(
        rm.open_resource(address_dict['vacgauge'])
    )

    psu_anode = PSW25045(
        rm.open_resource(address_dict['psu_anode'])
    )

    psu_disp = GPP3610H(
        rm.open_resource(address_dict['psu_disp'])
    )

    #-----------------------------------------------

    # Get DMM readings
    print(dmm_hpa.get_idn())

    dmm_hpa.set_function('DCI')

    n = 5
    readings = np.zeros(n)
    for i in range(n):
        pressure = gauge.get_pressure()
        print(f'Pressure: {pressure} +- {0.1*pressure} Torr')
        print(f'MCOUNT: {dmm_hpa.get_memory_count()}')
        readings[i] = dmm_hpa.get_reading()
    print(f'HP4357A reading: {readings.mean()} +- {readings.std(ddof=1)} V')

    plt.hist(readings)
    plt.show()

    print(f'BK5491 reading: {dmm_bk.getAdc()} A')
    print(f'HP34401 reading: {dmm_hp.getAdc()} A')
    print(f'MKS390 reading: {gauge.get_pressure()} Torr')

    print(psu_anode.get_idn())
    psu_anode.set_voltage(1)
    psu_anode.enable_output(True)
    print(psu_anode.get_out_voltage())
    print(psu_anode.get_out_current())
    time.sleep(1)
    psu_anode.enable_output(False)

    print(psu_disp.get_idn())
    psu_disp.set_voltage(1.0)
    psu_disp.set_current(0.1)
    psu_disp.enable_output(True)
    time.sleep(1)
    print(psu_disp.get_out_current())
    print(psu_disp.get_out_voltage())
    time.sleep(1)
    psu_disp.enable_output(False)

    #trace = scope.getTrace(4)
    #trace.plot()
    plt.show()