# Module containing oscilloscope and awg classes

# supported_instrument.txt must be in the current working directory (cwd).
# visa_devices.py and supported_instruments.py must in in the cwd or in a path in the sys.path list

try:
	import supported_instruments as si
except:
	print('\nsupported_instruments.py not found.')
	exit(0)

import numpy as np
import time
import string as s
import struct
import pylab as pl

class TekScope :
	def __init__(self, sid, ris) :
		# sid is a scope id key used in the recognized instruments dictionary ris.
		self.name = sid
		self.info = ris[sid]
		self.pipe = self.info['handle']		# This is a pyvisa instrument object.
		self.pipe.write('*IDN?')				#  A pyvisa instrument object has write and read methods.
		self.id = self.pipe.read()
		self.type = ' and '.join(self.info['type'])
		self.time = 0   # Used to report time taken for read and write functions
		self.error = ''
		self.channel = [self.Channel(0), self.Channel(1), self.Channel(2)]	# Channel 0 is meaningless, but scope could have 4 channels.

	class Channel :
	    def __init__(self, channel) :
	        # Specify channel 1 or 2
	        self.channel = channel
	        self.error = ''
	        self.xunit = 0
	        self.xinc = 0
	        self.yunit = 0
	        self.ymult = 0
	        self.yoff = 0
	        self.yzero = 0

	def write(self, stuff) :
		self.pipe.write(stuff)

	def Write(self, stuff) :
		self.pipe.write(stuff)

	def read(self) :
		return self.pipe.read()

	def Read(self) :
		return self.pipe.read()

	def ClearQueue(self) :
		self.Write('*esr?')
		return self.Read()

	def GetScaleParameters(self, ch):
		# ch = channel object
		error = ''
		if ch.channel > 2 | ch.channel <1 :
			error = "Incorrect channel request."
			return error
		self.write('data:source ch'+str(ch.channel))
		try :
			self.write('wfmpre:xunit?')
			#print self.visa.read()
			ch.xunit = self.read().strip('"')  # Assuming header=off.
		except :
			error = 'Parameter read error.  Is channel ' + str(ch.channel) + ' currently displayed?'
			return error
		self.write('wfmpre:xincr?')
		ch.xinc = float(self.read())
		#print ch.xinc
		self.write('wfmpre:yunit?')
		ch.yunit = self.read().strip('"')
		self.write('wfmpre:ymult?')
		ch.ymult = float(self.read())
		self.write('wfmpre:yoff?')
		ch.yoff = float(self.read())
		self.write('wfmpre:yzero?')
		ch.yzero = float(self.read())
		# Make a dictionary of values
		ch.parameters = {}
		ch.parameters['xunit'] = ch.xunit
		ch.parameters['xinc'] = ch.xinc
		ch.parameters['yunit'] = ch.yunit
		ch.parameters['ymult'] = ch.ymult
		ch.parameters['yoff'] = ch.yoff
		ch.parameters['yzero'] = ch.yzero
		ch.error = error
		return error

	def Query(self) :
		# Query each scope for its idn and current parameters for each channel.
		# If a channel is not currently displayed on the scope, and error will be reported.
		error = ''
		self.Write('*IDN?')
		self.idn = self.Read()
		self.ClearQueue()
		self.GetScaleParameters(self.channel[1])
		self.GetScaleParameters(self.channel[2])
		message = self.name + '\n' + self.idn + '\n'
		message += 'Ch 1 :\n'
		if self.channel[1].error :
			message += 'Error.  Not currently displayed?\n'
		else :
			for key in self.channel[1].parameters :
				message += key + ' : ' + str(self.channel[1].parameters[key]) + '\n'
		message += 'Ch 2 :\n'
		if self.channel[2].error :
			message += 'Error.  Not currently displayed?\n'
		else :
			for key in self.channel[2].parameters :
				message += key + ' : ' + str(self.channel[2].parameters[key]) + '\n'
		return message

	def GetWaveform(self, ch, mode) :
		# Mode is binary or ascii.
		# ch = channel object
		error = ''
		if ch.channel > 2 | ch.channel <1 :
			error = "Incorrect channel request."
			return error, []
		self.write('data:source ch'+str(ch.channel))
		self.write('data:start 1')
		self.write('data:stop 2500')
		if s.lower(mode).find('asc') > -1 :
			# Ascii transfer takes 4 seconds.
			mode = 'a'
			self.write('data:encdg ascii')
			self.write('curve?')
			w = self.read().split(',')
			#print w
		elif s.lower(mode).find('bin') > -1 :
			# Binary transfer takes a fraction of a second.
			mode = 'b'
			self.write('data:encdg ribinary')  # Request signed 16 bit integer, big endian
			self.write('data:width 2')
			self.write('curve?')
			data_block = self.read()
			#print(data_block[0:13]) # This yields ':CURVE #45000' which is correct when header=on and verbose=on.
			print(data_block[0:6]) # This yields '#45000' when header=off.
			# Notice that the slice function yields the characters from 0 to 13-1=12.
			#raw_data = data_block[13:]  # For header=on, strip away the first 13 characters.  This seems to be right.
			raw_data = data_block[6:]  # For header=off, strip away the first 6 characters.  This seems to be right.
			print(len(raw_data))
			w = struct.unpack('>2500h', raw_data) # A buffer of 2500 signed short (16 bit) integers in the big endian order, not the native little endian.
			#print w
		else :
			error = 'Only ascii or binary modes are possible.'
			return error, [], []
		x, y = self.ScaleWaveform(w, ch)
		return error, np.array(x), np.array(y)

	def ScaleWaveform(self, w, ch) :
		y = []
		x = []
		delx = 0
		#yoff = p['yoff']
		for z in w:
			y.append((float(z)-ch.yoff)*ch.ymult + ch.yzero)
			x.append(delx)
			delx += ch.xinc
		return x, y

