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

# E, top, hollow, bridge, total
LO = np.array(
#[[89.2, 5., 3., 3., 11.],
[[89.2, 7., 1., 3., 11.],
[102.4, 11., 0., 11., 22.],
[111.8, 17., 5., 12., 34.],
[120.0, 11., 5., 17., 33.]]
)

# E, top, hollow, bridge, total
v1 = np.array(
[[89.2, 14., 7., 6., 27.],
[97.4, 30., 2., 18., 50.],
[102.4, 53., 16., 46., 115.],
[111.8, 39., 24., 52., 114.],
[120.0, 31., 15., 35., 81.]]
)

pt_LO_old = np.array(
[[81.7, 4, 6, 13, 23],
[89.3, 4, 8, 18, 30],
[97.4, 10, 5, 15, 30],
[102.5, 11, 7, 25, 43],
[111.9, 12, 13, 22, 47],
[120.1, 23, 16, 31, 70]]
)

pt_LO = np.array(
[[81.7, 12 ,  0 ,  11 ,  23],
[89.3, 19 ,  0 ,  11 ,  30],
[97.4, 20 ,  4 ,  6 ,  30],
[102.5, 31 ,  5 ,  7 ,  43],
[111.9, 24 ,  4 ,  19 ,  47],
[120.1, 41 ,  5 ,  24 ,  70]]
)

pt_v1_old = np.array(
[[60.7, 6, 8, 11, 25],
[71.4, 9, 5, 10, 24],
[81.9, 12, 6, 15, 33],
[92.2, 9, 12, 21, 42],
[104.6, 13, 11, 9, 36]]
)

pt_v1 = np.array(
[[60.7, 13 ,  3 ,  9 ,  25],
[71.4, 17 ,  1 ,  6 ,  24],
[81.9, 20 ,  4 ,  9 ,  33],
[92.2, 22 ,  3 ,  17 ,  42],
[104.6, 44 ,  2 ,  23 ,  69]]
)

ni_LO = np.array(
[[101.1, 10, 2, 2, 14],
[112.3, 10, 7, 11, 28],
[121.2, 24, 7, 9, 40]]
#[130.7, 32, 24, 52, 108],
#[136.4, 21, 12, 29, 62],
#[146.6, 24, 28, 54, 106],
#[160.4, 42, 37, 61, 140]]
)

ni_v1 = np.array(
[[71.4, 8, 1, 2, 11],
[89.2, 18, 4, 8, 30],
[101.1, 15, 8, 22, 45],
[112.3, 16, 12, 24, 52],
[121.2, 30, 15, 39, 84],
[130.7, 25, 28, 49, 102],
[136.4, 27, 32, 65, 124],
[146.6, 41, 43, 77, 161],
[160.4, 46, 56, 103, 205]]
)

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

        return SE

# E, S, error
LO_ = np.zeros((len(LO), 7))
for i in range(len(LO)):
	LO_[i][0] = LO[i][0]

	LO_[i][1] = LO[i][1] / LO[i][4]
	LO_[i][2] = confidence( LO_[i][1], LO[i][4] )

	LO_[i][3] = LO[i][2] / LO[i][4]
	LO_[i][4] = confidence( LO_[i][3], LO[i][4] )

	LO_[i][5] = LO[i][3] / LO[i][4]
	LO_[i][6] = confidence( LO_[i][5], LO[i][4] )

v1_ = np.zeros((len(v1), 7))
for i in range(len(v1)):
	v1_[i][0] = v1[i][0]

	v1_[i][1] = v1[i][1] / v1[i][4]
	v1_[i][2] = confidence( v1_[i][1], v1[i][4] )

	v1_[i][3] = v1[i][2] / v1[i][4]
	v1_[i][4] = confidence( v1_[i][3], v1[i][4] )

	v1_[i][5] = v1[i][3] / v1[i][4]
	v1_[i][6] = confidence( v1_[i][5], v1[i][4] )

pt_LO_ = np.zeros((len(pt_LO), 7))
for i in range(len(pt_LO)):
	pt_LO_[i][0] = pt_LO[i][0]

	pt_LO_[i][1] = pt_LO[i][1] / pt_LO[i][4]
	pt_LO_[i][2] = confidence( pt_LO_[i][1], pt_LO[i][4] )

	pt_LO_[i][3] = pt_LO[i][2] / pt_LO[i][4]
	pt_LO_[i][4] = confidence( pt_LO_[i][3], pt_LO[i][4] )

	pt_LO_[i][5] = pt_LO[i][3] / pt_LO[i][4]
	pt_LO_[i][6] = confidence( pt_LO_[i][5], pt_LO[i][4] )

