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

# J, nu=0, 1, 2, 3
Erovib = np.array(
[[0,	0.1807216891,	0.5326696681,	0.8725158178,	1.2005140277],
[1,	0.1832914903,	0.5351683215,	0.8749440669,	1.202872624],
[2,	0.1884295011,	0.5401640537,	0.8797990086,	1.2075882788],
[3,	0.1961325398,	0.5476537166,	0.8870775318,	1.2146579181],
[4,	0.2063958383,	0.557632593,	0.8967749738,	1.2240769354],
[5,	0.2192130471,	0.5700944009,	0.908885126,	1.2358391961],
[6,	0.2345762428,	0.5850313011,	0.9234002402,	1.2499370449],
[7,	0.252475936,	0.6024339051,	0.9403110371,	1.2663613136],
[8,	0.2729010819,	0.6222912858,	0.9596067165,	1.285101331],
[9,	0.2958390924,	0.6445909892,	0.9812749689,	1.3061449348],
[10,	0.3212758494,	0.6693190474,	1.0053019886,	1.3294784842],
[11,	0.3491957198,	0.6964599943,	1.0316724887,	1.3550868744],
[12,	0.3795815724,	0.7259968814,	1.0603697168,	1.3829535529],
[13,	0.4124147952,	0.7579112957,	1.0913754729,	1.4130605364],
[14,	0.4476753156,	0.7921833794,	1.1246701285,	1.44538843],
[15,	0.4853416204,	0.82879185,	1.1602326465,	1.4799164465],
[16,	0.5253907783,	0.8677140221,	1.1980406029,	1.5166224283],
[17,	0.5677984625,	0.9089258311,	1.2380702093,	1.555482869],
[18,	0.6125389754,	0.9524018562,	1.2802963367,	1.5964729375],
[19,	0.6595852736,	0.9981153467,	1.3246925405,	1.639566502],
[20,	0.7089089943,	1.0460382473,	1.3712310857,	1.6847361553],
[21,	0.7604804825,	1.0961412256,	1.4198829742,	1.7319532416],
[22,	0.8142688186,	1.1483936997,	1.4706179714,	1.7811878828],
[23,	0.8702418476,	1.2027638669,	1.523404635,	1.8324090065],
[24,	0.9283662077,	1.2592187326,	1.5782103432,	1.8855843745],
[25,	0.9886073609,	1.3177241405,	1.6350013243,	1.9406806116],
[26,	1.050929623,	1.3782448023,	1.6937426864,	1.997663235],
[27,	1.115296194,	1.4407443285,	1.7543984476,	2.0564966841],
[28,	1.1816691898,	1.5051852594,	1.816931566,	2.1171443512],
[29,	1.2500096728,	1.5715290957,	1.8813039711,	2.1795686118]]
)

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

plt.subplot(nrows,ncols,1)

T = np.linspace(0.1, 600., 100)
fj = []
for i in range( len(T) ):
	tmp = []
	for j in range( 30 ):
		tmp.append( (2*j + 1) * np.exp( -(Erovib[j][1] - Erovib[0][1]) / (8.617333262145*10**(-5)*T[i]) ) )
	tmp = np.array( tmp ) / sum( tmp )
	fj.append( tmp )
fj = np.array( fj )

for j in range( 9 ):
	plt.plot( T, fj[:,j], label='J={:d}'.format(j) )

plt.legend(fontsize=8)
plt.xlim(0.,600.)
plt.ylim(0.,1.)
plt.tick_params(labelbottom=True, length=6, width=1, direction='in', top=True, right=True)
plt.ylabel('Rotational population')
plt.xlabel(r'$T_\textrm{rot}$ (K)')
plt.annotate('(a)', xy=(0.1, 0.8), xycoords='axes fraction')

ax2 = plt.subplot(nrows,ncols,2)

T = np.linspace(300., 1200., 100)
fv = []
for i in range( len(T) ):
	tmp = []
	for j in range( 30 ):
		tmp.append( (2*j + 1) * np.exp( -(Erovib[j][1] - Erovib[0][1]) / (8.617333262145*10**(-5)*T[i]) ) )
	Z0 = sum( tmp )
	tmp = []
	for j in range( 30 ):
		tmp.append( (2*j + 1) * np.exp( -(Erovib[j][2] - Erovib[0][1]) / (8.617333262145*10**(-5)*T[i]) ) )
	Z1 = sum( tmp )
	tmp = []
	for j in range( 30 ):
		tmp.append( (2*j + 1) * np.exp( -(Erovib[j][3] - Erovib[0][1]) / (8.617333262145*10**(-5)*T[i]) ) )
	Z2 = sum( tmp )
	fv.append( [Z0/(Z0+Z1+Z2), Z1/(Z0+Z1+Z2), Z2/(Z0+Z1+Z2)] )
fv = np.array( fv )

plt.plot( T, fv[:,0], label=r'$\nu=0$' )
plt.plot( T, fv[:,1], label=r'$\nu=1$' )
plt.plot( T, fv[:,2], label=r'$\nu=2$' )

plt.annotate('(b)', xy=(0.1, 0.8), xycoords='axes fraction')

plt.legend(fontsize=8)
plt.xlim(300.,1200.)
plt.ylim(10**-7,2.)
plt.yscale('log')
ax2.yaxis.set_label_position('right')
plt.xlabel(r'$T_\textrm{vib}$ (K)')
plt.ylabel('Vibrational population')
plt.tick_params(length=6, width=1, direction='in', top=True, right=True, labelleft=False, labelright=True)
plt.tick_params(axis='y', which='minor', direction='in', right=True)

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