class TekAWG :
	def __init__(self, sid, ris) :
		# sid is a scope id key used in the recognized instruments dictionary ris.
		self.name = sid
		self.info = ris[sid]
		self.pipe = self.info['handle']		# This is a pyvisa instrument object.
		self.pipe.write('*IDN?')				#  A pyvisa instrument object has write and read methods.
		self.id = self.pipe.read()
		self.type = ' and '.join(self.info['type'])
		self.time = 0   # Used to report time taken for read and write functions
		self.error = ''

	def Write(self, stuff) :
		try:
			self.pipe.write(stuff)
		except :
			self.error = self.name + ' write operation probably timed out.'

	def write(self, stuff) :
		self.Write(stuff)

	def Read(self) :
		try :
			q = self.pipe.read()
		except:
			self.error = self.name + ' read operation probably timed out.'
			q = self.error
		return q

	def read(self) :
		return self.Read()

	def ClearQueue(self) :
		self.Write('*esr?')
		return self.Read()


def CreateOscilloscopeObjects(ris, print_messages=False) :
	# Creates a dictionary of available oscilloscope objects.
	#ris = si.RecognizedDevices()
	scopes = {}
	for key in ris :
		if (ris[key]['type'][0].lower()=='osc') and (ris[key]['vendor'].lower()=='tektronix') :
			scopes[key] = TekScope(key, ris)
	if print_messages :
		print('\nDictionary of oscilloscope objects created:')
		for key in scopes :
			print('\t' + key + ' : ' + str(scopes[key]))
			print('\t\tInstrument ID = ' + scopes[key].id)
	return scopes

def CreateAWGObjects(ris, print_messages=False) :
	# ris is the dictionary of recognized devices.
	# Creates a dictionary of available oscilloscope objects.
	awgs = {}
	for key in ris :
		if (ris[key]['type'][0].lower()=='wave') and (ris[key]['vendor'].lower()=='tektronix') :
			awgs[key] = TekAWG(key, ris)
	if print_messages :
		print('\nDictionary of awg objects created:')
		for key in awgs :
			print('\t' + key + ' : ' + str(awgs[key]))
			print('\t\tInstrument ID = ' + awgs[key].id)
	return awgs

