#!/usr/bin/env python
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np

# E, top, hollow, bridge, total # at the moment of reaction
nNN = np.array(
[[114, 64, 77, 246, 387, 20000.],
[124, 73, 160, 430, 663, 20000.],
[150, 220, 595, 1303, 2118, 20000.],
[174, 349, 1039, 2050, 3438, 20000.],
[205, 501, 1505, 2894, 4900, 20000.],
[247, 770, 2046, 3886, 6702, 20000.]]
)

nNN_FS = np.array(
[[114, 69, 94, 345, 508, 20000.],
[124, 95, 146, 545, 786, 20000.],
[150, 240, 586, 1467, 2293, 20000.],
[174, 352, 1030, 2286, 3668, 20000.],
[205, 516, 1526, 3082, 5124, 20000.],
[247, 731, 2006, 3987, 6724, 20000.]]
)

nNN_v2 = np.array(
[[43, 1361, 626, 1655, 3642, 96070.],
[50, 2606, 1243, 3232, 7081, 95555.],
[75, 4780, 4257, 8900, 17937, 94401.],
[94, 5828, 5566, 11039, 22433, 89828.],
[114, 6947, 6555, 13501, 27003, 92873.],
[124, 7150, 6848, 13821, 27819, 88917.]]
)

dErovib = 66.9822596142
nNN_v2[:,0] += dErovib

def confidence(P,T):
	'''compute the 1 sigma confidence given a probability P and the total number of events T'''
	if P == 0:
		sigma = 1 - 0.68**(1./T)
	else:
		sigma = np.sqrt( P * (1-P) / T )

	return sigma

def site( nsites ):
	psites = np.zeros((len(nsites), 7))

	for i in range(len(nsites)):
		psites[i][0] = float(nsites[i][0])

		psites[i][1] = float(nsites[i][1]) / float(nsites[i][4])
		psites[i][2] = confidence( psites[i][1], float(nsites[i][4]) )

		psites[i][3] = float(nsites[i][2]) / float(nsites[i][4])
		psites[i][4] = confidence( psites[i][3], float(nsites[i][4]) )

		psites[i][5] = float(nsites[i][3]) / float(nsites[i][4])
		psites[i][6] = confidence( psites[i][5], float(nsites[i][4]) )

	return psites

def prob( nsites ):
	psites = np.zeros((len(nsites), 7))

	for i in range(len(nsites)):
		psites[i][0] = nsites[i][0]

		psites[i][1] = float(nsites[i][1]) / float(nsites[i][5])
		psites[i][2] = confidence( psites[i][1], float(nsites[i][5]) )

		psites[i][3] = float(nsites[i][2]) / float(nsites[i][5])
		psites[i][4] = confidence( psites[i][3], float(nsites[i][5]) )

		psites[i][5] = float(nsites[i][3]) / float(nsites[i][5])
		psites[i][6] = confidence( psites[i][5], float(nsites[i][5]) )

	return psites

site_NN = site( nNN )
p_NN = prob( nNN )
#site_NN_FS = site( nNN_FS )
#p_NN_FS = prob( nNN_FS )
site_NN_v2 = site( nNN_v2 )
p_NN_v2 = prob( nNN_v2 )

mpl.rc("text", usetex=True)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(3.69,5.))

plt.subplot(2,1,1)

plt.errorbar( site_NN[:,0], site_NN[:,1], yerr=site_NN[:,2], marker='o', label='Top', color='b', capsize=4)
plt.errorbar( site_NN[:,0], site_NN[:,3], yerr=site_NN[:,4], marker='d', label='Hollow', color='r', capsize=4)
plt.errorbar( site_NN[:,0], site_NN[:,5], yerr=site_NN[:,6], marker='v', label='Bridge', color='g', capsize=4)

#plt.errorbar( site_NN_FS[:,0], site_NN_FS[:,1], yerr=site_NN_FS[:,2], marker='o', label='Top', color='b', linestyle='--', markerfacecolor='none', capsize=4)
#plt.errorbar( site_NN_FS[:,0], site_NN_FS[:,3], yerr=site_NN_FS[:,4], marker='d', label='Hollow', color='r', linestyle='--', markerfacecolor='none', capsize=4)
#plt.errorbar( site_NN_FS[:,0], site_NN_FS[:,5], yerr=site_NN_FS[:,6], marker='v', label='Bridge', color='g', linestyle='--', markerfacecolor='none', capsize=4)

