#!/usr/bin/python -u

#
# @file process.py
# @purpose Process TCS data
#		Takes S(E) = dI/dE
# @author Sam Moore
# @date August 2012
#

import sys
import os
import re # Regular expressions - for removing comments
import odict #ordered dictionary
import copy

import Gnuplot, Gnuplot.funcutils
import string

gnuplot = Gnuplot.Gnuplot()

def Reset():
	gnuplot = Gnuplot.Gnuplot()

def FindDataFiles(directory=".", depth=1, result=None):
	if result == None:
		result = []
	
	for f in os.listdir(directory):
		if os.path.isdir(directory+"/"+str(f)):
			if depth > 1:
				result += FindDataFiles(directory+"/"+str(f), depth-1, result)
			continue
		s = f.split(".")
		if (len(s) == 2 and s[1] == "dat"):
			result.append(directory+"/"+str(f))

	return result

def BaseName(f):
	a = f.split("/")
	return a[len(a)-1]
	

def DirectoryName(f, start=0,back=1):
	a = f.split("/")
	return string.join(a[start:(len(a)-back)], "/")

def GetData(filename):
	input_file = open(filename, "r")
	data = []
	for line in input_file:
		line = re.sub("#.*", "", line).strip("\r\n ")
		if len(line) == 0:
			continue
		data.append(map(lambda e : float(e), line.split("\t")))
	return data

def DoNothing(data):
	return data

def AverageAllDataSets(directory=".", function=DoNothing):
	dirs = {}
	for f in os.listdir(directory):
		if os.path.isdir(directory+"/"+str(f)) == True:
			data_set = []
			for datafile in os.listdir(directory+"/"+str(f)):
				if datafile.split(".")[1] == "dat":
					data_set.append(GetData(f))

			avg = Average(data_set)
			dirs.update({f : avg})
	return dirs

def GetDataSets(directory="."):
	data_sets = []
	for f in os.listdir(directory):
		if os.path.isdir(directory+"/"+str(f)) == False:
			if (len(f.split(".")) > 1 and f.split(".")[1] == "dat"):
				d = GetData(directory+"/"+str(f))
				if len(d) > 0:
					data_sets.append(d)
	return data_sets	



def Derivative(data, a=1, b=2, sigma=None,step=1):
	result = [[]]
	n = 0
	dI = [0,0]
	dE = [0,0]
	
	for i in range(0, len(data),step):
		result[len(result)-1] = [d for d in data[i]]
		if (i >= step):
			dE[0] = data[i][a] - data[i-step][a]
			dI[0] = data[i][b] - data[i-step][b]
		else:
			dI[0] = None

		if (i < len(data)-step):
			dE[1] = data[i+step][a] - data[i][a]
			dI[1] = data[i+step][b] - data[i][b]
		else:
			dI[1] = None		

		if sigma != None:
			#print str(data[i]) + " ["+str(sigma)+"] = " + str(data[i][int(abs(sigma))])
			if sigma < 0:
				if dI[0] != None: dI[0] -= 0.5*data[i][int(abs(sigma))]
				if dI[1] != None: dI[1] -= 0.5*data[i][int(abs(sigma))]
			else:
				if dI[0] != None: dI[0] += 0.5*data[i][int(abs(sigma))]
				if dI[1] != None: dI[1] += 0.5*data[i][int(abs(sigma))]

		deltaE = 0.0
		deltaI = 0.0
		count = 0
		if dI[0] != None:
			deltaE += dE[0]
			deltaI += dI[0]
			count += 1
		if dI[1] != None:
			deltaE += dE[1]
			deltaI += dI[1]
			count += 1

		deltaI /= float(count)
		deltaE /= float(count)


		if (deltaE != 0):	
			result[len(result)-1][b] = (deltaI / deltaE)
		else:
			result[len(result)-1][b] = 0.0
		result.append([])
			
	return result[0:len(result)-1]

def MaxNormalise(data, u=1):	
	result = copy.deepcopy(data)
	if (len(data) <= 0):
		return result
	maxval = max(data, key = lambda e : e[u])[u]

	for d in result:
		d[u] = d[u] / maxval
		
	return result
	
def Average(data_sets, u=1): 
	avg = odict.odict()
	for t in data_sets:
		for p in t:
			if p[u] in avg:
				#print "Already have " + str(p[u])
				avg[p[u]][1] += 1
				for i in range(0, len(p)):
					avg[p[u]][0][i] += p[i]
			else:
				#print "Create key for " + str(p[u])
				avg.update({p[u] : [p, 1]})

	for a in avg.keys():
		for i in range(0, len(avg[a][0])):
			avg[a][0][i] /= float(avg[a][1])
	return map(lambda e : e[1][0], sorted(avg.items(), key = lambda e : e[0]))

def FullWidthAtHalfMax(data, u=1):
	maxval = max(data, key = lambda e : e[u])
	peak = data.index(maxval)
	maxval = maxval[0]
	lhs = None
	rhs = None
	for i in range(1, len(data)/2):
		if lhs == None:
			if (peak-i > 0 and data[peak-i] < 0.50*maxval):
				lhs = data[peak-i][u]
		if rhs == None:
			if (peak+i < len(data) and data[peak+i] < 0.50*maxval):
				rhs = peak+i
		if lhs != None and rhs != None:
			break
	if rhs == None or lhs == None:
		return abs(data[len(data)-1][0] - data[0][0])
	else:
		return abs(rhs - lhs)

