# install pip install pyvisa pyvisa-py numpy
import pyvisa
import numpy as np
import csv

IP = "10.233.17.204"
RESOURCE = f"TCPIP0::{IP}::INSTR"

rm = pyvisa.ResourceManager('@py')
scope = rm.open_resource(RESOURCE)

scope.timeout = 10000

# Identify instrument
print(scope.query("*IDN?"))

# Stop acquisition for stable data
scope.write(":STOP")

# Select channel
scope.write(":WAV:SOUR CHAN1")

# Set waveform mode and format
scope.write(":WAV:MODE NORM")
scope.write(":WAV:FORM BYTE")

# Get scaling parameters
yinc = float(scope.query(":WAV:YINC?"))
yorig = float(scope.query(":WAV:YOR?"))
yref = float(scope.query(":WAV:YREF?"))

xinc = float(scope.query(":WAV:XINC?"))
xorig = float(scope.query(":WAV:XOR?"))
xref = float(scope.query(":WAV:XREF?"))

# Request waveform data
raw = scope.query_binary_values(":WAV:DATA?", datatype='B', container=np.array)

# Convert to voltage
voltage = (raw - yref) * yinc + yorig

# Create time axis
time = (np.arange(len(voltage)) - xref) * xinc + xorig

# Save to CSV
with open("waveform.csv", "w", newline="") as f:
    writer = csv.writer(f)
    writer.writerow(["Time (s)", "Voltage (V)"])
    for t, v in zip(time, voltage):
        writer.writerow([t, v])

print("Waveform saved to waveform.csv")

scope.close()
