import pyvisa
import numpy as np
import matplotlib.pyplot as plt
import time
import math

def main():
    print("Starting")

    k = 0.00176727
    m = 0.3595224
    roomtemp = 23

    def get_temp(r):
        return (r-m)/k


    rm = pyvisa.ResourceManager()
    print(rm.list_resources())
    VOLT = "GPIB0::4::INSTR"
    AMP = "GPIB::2::INSTR"

    volt = rm.open_resource(VOLT)
    amp = rm.open_resource(AMP)

    volt.write("*RST")
    amp.write("*RST")

    volt.write("CONF:VOLT:DC 10")
    volt.write("INP:IMP:AUTO ON")
    amp.write("CONF:CURR:DC 10")

    # Do measurement to zero before zeroing is turned off
    amp.query("READ?")
    volt.query("READ?")

    amp.write("SENS:CURR:DC:ZERO:AUTO OFF")
    amp.write("TRIG:DEL 0")
    amp.write("SENS:CURR:DC:NPLC 0.02")

    volt.write("VOLT:DC:NPLC 0.02")
    volt.write("TRIG:DEL 0")
    volt.write("TRIG:SOUR IMM")
    volt.write("ZERO:AUTO OFF")


    d_i = []
    d_u = []
    d_t = []

    
    t_start = time.time_ns()

    try:
        while True:
            t = (time.time_ns() - t_start) * 10**-9
            u = float(volt.query("READ?"))
            i = float(amp.query("READ?"))

            r = u/i

            print(f"{r} t = {get_temp(r)}")
            
            d_i.append(i)
            d_u.append(u)
            d_t.append(t)

    except KeyboardInterrupt:
        print("Stopping!")

        d_i = np.array(d_i[1:])
        d_u = np.array(d_u[1:])
        d_t = np.array(d_t[1:])

        p = d_i * d_u
        r = d_u / d_i

        plt.plot(d_t, p, label="Power")
        plt.plot(d_t, r, label="Resistance")
        plt.plot(d_t, get_temp(r), label="Temperature")
        plt.legend()
        plt.show()

        filename = input("Enter filename:")
        data_out  = np.array([d_t, d_u, d_i])
        np.savetxt(filename, data_out, delimiter=",")


    

if __name__ == "__main__":
    main()

#  !uv run stek.py