import numpy as np
from astropy.io import fits
from astropy.table import Table
from matplotlib.colors import LogNorm

# Set up matplotlib
import matplotlib.pyplot as plt
%matplotlib inline

"""
Explanation: Viewing and manipulating data from FITS tables
Authors
Lia Corrales, Kris Stern
Learning Goals

Download a FITS table file from a URL 
Open a FITS table file and view table contents
Make a 2D histogram with the table data
Close the FITS file after use

Keywords
FITS, file input/output, table, numpy, matplotlib, histogram
Summary
This tutorial demonstrates the use of astropy.utils.data to download a data file, then uses astropy.io.fits and astropy.table to open the file. Lastly, matplotlib is used to visualize the data as a histogram.
End of explanation
"""


from astropy.utils.data import download_file

"""
Explanation: The following line is needed to download the example FITS files used in this tutorial.
End of explanation
"""


event_filename = download_file('http://data.astropy.org/tutorials/FITS-tables/chandra_events.fits', 
                               cache=True)

"""
Explanation: FITS files often contain large amounts of multi-dimensional data and tables.  
In this particular example, we'll open a FITS file from a Chandra observation of the Galactic Center.  The file contains a list of events with x and y coordinates, energy, and various other pieces of information.
End of explanation
"""


hdu_list = fits.open(event_filename, memmap=True)

hdu_list.info()

"""
Explanation: Opening the FITS file and viewing table contents
Since the file is big, let's open it with memmap=True to prevent RAM storage issues.
End of explanation
"""


print(hdu_list[1].columns)

"""
Explanation: In this case, we're interested in reading EVENTS, which contains information about each X-ray photon that hit the detector.
To find out what information the table contains, let's print the column names.
End of explanation
"""


evt_data = Table(hdu_list[1].data)

"""
Explanation: Now we'll take this data and convert it into an astropy table. While it's possible to access FITS tables directly from the .data attribute, using Table tends to make a variety of common tasks more convenient.
End of explanation
"""


evt_data

"""
Explanation: For example, a preview of the table is easily viewed by simply running a cell with the table as the last line:
End of explanation
"""


energy_hist = plt.hist(evt_data['energy'], bins='auto')

"""
Explanation: We can extract data from the table by referencing the column name. Let's try making a histogram for the energy of each photon, which will give us a sense for the spectrum (folded with the detector's efficiency).
End of explanation
"""


ii = np.in1d(evt_data['ccd_id'], [0, 1, 2, 3])
np.sum(ii)

"""
Explanation: Making a 2D histogram with some table data
We'll make an image by binning the x and y coordinates of the events into a 2D histogram.
This particular observation spans five CCD chips.  First, we determine the events that only fell on the main (ACIS-I) chips, which have number ids 0, 1, 2, and 3.
End of explanation
"""


NBINS = (100,100)

img_zero, yedges, xedges = np.histogram2d(evt_data['x'][ii], evt_data['y'][ii], NBINS)

extent = [xedges[0], xedges[-1], yedges[0], yedges[-1]]

plt.imshow(img_zero, extent=extent, interpolation='nearest', cmap='gist_yarg', origin='lower')

plt.xlabel('x')
plt.ylabel('y')

# To see more color maps
# http://wiki.scipy.org/Cookbook/Matplotlib/Show_colormaps

"""
Explanation: Method 1: Use numpy to make a 2D histogram and imshow to display it
This method allows us to create an image without stretching:
End of explanation
"""


NBINS = (100,100)
img_zero_mpl = plt.hist2d(evt_data['x'][ii], evt_data['y'][ii], NBINS, 
                          cmap='viridis', norm=LogNorm())

cbar = plt.colorbar(ticks=[1.0,3.0,6.0])
cbar.ax.set_yticklabels(['1','3','6'])

plt.xlabel('x')
plt.ylabel('y')

"""
Explanation: Method 2: Use hist2d with a log-normal color scheme
End of explanation
"""


hdu_list.close()

"""
Explanation: Close the FITS file
When you're done using a FITS file, it's often a good idea to close it.  That way you can be sure it won't continue using up excess memory or file handles on your computer.  (This happens automatically when you close Python, but you never know how long that might be...)
End of explanation
"""


