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


def main():
    measureseqs = [
        [120, 8],
        [120, 12],
    ]

    for seq in measureseqs:
        length = seq[0]
        period = seq[1]
        filename = f"measurement_L{length}min_P{period}min.csv"
        print("Measuring " + filename)
        do_measurement(filename, length * 60, period * 60)
        print("Done")
   

    # 
    # plt.plot(d_t, d_1, label="1")
    # plt.plot(d_t, d_2, label="2")
    # # plt.plot(d_t, d_v, label="V")
    # plt.legend()
    # plt.show()

def do_measurement(filename, MEAS_LENGTH, PERIOD):
    print("Starting")
    MULT1= "GPIB0::9::INSTR"
    MULT2= "GPIB0::12::INSTR"
    PSU="GPIB0::10::INSTR"
    rm = pyvisa.ResourceManager()
    print(rm.list_resources())

    print("Opening")
    mult1 = rm.open_resource(MULT1)
    mult2 = rm.open_resource(MULT2)
    psu = rm.open_resource(PSU)

    print("Configuring...")
    for m in [mult1, mult2]:
        m.write("*RST")
        m.write("CONF:VOLT:DC 0.001")
        # m.write("*RST")
        m.write("INP:IMP:AUTO ON")

    psu.write("*RST")
    psu.write("VOLT:RANG P8V")
    psu.write("VOLT 2")
    psu.write("OUTP ON")

    print("Querying")
    print(mult1.query("READ?"))
    print(mult2.query("READ?"))

    d_1 = []
    d_2 = []
    d_v = []
    d_t = []
    

    t_start = time.time_ns()

    t = 0

    AMP = 6
    FREQ = 1/PERIOD

    while t < MEAS_LENGTH:
        t = (time.time_ns() - t_start) * 10**-9
        remain = np.floor((MEAS_LENGTH - t) / 60)
        print(f"{t/MEAS_LENGTH*100:.3g}%  {remain}m")

        volt = ((math.sin(t * 2 * np.pi * FREQ) + 1) / 2) * AMP

        psu.write(f"VOLT {volt:.3g}")
        d_1.append(float(mult1.query("READ?")))
        d_2.append(float(mult2.query("READ?")))
        d_t.append(t)
        d_v.append(volt)

    # psu.write("OUTP OFF")
    print("Setting idle power")
    psu.write("VOLT 3.0")

    data_out  = np.array([d_t, d_1, d_2, d_v])
    np.savetxt(filename, data_out, delimiter=",")

    return (d_t, d_1, d_2, d_v)

if __name__ == "__main__":
    main()

#  !uv run main.py