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

exp_gs = np.array(
[[0.020651995821983293, 0.013771069285436477],
[0.03963227312757267, 0.05820703709013464],
[0.07131300305159005, 0.07718398631904488],
[0.09050832531181463, 0.09833058553669072],
[0.10925307719089816, 0.16827395703197023],
[0.10936571978618437, 0.1560747639625617],
[0.11037182296680093, 0.19711378950171032],
[0.1292906588567797, 0.24820386261699467],
[0.14295625371208542, 0.268219939787413],
[0.14684498330841533, 0.29707052450488436],
[0.16013168943412426, 0.358120251090585],
[0.16280951113113606, 0.368112161304196],
[0.18136993876338914, 0.45801784873123474],
[0.18149282159460967, 0.4447096381100619],
[0.1817385872570505, 0.4180932168677165],
[0.19244987404509795, 0.4580608577221619],
[0.22077692670039095, 0.5402410551539106],
[0.2316827779712045, 0.5591373625248324],
[0.23195926434145053, 0.5291938886271939],
[0.42224076842730485, 0.7717069961291909],
[0.42244557314600595, 0.749526645093903],
[0.5695337620578784, 0.8698757859381078],
[0.7134909988326137, 0.879307043234276],
[0.8560632436971356, 0.8887329244065781],
[0.8572229504167783, 0.9131366866692607],
[0.8576325598541801, 0.8687759845986851],
[0.9542850267270158, 0.901313822270465],
[1.0016512380445257, 0.8715531365842669]]
)

exp_excited = np.array(
[[0.054590697769676844, 0.08820964835029743],
[0.08283582853748928, 0.17926198619616174],
[0.09679071005795992, 0.3179483175292359],
[0.1049777786880206, 0.33128878489360347],
[0.2930857926966641, 0.6591908677575932],
[1.2883548036946784, 0.8715569766727422],
[1.5843949044586, 0.9104140639400331],
[2.0205956745243445, 0.9198706658201402]]
)

gs_carter = np.array(
[[0.05024492740782016, 0.006297229219143663],
[0.10047593129374843, 0.026448362720403296],
[0.14959052187891608, 0.15743073047858935],
[0.1997202637874511, 0.27833753148614626],
[0.24982722175107275, 0.42191435768261976],
[0.29893928078680565, 0.5554156171284635],
[0.3490740857942104, 0.6712846347607054],
[0.39923040897181117, 0.7657430730478592],
[0.4494069845448908, 0.8400503778337532],
[0.4986101793602774, 0.8828715365239295],
[0.5990645924838299, 0.9307304785894207],
[0.6985443590749721, 0.9483627204030227],
[0.8000430363403924, 0.9571788413098237],
[0.8975330050757567, 0.9546599496221664],
[0.9980253914408322, 0.964735516372796]]
)

gs_MS_RPBEl = np.array(
[[0.025, 0.33100],
[0.05, 0.77400],
[0.075, 0.95000],
[0.1, 0.99300]]
)

TN_Trot_MS_RPBEl = np.array(
[[0.01, 0.00100],
[0.025, 0.30133],
[0.05, 0.76467],
[0.075, 0.94333],
[0.1, 0.99100]]
)

gs_HSE03_13X = np.array(
[[0.05, 0.03100],
[0.075, 0.15300],
[0.1, 0.30300],
[0.15, 0.61600],
[0.2, 0.880],
[0.3, 0.99000]]
)

# Nozzle temperature is 300 K and Trot = 0.8*TN
#TN_HSE03_13X = np.array(
#[[0.05, 0.01900],
#[0.075, 0.10400],
#[0.1, 0.24800],
#[0.15, 0.61900],
#[0.2, 0.89800],
#[0.3, 0.99400]]
#)

# Nozzle temperature is 300 K and Trot = 0.05*TN
#TN_Trot_HSE03_13X = np.array(
#[[0.05, 0.02200],
#[0.075, 0.11600],
#[0.1, 0.26800],
#[0.15, 0.62400],
#[0.2, 0.89400],
#[0.3, 0.995]]
#)

# Nozzle temperature is 300 K and Trot = 0.03*TN = 9 K
TN_Trot_HSE03_13X = np.array(
[[0.05, 0.02200],
[0.075, 0.11550],
[0.1, 0.25250],
[0.125, 0.44350],
[0.15, 0.61800],
[0.175, 0.77467],
[0.2, 0.89600],
[0.225, 0.96000],
[0.25, 0.98100]]
)