## Create oscilloscope and awg instruments objects.
def CreateInstrumentObjects(print_messages=False) :
	# ris is the dictionary of recognized devices.
	ris = si.RecognizedDevices()
	instruments = {}
	for key in ris :
		if (ris[key]['type'][0].lower()=='wave') and (ris[key]['vendor'].lower()=='tektronix') :
			instruments[key] = TekAWG(key, ris)
		elif (ris[key]['type'][0].lower()=='osc') and (ris[key]['vendor'].lower()=='tektronix') :
			instruments[key] = TekScope(key, ris)
	if print_messages :
		print('\nDictionary of instrument objects created:')
		for key in instruments :
			print('\t' + key + ' : ' + str(instruments[key]))
			print('\t\tInstrument ID = ' + instruments[key].id + ', Type = ' + instruments[key].type )
	# Returns a dictionary of available oscilloscope and awg objects.
	return instruments

## Select one instrument of each type for use.
def SelectOneInstrumentOfEachType(instruments, osc, wg, print_messages=False) :
	# Returns the first instrument object of each type or single, specified instruments of each type.
	# instruments is a dictionary of active instruments.
	# osc is the identifier for a scope of the form model-sn.
	#    If sn is not present or if osc is empty or None, then the first scope of this model number will be used.
	# wg is the identifier for a wg of the form model-sn.
	#    If sn is not present or if wg is empty or None, then the first wg of this model number will be used.
	osc_tags= osc.split('-')
	osc_model = osc_tags[0].strip()
	if len(osc_tags) > 1 :
		osc_sn = osc_tags[1].strip()
	else :
		osc_sn = ''
	wg_tags= wg.split('-')
	wg_model = wg_tags[0].strip()
	if len(wg_tags) > 1 :
		wg_sn = wg_tags[1].strip()
	else :
		wg_sn = ''
	scope = None
	awg = None
	# Select a oscilloscope.
	for key in instruments :
		if (instruments[key].type.find('osc') > -1) :
			# Match model number and sn.
			if (osc == key) :
				scope = instruments[key]
				break
			# Or match just the model number if sn is not present.
			elif (osc_sn == '') and (key.find(osc_model) > -1) :
				scope = instruments[key]
				break
			# Or just select the first scope encountered is model number is not present.
			elif osc_model == '' :
				scope = instruments[key]
				break
	# Select an awg.
	for key in instruments :
		if (instruments[key].type.find('wave') > -1) :
			# Match model number and sn.
			if (wg == key) :
				awg = instruments[key]
				break
			# Or match just the model number if sn is not present.
			elif (wg_sn == '') and (key.find(wg_model) > -1) :
				awg = instruments[key]
				break
			# Or just select the first awg encountered is model number is not present.
			elif wg_model == '' :
				awg = instruments[key]
				break
	# Return an instrument object or None for each type of device.
	return scope, awg

def Test():
	instruments = CreateInstrumentObjects(print_messages=True)
	#wg = 'AFG3021B-C031285'
	#wg = 'AFG3021B'
	wg = ''
	#osc = 'TDS1012B-C040717'
	#osc = 'TDS1012B'
	osc = ''
	scope, awg = SelectOneInstrumentOfEachType(instruments, osc, wg, print_messages=True)
	print 'Selected scope: ', scope
	print 'Selected wg:', awg
	# Change the awg frequency.
	awg.write('function sin;frequency 2.0e3')
	# Test the ClearQueue() function.
	print awg.ClearQueue()
	# scope.channel[1] should be the channel 1 object.
	print scope.channel[1]
	scope.GetScaleParameters(scope.channel[1])
	print scope.channel[1].parameters
	print scope.Query()
	# Acquire the channel 1 waveform, with x = list of times, y = list of signals at these times.
	error, x, y = scope.GetWaveform(scope.channel[1], 'ascii')		# Change to 'binary' for much faster data transfer.
	# Save data in a csv file as (x,y) pairs.
	f = open('data.csv', 'w')
	for i in range(len(x)) :
		f.write('\n' + str(x[i]) + ',' + str(y[i]))
	f.close()
	#Plot y versus x.
	pl.plot(x,y)
	pl.show()

#---------------------------------------------------------------------------------

if __name__=='__main__' :
	Test()