# Minimal demonstration of using AstroPy specutils to read a plot reduced
# spectra from Liverpool Telescope FRODO level 2 reduced FITS files

# Plot the final extraction, using the pipeline's automated choice of which IFU
# fibres to include in the stack.

# Tested using
#   astropy                       4.0.1.post1        
#   specutils                     0.6             
#   matplotlib                    3.3.0              
#   numpy                         1.19.1             
# Specific version requirements are not well defined

import argparse
from astropy.io import fits
from astropy import units as u
from astropy.visualization import quantity_support
from matplotlib import pyplot as plt

import astropy.wcs as fitswcs
import specutils as sp

quantity_support()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Plot a FRODO spectrum')

    parser.add_argument('infile', action='store', help='Reduced FRODO FITS to display. Filename ends _2.fits, mandatory')
    parser.add_argument('-x', dest='ext_str', choices=['5','SPEC_NONSS','6','SPEC_SS'], default='6', help='Extension to display (default: 6, SPEC_SS)')
    parser.add_argument('-v', dest='verbose', action='store_true', help='Turn on verbose mode (default: Off)')

    args = parser.parse_args()

    # Parse the string names of the extensions into extension numbers
    if args.ext_str in ['5','SPEC_NONSS']: 
      extension = 5
      extName = "SPEC_NONSS"
      unitName = "adu"
    elif args.ext_str in ['6','SPEC_SS']: 
      extension = 6
      extName = "SPEC_SS"
      unitName = "adu"
    

if args.verbose:
  print (args)

f=fits.open(args.infile)

# Spectra are stored as 2D NAXIS1 x 1 arrays. Read and convert to a 1D NAXIS vector. 
specarray=f[extension].data  
specdata = specarray[0] 

specheader=f[extension].header 

flux = specdata * u.Unit(unitName)

# Make WCS Wavelength Calibration from fits header 
my_wcs = fitswcs.WCS(specheader)

# Make Spectrum1D object 
sp1d = sp.Spectrum1D(flux=flux,wcs=my_wcs)
 
plt.plot(sp1d.spectral_axis,sp1d.flux)
plt.title(specheader['OBJECT'] + " " + args.infile + "[" + str(extName) + "]")
plt.show()

f.close()