'''A program to analyse and plot the data from an interferometer experiment with TK power meter.'''
import os.path
import matplotlib.pyplot as plt	
import string
import glob
import matplotlib
 
font = {'family' : 'normal',
        'weight' : 'bold',
        'size'   : 22}
 
matplotlib.rc('font', **font)
 
def generateListIncre(Min, Max, Increment):
	'''a function to generate and return a list of float values from 3 float inputs. 
	Increments between Min and Max by Increment.'''
	newList = []
	i = Min
	while i!=Max:
		newList.append(i)
		i = i + Increment
		if i > Max:
			print 'Error in incrementing function in generateListIncre - Line 45'
			return(newList)
	newList.append(i)
	return(newList)
 
 
def tryFloat(rawInputString):
	'''is an exception handler - tries to convert to float - if fails repeats input request
	takes a raw_input() as an argument and returns a float'''
	try:
		floatArg = float(raw_input(rawInputString))
	except ValueError:
		print '\nYour input produced a ValueError. \n Please ensure your input is an integer or decimal - no alphabet characters.\n'
		floatArg = tryFloat(rawInputString)
	return(floatArg)
 
 
def getList(axis):
	'''a function to get the xList from user input. The axis argument is used 
	to construct the user interface and represents the x, y or z axis'''
	optionError = True
	listArg = []
	while optionError == True:
		incrOpt = raw_input('\nTo enter OPL (optical path length) as MAX,MIN and INCREMENT, enter 1. Else,'
		+ ' to enter as a list enter 2: ')
		if incrOpt == '1':
			optionError = False
			minMaxValid = False
			while minMaxValid == False:
				MIN  = tryFloat('Please enter the OPL minimum: ')
				MAX = tryFloat('Please enter the OPL maximum: ')
				if MIN == MAX or MIN > MAX:
					print 'Your values for MIN and MAX are invalid.'
				else:
					minMaxValid = True
			INCRValid = False
			while INCRValid == False:
				INCR = tryFloat('Please enter the increment: ')
				if INCR > (MAX-MIN):
					print 'Your increment is too large.'
				elif INCR == 0:
					print 'Cannot have a Zero increment.'
				else:
					INCRValid = True
			listArg = generateListIncre(MIN, MAX, INCR)
		elif incrOpt == '2':
			optionError = False
			finished = False
			print('Please input the list of OPL positions 1 at a time. To end the list, input end')
			i = 0
			while finished == False:
				i += 1
				tempInput = raw_input('Input OPL' + str(i) + ': ')
				if tempInput.lower() == 'end':
					finished = True
				else:
					try:
						listArg.append(float(tempInput))			
					except ValueError:
						print 'Please ensure the input is either integer or decimal. No alphabet characters.'
						i -= 1
		else:
			print 'Your input produced an error, please enter either 1 or 2'
			optionError = True
	return(listArg)
 
def processDataNoInput(freq,rootFolderName):
	'''a function to process all the data with no need for user input more than root dir'''
	outFileName = rootFolderName+str(freq)+'GHz_mean_sdev.txt'
	outFile = open(outFileName, 'w')
	outFile.write('Interferometer Measurements\n Nominal Frequency = ' + str(freq) + 'GHz\n')
	outFile.write('x(mm)\tMean (W)\tError (W)\n')
	dirList = glob.glob(rootFolderName+str(freq)+'GHz/*')
	for i in range(len(dirList)):
		if os.path.isdir(dirList[i]):
			dirList.pop(i);
	sortDirList(dirList)
	for i in dirList:
		print 'i = ', i
		inputFile = open(i, 'r')
		dataList = []
		dataList = inputFile.readlines()
		dataList.pop(0)
		dataList.pop(0)
		dataList.pop(0)
		dataList.pop(0)
		for k in range(len(dataList)):
			dataList[k] = string.split(dataList[k])
		dataMean = mean(dataList)
		dataSDev = sDev(dataList)
		outFile.write(str(i[findStartNum(i):]))
		outFile.write('\t'+str(dataMean) + '\t'+ str(dataSDev)+'\n')
	print '>>> data successfully written to ', outFileName
	return(outFileName)
 
def mean(listOfData):
	'''takes a list(input) of data from TK Power Meter and finds the mean of all the voltages(returns float)'''
	sum = 0.00
	total = 0.00
	for i in range(len(listOfData)):
		tempList = listOfData[i]
		sum += float(tempList[1])
		total += 1.00
	mean = sum/total
	return mean
 
def sDev(listOfData):
	'''takes a list(input) of data from TK Power Meter and finds the standard deviation 
	of all the voltages(returns float)'''
	sDevMean = mean(listOfData)
	N = len(listOfData)
	sum = 0.00
	for i in range(N):
		tempList = listOfData[i]
		operand = float(tempList[1]) - sDevMean
		sum+= operand**2
	standardDeviation = operand/float(N)
	return standardDeviation
 
