import numpy as np 
import matplotlib.pyplot as plt
import argparse 
import glob
import re

parser = argparse.ArgumentParser(description=' script ... ')
parser.add_argument('-i','--input', help='Input file dir',required=False,default='')#parsing for the input file
parser.add_argument('-p','--plot', help='plot data and fit; 0: plot; 1: no plot', type=int,default = 0)# plot yes or no
#parser.add_argument('-u','--unit', help='what is the unit?', type=int,default = 0)# plot yes or no

# read in parser arguments and input file
args = parser.parse_args()
#data_fromfile = np.genfromtxt(args.input)

filelist =glob.glob(args.input+"*.dat")
filelist.sort()

print "#freq bright dark fid fiderr loadings atombin fano fanodark"

for file in filelist:
	data_fromfile = np.genfromtxt(file)
	#print(re.findall('\d+', file))
	variable = re.findall('\d+', file)[3]

	blue_atom = data_fromfile[:,0]
	blue_bg = data_fromfile[:,1]
	red_atom = data_fromfile[:,2]
	red_bg = data_fromfile[:,3]
	sum_atom = blue_atom+red_atom
	sum_bg = blue_bg+red_bg

	###  histogram TRANSMISSION
	binwidth = 1
	binrange=  np.arange(0,200,binwidth)
	hist_blue_atom, bins_blue_atom = np.histogram(blue_atom, binrange)
	hist_blue_bg, bins_blue_bg = np.histogram(blue_bg, binrange)

	width_blue = binwidth*0.7
	center_blue = (bins_blue_atom[:-1] + bins_blue_atom[1:]) / 2


	###  histogram REFELCTION
	hist_red_atom, bins_red_atom = np.histogram(red_atom,binrange)
	hist_red_bg, bins_red_bg = np.histogram(red_bg, binrange)

	width_red = 0.7 * (bins_red_atom[1] - bins_red_atom[0])
	center_red = (bins_red_atom[:-1] + bins_red_atom[1:]) / 2

	hist_total_atom, bins = np.histogram(sum_atom, binrange)
	hist_total_bg, bins = np.histogram(sum_bg, binrange)

	### determine fidelity
	steps=7
	fidelities = np.zeros(steps)
	for i in range(1,steps):
		threshold=i
		err_atom = sum(hist_total_atom[:threshold])*1.0/sum(hist_total_atom)
		err_bg = sum(hist_total_bg[threshold:])*1.0/sum(hist_total_bg)
		fidelities[i]= 100* (1- 0.5* (err_atom+err_bg) )
	threshold = fidelities.argmax()
	err_atom = sum(hist_total_atom[:threshold])*1.0/sum(hist_total_atom)
	err_bg = sum(hist_total_bg[threshold:])*1.0/sum(hist_total_bg)
	fidelity_max = 100* (1- 0.5* (err_atom+err_bg) )
	fidelity_err = 100* np.sqrt( (np.sqrt(sum(hist_total_atom[:threshold]))/sum(hist_total_atom)/2)**2 
		+ (np.sqrt(sum(hist_total_bg[threshold:]))/sum(hist_total_bg)/2)**2 ) 

	fano_factor = sum_atom.std()**2/sum_atom.mean()
	fano_factor_bg = sum_bg.std()**2/sum_bg.mean()
	
	print(variable+' %.2f'%sum_atom.mean() +' %.3f'%(sum_atom.std()/np.sqrt(len(sum_atom)))
		+' %.2f'%sum_bg.mean() +' %.3f'%(sum_bg.std()/np.sqrt(len(sum_bg)))
		+ ' %.3f'%fidelity_max + ' %.3f'%fidelity_err 
		+ ' %d'%len(data_fromfile[:,0]) + ' %d'%(red_atom+blue_atom).sum()
		+ ' %.1f'%fano_factor + ' %.1f'%fano_factor_bg + ' %d'%threshold)

	print 'error of bright mean',sum_atom.std()/np.sqrt(len(sum_atom))
	print 'error in bright/dark',err_atom, err_bg

	np.savetxt(file+'.histo', np.array([bins[0:bins.size-1],hist_total_atom*1.0/hist_total_atom.sum(),hist_total_bg*1.0/hist_total_bg.sum()]).T , fmt='%.3f')

	if (args.plot):
		#blue_fig = plt.figure(1)
		#plt.bar(center_blue, hist_blue_atom, align='center', width=width_blue,alpha = 0.5)
		#plt.bar(center_blue, hist_blue_bg, align='center', width=width_blue, color='red', alpha = 0.5)
		#plt.xlabel('counts')
		#plt.ylabel('# of occurances')
		#plt.title('BLUE detector')
		#blue_fig.show()

		#red_fig = plt.figure(2)
		#plt.xlabel('counts')
		#plt.ylabel('# of occurances')
		#plt.title('RED detector')
		#plt.bar(center_red, hist_red_atom, align='center', width=width_red,alpha = 0.5)
		#plt.bar(center_red, hist_red_bg, align='center', width=width_red, color='red',alpha = 0.5)
		#red_fig.show()

		tot_fig = plt.figure(3)
		plt.xlabel('counts')
		plt.ylabel('# of occurances')
		plt.title('RED + BLUE')
		plt.bar(center_red, hist_total_atom, align='center', width=width_red,alpha = 0.5)
		plt.bar(center_red, hist_total_bg, align='center', width=width_red, color='red',alpha = 0.5)
		plt.ylim(ymax = hist_total_atom.max()*1.1, ymin = 0)
		#tot_fig.savefig('tmp.pdf', bbox_inches='tight')
		#tot_fig.show()
		plt.show()