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

# Energy, sticking, error
# For LO the 4th column is average total energy due to vibrationally excited states
LO = np.array(
[[160.3, 0.00026, 0.00007, 179.0],
[172.3, 0.00105, 0.00023, 193.5],
[181.3, 0.00216, 0.00023, 205.3]]
)

v0 = np.array(
[[172.3, 0.000106, 0.000021],
[181.3, 0.000301, 0.000045]]
)

v1 = np.array(
[[143.9, 0.000140, 0.000022],
[160.3, 0.00070, 0.00019],
[172.3, 0.00171, 0.00029],
[181.3, 0.00251, 0.00050]]
)

v2 = np.array(
[[83.4, 0.000064, 0.000024],
[97.0, 0.000525, 0.000046],
[111.6, 0.00165, 0.00029],
[127.3, 0.00441, 0.00066],
[143.9, 0.01374, 0.00117],
[160.4, 0.02460, 0.00155],
[172.3, 0.03567, 0.00186],
[181.3, 0.04858, 0.00216]]
)

v2fs = np.array(
[[83.4, 0.00015, 0.00005],
[97.0, 0.00115, 0.00024],
[111.6, 0.00380, 0.00044],
[127.3, 0.01291, 0.00113],
[143.9, 0.03483, 0.00183],
[160.4, 0.05393, 0.00226],
[172.3, 0.06906, 0.00254],
[181.3, 0.08285, 0.00277]]
)

v2cs = np.array(
[[83.4, 0.00025, 0.00006],
[97.0, 0.00180, 0.00030],
[111.6, 0.00600, 0.00055],
[127.3, 0.01431, 0.00119],
[143.9, 0.03571, 0.00186],
[160.4, 0.05168, 0.00222],
[172.3, 0.07635, 0.00267],
[181.3, 0.08415, 0.00279]]
)

# AIMD results
v1_AIMD = np.array(
[[160.4, 0.00, 0.001],
[181.3, 0.005, 0.002]]
)

v2_AIMD = np.array(
[[160.4, 0.024, 0.005],
[181.3, 0.048, 0.007]]
)

mpl.rc("font", size=10)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(3.69,9))

ax = plt.subplot(3,1,1)

#x_axis = np.linspace(120.,200.,160)
#fit_1100_S = []
#A = 0.4
#E0 = 150.
#W = 80.
#for i in range(len(x_axis)):
#	fit_1100_S.append( A/2. * ( 1 + erf((x_axis[i] - E0)/W) ) )
#plt.plot(x_axis, fit_1100_S, color='r', label='Exp. 1100 K')

AIMD_key = plt.errorbar( v1_AIMD[:,0], v1_AIMD[:,1], yerr=v1_AIMD[:,2], marker='D', color='blue', linestyle='None', markersize=8, fillstyle='left', markerfacecoloralt='red' )
plt.errorbar( v1_AIMD[:,0], v1_AIMD[:,1], yerr=v1_AIMD[:,2], marker='D', color='blue', linestyle='None', markersize=8 )
plt.errorbar( v2_AIMD[:,0], v2_AIMD[:,1], yerr=v2_AIMD[:,2], marker='D', color='red', linestyle='None', markersize=8 )
plt.errorbar( LO[:,0], LO[:,1], yerr=LO[:,2], marker='o', label='Laser-off', color='black' )
plt.errorbar( v0[:,0], v0[:,1], yerr=v0[:,2], marker='o', label='Ground state', color='green' )
plt.errorbar( v1[:,0], v1[:,1], yerr=v1[:,2], marker='o', label=r'$v_1=1$', color='blue' )
plt.errorbar( v2[:,0], v2[:,1], yerr=v2[:,2], marker='o', label=r'$v_1=2$', color='red' )
#plt.errorbar( v2fs[:,0], v2fs[:,1], yerr=v2fs[:,2], marker='o', color='red', linestyle='--', fillstyle='none' )
#plt.errorbar( v2cs[:,0], v2cs[:,1], yerr=v2cs[:,2], marker='o', color='red', linestyle=':', fillstyle='top' )

