'''
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
file_spec = pyfits.open('spSpec-51792-0354-362.fit')
data1 = file_spec[0].data

#Extract fitsfile header data, such as Lamda, redshift, etc.
lamda0 = file_spec[0].header['coeff0']    # Center wavele (log10) of first pixel
interval = file_spec[0].header['coeff1'] # Log10 dispersion per pixel
index_pix = np.arange(file_spec[0].header['naxis1']) # col=3865,row=5

zz = file_spec[0].header['z']
zz_err = file_spec[0].header['z_err']

lamda=lamda0+interval*index_pix
lamda1=10.0**lamda

flux1 = data1[0] # SPECTRUM
flux2 = data1[1] # CONTINUUM-SUBTRACTED SPECTRUM
flux3 = data1[2] # ERROR
flux4 = data1[4] # MASK

print(lamda1,flux1,'z = ',zz,'z_error =',zz_err)

#create plot
fig, ax = plt.subplots(1,1)

# Plots data scaled to 1.0e13
#ax.plot(lamda1, (flux1*1.0e13), 'b')

plt.plot(lamda1, flux1,'k')
plt.plot(lamda1, flux2,'r')
plt.plot(lamda1, flux3,'g')
#plt.plot(lamda1, flux4,'c')

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

# Define limits
plt.xlim(3700,9200)
plt.ylim(-20,50)

mz=1   # Set thickness of the tick marks
lz=6   # Set length of the tick marks
## Make tick lines thicker
for l in ax.get_xticklines():
 l.set_markersize(lz)
 l.set_markeredgewidth(mz)
for l in ax.get_yticklines():
 l.set_markersize(lz)
 l.set_markeredgewidth(mz)

## Make figure box thicker
for s in ax.spines.values():
 s.set_linewidth(mz)

# Overplot lines for the locations of the Balmer lines
xline = np.array([1.0,1.0])
yline17 = np.array([1.7,1.0])
yline35 = np.array([3.5,1.0])
yline4 = np.array([4,1.0])
yline45 = np.array([4.5,1.0])
ax.plot(6562.85*(1+zz)*xline, yline17, 'k')
ax.plot(4861.36*(1+zz)*xline, yline35, 'k')
ax.plot(4340.49*(1+zz)*xline, yline4, 'k')
ax.plot(4101.77*(1+zz)*xline, yline4, 'k')
ax.plot(3970.07*(1+zz)*xline, yline4, 'k')
ax.plot(3889.05*(1+zz)*xline, yline45, 'k')
ax.plot(3835.38*(1+zz)*xline, yline45, 'k')
ax.plot(3797.90*(1+zz)*xline, yline45, 'k')
ax.plot(3770.63*(1+zz)*xline, yline45, 'k')
ax.plot([3770.63*(1+zz), 6562.85*(1+zz)], xline,'k')

# Add Text to the figure at position 5000A and Flux=1.3
plt.text(5000*(1+zz),-13.0, 'Hydrogen Balmer Series', fontsize = 11)

# Plot an arrow to mark H Balmer Series
#plt.annotate(r'$\mathrm{H}_\alpha$',xy=(6562.85*(1+zz),0.),xytext=(6562.85*(1+zz),-10.),arrowprops=dict(arrowstyle='->',connectionstyle='arc3'))
#plt.annotate(r'$\mathrm{H}_\beta$',xy=(4861.36*(1+zz),0.),xytext=(4861.36*(1+zz),-10.),arrowprops=dict(arrowstyle='->',connectionstyle='arc3'))
#plt.annotate(r'$\mathrm{H}_\gamma$',xy=(4340.49*(1+zz),0.),xytext=(4340.49*(1+zz),-10.),arrowprops=dict(arrowstyle='->',connectionstyle='arc3'))

#OR

#Plot a line to mark H Balmer Series
y = np.linspace(-5.,0.)
plt.plot(6562.85*(1+zz)+0.0*y,y,'k-')
plt.text(6562.85*(1+zz)-80,-8.,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,-8.,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,-8.,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,-8.,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,-8.,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