# NumPy is a library for the Python programming language, adding support for large, multi-dimensional arrays and matrices, along with a large collection of high-level mathematical functions to operate on these arrays
import numpy as np

# Matplotlib is a plotting library for the Python programming language and its numerical mathematics extension NumPy
import matplotlib.pyplot as plt

# Astropy is a collection of software packages written in the Python programming language and designed for use in astronomy
# In general, we do not "from astropy import *". Only import the packages you need.
from astropy.io import fits
from astropy import wcs
from astropy.modeling import models, fitting

# The math module provides access to the mathematical functions defined by the C standard
import math

# The sys module is a set of functions which provide crucial information about how your Python script is interacting with the host system
import sys


# Get the cube data
# rsync -avz rsync://data.sdss.org/dr13/manga/spectro/redux/v1_5_4/7443/stack/manga-7443-12703-LOGCUBE.fits.gz .
# It is possible to read in a zipped tar file and have python unzip it
# But that takes much longer


#
# Section 1: reading in fits files
#

# fits.open is the factory function to open a FITS file and return an `HDUList` object
# import the fits image with fits.open command:


# For reference, to get help with a function in python, use
# help(fits.open) or fits.open? if using ipython

# We can look at the header info. This is the global header:


# We can call information from the header from the keywords like the RA and DEC:


# An HDU (Header Data Unit) is the highest level component of the
# FITS file structure, consisting of a header and (typically) a data
# array or table
# List the HDUList and see what the fits file contains:
cube.info()

# There are 16 different HDUs (data tables) in this file
# We can call them by either their HDU number or their name
# HDU1 is the flux data
# Call the flux data array with both the HDU number and the HDU name:


# If an HDU's data is an image, the data attribute of the HDU object will
# return a numpy ndarray object
# Check that flux is an array:


# What is the size of the image along the X, Y, and velocity axis:


# Pull out the header info for the data array we want (flux):


# FLUX is the 3D rectified cube in units of ..?
# We know the units from BUNIT in the flux header. Look it up:


# Make an array for flux, mask, and wave, the three parameters we care about
# Note, MANGA 3D data structures are in the order (wavelength, DEC, RA)
# I prefer to seeing coords as RA (x-coord) & DEC (y-coord)
# Re-order FLUX and MASK arrays from (wavelength, DEC, RA) to (RA, DEC, wavelength)
flux = np.transpose(cube['FLUX'].data, axes=(2, 1, 0))


# Note the different dimension of the wave HDU compared the flux HDU



# Don't continue onto section 2 yet
# remove (comment out) this when we start section 2
sys.exit(0)



#
# Section 2: Sum over the velocity axis and plot a white light image
#

# Ignore pixels that are masked out for us in the mask array
# In MANGA data, valid pixels have a mask value of 0
# Therefore do not use flux values if their mask array is not zero
do_not_use = (mask) != 0

# np.ma.array: An array class with possibly masked values
# Masked values of True exclude the corresponding element from any computation
flux_m = np.ma.array(flux, mask=do_not_use)

# To make the image, sum over the velocity axis:
image_data = flux_m.sum(axis=2)

# Note: (x=0, y=0) corresponds to the upper left if North is up and East is left
# Invert the array (image_data.T):
imd = image_data.T

# Now, let's plot the figure!
plt.figure()

# plt.imshow allows us to display an image on the axes
plt.imshow(imd, cmap='magma_r', origin='lower', interpolation='none')

plt.title('Whitelight')
plt.xlabel('x pixel')
plt.ylabel('y pixel')

plt.show()








#
# Plot an H-alpha narrow band image 6564 A 
#

# Need to correct the observed wavelengths to get rest-frame wavelengths
redshift = 0.0404448

# Get the indices for the wavelengths +/- around H-alpha
# np.where can help us with this condition:


# Sum over the maxed flux values only around the H-alpha line:


# Transpose the image:



# Plot the image
# In imshow, we can specify the max and min values to normalize the flux values
# Try a few values for vmin= and vmax=

# plot image with imshow


# plot title


# plot colorbar


# plot x and y labels



# show plot






# Don't continue onto section 3 yet
# remove this when we start section 3
sys.exit(0)



#
# Section 3: Plot the central spectrum
#

# We specificed wave and flux earlier. Look at the shapes of wave and flux:


# We need to specify what pixels we want the spectrum at. Pick the center pixels.
# We can get the center pixels CRPIX in the flux header. Call and define those:


# Note, the reference pixel is 1-indexed
# Python indexing begins at 0, so we need to zero-index it

# Plot the spectrum with plt.plot


# Set the x, y labels, the title, and show the plot.
# Rember, the flux value can be taken from the flux_header!





# Compare the spectrum at three different regions in the galaxy

xa, ya = [34, 24] # random spot in the arm
xo, yo = [22, 60] # random spot in the outskirt

# Let's do a subplot with two plots, showing the galaxy image and spectrum
fig = plt.figure(figsize=(11, 5))

# figure 1, the spectrum
plt.subplot(1, 2, 1)
# plot the center value

# plot the center value


# plot the outskirt value



# figure 2, the image
plt.subplot(1, 2, 2)

# plt.imshow to show the galaxy


# use plt.scatter to show where these three regions are in the galaxy






# Don't continue onto section 4 yet
# remove this when we start section 4
sys.exit(0)



#
# Section 4: Fit a Gaussian to Hbeta 4861A
#

# Mask out wavelengths not around the line:


# Create a composite model with some initial parameters close to
# the desired result (usually sufficient to guess from looking at
# a plot of the data). Several iterations are required for a good fit:
bg = models.Const1D(amplitude=1.)
gs = models.Gaussian1D(amplitude=10., mean=5055, stddev=1)
init_model = bg - gs
fitter = fitting.LevMarLSQFitter()
fit_model = fitter(init_model, wave[ind_wave_fit], flux[x_center, y_center][ind_wave_fit])
y_fit = fit_model(wave[ind_wave_fit])
print(fit_model)

# plot the data


#plot the Gaussian fit



# Area under the curve
area = (math.sqrt(2 * math.pi) * abs(fit_model.amplitude_1) *
    fit_model.stddev_1)

# Continuum level
height = fit_model.amplitude_0 
print('EW = {:.2f} Angstrom'.format(area / height))

# For help with positional formatting:
# https://www.programiz.com/python-programming/methods/string/format
# and https://pyformat.info