# pylab is a nice interface to all that is matplotlib
import matplotlib.pyplot as plt
import numpy as np

# Should plots pop up as you type? Use interactive on/off plt.ion() or plt.ioff()
plt.ion()

# if you want the plots to appear in the notebook
%matplotlib inline

from IPython.display import Image
Image(url='http://matplotlib.org/_images/fig_map.png')

# basic
fig = plt.figure()
ax = plt.axes()
# plt.figure?
# plt.axes?

# my favorite:
fig, ax = plt.subplots()
# plt.subplots?
# fig, ax = plt.subplots(ncols=2)
# fig, ax = plt.subplots(nrows=2, ncols=2, figsize=(8,8))

# clear them:
plt.close()  # close the most recent window
plt.close('all')

# define something to plot
y = np.sin(np.linspace(0, 2*np.pi, 20))
x = np.arange(20)

fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(14, 14))
# plt.plot?
ax[0, 0].plot(x, y)
ax[0, 1].plot(x, y, ':')
ax[0, 2].plot(x, y, linestyle=':')
ax[1, 0].plot(x, y, linestyle='steps')
ax[1, 1].plot(x, y, linestyle='steps-pre:')
ax[1, 2].plot(x, y, linestyle='steps-pre-')
ax[1, 2].plot(x, y, linestyle='steps-mid-')
ax[1, 2].plot(x, y, linestyle='steps-post-')

fig, ax = plt.subplots(ncols=2, figsize=(14, 4))
ax[0].plot(x, y, 'o')
ax[1].plot(x, y, 'o', ms=30, mec='white', alpha=0.3)

# A scatter plot that uses colors picked by the y value (you must set the color map)
fig, ax = plt.subplots(ncols=2, figsize=(14,4))

ax[0].scatter(x, y, c=y, cmap=plt.cm.RdYlBu, s=150)

# Same as above, with the marker size set by the array z
z = np.random.random(20) * 1000
ax[1].scatter(x, y, c=y, cmap=plt.cm.RdYlBu, s=z)

xerr = np.sqrt(x)
yerr = 0.1
fig, ax = plt.subplots(ncols=3, figsize=(16, 4))
ax[0].errorbar(x, y, xerr=xerr, yerr=yerr)
# no error caps, color the error bars differently than the data:
ax[1].errorbar(x, y, xerr=xerr, yerr=yerr, capsize=0, ecolor='gray')
# don't plot values? Just error bars? set fmt='none'
ax[2].errorbar(x, y, xerr=xerr, yerr=yerr, capsize=0, ecolor='gray', fmt='none')

x = np.random.normal(size=10000)
y = np.random.normal(size=10000)


fig, ax = plt.subplots(ncols=2, figsize=(14, 4))
h, bins, p = ax[0].hist(x, bins=20)

h, xe, ye, p = ax[1].hist2d(x, y, bins=20, cmap=plt.cm.RdYlBu)

fig, ax = plt.subplots()
h, bins = np.histogram(x, bins=20)
ax.plot(bins[1:], h, linestyle='steps-pre')
plt.draw()

from astropy.io import fits
# replace this with the file you downloaded
fitsfile = 'hlsp_angst_hst_acs-wfc_10605-ugc-5336_f555w-f814w_v1_gst.fits'

hdu = fits.open(fitsfile)
data = hdu[1].data
photsys = hdu[0].header['CAMERA']  # survey used ACS and WFPC2

# the magnitude fields in the fits file are named MAG[1 or 2]_[ACS or WFPC2] 
mag1 = data['MAG1_%s' % photsys]
mag2 = data['MAG2_%s' % photsys]

mag1_err = data['MAG1_ERR']
mag2_err = data['MAG2_ERR']

color = mag1 - mag2
color_err = np.sqrt(data['MAG1_ERR'] ** 2 + data['MAG2_ERR'] ** 2)

# these fits file contains stars that are recovered in only one filter
# stars not recovered are given values >= 90. No need to plot them, messes up the autoscaling.
good, = np.nonzero((np.abs(color) < 30) & (np.abs(mag2) < 30))

fig, ax = plt.subplots(figsize=(8, 8))
ax.plot(color[good], mag2[good], '.', color='black', ms=3)

# CMDs have yaxis reversed, a simple way to reverse axis depend on autoscaling
ax.set_ylim(ax.get_ylim()[::-1])

fig, ax = plt.subplots(figsize=(8, 8))
ax.plot(color[good], mag2[good], '.', color='black', ms=3)

# CMDs have yaxis reversed, a simple way to reverse axis depend on autoscaling
ax.set_ylim(ax.get_ylim()[::-1])

