import pyvisa as visa
import time
import numpy as np

resources = visa.ResourceManager('@py')
device = resources.open_resource("TCPIP::192.168.1.145::INSTR", write_termination='\n', query_delay=0.25)
print(device.query("*IDN?"))

device.write("*RST")
time.sleep(5)

device.write("MSIZ 20M")
device.write("TIME_DIV 200MS")
device.write("C1:VDIV 2V")
device.write("C2:TRA ON")
device.write("C2:VDIV 2V")

time.sleep(5)

device.write("STOP")


for channel in range(1, 3):
    sample_num = int(float(device.query("SANU? C{}".format(channel))))
    print("C{}:{}".format(channel, sample_num))

    device.write("C{}:WF? DAT2".format(channel))
    response = device.read_raw()[18:]
    assert len(response) == sample_num

    array = np.frombuffer(response, dtype="u8")
    np.save("siglent_ch{}.npy".format(channel), array)

    time.sleep(1)