#!/usr/bin/env python
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from scipy import interpolate
from scipy.special import erf
from sticking_data import *

def S0_erf_plot( E0, W, A, E ):
	S0 = []
	for i in range( len(E) ):
		S0.append( A / 2. * ( 1 + erf( (E[i] - E0) / W ) ) )

	return np.array(S0)

# population of v=1 for DS1 according to Shirhatti 2016
DS1_v1_pop = np.array(
[[0.94, 4.90E-07],
[1.18, 2.30E-05],
[1.29, 1.70E-04],
[1.55, 9.40E-04],
[1.80, 3.00E-03],
[2.12, 8.70E-03],
[2.56, 1.70E-02]]
)
fDS1_v1_pop = interpolate.UnivariateSpline(DS1_v1_pop[:,0]*96.48530749926, DS1_v1_pop[:,1], s=0, k=2)

mpl.rc("text", usetex=True)
Nrows = 2
Ncols = 2
fig, ax = plt.subplots(nrows=Nrows, ncols=Ncols, figsize=(5.,4.))
#fig = plt.figure(figsize=(3.69,3))
#plt.rcParams.update({'font.size': 10})

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

plt.errorbar( Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,1], yerr=Shirhatti_DS1[:,2], marker='s', label='DS1 old', color='C0', capsize=4., fillstyle='none', linestyle='none' )
plt.scatter( Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,3], marker='^', label='DS1 new (lower limit)', color='C1' )
plt.scatter( Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,4], marker='v', label='DS1 new (upper limit)', color='C2' )
#plt.errorbar( SRP32vdW[:,0]*96.48530749926, SRP32vdW[:,1], yerr=SRP32vdW[:,2], marker='d', label='SRP32-vdW-DF1 DS1', color='orange', capsize=4., fillstyle='none' )
#plt.errorbar( NN_DS1[:,0]*96.48530749926, NN_DS1[:,1], yerr=NN_DS1[:,2], marker='o', label='MS-RPBEl DS1', color='r', capsize=4., fillstyle='full' )
#plt.fill_between(Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,3], Shirhatti_DS1[:,4], color='grey', alpha=0.5)

xaxis = np.linspace( 90., 250., 200 )
yaxis = (1.-fDS1_v1_pop( xaxis ))*S0_erf_plot( 385.9, 48.24, 1., xaxis ) + fDS1_v1_pop( xaxis )*S0_erf_plot( 164., 48.24, 1., xaxis )
plt.plot( xaxis, yaxis, color='C0', linestyle='-' )

yaxis = (1.-fDS1_v1_pop( xaxis ))*S0_erf_plot( 395.6, 106.1, 1., xaxis ) + fDS1_v1_pop( xaxis )*S0_erf_plot( 115.8, 9.649, 1., xaxis )
plt.plot( xaxis, yaxis, color='C1', linestyle='-' )

yaxis = (1.-fDS1_v1_pop( xaxis ))*S0_erf_plot( 299.1, 77.19, 1., xaxis ) + fDS1_v1_pop( xaxis )*S0_erf_plot( 96.49, 19.3, 1., xaxis )
plt.plot( xaxis, yaxis, color='C2', linestyle='-' )

plt.annotate('(a)', xy=(100, 0.05), size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(80., 260.)
plt.ylim(0.0, 0.2)
#plt.ylim(1e-5, 0.5)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
plt.ylabel('Reaction probability')

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

plt.errorbar( Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,1], yerr=Shirhatti_DS1[:,2], marker='s', label='DS1 old', color='C0', capsize=4., fillstyle='none', linestyle='none' )
plt.scatter( Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,3], marker='^', label='DS1 new (lower limit)', color='C1' )
plt.scatter( Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,4], marker='v', label='DS1 new (upper limit)', color='C2' )
#plt.errorbar( SRP32vdW[:,0]*96.48530749926, SRP32vdW[:,1], yerr=SRP32vdW[:,2], marker='d', label='SRP32-vdW-DF1 DS1', color='orange', capsize=4., fillstyle='none' )
#plt.errorbar( NN_DS1[:,0]*96.48530749926, NN_DS1[:,1], yerr=NN_DS1[:,2], marker='o', label='MS-RPBEl DS1', color='r', capsize=4., fillstyle='full' )
#plt.fill_between(Shirhatti_DS1[:,0]*96.48530749926, Shirhatti_DS1[:,3], Shirhatti_DS1[:,4], color='grey', alpha=0.5)

yaxis = (1.-fDS1_v1_pop( xaxis ))*S0_erf_plot( 385.9, 48.24, 1., xaxis ) + fDS1_v1_pop( xaxis )*S0_erf_plot( 164., 48.24, 1., xaxis )
plt.plot( xaxis, yaxis, color='C0', linestyle='-' )