def SaveData(filename, data):
	out = open(filename, "w", 0)
	for a in data:
		for i in range(0, len(a)):
			out.write(str(a[i]))
			if (i < len(a) - 1):
				out.write("\t")
			out.write("\n")

def CalibrateData(data, ammeter_scale=1e-6):
	for i in range(0, len(data)):
		data[i][1] = 16.8 * float(data[i][1]) / 4000.0
		data[i][2] = ammeter_scale * 0.170 * float(data[i][2]) / 268.0
		data[i][3] = ammeter_scale * 0.170 * float(data[i][3]) / 268.0
	return data

def ShowTCS(filename, calibrate=True, normalise=False, show_error=False, plot=gnuplot.plot,w="lp", step=1, output=None, title=""):
	if type(filename) == type(""): 
		data = GetData(filename)
	else:
		data = filename
		filename = "tcs data"

	if (title == ""):
		title = BaseName(filename)

	if (len(data) <= 0):
		return data

	if calibrate: 
		data = CalibrateData(data)
		units = ["V", "uA / V"]
	else:
		units = ["DAC counts", "ADC counts / DAC counts"]

	if not normalise:
		gnuplot("set ylabel \"dI(E)/dE ("+str(units[1])+")\"")
	else:
		data = MaxNormalise(data)
		gnuplot("set ylabel \"dI(E)/dE (normalised)\"")

	if (output != None and type(output) == type("")):
		gnuplot("set term png size 640,480")
		gnuplot("set output \""+str(output)+"\"")

	gnuplot("set title \"Total Current Spectrum S(E)\"")
	gnuplot("set xlabel \"U ("+str(units[0])+")\"")


	d = Derivative(data, 1, 2, step=step)
	
	plot(Gnuplot.Data(d, using="2:3", with_=w,title=title))
	if (show_error):
		error1 = Derivative(data, 1, 2, -3,step=step)
		error2 = Derivative(data, 1, 2, +3,step=step)
		gnuplot.replot(Gnuplot.Data(error1, using="2:3", with_=w,title="-sigma/2"))
		gnuplot.replot(Gnuplot.Data(error2, using="2:3", with_=w, title="+sigma/2"))

	gnuplot("set term wxt")
	return data

def ShowData(filename,calibrate=True, normalise=False, show_error=False, plot=gnuplot.plot,w="lp", step=1, output=None, title=""):
	if type(filename) == type(""): 
		data = GetData(filename)
	else:
		data = filename
		filename = "raw data"

	if (title == ""):
		title = BaseName(filename)

	if len(data) <= 0:
		return data
	if calibrate: 
		data = CalibrateData(data)
		units = ["V", "uA"]
	else:
		units = ["DAC counts", "ADC counts"]

	if not normalise:
		gnuplot("set ylabel \"I(E) ("+str(units[1])+")\"")
	else:
		data = MaxNormalise(data)
		gnuplot("set ylabel \"I(E) (normalised)\"")

	if (output != None and type(output) == type("")):
		gnuplot("set term png size 640,480")
		gnuplot("set output \""+str(output)+"\"")

	gnuplot("set title \"Sample Current I(E)\"")
	gnuplot("set xlabel \"U ("+str(units[0])+")\"")


	#d = Derivative(data, 1, 2, step=step)
	
	plot(Gnuplot.Data(data, using="2:3", with_=w,title=title))
	if (show_error):
		error1 = copy.deepcopy(data)
		error2 = copy.deepcopy(data)
		for i in range(len(data)):
			#print str(data[i])
			error1[i][2] -= 0.50*float(data[i][3])
			error2[i][2] += 0.50*float(data[i][3])
		gnuplot.replot(Gnuplot.Data(error1, using="2:3", with_=w,title="Error : Low bound"))
		gnuplot.replot(Gnuplot.Data(error2, using="2:3", with_=w, title="Error : Upper bound"))
		
	gnuplot("set term wxt")
	return data

def main():	
	if (len(sys.argv) < 2):
		sys.stderr.write(sys.argv[0] + " - Require arguments (filename)\n")
		return 1

	tcs = []
	gnuplot("set style data lp")
	gnuplot("set key outside right")
	#gnuplot("set title \"Au on Si (50min 3.5A 3-6 e-8mbar)\"")
	#gnuplot("set xlabel \"E (DAC Counts)\"")
	#gnuplot("set ylabel \"S(E) (ADC/DAC Counts)\"")
	#gnuplot("set term postscript colour")
	#gnuplot("set output \"test.eps\"")
	for i in range(1, len(sys.argv)):
		tcs.append(Derivative(GetData(sys.argv[i]), 1, 2))
		#tcs.append(GetTCS(GetData(sys.argv[i])))
		if (len(tcs[i-1]) > 0):
			gnuplot.replot(Gnuplot.Data(tcs[i-1], title=sys.argv[i], with_="lp"))

	# Now average the data
	

	
	avg = Average(tcs)
	for a in avg:
		sys.stdout.write(str(a[0]) + "\t" + str(a[1]) + "\t" + str(a[1]) + "\n")
	gnuplot.replot(Gnuplot.Data(avg, title="Average", with_="l lw 2"))
	
	sys.stdout.write("Save averaged data as (blank for no save): ")
	filename = sys.stdin.readline().strip(" \r\n\t")
	if (filename != ""):
		SaveData(filename, avg)
		
	return 0


if __name__ == "__main__":
	sys.exit(main())