#!/usr/bin/python3

import pyvisa as visa
import time

def Query(cmd : str) -> str :
  haha = osc.query(cmd).rstrip('\n')
  print(cmd + " : |" + haha + "|")
  return haha

def Write(cmd : str):
  osc.write(cmd)

def AskWaveParameter():
  haha = Query('WAVEFORM:Preamble?')
  result = haha.split(',')
  print(" Format :" + result[0]) 
  print("   Mode :" + result[1]) 
  print(" Points :" + result[2]) 
  print("Average :" + result[3]) 
  print(" X-Step :" + result[4] + " sec") 
  print("X-Start :" + result[5] + " sec") 
  print("  X-Ref :" + result[6] ) 
  print(" Y-Step :" + result[7] + " V" ) 
  print("Y-Start :" + result[8] + " V") 
  print("  Y-Ref :" + result[9] ) 
  return result


#=========================================

# Initialize visa resource manager
rm = visa.ResourceManager('@py')

osc = rm.open_resource('TCPIP::128.186.111.202::INSTR')
osc.timeout = 25000
Query('*IDN?')
# Query(':SYST:VERSION?')

#Write(':STOP')
#Write(':RUN')

#================ channel setting
#Write('CHAN1:OFFSET 0.0')
#Write('CHAN1:IMPEDANCE FIFTY') #set to 50 Ohm
#Write('CHAN1:IMPEDANCE OMEG') #set to 1Mohm
#Write('CHAN1:SCALE 0.2')  #set vertical scale is 0.1V/div

#================ mesurement
# Query('MEASURE:FREQ? CHAN1') #measure freqeuncy of ch1 in Hz
#Query('MEASURE:VMAX? CHAN1') #measure VMaX of ch1 in V
#Query('MEASURE:VMIN? CHAN1') #measure VMin of ch1 in V
#Query('MEASURE:VPP? CHAN1') #measure Vpp of ch1 in V
#Query('MEASURE:VRMS? CHAN1') #measure Vrms of ch1 in V

#================ time axis
#Query('TIMEBASE:OFFSET?') #offset in sec
#Write('TIMEBASE:OFFSET 6.0e-6') #offset in sec, pos to the right, neg to the left
#Write('TIMEBASE:SCALE 2.0e-6') # set time scale to be 2us/div

#================ trigger
#Write('TRIG:MODE EDGE') #set trigger mode to be EDGE; EDGE|PULSe|RUNT|NEDG|SLOPe|VIDeo|PATTern|RS232|IIC|SPI|CAN|FLEXray|USB
#Query('TRIG:STATUS?') #ask if trigger is ok, return TD = triggered
#Query('TRIG:EDGE:SLOPE?')
#Write('TRIG:EDGE:SLOPE POS') #set trigger at positive slope; NEG
#Write('TRIG:EDGE:SOURCE CHAN1') #set trigger for channel 1
#Query('TRIG:EDGE:LEVEL?') #ask trigger level

#=============== Waveform

#--------------- read waveform on the screen
Write('RUN')
Write('WAVEFORM:RESET')
Write('WAVEFORM:SOURCE CHAN1')
Write('WAVEFORM:FORMAT BYTE')
Write('WAVEFORM:MODE NORM')

para = AskWaveParameter()
wfSize = int(para[2])

Query('WAVEFORM:SOURCE?')
Query('WAVEFORM:MODE?')
Query('WAVEFORM:FORMAT?')

# print("----------------- get data ")

[status, wfLen] = Query('WAVEFORM:STATUS?').split(',')
if status == "IDLE" and int(wfLen) >= 0 :
  print("---- get waveform data")
  osc.write('WAVEFORM:DATA?')
  binary = osc.read_bytes(int(wfSize) +11 + 1) # 
  # print(binary)
  

#---------- read data from internal memory, 1.12 million points, take like 
# Write('STOP')
# Write('WAVEFORM:RESET')
# Write('WAVEFORM:SOURCE CHAN1')
# Write('WAVEFORM:MODE RAW')  #internal 
# Write('WAVEFORM:FORMAT BYTE')
# Write('WAVEFORM:BEGIN')

# para = AskWaveParameter()
# wfSize = int(para[2])
# Write('WAVEFORM:POINTS %d' % wfSize) 
# Query('WAVEFORM:POINTS?') 

# # print("----------------- get data ")
# count = 0
# while(count<10):
#   [status, wfLen] = Query('WAVEFORM:STATUS?').split(',')
#   if status == "IDLE" and int(wfLen) >= wfSize :
#     print("---- get waveform data")
#     osc.write('WAVEFORM:DATA?')
#     start_time = time.time()
#     binary = osc.read_bytes(int(wfLen)+11 + 1) # 
#     end_time = time.time()
#     print(binary)
#     break
  
#   time.sleep(1)
#   count = count + 1

#Write('WAVEFORM:END')
 
# duration = end_time - start_time
# print("Duration:", duration, "seconds")

#----------------- decode binary

# Extracting the header (first 11 bytes)
header = binary[:11]
data_points = binary[11:]
# Decode the 1-byte data points
decoded_data = [byte for byte in data_points]
print("Decoded data:", decoded_data)