# v=0, J=2, K=1, mJ sampled randomly
# Ei, sticking of all mk states, mK=-1, 0, 1, seed 11, 10000 trajectories
HSE03_13X_random = np.array(
[[0.05, 0.02160, 0.02370, 0.01680, 0.02430],
[0.075, 0.11428, 0.12418, 0.08308, 0.13557],
[0.1, 0.26647, 0.30539, 0.16168, 0.33234],#1000 trajectories
[0.125, 0.45907, 0.52231, 0.32208, 0.53281],#seed 8, 8000 trajectories
[0.15, 0.62275, 0.68862, 0.45210, 0.72754],#seed 10, 1000 trajectories
[0.175, 0.76462, 0.83058, 0.61169, 0.85157],#seed 10, 2000 trajectories
[0.2, 0.88922, 0.92216, 0.81138, 0.93413],#1000 trajectories
[0.225, 0.95908, 0.96707, 0.92814, 0.98204],#seed 8, 1000 trajectories
[0.25, 0.98703, 0.99401, 0.97605, 0.99102]]#1000 trajectories
)
HSE03_13X_random[:,0] *= 96.4853075

HSE03_13X_random_cartwheel = (3.*HSE03_13X_random[:,1]-HSE03_13X_random[:,4])/2.
HSE03_13X_random_perpendicular = 2.*HSE03_13X_random_cartwheel - HSE03_13X_random[:,4]

carter_helicopter = np.array(
[[0.05459, 0.02337],
[0.10924, 0.18114],
[0.14202, 0.26461],
[0.18574, 0.38898],
[0.22945, 0.51002],
[0.27275, 0.61185],
[0.31685, 0.68447],
[0.36015, 0.76127],
[0.40343, 0.82888],
[0.44712, 0.87062],
[0.49079, 0.89983]]
)
carter_helicopter[:,0] *= 96.4853075

carter_random = np.array(
[[0.05500, 0.02588],
[0.10882, 0.15359],
[0.14201, 0.23038],
[0.18571, 0.33139],
[0.22943, 0.44825],
[0.27313, 0.55175],
[0.31684, 0.64274],
[0.36054, 0.73038],
[0.40383, 0.79883],
[0.44751, 0.85142],
[0.49119, 0.88314]]
)
carter_random[:,0] *= 96.4853075

carter_perpendicular = np.array(
[[0.05459, 0.02421],
[0.10921, 0.10351],
[0.14238, 0.15860],
[0.18526, 0.21369],
[0.22938, 0.32554],
[0.27268, 0.42988],
[0.31640, 0.56093],
[0.36011, 0.66528],
[0.40340, 0.73957],
[0.44710, 0.81636],
[0.49118, 0.85476]]
)
carter_perpendicular[:,0] *= 96.4853075

exp_random = np.array(
[[0.10437, 0.12285],
[0.10572, 0.12943],
[0.11752, 0.13748],
[0.13606, 0.19890],
[0.17151, 0.29762],
[0.17153, 0.32029],
[0.18548, 0.34150],
[0.18737, 0.35101],
[0.20753, 0.43071],
[0.20780, 0.43876],
[0.25317, 0.52358],
[0.26015, 0.54918],
[0.27441, 0.62523],
[0.28192, 0.61865],
[0.32484, 0.66179],
[0.33450, 0.67788]]
)
exp_random[:,0] *= 96.4853075

exp_helicopter = np.array(
[[0.10413, 0.15868],
[0.10575, 0.17038],
[0.11675, 0.18501],
[0.13584, 0.26691],
[0.17157, 0.38172],
[0.17159, 0.41609],
[0.18500, 0.42048],
[0.18742, 0.43510],
[0.20758, 0.50603],
[0.20785, 0.51700],
[0.25348, 0.59305],
[0.25994, 0.62303],
[0.27392, 0.68227],
[0.28196, 0.67642],
[0.32460, 0.69835],
[0.33426, 0.71737]]
)
exp_helicopter[:,0] *= 96.4853075

exp_perpendicular = np.array(
[[0.10432, 0.04753],
[0.10540, 0.05192],
[0.11719, 0.04241],
[0.13570, 0.06801],
[0.17140, 0.13309],
[0.17142, 0.13400],
[0.18699, 0.18135],
[0.18564, 0.18062],
[0.20796, 0.28080],
[0.20798, 0.28100],
[0.25308, 0.38684],
[0.25979, 0.40219],
[0.27461, 0.51700],
[0.28184, 0.50823],
[0.32452, 0.58867],
[0.33419, 0.60256]]
)
exp_perpendicular[:,0] *= 96.4853075

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

ax2 = plt.subplot(Nrows,Ncols,1)