yaxis = (1.-fDS1_v1_pop( xaxis ))*S0_erf_plot( 395.6, 106.1, 1., xaxis ) + fDS1_v1_pop( xaxis )*S0_erf_plot( 115.8, 9.649, 1., xaxis )
plt.plot( xaxis, yaxis, color='C1', linestyle='-' )

yaxis = (1.-fDS1_v1_pop( xaxis ))*S0_erf_plot( 299.1, 77.19, 1., xaxis ) + fDS1_v1_pop( xaxis )*S0_erf_plot( 96.49, 19.3, 1., xaxis )
plt.plot( xaxis, yaxis, color='C2', linestyle='-' )

plt.annotate('(c)', xy=(100, 0.1), size=12)

#plt.legend(loc='best', numpoints=1, frameon=True)
plt.xlim(80., 260.)
#plt.ylim(0.0, 0.375)
plt.ylim(1e-6, 0.5)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
plt.tick_params(which='minor', direction='in', axis='y', right=True)
plt.ylabel('Reaction probability')
#ax2.yaxis.set_label_position("right")
plt.xlabel('Normal incidence energy (kJ/mol)')

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

plt.errorbar( Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,1], yerr=Shirhatti_DS2[:,2], marker='s', label='DS2 old', color='C0', capsize=4., fillstyle='none', linestyle='none' )
plt.scatter( Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,3], marker='^', label='DS2 new (lower limit)', color='C1' )
plt.scatter( Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,4], marker='v', label='DS2 new (upper limit)', color='C2' )
#plt.errorbar( NN_DS2[:,0]*96.48530749926, NN_DS2[:,1], yerr=NN_DS2[:,2], marker='o', label='MS-RPBEl DS2', color='r', capsize=4., fillstyle='full' )
#plt.fill_between(Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,3], Shirhatti_DS2[:,4], color='grey', alpha=0.5)

xaxis = np.linspace( 40., 250., 200 )
yaxis = 0.98*S0_erf_plot( 385.9, 48.24, 1., xaxis ) + 0.02*S0_erf_plot( 164., 48.24, 1., xaxis )
plt.plot( xaxis, yaxis, color='C0', linestyle='-' )

yaxis = 0.98*S0_erf_plot( 395.6, 106.1, 1., xaxis ) + 0.02*S0_erf_plot( 115.8, 9.649, 1., xaxis )
plt.plot( xaxis, yaxis, color='C1', linestyle='-' )

yaxis = 0.98*S0_erf_plot( 299.1, 77.19, 1., xaxis ) + 0.02*S0_erf_plot( 96.49, 19.3, 1., xaxis )
plt.plot( xaxis, yaxis, color='C2', linestyle='-' )

plt.annotate('(b)', xy=(60, 0.05), size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(30., 260.)
plt.xticks([50,100,150,200,250])
plt.ylim(0.0, 0.2)
#plt.ylim(1e-5, 0.5)
#plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False, labelleft=False)
#plt.ylabel('Reaction probability')

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

plt.errorbar( Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,1], yerr=Shirhatti_DS2[:,2], marker='s', label='DS2 old', color='C0', capsize=4., fillstyle='none', linestyle='none' )
plt.scatter( Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,3], marker='^', label='DS2 new (lower limit)', color='C1' )
plt.scatter( Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,4], marker='v', label='DS2 new (upper limit)', color='C2' )
#plt.errorbar( NN_DS2[:,0]*96.48530749926, NN_DS2[:,1], yerr=NN_DS2[:,2], marker='o', label='MS-RPBEl DS2', color='r', capsize=4., fillstyle='full' )
#plt.fill_between(Shirhatti_DS2[:,0]*96.48530749926, Shirhatti_DS2[:,3], Shirhatti_DS2[:,4], color='grey', alpha=0.5)

yaxis = 0.98*S0_erf_plot( 385.9, 48.24, 1., xaxis ) + 0.02*S0_erf_plot( 164., 48.24, 1., xaxis )
plt.plot( xaxis, yaxis, color='C0', linestyle='-' )

yaxis = 0.98*S0_erf_plot( 395.6, 106.1, 1., xaxis ) + 0.02*S0_erf_plot( 115.8, 9.649, 1., xaxis )
plt.plot( xaxis, yaxis, color='C1', linestyle='-' )

yaxis = 0.98*S0_erf_plot( 299.1, 77.19, 1., xaxis ) + 0.02*S0_erf_plot( 96.49, 19.3, 1., xaxis )
plt.plot( xaxis, yaxis, color='C2', linestyle='-' )

plt.annotate('(d)', xy=(60, 0.1), size=12)

#plt.legend(loc='best', numpoints=1, frameon=True)
plt.xlim(30., 260.)
plt.xticks([50,100,150,200,250])
#plt.ylim(0.0, 0.375)
plt.ylim(1e-6, 0.5)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelleft=False)
plt.tick_params(which='minor', direction='in', axis='y', right=True)
#plt.ylabel('Reaction probability')
#ax2.yaxis.set_label_position("right")
plt.xlabel('Normal incidence energy (kJ/mol)')

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