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

mpl.rc("text", usetex=True)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(3.69,5.5))
plt.rcParams.update({'font.size': 10})

#Ei, Reaction, error
Geweke_2016_12 = np.array(
[[0.64, 0.0017],
[0.67, 0.0035],
[0.94, 0.0082],
[0.99, 0.015],
[1.06, 0.017]]
)

Geweke_2016_01 = np.array(
[[0.99, 0.0006],
[0.67, 0.00015],
[1.37, 0.0018],
[1.12, 0.00081],
[0.86, 0.00054],
[0.59, 0.000068]]
)

#Ei, excitation probability, standard error
# dt = 0.4 fs
NN_12 = np.array(
[[0.28, 0.0020, 0.0002],
[0.32, 0.0031, 0.0003],
[0.45, 0.0087, 0.0004],
[0.52, 0.0133, 0.0005],
[0.78, 0.0287, 0.0008],
[0.94, 0.0342, 0.0009],
[1.27, 0.0275,  0.0009]]
)

# dt = 0.2 fs
NN_12 = np.array(
[[0.28, 0.000194, 0.000044],
[0.32, 0.000419, 0.000065],
[0.45, 0.002074, 0.000145],
[0.52, 0.005293, 0.000232],
[0.78, 0.0230, 0.0005],
[0.97, 0.0295, 0.0006],
[1.27, 0.0284, 0.0006]]
)

# angle is 10-20
NN_12_angle = np.array(
[[0.28, 0.000250, 0.000144],
[0.32, 0.000568, 0.000214],
[0.45, 0.004938, 0.000594],
[0.52, 0.013939, 0.000962],
[0.78, 0.056768, 0.001925],
[0.97, 0.078685, 0.002433],
[1.27, 0.063793, 0.002136]]
)

NN_01 = np.array(
[[0.52, 0.000082, 0.000029],
[0.78, 0.000776, 0.000089],
[0.97, 0.004738, 0.000220],
[1.27, 0.023531, 0.000492]]
)

# angle is 7.5-22.5
NN_01_angle = np.array(
[[0.52, 0.000107, 0.000076],
[0.78, 0.001353, 0.000251],
[0.97, 0.008894, 0.000590],
[1.27, 0.036086, 0.001014]]
)

Gernot_12 = np.array( [[0.94, 0.074, 0.010]] )

liu_01 = np.array(
[[0.99, 0.0133],
[1.12, 0.0247],
[1.37, 0.0500]]
)

plt.subplot(2,1,1)

plt.scatter( Geweke_2016_12[:,0]*96.48530749926, Geweke_2016_12[:,1], marker='s', label='Exp.', color='k' )
plt.errorbar( NN_12[:,0]*96.48530749926, NN_12[:,1], yerr=NN_12[:,2], marker='o', label='MS-RPBEl', color='r', capsize=4, linestyle='', fillstyle='full' )
plt.errorbar( NN_12_angle[:,0]*96.48530749926, NN_12_angle[:,1], yerr=NN_12_angle[:,2], marker='o', label='MS-RPBEl ($10^\circ$-$20^\circ$)', color='b', capsize=4, linestyle='', fillstyle='full' )
#plt.errorbar( NN_12_gauss_nu[:,0]*96.48530749926, NN_12_gauss_nu[:,1], yerr=NN_12_gauss_nu[:,2], marker='o', label=r'MS-RPBEl (GB $\nu$)', color='r', capsize=4, linestyle='', fillstyle='left' )
#plt.errorbar( NN_12_gauss_nuj[:,0]*96.48530749926, NN_12_gauss_nuj[:,1], yerr=NN_12_gauss_nuj[:,2], marker='o', label=r'MS-RPBEl (GB $\nu, j$)', color='r', capsize=4, linestyle='', fillstyle='right' )
#plt.errorbar( NN_12_1GB[:,0]*96.48530749926, NN_12_1GB[:,1], yerr=NN_12_1GB[:,2], marker='o', label=r'MS-RPBEl (1GB $\nu, j, an, am$)', color='r', capsize=4, linestyle='', fillstyle='full' )
#plt.errorbar( NN_12_gauss_nuj[:,0]*96.48530749926, NN_12_gauss_nuj[:,1], yerr=NN_12_gauss_nuj[:,2], marker='o', label=r'MS-RPBEl (GB $\nu, j$)', color='r', capsize=4, linestyle='' )
#plt.errorbar( Gernot_12[:,0]*96.48530749926, Gernot_12[:,1], yerr=Gernot_12[:,2], marker='d', label='SRP32-vdW-DF1 (LDFA)', color='k', capsize=4, linestyle='' )

plt.annotate(r'$\nu=1, j=1\rightarrow\nu=2$', xy=(120., 5*10**-3), size=12)
plt.annotate('(a)', xy=(240., 5*10**-2), size=12)

plt.legend(loc='best', numpoints=1, frameon=True)
plt.xlim(0., 270.)
#plt.ylim(0.0, 0.05)
plt.ylim(1e-4, 0.1)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
#plt.ylabel('Vibrational transition probability')
#plt.xlabel('Incidence energy (kJ/mol)')

plt.subplot(2,1,2)

plt.scatter( Geweke_2016_01[:,0]*96.48530749926, Geweke_2016_01[:,1], marker='s', label='Exp.', color='k' )
plt.errorbar( NN_01[:,0]*96.48530749926, NN_01[:,1], yerr= NN_01[:,2], marker='o', label='MS-RPBEl', color='r', capsize=4, linestyle='', fillstyle='full' )
plt.errorbar( NN_01_angle[:,0]*96.48530749926, NN_01_angle[:,1], yerr= NN_01_angle[:,2], marker='o', label='MS-RPBEl ($7.5^\circ$-$22.5^\circ$)', color='b', capsize=4, linestyle='', fillstyle='full' )
#plt.errorbar( NN_01_gauss[:,0]*96.48530749926, NN_01_gauss[:,1], yerr= NN_01_gauss[:,2], marker='o', label=r'MS-RPBEl (GB $\nu, j$)', color='r', capsize=4, linestyle='', fillstyle='full' )
#plt.errorbar( NN_01_gauss_nuj[:,0]*96.48530749926, NN_01_gauss_nuj[:,1], yerr= NN_01_gauss_nuj[:,2], marker='o', label=r'MS-RPBEl (GB $\nu, j$)', color='r', capsize=4, linestyle='' )
#plt.scatter( liu_01[:,0]*96.48530749926, liu_01[:,1], marker='v', label='RPBE', color='k' )

plt.annotate(r'$\nu=0, j=0\rightarrow\nu=1$', xy=(100., 10**-4), size=12)
plt.annotate('(b)', xy=(240., 2*10**-2), size=12)

plt.legend(loc='lower right', numpoints=1, frameon=True)
plt.xlim(0., 270.)
#plt.ylim(0.0, 0.05)
plt.ylim(1e-6, 0.09)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
#plt.ylabel('Vibrational transition probability')
plt.xlabel('Incidence energy (kJ/mol)')

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.subplots_adjust(left=0.15)

fig.add_subplot(111, frameon=False)
# hide tick and tick label of the big axis
plt.tick_params(labelcolor='none', top=False, bottom=False, left=False, right=False)
plt.ylabel('Transition probability', labelpad=6.)

plt.savefig('Tvjprobability_angle.pdf')
#plt.show()