import pyvisa
import csv

# -----------------------------
# Connect to the oscilloscope
# -----------------------------
rm = pyvisa.ResourceManager()
scope = rm.open_resource("USB0::0x0957::0x1734::MY12345678::INSTR")  # <-- change this

scope.timeout = 5000
scope.write_termination = '\n'
scope.read_termination = '\n'

# -----------------------------
# Configure waveform output
# -----------------------------
scope.write(":WAV:FORM ASC")
scope.write(":WAV:SOUR CHAN1")
scope.write(":WAV:STAR 1")
scope.write(":WAV:STOP 10000")

# -----------------------------
# Read preamble
# -----------------------------
preamble = scope.query(":WAV:PRE?")
fields = [float(x) for x in preamble.split(',')]

fmt, wtype, points, count, xinc, xorig, xref, yinc, yorig, yref = fields

# -----------------------------
# Read waveform data
# -----------------------------
raw = scope.query(":WAV:DATA?")
raw = raw[raw.find(',')+1:]          # remove #9xxxxxx header
samples = [float(x) for x in raw.split(',')]

# -----------------------------
# Build time + voltage arrays
# -----------------------------
time = []
volt = []

for n, d in enumerate(samples):
    t = xorig + (n - xref) * xinc
    v = (d - yref) * yinc + yorig
    time.append(t)
    volt.append(v)

# -----------------------------
# Write CSV
# -----------------------------
with open("waveform.csv", "w", newline='') as f:
    writer = csv.writer(f)
    writer.writerow(["Time", "Voltage"])
    writer.writerows(zip(time, volt))

print("CSV saved as waveform.csv")