#!/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.]]
)

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 )

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

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.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.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.legend(loc='upper left', numpoints=1, handletextpad=0.5, borderaxespad=0.2, frameon=False, fontsize=13)
plt.xlim(100.,260.)
plt.ylim(0.,0.25)
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.pdf')
#plt.show()