ax.errorbar(color[good], mag2[good], fmt='none', lw=1, xerr=color_err[good],
            yerr=mag2_err[good], capsize=0, ecolor='gray')

binsize = 0.1
mbin = np.arange(mag2.min(), mag2.max(), binsize) 
lf, bins = np.histogram(mag2, bins=mbin)
fig, ax = plt.subplots(ncols=2, figsize=(14, 4))

[ax[i].plot(bins[1:], lf, linestyle='steps-pre') for i in range(2)]

# maybe you want a log scale
for i, yscale in enumerate(['log', 'linear']):
    ax[i].set_yscale(yscale)
    ax[i].set_title(yscale)

cbin = np.arange(color.min(), color.max(), binsize)
hess, cbin, mbin = np.histogram2d(color, mag2, bins=[cbin, mbin])

# Set the extent of the image
extent = [np.min(cbin), np.max(cbin), np.max(mbin), np.min(mbin)]

fig, ax = plt.subplots(ncols=2, figsize=(14, 4))

im = ax[0].imshow(hess.T, cmap=plt.cm.gray, interpolation='nearest',
                  extent=extent, aspect='auto')

# Perhaps you want log bin counts
from matplotlib.colors import LogNorm
im = ax[1].imshow(hess.T, cmap=plt.cm.gray, interpolation='nearest',
                  extent=extent, aspect='auto', norm=LogNorm())

# some code for the next sections
import os

def directory_check(directory):
    """check if a directory exists, if not print the mkdir command to create it"""
    if not os.path.isdir(directory):
        print 'Need to make the directory.'
        print 'mkdir', directory
    else:
        print '%s exists' % directory

def plot_demo():
    """make a plot for a demo"""
    fig, ax = plt.subplots()
    ax.plot(np.sin(np.linspace(0, 2*np.pi)))
    ax.plot(np.sin(np.linspace(0, 2*np.pi)), 'o')
    return ax

print 'matplotlib data path:', mpl.get_data_path()
mplrc = os.path.join(mpl.get_data_path(), 'matplotlibrc')

# In terminal use 'cat' to see the contents of the matplotlibrc file, or open it.
print '$ cat %s' % mplrc

# Quick check if you may need to create the directory .matplotlib
user_mpl_path = os.path.join(os.environ['HOME'], '.matplotlib')
directory_check(user_mpl_path)

print 'cp %s %s/matplotlibrc' % (mplrc, user_mpl_path)

#plt.rcParams.keys()

print plt.style.available

# default parameters
plot_demo()

# set all future plots to use bmh
plt.style.use('bmh')
plot_demo()

# plot each of the available styles
for style in plt.style.available:
    with plt.style.context(style):
        ax = plot_demo()
        ax.set_title(style)

# back to bmh
plot_demo()

stylelib_path = os.path.join(mpl.get_data_path(), 'stylelib')
print 'mpl style path:', stylelib_path
print 'mpl style files:', os.listdir(stylelib_path)

# Quick check if you may need to create the directory stylelib
user_stylelib_path = os.path.join(os.environ['HOME'], '.matplotlib', 'stylelib')
directory_check(user_stylelib_path)

# for copy and paste ease, here are the full paths
for s in os.listdir(stylelib_path):
    print os.path.join(stylelib_path, s)

%%file presentation.mplstyle
# Author: Phil Rosenfield
# Adapted from: 
# Author: Cameron Davidson-Pilon, replicated styles from FiveThirtyEight.com
# See https://www.dataorigami.net/blogs/fivethirtyeight-mpl

lines.linewidth: 3
lines.solid_capstyle: butt

legend.frameon: False
legend.numpoints: 1
legend.fontsize: 16

axes.color_cycle: 000000, 30a2da, fc4f30, e5ae38, 6d904f, 8b8b8b
axes.axisbelow: true
axes.edgecolor: f0f0f0
axes.linewidth: 3.0
axes.titlesize: 24
axes.labelsize: 20

patch.edgecolor: f0f0f0
patch.linewidth: 0.5

font.size: 14.0

figure.figsize: 8, 8  # square figures (inches)
figure.subplot.left: 0.08
figure.subplot.right: 0.95 
figure.subplot.bottom: 0.07
figure.facecolor: ffffff


xtick.labelsize : 16
ytick.labelsize : 16

text.usetex: True
text.latex.unicode: False

image.aspect: 1.0
image.cmap: Blues

print 'cp %s %s/' % (os.path.join(os.getcwd(), 'presentation.mplstyle'), user_stylelib_path)

%%file data_plots.py
#!/usr/bin/env python
""" 
Script to make CMD, LF, or Hess diagram from a binary fits table
Written by: Phil Rosenfield
"""
import argparse
from astropy.io import fits
import matplotlib.pyplot as plt
import numpy as np
import sys

