'''
Written by Li Feng, Nov. 25, 2016. @NJNU
'''
# ------------------------ python code begins here

# Import necessary modules
import pyfits
import numpy as np
import matplotlib.pyplot as plt
 
#Read data
hdulist = pyfits.open('spec-0768-52281-0572.fits')

data1 = hdulist[1].data
data2 = hdulist[2].data

flux     = data1.field(0) # flux
lamda    = data1.field(1) # loglam
ivar     = data1.field(2) # ivar
and_mask = data1.field(3) # and_mask
or_mask  = data1.field(4) # or_mask
wdisp    = data1.field(5) # wdisp
sky      = data1.field(6) # sky
model    = data1.field(7) # model

print(flux,lamda)

lamda1 = 10**lamda

#Read redshift (zz) & redshift error (zz_err)
#redshift = hdulist[2].data['z']
#redshift_err = hdulist[2].data['z_err']
#zz = redshift[0]
#zz_err = redshift_err[0]
#OR
zz = data2.field(63)
zz_err = data2.field(64)

#print(redshift,redshift_err,zz,zz_err)
print(zz,zz_err)

plt.plot(lamda1, flux,'r')
plt.plot(lamda1, ivar,'g')
plt.plot(lamda1, model,'k')
#plt.plot(lamda1, sky,'c')

# Axis Lables
plt.ylabel('Flux (10 $^{-13}$ ergs cm$^{-2}$ s$^{-1}$ $\AA ^{-1}$)')
plt.xlabel('Wavelength ($\AA$)')

# Define limits
plt.xlim(3700,9300)
plt.ylim(-5,30)

x = np.linspace(3900.,4400.)
plt.plot(x,28.+0.0*x,'r-',x,26.+0.0*x,'g-',x,24.+0.0*x,'k-',linewidth=2)
plt.text(4450,27.5,'flux', color='red', size=14)
plt.text(4450,25.5,'ivar', color='green', size=14)
plt.text(4450,23.5,'model', color='black', size=14)
#plt.text(4230,21.5,' sky', color='cyan', size=12)

#Plot a line to mark H Balmer Series
y = np.linspace(-2.,10.)
plt.plot(6562.85*(1+zz)+0.0*y,y,'k-')
plt.text(6562.85*(1+zz)-80,-4.,r'$\mathrm{H}_\alpha$', color='black', size=16)
plt.plot(4861.36*(1+zz)+0.0*y,y,'k-')
plt.text(4861.36*(1+zz)-80,-4.,r'$\mathrm{H}_\beta$', color='black', size=16)
plt.plot(4340.49*(1+zz)+0.0*y,y,'k-')
plt.text(4340.49*(1+zz)-80,-4.,r'$\mathrm{H}_\gamma$', color='black', size=16)
plt.plot(4101.77*(1+zz)+0.0*y,y,'k-')
plt.text(4101.77*(1+zz)-80,-4.,r'$\mathrm{H}_\delta$', color='black', size=16)
plt.plot(3970.07*(1+zz)+0.0*y,y,'k-')
plt.text(3970.07*(1+zz)-80,-4.,r'$\mathrm{H}_\epsilon$', color='black', size=16)
#plt.plot(3889.05*(1+zz)+0.0*y,y,'k-')
#plt.text(3889.05*(1+zz)-80,-8.,r'$\mathrm{H}_\gamma$', color='black', size=16)
#plt.plot(3835.38*(1+zz)+0.0*y,y,'k-')
#plt.text(3835.38*(1+zz)-80,-8.,r'$\mathrm{H}_\gamma$', color='black', size=16)
#plt.plot(3797.90*(1+zz)+0.0*y,y,'k-')
#plt.text(3797.90*(1+zz)-80,-8.,r'$\mathrm{H}_\gamma$', color='black', size=16)
#plt.plot(3770.63*(1+zz)+0.0*y,y,'k-')
#plt.text(3770.63*(1+zz)-80,-8.,r'$\mathrm{H}_\gamma$', color='black', size=16)

plt.show()
# ------------------------ python code ends here