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

def baule(E, mass):
	mu = (35.453 + 1.008) / mass

	ET = E * 2.4 * mu / (1 + mu)**2
#	stderr = ( ET / 2.4 * 4. - ET )

	return ET

def baule_limit(E, mass):
	mu = (35.453 + 1.008) / mass

	ET = E * 4. * mu / (1 + mu)**2

	return ET

# Ei, mean, standard deviation and standard error
ET_NN = np.array(
[[91, 25.295827196871933, 11.296380037071488, 0.07994144803090951],
[114, 34.96397668577859, 14.703609490892957, 0.10510625133906507],
[124, 38.50627417110618, 15.044178775725934, 0.10835776205817946],
[150, 49.04502136931593, 17.418412724928636, 0.13085090119680326],
[174, 58.749627430131696, 19.101037441048504, 0.14942763736370507],
[205, 71.52297442491462, 21.693732586810068, 0.17850843280597178],
[247, 89.98461806034919, 23.889612329153486, 0.21025482993689487]]
)

#Ei, mean, standard deviation and standard error for v=2,j=1 -> v=1,j=5
ET_NN_v1 = np.array(
[[27.016, 12.660240797845475, 8.085872877264114, 0.3301043779051192],
[30.875, 15.222274607122172, 9.185295953211277, 0.36139063568456575],
[43.418, 21.11005385319224, 11.147964188558971, 0.4715086462487053],
[50.172, 24.51243303930573, 11.518152357655874, 0.5312927504384511],
[75.259, 33.45933513687, 14.763088009885825, 0.9086052572979781],
[90.696, 37.381221780895736, 16.418657957865864, 1.2980088817526252],
[113.85, 42.91385765658492, 18.079325381710486, 1.6042807063312763],
[124.47, 52.10184200784529, 19.488359136911093, 1.8016994395587198],
[149.55, 58.46898188549315, 20.0881033750611, 2.141399014894291],
[173.67, 73.76299044725727, 22.61746672608384, 2.8495329639760034],
[204.55, 77.61513424032324, 23.54990210663145, 3.1469880419671705],
[247., 101.1689015849009, 24.917281418680712, 3.756421511884474]]
)

#Ei, mean, standard deviation and standard error for v=2,j=1 -> v=2,j=5
ET_NN_v2 = np.array(
[[27.016, 5.536219356623555, 5.55974554775645, 0.06676717438097401],
[30.875, 6.654528378333897, 6.111670711741596, 0.07824551160716992],
[43.418, 9.394600440095749, 7.628031091696467, 0.11179077545861812],
[50.172, 10.62021145882067, 8.215167370067045, 0.12623788465670985],
[75.259, 16.86181975433683, 8.507432817325917, 0.1312881688034286],
[90.696, 22.258510867678797, 9.229748623639418, 0.13820458259128665],
[113.85, 30.738607153758217, 11.446611316962393, 0.17411434556416713],
[124.47, 33.837416627914344, 11.180100701519983, 0.17447621963117793],
[149.55, 44.40422093259904, 14.248465259683105, 0.21938897552819003],
[173.67, 54.55976617488515, 16.43621619251838, 0.2918795876751579],
[204.55, 68.00129851174219, 19.069801340071002, 0.34189768741000265],
[247., 87.21473277959602, 22.11274578447967, 0.41946757927689726]]
)

#Ei, mean translational energy of scattered molecule, standard deviation and standard error for v=2,j=1 -> v=1,j=5
Escat_NN_v1 = np.array(
[[27.016, -9.759368447667221, 9.512627705114026, 0.38835139984319816],
[30.875, -8.745132999916152, 10.420165453688119, 0.40997592635326624],
[43.418, -4.107310622198214, 12.052800196689127, 0.5097791316982717],
[50.172, -1.3750078152898093, 12.574124779900322, 0.580001126155369],
[75.259, 7.657729375100001, 15.459511998832209, 0.9514671908745704],
[90.696, 11.60993921532143, 16.843096901280262, 1.3315637264742455],
[113.85, 19.964468462972658, 17.44955817796358, 1.5483979035651163],
[124.47, 24.924362641913135, 19.149804042556767, 1.7704000100134876],
[149.55, 34.287625971764044, 21.740428618500147, 2.317537477671123],
[173.67, 49.26720710901563, 22.361325228888123, 2.8172621686415455],
[204.55, 56.0681200056579, 25.91276567019054, 3.4627389672003797],
[247., 75.66120519996667, 24.395574941981444, 3.677771301252212]]
)

#Ei, mean translational energy of scattered molecule, standard deviation and standard error for v=2,j=1 -> v=2,j=5
Escat_NN_v2 = np.array(
[[27.016, 7.581090606173253, 6.83035325135507, 0.0820259457379135],
[30.875, 8.698980605284907, 7.266519420086344, 0.09303062230361392],
[43.418, 11.680658011794343, 8.416411514575774, 0.12334469517533363],
[50.172, 13.05785822357631, 8.57166189021683, 0.13171593666565515],
[75.259, 19.892228822044522, 8.602967837761241, 0.13276248169650345],
[90.696, 25.40129175049967, 9.362540898371625, 0.14019299003866906],
[113.85, 33.889769020575876, 11.522219915963008, 0.17526442757267455],
[124.47, 37.002409087400174, 11.265465823826776, 0.1758084244320182],
[149.55, 47.6064659029536, 14.234745544116894, 0.21917772791044646],
[173.67, 57.79965121269577, 16.469656068344484, 0.292473423692749],
[204.55, 71.24023993581451, 19.08361146335012, 0.34214528564806374],
[247., 90.46858784513545, 22.124409712862576, 0.4196888380862429]]
)

