# 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)

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 = ''

	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()

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) :
		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 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

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

if __name__=='__main__' :
	Test()