def getXDistwErrors(dataList):
	'''generates a list of x values, a list of data points and a list of errors where the xvalues are dataList[i][0], data points are dataList[i][1] and errors are dataList[i][2]. Returns a nested list where nestedList = [xValues, dataPoints, errors]'''
	xValues = []
	dataPoints = []
	errors = []
	for i in dataList:
		xValues.append(float(i[0]))
		dataPoints.append(float(i[1]))
		errors.append(float(i[2]))
	return([xValues, dataPoints, errors])
 
def findStartNum(stringArg):
	'''given string with number at the end, returns the index where the number at the end begins.'''
	i = len(stringArg)-1
	while i > 0:
		if stringArg[i] in string.digits or stringArg[i] == '-' or stringArg[i] == '.' or stringArg[i] in string.whitespace:
			i-=1
		elif stringArg[i] == '=' or stringArg[i] == 'x' or stringArg[i] in string.ascii_letters or stringArg[i] in string.punctuation:
			return i+1
		else:
			i-=1		
	return i
 
def sortDirList(dirList):	
	'''sorts a listArg by list[i][0][30:]'''
	newList = []
	numStart = findStartNum(dirList[0])
	tempStr = dirList[0][:numStart]
	for i in dirList:
		newList.append([i[numStart:]])
	finished = False
	level = 0
	i = 0
	j = 0
	k = 0
	tempList = []
	checkLen = len(newList)
	while len(newList)!=1:
		while i < checkLen-1:
			finished = False
			while finished == False:
				if len(newList[i]) == j and len(newList[i+1]) == k:
					newList[i] = tempList[:]
					tempList = []
					newList.pop(i+1)
					i+=1
					j = 0
					k = 0
					checkLen = len(newList)
					finished = True
				elif len(newList[i]) == j:
					tempList.append(newList[i+1][k])
					k+=1
				elif len(newList[i+1])==k:
					tempList.append(newList[i][j])
					j+=1
				elif float(newList[i][j]) < float(newList[i+1][k]):
					tempList.append(newList[i][j])
					j+=1
				else:
					tempList.append(newList[i+1][k])
					k+=1
		if len(newList[-1]) == 1:
			j = 0
			k = 0
			done = False
			while done == False:
				if len(newList[-1]) == k and len(newList[-2]) == j:
					done = True
					j = 0
					k = 0
					newList[-2] = tempList[:]
					newList.pop(-1)
					tempList = []
				elif len(newList[-1]) == k:
					tempList.append(newList[-2][j])
					j +=1
				elif len(newList[-2]) == j:
					tempList.append(newList[-1][k])
					k += 1
				elif float(newList[-1][k])<float(newList[-2][j]):
					tempList.append(newList[-1][k])
					k+=1
				else:
					tempList.append(newList[-2][j])
					j+=1
		i = 0
	for i in range(len(newList[0])):
		tempStr2 = tempStr + str(newList[0][i])
		dirList[i] = tempStr2
 
finished=False
dirSame = False
print('\nThis program analyses the interferogram data from the TK power meter and then plots a graph of x-position versus intensity.\n')
while finished == False:
	dirFound = False
	if dirSame == False:
		rootFolderName = ''
		while dirFound == False:
			rootFolderName = raw_input('Please enter the data root folder: ')
			if rootFolderName[-1] != '/':
				rootFolderName+='/'
			dirFound = os.path.exists(rootFolderName)
			if dirFound == False:
				print 'Directory not found.'
	freq = raw_input('Please input the frequency the emitter is set to: ' )
	iFi = open(processDataNoInput(freq, rootFolderName))
	lines = iFi.readlines()
	lines.pop(0)
	freqString = lines.pop(0)
	unitsString = lines.pop(0)
	plotList=[]
	for i in range(len(lines)):
		plotList.append(string.split(lines[i]))
	plotList = getXDistwErrors(plotList)
	for i in range(len(plotList[1])):
		plotList[1][i]=plotList[1][i]*1000
	plt.figure()
	plt.errorbar(plotList[0],plotList[1],yerr=plotList[2])
	plt.xlabel('Position (mm)')
	plt.ylabel('Power ('+'mW)')
	plotTitle = str(freq) + 'GHz Interferogram'	
	plt.title(plotTitle)
	fin = raw_input('If you need to analyse more data in the same rootdir, enter "1"\nif you need a new root dir, enter "2" \nelse enter anything: ')
	if fin != '1' and fin != '2':
		finished = True
	elif fin == '1':
		dirSame = True
	else: 
		dirSame = False
 
 
plt.show()