#plt.errorbar( Energy_LO, Stick_LO2, yerr=Stick_LO2_error, marker='o', linestyle='--', color='blue', fillstyle='none' )
#plt.errorbar( Energy_LO, Stick_LO, yerr=Stick_LO_error, marker='o', label='Laser off', color='blue' )
#plt.errorbar( Energy_v1, Stick_v12, yerr=Stick_v12_error, marker='o', linestyle='--', color='red', fillstyle='none' )
#plt.errorbar( Energy_v1, Stick_v1, yerr=Stick_v1_error, marker='o', label=r'$\nu_1=1$', color='red' )
plt.scatter( exp_gs[:,0]*96.4853075, exp_gs[:,1], c='k', marker='d', label='Exp.')
plt.plot( [0.,150.], [1.,1.], c='g', label='(R)PBE')
plt.plot( TN_Trot_MS_RPBEl[:,0]*96.4853075, TN_Trot_MS_RPBEl[:,1], c='b', marker='o', label='MS-RPBEl' )
#plt.plot( gs_HSE03_13X[:,0], gs_HSE03_13X[:,1], c='g', marker='o', label='HSE03-1/3X' )
plt.plot( TN_Trot_HSE03_13X[:,0]*96.4853075, TN_Trot_HSE03_13X[:,1], c='r', marker='o', label='HSE03-1/3X' )
plt.plot( gs_carter[:,0]*96.4853075, gs_carter[:,1], c='orange', marker='s', label='ECW')
#plt.plot( gs_HSE03_13X[:,0], gs_HSE03_13X[:,1], c='r', marker='o', label='HSE03-1/3X GS' )
#plt.plot( TN_HSE03_13X[:,0], TN_HSE03_13X[:,1], c='orange', marker='o', label='HSE03-1/3X TN, Tr=240' )

plt.annotate( '(a)', xy=(45., 0.1), xycoords='data', fontsize=10)

plt.legend(loc='lower right', numpoints=1, scatterpoints=1)
plt.xlim(0.,100.)
plt.ylim(0.,1.1)
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
plt.ylabel('Sticking probability')
#plt.xlabel('Incidence energy (kJ/mol)')

ax2 = plt.subplot(Nrows,Ncols,2)

#exp helicopter
x_axis = np.arange(10., 33., 0.1)
fit = []
A = 0.63
E0 = 18.
W = 9.
offset = 0.1
for i in range(len(x_axis)):
	fit.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) + offset )
plt.plot(x_axis, fit, color='r', label='', linestyle='--')

#exp random
fit = []
A = 0.62
E0 = 18.
W = 9.
offset = 0.06
for i in range(len(x_axis)):
	fit.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) + offset )
plt.plot(x_axis, fit, color='k', label='', linestyle='--')

#exp perpendicular
fit = []
A = 0.62
E0 = 22.
W = 9.
offset = 0.02
for i in range(len(x_axis)):
	fit.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) + offset )
plt.plot(x_axis, fit, color='b', label='', linestyle='--')

#HSE03-1/3X helicopter
x_axis = np.arange(5., 25., 0.1)
fit = []
A = 1.
E0 = 12.
W = 6.
offset = 0.
for i in range(len(x_axis)):
	fit.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) + offset )
plt.plot(x_axis, fit, color='r', label='', linestyle='-')

#HSE03-1/3X random
x_axis = np.arange(5., 25., 0.1)
fit = []
A = 1.
E0 = 13.
W = 7.
offset = 0.
for i in range(len(x_axis)):
	fit.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) + offset )
plt.plot(x_axis, fit, color='k', label='', linestyle='-')

#HSE03-1/3X perpendicular
x_axis = np.arange(5., 25., 0.1)
fit = []
A = 1.
E0 = 15.
W = 7.
offset = 0.
for i in range(len(x_axis)):
	fit.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) + offset )
plt.plot(x_axis, fit, color='b', label='', linestyle='-')

#ax2 = plt.subplot(Nrows,Ncols,1)

plt.scatter( exp_helicopter[:,0], exp_helicopter[:,1], edgecolor='r', marker='v', facecolor='None')
plt.scatter( exp_random[:,0], exp_random[:,1], edgecolor='k', marker='o', facecolor='None')
plt.scatter( exp_perpendicular[:,0], exp_perpendicular[:,1], edgecolor='b', marker='^', facecolor='None')
#plt.scatter( carter_helicopter[:,0], carter_helicopter[:,1], edgecolor='r', marker='v', facecolor='k')
#plt.scatter( carter_random[:,0], carter_random[:,1], edgecolor='k', marker='o', facecolor='k')
#plt.scatter( carter_perpendicular[:,0], carter_perpendicular[:,1], edgecolor='b', marker='^', facecolor='k')
plt.scatter( HSE03_13X_random[:,0], HSE03_13X_random[:,4], c='r', marker='v', label='helicopter' )
plt.scatter( HSE03_13X_random[:,0], HSE03_13X_random[:,1], c='k', marker='o', label='random' )
plt.scatter( HSE03_13X_random[:,0], HSE03_13X_random_perpendicular[:], c='b', marker='^', label='perpendicular' )

plt.annotate( r'$\nu=0, J=2, mK=1$', xy=(50., 0.8), xycoords='data', fontsize=10)

plt.annotate( '(b)', xy=(45., 0.1), xycoords='data', fontsize=10)

plt.legend(loc='lower right', numpoints=1, scatterpoints=1)
plt.xlim(0.,100.)
plt.ylim(0.,1.)
plt.yticks([0.00, 0.25, 0.50, 0.75])
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
plt.ylabel('Sticking probability')
plt.xlabel('Incidence energy (kJ/mol)')

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