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

def EnergySticking( E, S, minS=None, maxS=None, Nsteps=None, targetS=None ):
	if targetS == None:
		yaxis = np.linspace(minS, maxS, Nsteps)
	else:
		yaxis = [ targetS ]

	E_S = []
	for i in yaxis:	# Here we determine the energy that S0=i
		yreduced = S - i
		freduced = interpolate.UnivariateSpline(E, yreduced, s=0)
		E_S.append( freduced.roots()[0] )
	E_S = np.array( E_S )
	
	return E_S

prop_cycle = plt.rcParams['axes.prop_cycle']
color_cycle = prop_cycle.by_key()['color']

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

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

minS = max( min(NN_DS2_v1_j2[:,1]), min(NN_DS2_v1_j8[:,1]) ) + 0.00001
maxS = min( max(NN_DS2_v1_j2[:,1]), max(NN_DS2_v1_j8[:,1]) ) - 0.00001
Nsteps = 30
E0 = EnergySticking( NN_DS2_v1_j2[:,0]*96.48530749926, NN_DS2_v1_j2[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
E1 = EnergySticking( NN_DS2_v1_j8[:,0]*96.48530749926, NN_DS2_v1_j8[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
Efficacy = (E0-E1) / (E_J[7][1] - E_J[1][1])
xaxis = np.linspace( min(E1), max(E1), len(E1) )
plt.plot( xaxis, Efficacy, label=r'$\chi_{J=2\rightarrow8}(E_\textrm{i},\nu=1)$', color=color_cycle[1] )

minS = max( min(NN_DS2_v2_j2[:,1]), min(NN_DS2_v2_j8[:,1]) ) + 0.00001
maxS = min( max(NN_DS2_v2_j2[:,1]), max(NN_DS2_v2_j8[:,1]) ) - 0.00001
Nsteps = 30
E0 = EnergySticking( NN_DS2_v2_j2[:,0]*96.48530749926, NN_DS2_v2_j2[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
E1 = EnergySticking( NN_DS2_v2_j8[:,0]*96.48530749926, NN_DS2_v2_j8[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
Efficacy = (E0-E1) / (E_J[7][2] - E_J[1][2])
xaxis = np.linspace( min(E1), max(E1), len(E1) )
plt.plot( xaxis, Efficacy, label=r'$\chi_{J=2\rightarrow8}(E_\textrm{i},\nu=2)$', color=color_cycle[2] )

plt.plot( [0.,400.], [1., 1.], ls='--', c='k' )

plt.annotate('(a)', xy=(0.05, 0.85), xycoords='axes fraction', size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(40., 270.)
#plt.ylim(0.0, 0.39)
#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('Rotational efficacy')

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

plt.errorbar( NN_DS2[:,0]*96.48530749926, NN_DS2[:,1], yerr=NN_DS2[:,2], marker='None', label=r'Laser-off', color='k', capsize=4., fillstyle='none' )
plt.errorbar( NN_DS2_v2_j2_laser[:,0]*96.48530749926, NN_DS2_v2_j2_laser[:,1], yerr=NN_DS2_v2_j2_laser[:,2], marker='None', label=r'Laser-on ($\nu=0, J=3 \rightarrow \nu=2, J=2$)', color='b', capsize=4., fillstyle='none' )
plt.errorbar( NN_DS2_v2_j8_laser[:,0]*96.48530749926, NN_DS2_v2_j8_laser[:,1], yerr=NN_DS2_v2_j8_laser[:,2], marker='None', label=r'Laser-on ($\nu=0, J=7 \rightarrow \nu=2, J=8$)', color='r', capsize=4., fillstyle='none' )

plt.annotate('(b)', xy=(0.05, 0.85), xycoords='axes fraction', size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(40., 270.)
#plt.ylim(0.0, 0.39)
plt.ylim(2e-5, 0.5)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
plt.tick_params(axis='y', which='minor', direction='in', top=True, right=True)
plt.ylabel('Sticking probability')
ax2.yaxis.set_label_coords(-0.15,0.)
#plt.xlabel('Normal incidence energy (kJ/mol)')

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

plt.errorbar( NN_DS2_v2_j2[:,0]*96.48530749926, NN_DS2_v2_j2[:,1], yerr=NN_DS2_v2_j2[:,2], marker='None', label=r'$\nu=2, J=2$', color='b', capsize=4., fillstyle='none' )
plt.errorbar( NN_DS2_v2_j8[:,0]*96.48530749926, NN_DS2_v2_j8[:,1], yerr=NN_DS2_v2_j8[:,2], marker='None', label=r'$\nu=2, J=8$', color='r', capsize=4., fillstyle='none' )

plt.annotate('(c)', xy=(0.05, 0.85), xycoords='axes fraction', size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(40., 270.)
#plt.ylim(0.0, 0.39)
plt.ylim(6e-2, 1.)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
plt.tick_params(axis='y', which='minor', direction='in', top=True, right=True)
#plt.ylabel('Reaction probability')
plt.xlabel('Normal incidence energy (kJ/mol)')

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

################ Generation of SI figure
plt.clf()

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

minS = max( min(NN_DS2_v1_j2[:,1]), min(NN_DS2_v1_j8[:,1]) ) + 0.00001
maxS = min( max(NN_DS2_v1_j2[:,1]), max(NN_DS2_v1_j8[:,1]) ) - 0.00001
Nsteps = 30
E0 = EnergySticking( NN_DS2_v1_j2[:,0]*96.48530749926, NN_DS2_v1_j2[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
E1 = EnergySticking( NN_DS2_v1_j8[:,0]*96.48530749926, NN_DS2_v1_j8[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
Efficacy = (E0-E1) / (E_J[7][1] - E_J[1][1])
xaxis = np.linspace( min(E1), max(E1), len(E1) )
plt.plot( xaxis, Efficacy, label=r'$\chi_{J=2\rightarrow8}(E_\textrm{i},\nu=1)$', color=color_cycle[1] )

minS = max( min(NN_DS2_v2_j2[:,1]), min(NN_DS2_v2_j8[:,1]) ) + 0.00001
maxS = min( max(NN_DS2_v2_j2[:,1]), max(NN_DS2_v2_j8[:,1]) ) - 0.00001
Nsteps = 30
E0 = EnergySticking( NN_DS2_v2_j2[:,0]*96.48530749926, NN_DS2_v2_j2[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
E1 = EnergySticking( NN_DS2_v2_j8[:,0]*96.48530749926, NN_DS2_v2_j8[:,1], minS=minS, maxS=maxS, Nsteps=Nsteps )
Efficacy = (E0-E1) / (E_J[7][2] - E_J[1][2])
xaxis = np.linspace( min(E1), max(E1), len(E1) )
plt.plot( xaxis, Efficacy, label=r'$\chi_{J=2\rightarrow8}(E_\textrm{i},\nu=2)$', color=color_cycle[2] )

plt.plot( [0.,400.], [1., 1.], ls='--', c='k' )

plt.annotate('(a)', xy=(0.05, 0.85), xycoords='axes fraction', size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(40., 270.)
#plt.ylim(0.0, 0.39)
#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('Rotational efficacy')

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

plt.errorbar( NN_DS2[:,0]*96.48530749926, NN_DS2[:,1], yerr=NN_DS2[:,2], marker='None', label=r'Laser-off', color='k', capsize=4., fillstyle='none' )
plt.errorbar( NN_DS2_v1_j2_laser[:,0]*96.48530749926, NN_DS2_v1_j2_laser[:,1], yerr=NN_DS2_v1_j2_laser[:,2], marker='None', label=r'Laser-on ($\nu=0, J=3 \rightarrow \nu=1, J=2$)', color='b', capsize=4., fillstyle='none' )
plt.errorbar( NN_DS2_v1_j8_laser[:,0]*96.48530749926, NN_DS2_v1_j8_laser[:,1], yerr=NN_DS2_v1_j8_laser[:,2], marker='None', label=r'Laser-on ($\nu=0, J=7 \rightarrow \nu=1, J=8$)', color='r', capsize=4., fillstyle='none' )

plt.annotate('(b)', xy=(0.05, 0.85), xycoords='axes fraction', size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(40., 270.)
#plt.ylim(0.0, 0.39)
plt.ylim(1e-4, 0.5)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
plt.tick_params(axis='y', which='minor', direction='in', top=True, right=True)
plt.ylabel('Sticking probability')
ax2.yaxis.set_label_coords(-0.15,0.)
#plt.xlabel('Normal incidence energy (kJ/mol)')

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

plt.errorbar( NN_DS2_v1_j2[:,0]*96.48530749926, NN_DS2_v1_j2[:,1], yerr=NN_DS2_v1_j2[:,2], marker='None', label=r'$\nu=1, J=2$', color='b', capsize=4., fillstyle='none' )
plt.errorbar( NN_DS2_v1_j8[:,0]*96.48530749926, NN_DS2_v1_j8[:,1], yerr=NN_DS2_v1_j8[:,2], marker='None', label=r'$\nu=1, J=8$', color='r', capsize=4., fillstyle='none' )

plt.annotate('(c)', xy=(0.05, 0.85), xycoords='axes fraction', size=12)

plt.legend(loc='best', numpoints=1, frameon=True, fontsize=8)
plt.xlim(40., 270.)
#plt.ylim(0.0, 0.39)
plt.ylim(1e-4, 0.8)
plt.yscale('log')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
plt.tick_params(axis='y', which='minor', direction='in', top=True, right=True)
#plt.ylabel('Reaction probability')
plt.xlabel('Normal incidence energy (kJ/mol)')

plt.tight_layout()
plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig('reactionprobability_laser_SI.pdf')