pt_v1_ = np.zeros((len(pt_v1), 7))
for i in range(len(pt_v1)):
	pt_v1_[i][0] = pt_v1[i][0]

	pt_v1_[i][1] = pt_v1[i][1] / pt_v1[i][4]
	pt_v1_[i][2] = confidence( pt_v1_[i][1], pt_v1[i][4] )

	pt_v1_[i][3] = pt_v1[i][2] / pt_v1[i][4]
	pt_v1_[i][4] = confidence( pt_v1_[i][3], pt_v1[i][4] )

	pt_v1_[i][5] = pt_v1[i][3] / pt_v1[i][4]
	pt_v1_[i][6] = confidence( pt_v1_[i][5], pt_v1[i][4] )

ni_LO_ = np.zeros((len(ni_LO), 7))
for i in range(len(ni_LO)):
	ni_LO_[i][0] = ni_LO[i][0]

	ni_LO_[i][1] = ni_LO[i][1] / ni_LO[i][4]
	ni_LO_[i][2] = confidence( ni_LO_[i][1], ni_LO[i][4] )

	ni_LO_[i][3] = ni_LO[i][2] / ni_LO[i][4]
	ni_LO_[i][4] = confidence( ni_LO_[i][3], ni_LO[i][4] )

	ni_LO_[i][5] = ni_LO[i][3] / ni_LO[i][4]
	ni_LO_[i][6] = confidence( ni_LO_[i][5], ni_LO[i][4] )

ni_v1_ = np.zeros((len(ni_v1), 7))
for i in range(len(ni_v1)):
	ni_v1_[i][0] = ni_v1[i][0]

	ni_v1_[i][1] = ni_v1[i][1] / ni_v1[i][4]
	ni_v1_[i][2] = confidence( ni_v1_[i][1], ni_v1[i][4] )

	ni_v1_[i][3] = ni_v1[i][2] / ni_v1[i][4]
	ni_v1_[i][4] = confidence( ni_v1_[i][3], ni_v1[i][4] )

	ni_v1_[i][5] = ni_v1[i][3] / ni_v1[i][4]
	ni_v1_[i][6] = confidence( ni_v1_[i][5], ni_v1[i][4] )

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

def plot(ni, pd, pt, label, idx):
	plt.subplot(nrows,ncols,idx)

	plt.errorbar( ni[:,0], ni[:,1], yerr=ni[:,2], marker='o', label='Ni(111)', color='b')
	plt.errorbar( pd[:,0], pd[:,1], yerr=pd[:,2], marker='o', label='Pd(111)', color='k' )
	plt.errorbar( pt[:,0], pt[:,1], yerr=pt[:,2], marker='o', label='Pt(111)', color='r')

	if idx < 5:
		y = 0.25
	else:
		y = 0.5
	plt.plot([0., 400.], [y,y], linestyle=':', color='k', label='')

	if idx == 3:
		plt.legend(loc='upper left', numpoints=1, handletextpad=0.5, borderaxespad=0.2, frameon=False)
	if idx % 2 == 0:
		plt.xlim(50.,170.)
		plt.annotate(label, xy=(135, 0.83))
	else:
		plt.xlim(75.,125.)
		plt.annotate(label, xy=(110, 0.83))
	plt.ylim(0.,0.99)
	plt.tick_params(length=6, width=1)
	if idx % 2 == 0:
		plt.tick_params(labelleft='off')
	if idx < 5:
		plt.tick_params(labelbottom='off')
	if idx > 4:
		plt.xlabel('Incidence energy (kJ/mol)')
	if idx == 3:
		plt.ylabel('Probability of closest site')


plot( ni_LO_[:,0:3], LO_[:,0:3], pt_LO_[:,0:3], 'Top', 1 )
plot( ni_v1_[:,0:3], v1_[:,0:3], pt_v1_[:,0:3], 'Top', 2 )
plot( ni_LO_[:,[0,3,4]], LO_[:,[0,3,4]], pt_LO_[:,[0,3,4]], 'Hollow', 3 )
plot( ni_v1_[:,[0,3,4]], v1_[:,[0,3,4]], pt_v1_[:,[0,3,4]], 'Hollow', 4 )
plot( ni_LO_[:,[0,5,6]], LO_[:,[0,5,6]], pt_LO_[:,[0,5,6]], 'Bridge', 5 )
plot( ni_v1_[:,[0,5,6]], v1_[:,[0,5,6]], pt_v1_[:,[0,5,6]], 'Bridge', 6 )

plt.gcf().text(0.25, 0.95, 'Laser-off', fontsize=12)
plt.gcf().text(0.7, 0.95, r'$\nu_1=1$', fontsize=12)

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.subplots_adjust(top=0.93)
plt.savefig('probabilityreactionsite_comparison.pdf')