# Python program to load a FITS spectrum and display it

# import needed extensions
from numpy import *
import pyfits
import matplotlib.pyplot as plt # plotting package
import matplotlib.cm as cm # colormaps
from matplotlib.colors import LogNorm
from scipy.interpolate import interp1d


# calibration
centralp, centralw, slope = 1000.00, 480.00, 1.00000


plt.ion() # do plots in interactive mode

# read in the spectrum image file for the calibration star
# change the file names as appropriate
stdfile = pyfits.open('std1.fits')
# copy the image data into a numpy (numerical python) array
stdimg = stdfile[0].data

# calculate a spectrum
# y0 is the center of the band over which the spectrum is extracted
y0 = 94
# dy sets the width of the band (y0-dy to y0+dy)
dy = 7
# figure out dimensions of spectrum image
(ny, nx) = shape(stdimg)
# find a 1-d spectrum by integrating across the band
std = zeros(nx)
for i in range(nx):
  std[i] = sum(stdimg[(y0-dy):(y0+dy+1), i])

# calculate the wavelength for each bin
p = 1+arange(len(std)) # pixel numbers
w = centralw + slope*(p-centralp) # wavelengths
# change the wavelength to Angstroms
w = w*10.0

# keep only data at wavelengths 3500-6000 Ang
q = where((w >= 3500.0) & (w < 6000.0))
w, std = w[q], std[q]

# plot the spectrum versus wavelength
plt.figure(1)
plt.clf()
plt.xlabel('Wavelength (Ang)')
plt.ylabel('Counts')
plt.yscale('log')
#plt.axis([0, 750, 0, 1E5])
plt.plot(w, std, '-b')

# read in the reference spectrum from Hamuy 1992
# variables are wavelength, flux, flux in Janksy, bin width
rw, rf, rfj, rb = transpose(genfromtxt('feg274.dat'))

# plot the reference spectrum for comparison
# first normalize the reference spectrum
norm = max(std)/max(rf)
plt.plot(rw, norm*rf, '--r')

# now we wish to use the Hamuy reference spectrum to calculate a 
# conversion factor to apply to our measured spectra in order to
# convert them to physical units
# create a function that returns the reference flux at any wavelength
rfi = interp1d(rw, rf) 
# find reference flux in the bins of the measured spectrum
rfb = rfi(w)
# find the conversion factor
conv = rfb/std

# plot the conversion factor versus wavelength
plt.figure(2)
plt.clf()
plt.xlabel('Wavelength (Ang)')
plt.ylabel('Conversion')
#plt.axis([0, 750, 0, 1E5])
plt.plot(w, conv, '-b')


# read in the spectrum image file for the target star
# change the file names as appropriate
stdfile = pyfits.open('s.fits')
# copy the image data into a numpy (numerical python) array
simg = stdfile[0].data

# calculate a spectrum of the target
# y0 is the center of the band over which the spectrum is extracted
y0 = 61
# dy sets the width of the band (y0-dy to y0+dy)
dy = 7
# figure out dimensions of spectrum image
(ny, nx) = shape(simg)
# find a 1-d spectrum by integrating across the band
s = zeros(nx)
for i in range(nx):
  s[i] = sum(simg[(y0-dy):(y0+dy+1), i])


# calculate a sky background spectrum
# y0 is the center of the band over which the spectrum is extracted
y0 = 81
# dy sets the width of the band (y0-dy to y0+dy)
dy = 7
# figure out dimensions of spectrum image
(ny, nx) = shape(simg)
# find a 1-d spectrum by integrating across the band
b = zeros(nx)
for i in range(nx):
  b[i] = sum(simg[(y0-dy):(y0+dy+1), i])

# subtract the sky background from the spectrum
ss = s-b

# keep only data at wavelengths 3500-6000 Ang
q = where((w >= 3500.0) & (w < 6000.0))
w, s, b, ss = w[q], s[q], b[q], ss[q]

# plot the spectrum versus wavelength
plt.figure(3)
plt.clf()
plt.xlabel('Wavelength (Ang)')
plt.ylabel('Counts')
plt.yscale('log')
#plt.axis([0, 750, 0, 1E5])
plt.plot(w, s, '-b') # source spectrum
plt.plot(w, b, '-r') # background spectrum
plt.plot(w, ss, '-k') # source with sky subtracted


# convert the spectrum into physical units
sc = ss*conv
# plot the corrected spectrum versus wavelength
plt.figure(4)
plt.clf()
plt.xlabel('Wavelength (Ang)')
plt.ylabel('Flux ()')
#plt.yscale('log')
#plt.axis([0, 750, 0, 1E5])
plt.plot(w, sc, '-b')