handles, labels = ax.get_legend_handles_labels()
handles.append( AIMD_key )
labels.append( 'AIMD' )
plt.legend(handles, labels, loc='upper left', numpoints=1, frameon=False)

#plt.legend(loc='best', numpoints=1, frameon=False)
#plt.legend(loc='upper left', numpoints=1, frameon=False)
plt.xlim(60,190.)
plt.ylim(5*10**-5,1.0)
plt.yscale("log")
plt.tick_params(length=6, width=1)
plt.xlabel('Incidence energy (kJ/mol)')
plt.ylabel('Reaction probability')
plt.annotate('(a)', xy=(175, 0.25), size=12)

plt.subplot(3,1,2)

plt.errorbar( LO[:,3], LO[:,1], yerr=LO[:,2], marker='o', label='Laser-off', color='black' )
plt.errorbar( v0[:,0], v0[:,1], yerr=v0[:,2], marker='o', label='Ground state', color='green' )
plt.errorbar( v1[:,0]+36., v1[:,1], yerr=v1[:,2], marker='o', label=r'$v_1=1$', color='blue' )
plt.errorbar( v2[:,0]+72., v2[:,1], yerr=v2[:,2], marker='o', label=r'$v_1=2$', color='red' )
#plt.errorbar( v2fs[:,0], v2fs[:,1], yerr=v2fs[:,2], marker='s', color='red', label='Frozen surface', linestyle='--' )
#plt.errorbar( v2cs[:,0], v2cs[:,1], yerr=v2cs[:,2], marker='d', color='red', label='Distorted surface', linestyle=':' )

plt.legend(loc='upper left', numpoints=1, frameon=False)
plt.xlim(140,260.)
plt.ylim(5*10**-5,1.0)
plt.yscale("log")
plt.tick_params(length=6, width=1)
plt.ylabel('Reaction probability')
plt.xlabel('Total energy (kJ/mol)')
plt.annotate('(b)', xy=(245, 0.25), size=12)

plt.subplot(3,1,3)

#plt.errorbar( LO[:,3], LO[:,1], yerr=LO[:,2], marker='o', label='Laser-off', color='black' )
#plt.errorbar( v0[:,0], v0[:,1], yerr=v0[:,2], marker='o', label='Ground state', color='green' )
#plt.errorbar( v1[:,0]+36., v1[:,1], yerr=v1[:,2], marker='o', label=r'$v_1=1$', color='blue' )
plt.errorbar( v2[:,0], v2[:,1], yerr=v2[:,2], marker='o', label='Surface motion', color='red' )
plt.errorbar( v2fs[:,0], v2fs[:,1], yerr=v2fs[:,2], marker='s', color='red', label='Frozen surface', linestyle='--' )
plt.errorbar( v2cs[:,0], v2cs[:,1], yerr=v2cs[:,2], marker='d', color='red', label='Distorted surface', linestyle=':' )
#plt.errorbar( v1_AIMD[:,0]+36., v1_AIMD[:,1], yerr=v1_AIMD[:,2], marker='d', color='blue', linestyle='None' )
#plt.errorbar( v2_AIMD[:,0], v2_AIMD[:,1], yerr=v2_AIMD[:,2], marker='v', color='red', linestyle='None' )

plt.legend(loc='upper left', numpoints=1, frameon=False, borderaxespad=0.2)
plt.xlim(50,190.)
plt.ylim(5*10**-5,1.0)
plt.yscale("log")
plt.tick_params(length=6, width=1)
plt.ylabel('Reaction probability')
plt.xlabel('Incidence energy (kJ/mol)')
plt.annotate('(c)', xy=(175, 0.25), size=12)

plt.tight_layout()
#plt.subplots_adjust(hspace = 0.0, wspace = 0.0)
plt.savefig('reactionprobabilitymulti.pdf')
#plt.show()