#!/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 Obtained from the fits
# 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]]
# )

Geweke_2016_12 = np.array(
[[0.59, 0.00137049, 0.000662441],
[0.64, 0.00270825, 0.000648565],
[0.92, 0.00738576, 0.00122523],
[0.94, 0.0149115, 0.00244504],
[1.04, 0.0184508, 0.00312071]]
)

Geweke_2016_01 = np.array(
[[0.59, 0.000041764, 0.0000312661],
[0.64, 0.0000924593, 0.0000164302],
[0.86, 0.000639136, 0.000217479],
[0.94, 0.000577473, 0.0000897007],
[1.12, 0.00087208, 0.000276175],
[1.37, 0.00197916, 0.000549222]]
)

#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
#0.000250 0.000144
#0.000568 0.000214
#0.004938 0.000594
#0.013939 0.000962
#0.056768 0.001925
#0.078685 0.002433
#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 5-25
#0.000118 0.000053
#0.001038 0.000150
#0.006444 0.000352
#0.028332 0.000679
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.errorbar( Geweke_2016_12[:,0]*96.48530749926, Geweke_2016_12[:,1], yerr=Geweke_2016_12[:,2], capsize=4, linestyle='', marker='s', label='Exp.', color='g' )
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_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.errorbar( Geweke_2016_01[:,0]*96.48530749926, Geweke_2016_01[:,1], yerr=Geweke_2016_01[:,2], capsize=4, linestyle='', marker='s', label='Exp.', color='g' )
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_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.pdf')
#plt.show()