def plot_cmd(color, mag, color_err=None, mag_err=None, ax=None):
    '''
    Plot a Color Magnitude diagram with uncertainties
    
    Parameters
    ----------
    color, mag : color and magnitude arrays
    color_err, mag_err: uncertainties in color and mag
    
    ax : axes instance
    
    Returns
    -------
    ax : axes instance
    '''

    if ax is None:
        fig, ax = plt.subplots(figsize=(8, 8))

    ax.plot(color, mag, '.', ms=3)

    if color_err is not None and mag_err is not None:
        ax.errorbar(color, mag, fmt='none', lw=1, xerr=color_err,
                    yerr=mag_err, capsize=0, ecolor='gray')

    # reverse yaxis
    ax.set_ylim(ax.get_ylim()[::-1])
    return ax


def plot_lf(mag, binsize, mbin=None, yscale='log'):
    """
    Make a Luminosty function (binned magnitude) plot
    
    Parameters
    ----------
    mag : array
        magnitude array to be binned
    binsize : float
        width of magnitude bins
    mbin : array
        right edges of magnitude bins
    yscale : str
        plt.set_yscale option

    Returns
    -------
    ax : axes instance
    """
    if mbin is None:
        mbin = np.arange(mag.min(), mag.max(), binsize)
    
    lf, bins = np.histogram(mag, bins=mbin)
    fig, ax = plt.subplots()
    
    ax.plot(bins[1:], lf, linestyle='steps-pre')
    ax.set_yscale(yscale)
    return ax
    
def make_hess(color, mag, binsize, cbinsize=None, mbin=None, cbin=None):
    """
    Compute a hess diagram (surface-density CMD) on photometry data.

    Parameters
    ---------
    color : array
        color values

    mag : array
        magnitude values

    binsize, cbinsize: float, float
        width of mag, color bins in magnitudes

    cbin : array
        the right edges of the color bins

    mbin : array
        the right edges of the magnitude bins

    Returns
    -------
    cbin : array
        the centers of the color bins

    mbin : array
        the centers of the magnitude bins

    hess : 2d array
        The Hess diagram values
    """

    if mbin is None:
        mbin = np.arange(mag.min(), mag.max(), binsize)

    if cbin is None:
        if cbinsize is None:
            cbinsize = binsize
        cbin = np.arange(color.min(), color.max(), cbinsize)

    hess, cbin, mbin = np.histogram2d(color, mag, bins=[cbin, mbin])
    return hess, cbin, mbin


def plot_hess(color, mag, binsize, ax=None, colorbar=False,
              vmin=None, vmax=None, cbinsize=None, im_kwargs={}):
    """
    Plot a hess diagram with imshow.
    
    Parameters
    ----------
    color : array
        color array to be binned
    mag : array
        magnitude array to be binned

    binsize, cbinsize : float, float
        width of magnitude, color bins
    
    colorbar : bool
        option to also plot the colorbar
    
    vmin, vmax : float, float or None, None
        lower and upper limits sent to matplotlib.colors.LogNorm
    
    im_kwargs : dict
        dictionary passed to imshow by default:
        defaults = {'norm': LogNorm(vmin=vmin, vmax=vmax),
                    'cmap': plt.cm.gray,
                    'interpolation': 'nearest',
                    'extent' [limits of mag and color]
                    'aspect': 'auto'}
    Returns
    -------
    ax : axes instance
    """
    from matplotlib.colors import LogNorm
    if ax is None:
        fig, ax = plt.subplots()

    hess, cbin, mbin = make_hess(color, mag, binsize, cbinsize=cbinsize)
    extent = [np.min(cbin), np.max(cbin), np.max(mbin), np.min(mbin)]
    vmax = vmax or hess.max()

    defaults = {'norm': LogNorm(vmin=vmin, vmax=vmax),
                'cmap': plt.cm.gray,
                'interpolation': 'nearest',
                'extent': extent,
                'aspect': 'auto'}

    kwargs = dict(defaults.items() + im_kwargs.items())

    im = ax.imshow(hess.T, **kwargs)

    if colorbar is True:
        plt.colorbar(im)

    return ax