plt.errorbar( site_NN_v2[:,0], site_NN_v2[:,1], yerr=site_NN_v2[:,2], marker='o', label='Top', color='b', linestyle='--', markerfacecolor='none', capsize=4)
plt.errorbar( site_NN_v2[:,0], site_NN_v2[:,3], yerr=site_NN_v2[:,4], marker='d', label='Hollow', color='r', linestyle='--', markerfacecolor='none', capsize=4)
plt.errorbar( site_NN_v2[:,0], site_NN_v2[:,5], yerr=site_NN_v2[:,6], marker='v', label='Bridge', color='g', linestyle='--', markerfacecolor='none', capsize=4)

plt.plot([0., 400.], [0.25,0.25], linestyle='-', color='b', label='')
plt.plot([0., 400.], [0.25,0.25], linestyle='--', color='r', label='')
plt.plot([0., 400.], [0.5,0.5], linestyle=':', color='g', label='')

#plt.legend(loc='upper left', numpoints=1, handletextpad=0.5, borderaxespad=0.2, frameon=True)
#plt.xlim(100.,260.)
plt.xlim(100.,200.)
plt.ylim(0.,0.8)
plt.tick_params(labelbottom=False, length=6, width=1, direction='in', top=True, right=True)
#plt.xlabel('Incidence energy (kJ/mol)')
plt.ylabel('Probability of closest site')
plt.annotate('(a)', xy=(230, 0.65), size=12)

ax2 = plt.subplot(2,1,2)

plt.errorbar( p_NN[:,0], p_NN[:,1], yerr=p_NN[:,2], marker='o', label='Top', color='b', capsize=4)
plt.errorbar( p_NN[:,0], p_NN[:,3], yerr=p_NN[:,4], marker='d', label='Hollow', color='r', capsize=4)
plt.errorbar( p_NN[:,0], p_NN[:,5], yerr=p_NN[:,6], marker='v', label='Bridge', color='g', capsize=4)

#plt.errorbar( p_NN_FS[:,0], p_NN_FS[:,1], yerr=p_NN_FS[:,2], marker='o', color='b', linestyle='--', markerfacecolor='none', capsize=4)
#plt.errorbar( p_NN_FS[:,0], p_NN_FS[:,3], yerr=p_NN_FS[:,4], marker='d', color='r', linestyle='--', markerfacecolor='none', capsize=4)
#plt.errorbar( p_NN_FS[:,0], p_NN_FS[:,5], yerr=p_NN_FS[:,6], marker='v', color='g', linestyle='--', markerfacecolor='none', capsize=4)

plt.errorbar( p_NN_v2[:,0], p_NN_v2[:,1], yerr=p_NN_v2[:,2], marker='o', color='b', linestyle='--', markerfacecolor='none', capsize=4)
plt.errorbar( p_NN_v2[:,0], p_NN_v2[:,3], yerr=p_NN_v2[:,4], marker='d', color='r', linestyle='--', markerfacecolor='none', capsize=4)
plt.errorbar( p_NN_v2[:,0], p_NN_v2[:,5], yerr=p_NN_v2[:,6], marker='v', color='g', linestyle='--', markerfacecolor='none', capsize=4)

plt.legend(loc='upper left', numpoints=1, handletextpad=0.5, borderaxespad=0.2, frameon=False, fontsize=13)
#plt.xlim(100.,260.)
plt.xlim(100.,200.)
plt.ylim(0.,0.2)
plt.tick_params(labelleft=False, labelright=True, length=6, width=1, direction='in', top=True, right=True)
ax2.yaxis.set_label_position("right")
plt.xlabel('Incidence energy (kJ/mol)')
plt.ylabel('Site specific sticking probability')
plt.annotate('(b)', xy=(230, 0.21), size=12)

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig('reactionsite_p_v2.pdf')
#plt.show()