# 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

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

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'])

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

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


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'])

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

	def read(self) :
		return self.pipe.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

def CreateInstrumentObjects(ris, print_messages=False) :
	# ris is the dictionary of recognized devices.
	# Creates a dictionary of available oscilloscope objects.
	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 )
	return instruments

def SelectInstruments(instruments, print_messages=False) :
	# Returns the first instrument object of each type.
	scope = None
	awg = None
	for key in instruments :
		if instruments[key].type.find('osc') > -1 :
			if scope :
				continue
			else :
				scope = instruments[key]
		if instruments[key].type.find('wave') > -1 :
			if awg :
				continue
			else :
				awg = instruments[key]
	return scope, awg

def Test():
	ris = si.RecognizedDevices()
	#scopes = CreateOscilloscopeObjects(ris, print_messages=True)
	#awgs = CreateAWGObjects(ris, print_messages=True)
	instruments = CreateInstrumentObjects(ris, print_messages=True)
	scope, awg = SelectInstruments(instruments, print_messages=True)
	print scope, awg

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

if __name__=='__main__' :
	Test()