def load_data(fitsfile, yfilt='I'):
    """
    Load color, magnitude and uncertainties from binary fits table
    
    Parameters
    ----------
    fitsfile : string or file object
        path to binary fits table or object to be read by astropy.io.fits
    
    yfilt : string
        filter to use as mag (V or I)

    Returns
    -------
    color, mag : arrays of color and magnitude
    color_error, mag_err : arrays of summed quadriture uncertainies and magnitude uncertainties
    """

    hdu = fits.open(fitsfile)
    data = hdu[1].data
    photsys = hdu[0].header['CAMERA']

    # the magnitude fields in the fits file are named MAG{1,2}_[photsys]
    mag1 = data['MAG1_%s' % photsys]
    mag2 = data['MAG2_%s' % photsys]
    color = mag1 - mag2

    color_err = np.sqrt(data['MAG1_ERR'] ** 2 + data['MAG2_ERR'] ** 2)

    # choose what gets the yaxis V or I
    if yfilt.upper() == 'I':
        mag = mag2
        ymag = 'MAG2'
    else:
        mag = mag1
        ymag = 'MAG1'

    # the error fields in the fits file are named [MAG]_ERR
    mag_err = data['%s_ERR' % ymag]

    return color, mag, color_err, mag_err

    
def main(argv):
    parser = argparse.ArgumentParser(description="Generate a plot of a fits file")

    parser.add_argument('-p', '--plottype', type=str, default='cmd',
                        help='which plot to make: CMD, hess, or LF')

    parser.add_argument('-f', '--filters', type=str, default=None,
                        help='comma separated V and I filter names for plot labels')

    parser.add_argument('-y', '--yfilter', type=str, default='I',
                        help='plot V or I for LF or on y axis for cmd, Hess')

    parser.add_argument('-m', '--binsize', type=float, default=0.05,
                        help='hess diagram or LF mag binsize')

    parser.add_argument('-c', '--cbinsize', type=float, default=0.1,
                        help='hess diagram color binsize')

    parser.add_argument('-xlim', '--xlim', type=list, default=None,
                        help='comma separated x axis min, max')

    parser.add_argument('-ylim', '--ylim', type=list, default=None,
                        help='comma separated y axis min, max')

    parser.add_argument('-yscale', '--yscale', type=str, default='log',
                        help='y axis scale for LF')

    parser.add_argument('-x', '--colorbar', action='store_true',
                        help='add the hess diagram colorbar')

    parser.add_argument('-outfile', '--outfile', type=str, default='data_plot.png',
                        help='the name of the output file')

    parser.add_argument('-style', '--style', type=str, default='ggplot',
                        choices=plt.style.available,
                        help='the name of the matplotlib style')

    parser.add_argument('file', type=argparse.FileType('r'),
                        help='the name of the fits file')

    args = parser.parse_args(argv)

    # set the plot style
    plt.style.use(args.style)

    if args.filters is not None:
        filter1, filter2 = args.filters.split(',')
    else:
        print('warning: using V, I as default filter names')
        filter1 = 'V'
        filter2 = 'I'
    
    yfilt = args.yfilter
    
    color, mag, color_err, mag_err = load_data(args.file, yfilt=yfilt)
    # the fits file contains stars that are recovered in only one filter
    # stars not recovered are given values >= 90. No need to plot em.
    good, = np.nonzero((np.abs(color) < 30) & (np.abs(mag) < 30))    
    
    if args.plottype.lower() == 'cmd':    
        ax = plot_cmd(color[good], mag[good], color_err=color_err[good], mag_err=mag_err[good])

    if args.plottype.lower() == 'hess':
        ax = plot_hess(color[good], mag[good], colorbar=args.colorbar,
                       binsize=args.binsize, cbinsize=args.cbinsize)

    if args.plottype.lower() == 'lf':
        ax = plot_lf(mag[good], args.binsize, yscale=args.yscale)

        # make axis labels
        ax.set_xlabel(r'$%s$' % yfilt)
        ax.set_ylabel(r'$\#$')
    else:
        # make axis labels for cmd, hess
        if args.filters is not None:
            if yfilt == 'I':
                yfilt = filter2
            else:
                yfilt = filter1

        ax.set_ylabel(r'$%s$' % yfilt)
        ax.set_xlabel(r'$%s-%s$' % (filter1, filter2))

    if args.ylim is not None:
        ylim = np.array(''.join(args.ylim).split(','), dtype=float)
        ax.set_ylim(ylim)
    
    if args.xlim is not None:
        xlim = np.array(''.join(args.xlim).split(','), dtype=float)
        ax.set_xlim(xlim)

    plt.savefig(args.outfile)
    print('wrote %s' % args.outfile)

if __name__ == "__main__":
    main(sys.argv[1:])

! chmod u+x data_plots.py

!./data_plots.py -h

! ./data_plots.py -x -p hess -style presentation -f F555W,F814W -outfile hess_presentation.png hlsp_angst_hst_acs-wfc_10605-ugc-5336_f555w-f814w_v1_gst.fits

import matplotlib as mpl
mpl.use('Agg')