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

def baule(E, mass):
	mu = (14.001 + 1.008 * 3) / mass

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

	return ET

ind_1100 = np.arange(4)

ET_1100_phon = np.array(
[[74.0, 6.1400002847864155, 3.3196781753624123],
[87.8, 16.48687275900334, 4.073037703822967],
[102.9, 12.785133589851199, 4.065199290597144],
[119.5, 16.63999595604884, 6.504747097188525]]
)

ET_1100_para = np.array(
[[74.0, 16.547062508638746, 1.5694088380754976],
[87.8, 13.881005880664839, 1.6722765307877596],
[102.9, 13.368512635235232, 1.7923027215767877],
[119.5, 16.801377966630877, 2.729706881582817]]
)

ET_1100_rovib = np.array(
[[74.0, 35.71412019172484, 3.0367353709534703],
[87.8, 41.504714223694194, 4.2825182221367655],
[102.9, 57.42371057745147, 5.1455255081886815],
[119.5, 69.57536941469485, 6.593064427612765]]
)

ET_1100_Z = np.array(
[[74.0, 6.00112410193096, 0.992729252528557],
[87.8, 4.82928082780121, 0.7722055668060146],
[102.9, 6.047941502164149, 1.3906309086418152],
[119.5, 7.7160027161440805, 1.9297698611308314]]
)

#print ET_1100_rovib[:,1] / (ET_1100_rovib[:,1] + ET_1100_phon[:,1] + ET_1100_para[:,1] + ET_1100_Z[:,1])

ET_475_phon = np.array(
[[102.9, 18.95875304467068, 2.7564673643243722],
[119.5, 19.46802603047079, 3.7067534181590425]]
)

ET_475_para = np.array(
[[102.9, 12.993486475881705, 1.1069743685402131],
[119.5, 13.574742765833907, 1.3729931726382976]]
)

mpl.rc("text", usetex=True)
fig = plt.figure(figsize=(3.69,3.5))

width = 0.35

#plt.plot( [0., 130.], [baule(0., 101.07), baule(130., 101.07)], color='black', linestyle='-' )
#plt.errorbar( ET_475K[:,0], ET_475K[:,1], yerr=ET_475K[:,2], marker='o', color='b', label='475 K' )
#plt.errorbar( ET_1100K[:,0], ET_1100K[:,1], yerr=ET_1100K[:,2], marker='o', color='r', label='1100 K' )
p1 = plt.bar(ind_1100, ET_1100_phon[:,1], width, yerr=ET_1100_phon[:,2], color='r', align='center', label='Phonons', ecolor='k')
p2 = plt.bar(ind_1100, ET_1100_para[:,1], width, bottom=ET_1100_phon[:,1], yerr=ET_1100_para[:,2], color='b', align='center', label='Parallel', ecolor='k')
p3 = plt.bar(ind_1100, ET_1100_rovib[:,1], width, bottom=ET_1100_phon[:,1]+ET_1100_para[:,1], yerr=ET_1100_rovib[:,2], color='g', align='center', label='Rovibrational', ecolor='k')
p3 = plt.bar(ind_1100, ET_1100_Z[:,1], width, bottom=ET_1100_phon[:,1]+ET_1100_para[:,1]+ET_1100_rovib[:,1], yerr=ET_1100_Z[:,2], color='purple', align='center', label='Perpendicular', ecolor='k')

plt.xticks(ind_1100, ('74', '88', '103', '120'))
plt.legend(loc='upper left', numpoints=1, frameon=False)
plt.xlim(-0.5,3.5)
plt.ylim(0.,180.)
plt.tick_params(length=6, width=1)
plt.ylabel('Energy (kJ/mol)')
plt.xlabel('Incidence energy (kJ/mol)')

plt.tight_layout()
plt.savefig('energytrapping.pdf')
#plt.show()