# Minimal demonstration of using AstroPy specutils to read a plot reduced
# spectra from Liverpool Telescope FROSOSPEC level 2 reduced FITS files.
#
# Reading data from the data cube extensions

# Select the fibres to plot from FITS extension 7[COLCUBE] 
# Pixel coordinates following the typical FITS convention of a  
# raster-scan from bottom-left to top-right 
#   1,1  is bottom-left pixel in the COLCUBE image 
#   12,1 is bottom-right pixel 
#   1,12 is top-left pixel 
# In contrast, fibres map from IFU to their sequence on the slit 
# in boustrophedon manner from bottom-left to top-right. 

# 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 ast import literal_eval

import numpy as np
from astropy.io import fits
from astropy import units as u
from astropy.visualization import quantity_support
import astropy.wcs as fitswcs
import specutils as sp
from matplotlib import pyplot as plt

quantity_support()

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

    parser.add_argument('infile', action='store', help='Reduced FRODO FITS to display. Filename ends _2.fits, mandatory')
    parser.add_argument('fibre_list', action='store', type=literal_eval, help='Comma separated list of (x,y) tuples of the IFU fibres to read. E.g., "(1,1),(1,12)" to get bottom-left and top-left pixels of COLCUBE extension. Mandatory')
    parser.add_argument('-x', dest='ext_str', choices=['4','CUBE_SS','2','CUBE_NONSS'], default='4', help='Extension to display (default: 4, CUBE_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 ['2','CUBE_NONSS']: 
      extension = 2
      unitName = "adu"
    elif args.ext_str in ['4','CUBE_SS']: 
      extension = 4
      unitName = "adu"



if args.verbose:
  print (args)

countFibres = len(args.fibre_list)

# Read the FITS extension 
f = fits.open(args.infile)
specarray = f[extension].data
specheader = f[extension].header
# Some FRODO data erroneously have the units in the FITS header as 'Angstroms' instead of the correct 'Angstrom'
specheader['CUNIT3'] = 'Angstrom'
my_wcs = fitswcs.WCS(specheader)
f.close()


# Retain just the listed fibres to plot. 
# Python array indices are completely different from FITS conventions meaning the tuples in args.fibre_list
# need reversing and decrementing off-by-one to give zero based indices.
dataToPlot = np.zeros((countFibres,specarray.shape[0]))
for ii in range(0, countFibres) :
  dataToPlot[ii,:] = specarray[:,args.fibre_list[ii][1]-1,args.fibre_list[ii][0]-1]

# Create a Spectrum1D object containing the five separate fibres
dataToPlot = dataToPlot * u.Unit(unitName)
sp1d = sp.Spectrum1D(flux=dataToPlot,wcs=my_wcs)

fig = plt.figure(figsize=(12, 4), dpi=80, facecolor='w', edgecolor='k')
ax = plt.axes()
for ii in range(0, countFibres):
  print(args.fibre_list[ii])
  plt.plot(sp1d.spectral_axis,sp1d.flux[ii],label='Fibre '+str(args.fibre_list[ii]))

# Coadd the five fibres into one total and plot the result 
sp1d = sp.Spectrum1D( flux=np.nansum(dataToPlot,axis=0), wcs=my_wcs)
plt.plot(sp1d.spectral_axis, sp1d.flux, label='Coadded')
plt.legend()
plt.show()