#Ei, Escattered for v=2,j=1 -> v=1,j=5
ET_exp_v1 = np.array(
[[26.5335, 18.508777],
[30.5858, 21.093618],
[43.1289, 23.70451],
[50.1724, 26.89142],
[75.1621, 45.348],
[93.9767, 50.671189],
[122.343, 59.312413]]
)

#Ei, Escattered for v=2,j=1 -> v=2,j=5
ET_exp_v2 = np.array(
[[26.5335, 10.86425],
[30.5858, 13.011044],
[43.1289, 16.528898],
[50.1724, 19.548888],
[75.1621, 35.764209],
[93.9767, 42.874211],
[122.343, 51.653409]]
)

#Gernot 2016 paper
ET_gernot_PBE_298_total = np.array(
[[135.1, 41.1027],
[154.4, 50.0759],
[193.0, 65.5135]]
)

#Gernot 2019 paper, note that this also includes LDFA
ET_gernot_SRP32_total = np.array(
[[90.7, 45.3],
[102.0, 51.6]]
)

#print(ET_NN[:,1]/ET_NN[:,0])

mpl.rc("text", usetex=True)
fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(3.69,6.))

plt.subplot(2,1,1)
plt.plot( [0., 270.], [baule(0., 196.966), baule(270., 196.966)], color='black', linestyle='-', label='refined Baule model' )
plt.scatter( ET_NN[:,0], ET_NN[:,1], marker='o', color='r', label='HD-NNP' )
plt.scatter( ET_NN_v1[:,0], ET_NN_v1[:,1], marker='o', color='orange', label=r'HD-NNP $\nu=2,j=1 \rightarrow \nu=1,j=5$' )
plt.scatter( ET_NN_v2[:,0], ET_NN_v2[:,1], marker='o', color='blue', label=r'HD-NNP $\nu=2,j=1 \rightarrow \nu=2,j=5$' )

plt.scatter( ET_gernot_PBE_298_total[:,0], ET_gernot_PBE_298_total[:,1], marker='d', color='green', label='PBE' )

plt.annotate('(a)', xy=(220., 25), size=12)

plt.legend(loc='upper left', numpoints=1)
plt.xlim(0.,270.)
plt.ylim(0.,200.)
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelbottom=False)
plt.ylabel('Energy transfer (kJ/mol)')
plt.xlabel('Incidence energy (kJ/mol)')

plt.subplot(2,1,2)

plt.plot( [0., 270.], [baule(0., 196.966), baule(270., 196.966)], color='black', linestyle='-', label='' )
plt.plot( [0., 270.], [baule_limit(0., 196.966), baule_limit(270., 196.966)], color='black', linestyle='--', label='Baule limit' )
#plt.plot( ET_NN_v1[:,0], ET_NN_v1[:,1]-0.26*32.790160285, marker='o', color='orange', linestyle='-.', label='HD-NNP' )
#plt.plot( ET_NN_v2[:,0], ET_NN_v2[:,1], marker='o', color='blue', linestyle='-.', label='' )
plt.scatter( ET_gernot_SRP32_total[:,0], ET_gernot_SRP32_total[:,1], marker='^', color='grey', label=r'SRP32-vdW-DF1 ($\nu=1, J=1$)', zorder=10 )
plt.scatter( Escat_NN_v1[:,0], Escat_NN_v1[:,1], marker='o', color='orange', label=r'HD-NNP $\nu=2,j=1 \rightarrow \nu=1,j=5$' )
plt.scatter( Escat_NN_v2[:,0], Escat_NN_v2[:,1], marker='o', color='blue', label=r'HD-NNP $\nu=2,j=1 \rightarrow \nu=2,j=5$' )
#plt.plot( Escat_NN_v1[:,0], (Escat_NN_v1[:,0]-Escat_NN_v1[:,1])/0.3*0.5, marker='o', color='orange', linestyle='-.', label='HD-NNP' )
#plt.plot( Escat_NN_v2[:,0], (Escat_NN_v2[:,0]-Escat_NN_v2[:,1])/0.3*0.5, marker='o', color='blue', linestyle='-.', label='' )
plt.scatter( ET_exp_v1[:,0], ET_exp_v1[:,0]-ET_exp_v1[:,1], marker='s', color='orange', label=r'Exp. $\nu=2,j=1 \rightarrow \nu=1,j=5$' )
plt.scatter( ET_exp_v2[:,0], ET_exp_v2[:,0]-ET_exp_v2[:,1], marker='s', color='b', label=r'Exp. $\nu=2,j=1 \rightarrow \nu=2,j=5$' )

plt.annotate('(b)', xy=(220., 20), size=12)

plt.legend(loc='upper left', numpoints=1)
plt.xlim(0.,270.)
plt.ylim(-25.,199.)
plt.tick_params(length=6, width=1, direction='in', top=True, right=True)
plt.ylabel(r'Translational energy loss (kJ/mol)')
plt.xlabel('Incidence energy (kJ/mol)')

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