# lib.py: Module containing various plotting functions
#
#     Authors: William Cleveland (USRA),
#              Adam Goldstein (USRA) and
#              Daniel Kocevski (NASA)
#
#     Portions of the code are Copyright 2020 William Cleveland and
#     Adam Goldstein, Universities Space Research Association
#     All rights reserved.
#
#     Written for the Fermi Gamma-ray Burst Monitor (Fermi-GBM)
#
#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU General Public License as published by
#     the Free Software Foundation, either version 3 of the License, or
#     (at your option) any later version.
#
#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU General Public License for more details.
#
#     You should have received a copy of the GNU General Public License
#     along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
import matplotlib.pyplot as plt
from astropy.coordinates import SkyCoord
from matplotlib.colors import colorConverter
from matplotlib.patches import Polygon

from .globals import *
from .lal_post_subs import *
from ..coords import calc_mcilwain_l, saa_boundary, radec_to_spacecraft


# ---------- Lightcurve and Spectra ----------#
def histo(bins, ax, color='C0', edges_to_zero=False, **kwargs):
    """Plot a rate histogram either lightcurves or count spectra
    
    Args:
        bins (:class:`gbm.data.primitives.TimeBins` or \
               :class:`gbm.data.primitives.EnergyBins`)
            The lightcurve or count spectrum histograms
        ax (:class:`matplotlib.axes`): The axis on which to plot
        color (str, optional): The color of the histogram. Default is 'C0'
        edges_to_zero (bool, optional):
            If True, then the farthest edges of the histogram will drop to zero.
            Default is True.
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the plot objects
    """
    bin_segs = bins.contiguous_bins()
    refs = []
    for seg in bin_segs:
        edges = np.concatenate(
            ([seg.lo_edges[0]], seg.lo_edges, [seg.hi_edges[-1]]))
        if edges_to_zero:
            rates = np.concatenate(([0.0], seg.rates, [0.0]))
        else:
            rates = np.concatenate(
                ([seg.rates[0]], seg.rates, [seg.rates[-1]]))

        p = ax.step(edges, rates, where='post', color=color, **kwargs)
        refs.append(p)
    return refs


def histo_errorbars(bins, ax, color='C0', **kwargs):
    """Plot errorbars for lightcurves or count spectra
    
    Args:
        bins (:class:`gbm.data.primitives.TimeBins` or \
               :class:`gbm.data.primitives.EnergyBins`)
            The lightcurve or count spectrum histograms
        ax (:class:`matplotlib.axes`): The axis on which to plot
        color (str, optional): The color of the errorbars. Default is 'C0'
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the plot objects
    """
    bin_segs = bins.contiguous_bins()
    refs = []
    for seg in bin_segs:
        p = ax.errorbar(seg.centroids, seg.rates, seg.rate_uncertainty,
                        capsize=0, fmt='none', color=color, **kwargs)
        refs.append(p)
    return refs


def histo_filled(bins, ax, color=DATA_SELECTED_COLOR,
                 fill_alpha=DATA_SELECTED_ALPHA, **kwargs):
    """Plot a filled histogram
    
    Args:
        bins (:class:`gbm.data.primitives.TimeBins` or \
               :class:`gbm.data.primitives.EnergyBins`)
            The lightcurve or count spectrum histograms
        ax (:class:`matplotlib.axes`): The axis on which to plot
        color (str, optional): The color of the filled histogram
        fill_alpha (float, optional): The alpha of the fill
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the plot objects
    """
    h = histo(bins, ax, color=color, zorder=3, **kwargs)
    zeros = np.zeros(bins.size + 1)
    rates = np.append(bins.rates, bins.rates[-1])
    edges = np.append(bins.lo_edges, bins.hi_edges[-1])
    b1 = ax.plot((edges[0], edges[0]), (zeros[0], rates[0]), color=color,
                 zorder=2)
    b2 = ax.plot((edges[-1], edges[-1]), (zeros[-1], rates[-1]), color=color,
                 zorder=2)
    f = ax.fill_between(edges, zeros, rates, color=color, step='post',
                        alpha=fill_alpha, zorder=4, **kwargs)
    refs = [h, b1, b2, f]
    return refs


def selection_line(xpos, ax, **kwargs):
    """Plot a selection line
    
    Args:
        xpos (float): The position of the selection line
        ax (:class:`matplotlib.axes`): The axis on which to plot
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the plot objects
    """
    ylim = ax.get_ylim()
    ref = ax.plot([xpos, xpos], ylim, **kwargs)
    return ref


def selections(bounds, ax, **kwargs):
    """Plot selection bounds
    
    Args:
        bounds (list of tuples): List of selection bounds
        ax (:class:`matplotlib.axes`): The axis on which to plot
        **kwargs: Other plotting options
    
    Returns:
        tuple: The reference to the lower and upper selection
    """
    refs1 = []
    refs2 = []
    for bound in bounds:
        p = selection_line(bound[0], ax, **kwargs)
        refs1.append(p[0])
        p = selection_line(bound[1], ax, **kwargs)
        refs2.append(p[0])

    return (refs1, refs2)


def errorband(x, y_upper, y_lower, ax, **kwargs):
    """Plot an error band
    
    Args:
        x (np.array): The x values
        y_upper (np.array): The upper y values of the error band
        y_lower (np.array): The lower y values of the error band
        ax (:class:`matplotlib.axes`): The axis on which to plot
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the lower and upper selection    
    """
    refs = ax.fill_between(x, y_upper.squeeze(), y_lower.squeeze(), **kwargs)
    return refs


def lightcurve_background(backrates, ax, cent_color=None, err_color=None,
                          cent_alpha=None, err_alpha=None, **kwargs):
    """Plot a lightcurve background model with an error band
    
    Args:
        backrates (:class:`~gbm.background.BackgroundRates`):
            The background rates object integrated over energy. If there is more
            than one remaining energy channel, the background will be integrated
            over the remaining energy channels.
        ax (:class:`matplotlib.axes`): The axis on which to plot
        cent_color (str): Color of the centroid line
        err_color (str): Color of the errorband
        cent_alpha (float): Alpha of the centroid line
        err_alpha (fl): Alpha of the errorband
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the lower and upper selection    
    """
    if backrates.numchans > 1:
        backrates = backrates.integrate_energy()
    times = backrates.time_centroids
    rates = backrates.rates
    uncert = backrates.rate_uncertainty
    p2 = errorband(times, rates + uncert, rates - uncert, ax, alpha=err_alpha,
                   color=err_color, linestyle='-', **kwargs)
    p1 = ax.plot(times, rates, color=cent_color, alpha=cent_alpha,
                 **kwargs)
    refs = [p1, p2]
    return refs


def spectrum_background(backspec, ax, cent_color=None, err_color=None,
                        cent_alpha=None, err_alpha=None, **kwargs):
    """Plot a count spectrum background model with an error band
    
    Args:
        backspec (:class:`~gbm.background.BackgroundSpectrum`):
            The background rates object integrated over energy. If there is more
            than one remaining energy channel, the background will be integrated
            over the remaining energy channels.
        ax (:class:`matplotlib.axes`): The axis on which to plot
        cent_color (str): Color of the centroid line
        err_color (str): Color of the errorband
        cent_alpha (float): Alpha of the centroid line
        err_alpha (fl): Alpha of the errorband
        **kwargs: Other plotting options
    
    Returns:
        list: The reference to the lower and upper selection    
    """
    rates = backspec.rates / backspec.widths
    uncert = backspec.rate_uncertainty / backspec.widths
    edges = np.append(backspec.lo_edges, backspec.hi_edges[-1])
    # plot the centroid of the model
    p1 = ax.step(edges, np.append(rates, rates[-1]), where='post',
                 color=cent_color, alpha=cent_alpha, zorder=1, **kwargs)

    # construct the stepped errorband to fill between
    energies = np.array(
        (backspec.lo_edges, backspec.hi_edges)).T.flatten()  # .tolist()
    upper = np.array((rates + uncert, rates + uncert)).T.flatten()  # .tolist()
    lower = np.array((rates - uncert, rates - uncert)).T.flatten()  # .tolist()
    p2 = errorband(energies, upper, lower, ax, color=err_color,
                   alpha=err_alpha,
                   linestyle='-', **kwargs)
    refs = [p1, p2]
    return refs


# ---------- DRM ----------#
def response_matrix(phot_energies, chan_energies, matrix, ax, cmap='Greens',
                    num_contours=100, norm=None, **kwargs):
    """Make a filled contour plot of a response matrix
    
    Parameters:
    -----------
    phot_energies: np.array
        The incident photon energy bin centroids
    chan_energies: np.array
        The recorded energy channel centroids
    matrix: np.array
        The effective area matrix corresponding to the photon bin and 
        energy channels
    ax: matplotlib.axes
        The axis on which to plot
    cmap: str, optional
        The color map to use. Default is 'Greens'
    num_contours: int, optional
        The number of contours to draw.  These will be equally spaced in 
        log-space. Default is 100
    norm: matplotlib.colors.Normalize or similar, optional
        The normalization used to scale the colormapping to the heatmap values.
        This can be initialized by Normalize, LogNorm, SymLogNorm, PowerNorm,
        or some custom normalization.
    **kwargs: optional
        Other keyword arguments to be passed to matplotlib.pyplot.contourf
    
    Returns:
    -----------
    image:  matplotlib.collections.QuadMesh
        The reference to the plot object
    """
    mask = (matrix > 0.0)
    levels = np.logspace(np.log10(matrix[mask].min()),
                         np.log10(matrix.max()), num_contours)

    image = ax.contourf(phot_energies, chan_energies, matrix, levels=levels,
                        cmap=cmap, norm=norm)
    return image


def effective_area(bins, ax, color='C0', orientation='vertical', **kwargs):
    """Plot a histogram of the effective area of an instrument response
    
    Parameters:
    -----------
    bins: Bins
        The histogram of effective area
    ax: matplotlib.axes
        The axis on which to plot
    **kwargs: optional
        Other plotting options
    
    Returns:
    -----------
    refs: list
        The reference to the plot objects
    """
    bin_segs = bins.contiguous_bins()
    refs = []
    for seg in bin_segs:
        edges = np.concatenate(
            ([seg.lo_edges[0]], seg.lo_edges, [seg.lo_edges[-1]]))
        counts = np.concatenate(([0.0], seg.counts, [0.0]))
        if orientation == 'horizontal':
            p = ax.step(counts, edges, where='post', color=color, **kwargs)
        else:
            p = ax.step(edges, counts, where='post', color=color, **kwargs)
        refs.append(p)
    return refs


# ---------- Earth and Orbital ----------#
def saa_polygon(m, color='darkred', alpha=0.4, **kwargs):
    """Plot the SAA polygon on Basemap
    
    Parameters:
    -----------
    m: Basemap
        The basemap references
    color: str, optional
        The color of the polygon
    alpha: float, optional
        The alpha opacity of the interior of the polygon
    kwargs: optional
        Other plotting keywordss
    
    Returns:
    -----------
    poly: Polygon
        The SAA polygon object
    """
    # Define the SAA boundaries
    lat_saa, lon_saa = saa_boundary()

    edge = colorConverter.to_rgba(color, alpha=1.0)
    face = colorConverter.to_rgba(color, alpha=alpha)

    # plot the polygon
    x, y = m(lon_saa, lat_saa)
    xy = list(zip(x, y))
    poly = Polygon(xy, edgecolor=edge, facecolor=face, **kwargs)
    return poly


def mcilwain_map(lat_range, lon_range, m, ax, saa_mask=False, color=None,
                 alpha=0.5, **kwargs):
    """Plot the McIlwain L heatmap on a Basemap
    
    Parameters:
    -----------
    lat_range: (float, float)
        The latitude range
    lon_range: (float, float)
        The longitude range
    m: Basemap
        The basemap references
    ax: matplotlib.axes
        The plot axes references
    saa_mask: bool, optional
        If True, mask out the SAA from the heatmap.  Default is False.
    color: str, optional
        The color of the heatmap
    alpha: float, optional
        The alpha opacity of the heatmap
    kwargs: optional
        Other plotting keywords
    
    Returns:
    -----------
    image: QuadContourSet
        The heatmap plot object
    """
    # do create an array on the earth
    lat_array = np.linspace(*lat_range, 108)
    lon_array = np.linspace(*lon_range, 720)
    LAT, LON = np.meshgrid(lat_array, lon_array)
    # convert to projection coordinates
    mLON, mLAT = m(LON, LAT)
    # mcilwain l over the grid
    mcl = calc_mcilwain_l(LAT, LON)

    # if we want to mask out the SAA
    if saa_mask:
        saa_path = saa_polygon(m).get_path()
        mask = saa_path.contains_points(
            np.array((mLON.ravel(), mLAT.ravel())).T)
        mcl[mask.reshape(mcl.shape)] = 0.0

    # do the plot  
    levels = [0.9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7]
    image = ax.contourf(mLON, mLAT, mcl, levels=levels, alpha=alpha, **kwargs)
    return image


def earth_line(lat, lon, m, color='black', alpha=0.4, **kwargs):
    """Plot a line on the Earth (e.g. orbit)
    
    Parameters:
    -----------
    lat: np.array
        Array of latitudes
    lon: np.array
        Array of longitudes
    m: Basemap
        The basemap references
    color: str, optional
        The color of the lines
    alpha: float, optional
        The alpha opacity of line
    kwargs: optional
        Other plotting keywords
    
    Returns:
    -----------
    refs: list
        The list of line plot object references
    """
    lon[(lon > 180.0)] -= 360.0
    path = np.vstack((lon, lat))
    isplit = np.nonzero(np.abs(np.diff(path[0])) > 5.0)[0]
    segments = np.split(path, isplit + 1, axis=1)

    refs = []
    for segment in segments:
        x, y = m(segment[0], segment[1])
        refs.append(m.plot(x, y, color=color, alpha=alpha, **kwargs))
    return refs


def earth_points(lat, lon, m, color='black', alpha=1.0, **kwargs):
    """Plot a point or points on the Earth
    
    Parameters:
    -----------
    lat: np.array
        Array of latitudes
    lon: np.array
        Array of longitudes
    m: Basemap
        The basemap references
    color: str, optional
        The color of the lines
    alpha: float, optional
        The alpha opacity of line
    kwargs: optional
        Other plotting keywords
    
    Returns:
    -----------
    ref: 
        The scatter plot object reference
    """
    lon[(lon > 180.0)] -= 360.0
    x, y = m(lon, lat)
    if 's' not in kwargs.keys():
        kwargs['s'] = 1000
    ref = m.scatter(x, y, color=color, alpha=alpha, **kwargs)
    return [ref]


# ---------- Sky and Fermi Inertial Coordinates ----------#
def sky_point(x, y, ax, flipped=True, fermi=False, **kwargs):
    """Plot a point on the sky defined by the RA and Dec
        
    Inputs:
    -----------
    x: float
        The azimuthal coordinate (RA or Fermi azimuth), in degrees
    y: float
        The polar coordinate (Dec or Fermi zenith), in degrees
    ax: matplotlib.axes
        Plot axes object
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    fermi: bool, optional
        If True, plot in Fermi spacecraft coordinates, else plot in equatorial.
        Default is False.
    **kwargs: optional
        Other plotting options
    
    Returns:
    ----------
    point: matplotlib.collections.PathCollection
        The plot object
    """
    theta = np.array(np.deg2rad(y))
    phi = np.array(np.deg2rad(x - 180.0))
    if fermi:
        flipped = False
        theta = np.pi / 2.0 - theta
        phi -= np.pi
        phi[phi < -np.pi] += 2 * np.pi

    if flipped:
        phi = -phi
    point = ax.scatter(phi, theta, **kwargs)
    return point


def sky_line(x, y, ax, flipped=True, fermi=False, **kwargs):
    """Plot a line on a sky map, wrapping at the meridian
        
    Inputs:
    -----------
    x: np.array
        The azimuthal coordinates (RA or Fermi azimuth), in degrees
    y: np.array
        The polar coordinates (Dec or Fermi zenith), in degrees
    ax: matplotlib.axes
        Plot axes object
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    fermi: bool, optional
        If True, plot in Fermi spacecraft coordinates, else plot in equatorial.
        Default is False.
    **kwargs: optional
        Other plotting options
    """
    theta = np.deg2rad(y)
    phi = np.deg2rad(x - 180.0)
    if fermi:
        flipped = False
        theta = np.pi / 2.0 - theta
        phi -= np.pi
        phi[phi < -np.pi] += 2 * np.pi

    if flipped:
        phi = -phi
    seg = np.vstack((phi, theta))

    # here is where we split the segments at the meridian
    isplit = np.nonzero(np.abs(np.diff(seg[0])) > np.pi / 16.0)[0]
    subsegs = np.split(seg, isplit + 1, axis=1)

    # plot each path segment
    segrefs = []
    for seg in subsegs:
        ref = ax.plot(seg[0], seg[1], **kwargs)
        segrefs.append(ref)

    return segrefs


def sky_circle(center_x, center_y, radius, ax, flipped=True, fermi=False,
               face_color=None, face_alpha=None, edge_color=None,
               edge_alpha=None, **kwargs):
    """Plot a circle on the sky
    
    Inputs:
    -----------
    center_x: float
        The azimuthal center (RA or Fermi azimuth), in degrees
    center_y: float
        The polar center (Dec or Fermi zenith), in degrees
    radius: float
        The ROI radius in degrees
    color: str
        The color of the ROI
    ax: matplotlib.axes
        Plot axes object
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    fermi: bool, optional
        If True, plot in Fermi spacecraft coordinates, else plot in equatorial.
        Default is False.
    face_color: str, optional
        The color of the circle fill
    face_alpha: float, optional
        The alpha of the circle fill
    edge_color: str, optional
        The color of the circle edge
    edge_alpha: float, optional
        The alpha of the circle edge
    **kwargs: optional
        Other plotting options

    Returns:
    -----------
    patches: list of matplotlib.patches.Polygon
        The circle polygons
    """
    try:
        import mpl_toolkits.basemap
    except:
        raise ImportError('Cannot execute sky_circle due to missing Basemap.')
    theta = np.deg2rad(90.0 - center_y)
    phi = np.deg2rad(center_x)
    if fermi:
        flipped = False
        theta = np.pi / 2.0 - theta
        phi -= np.pi
        if phi < -np.pi:
            phi += 2 * np.pi

    rad = np.deg2rad(radius)

    # Call Leo's lalinference plotting helper functions
    # The native matplotlib functions don't cut and display the circle polygon
    # correctly on map projections
    roi = make_circle_poly(rad, theta, phi, 100)
    roi = cut_prime_meridian(roi)

    # plot each polygon section
    edge = colorConverter.to_rgba(edge_color, alpha=edge_alpha)
    face = colorConverter.to_rgba(face_color, alpha=face_alpha)
    patches = []
    for section in roi:
        section[:, 0] -= np.pi
        if flipped:
            section[:, 0] *= -1.0
        patch = ax.add_patch(
            plt.Polygon(section, facecolor=face, edgecolor=edge, \
                        **kwargs))
        patches.append(patch)
    return patches


def sky_annulus(center_x, center_y, radius, width, ax, color='black',
                alpha=0.3,
                fermi=False, flipped=True, **kwargs):
    """Plot an annulus on the sky defined by its center, radius, and width
    
    Inputs:
    -----------
    center_x: float
        The azimuthal center (RA or Fermi azimuth), in degrees
    center_y: float
        The polar center (Dec or Fermi zenith), in degrees
    radius: float
        The radius in degrees, defined as the angular distance from the center
        to the middle of the width of the annulus band
    width: float
        The width of the annulus in degrees
    ax: matplotlib.axes
        Plot axes object
    color: string, optional
        The color of the annulus. Default is black.
    alpha: float, optional
        The opacity of the annulus. Default is 0.3
    fermi: bool, optional
        If True, plot in Fermi spacecraft coordinates, else plot in equatorial.
        Default is False.
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    **kwargs: optional
        Other plotting options
    """
    try:
        import mpl_toolkits.basemap
    except:
        raise ImportError('Cannot execute sky_annulus due to missing Basemap.')
    edge = colorConverter.to_rgba(color, alpha=1.0)
    face = colorConverter.to_rgba(color, alpha=alpha)

    inner_radius = np.deg2rad(radius - width / 2.0)
    outer_radius = np.deg2rad(radius + width / 2.0)
    center_theta = np.deg2rad(90.0 - center_y)
    center_phi = np.deg2rad(center_x)
    if fermi:
        flipped = False
        center_theta = np.pi / 2.0 - center_theta
        center_phi = np.deg2rad(center_x + 180.0)

    # get the plot points for the inner and outer circles
    inner = make_circle_poly(inner_radius, center_theta, center_phi, 100)
    inner = cut_prime_meridian(inner)
    outer = make_circle_poly(outer_radius, center_theta, center_phi, 100)
    outer = cut_prime_meridian(outer)

    x1 = []
    y1 = []
    x2 = []
    y2 = []
    polys = []
    # plot the inner circle
    for section in inner:
        section[:, 0] -= np.pi
        if flipped:
            section[:, 0] *= -1.0
        polys.append(plt.Polygon(section, ec=edge, fill=False, **kwargs))
        ax.add_patch(polys[-1])
        x1.extend(section[:, 0])
        y1.extend(section[:, 1])

    # plot the outer circle
    for section in outer:
        section[:, 0] -= np.pi
        if flipped:
            section[:, 0] *= -1.0
        polys.append(plt.Polygon(section, ec=edge, fill=False, **kwargs))
        ax.add_patch(polys[-1])
        x2.extend(section[:, 0])
        y2.extend(section[:, 1])

    # organize and plot the fill between the circles
    # organize and plot the fill between the circles
    x1.append(x1[0])
    y1.append(y1[0])
    x2.append(x2[0])
    y2.append(y2[0])
    x2 = x2[::-1]
    y2 = y2[::-1]
    xs = np.concatenate((x1, x2))
    ys = np.concatenate((y1, y2))
    f = ax.fill(np.ravel(xs), np.ravel(ys), facecolor=face, zorder=1000)

    refs = polys
    refs.append(f)
    return refs


def sky_polygon(x, y, ax, face_color=None, edge_color=None, edge_alpha=1.0,
                face_alpha=0.3, flipped=True, fermi=False, **kwargs):
    """Plot single polygon on a sky map, wrapping at the meridian
        
    Inputs:
    -----------
    paths: matplotlib.path.Path
        Object containing the contour path
    ax: matplotlib.axes
        Plot axes object
    face_color: str, optional
        The color of the polygon fill
    face_alpha: float, optional
        The alpha of the polygon fill
    edge_color: str, optional
        The color of the polygon edge
    edge_alpha: float, optional
        The alpha of the polygon edge
    fermi: bool, optional
        If True, plot in Fermi spacecraft coordinates, else plot in equatorial.
        Default is False.
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    **kwargs: optional
        Other plotting options
    """
    line = sky_line(x, y, ax, color=edge_color, alpha=edge_alpha,
                    flipped=flipped,
                    fermi=fermi, **kwargs)
    refs = line

    theta = np.deg2rad(y)
    phi = np.deg2rad(x - 180.0)
    if fermi:
        flipped = False
        theta = np.pi / 2.0 - theta
        phi -= np.pi
        phi[phi < -np.pi] += 2 * np.pi

    if flipped:
        phi = -phi

    # this is needed to determine when contour spans the full ra range
    if (np.min(phi) == -np.pi) and (np.max(phi) == np.pi):
        if np.mean(theta) > 0.0:  # fill in positive dec
            theta = np.insert(theta, 0, np.pi / 2.0)
            theta = np.append(theta, np.pi / 2.0)
            phi = np.insert(phi, 0, -np.pi)
            phi = np.append(phi, np.pi)
        else:  # fill in negative dec
            theta = np.insert(theta, 0, -np.pi / 2.0)
            theta = np.append(theta, -np.pi / 2.0)
            phi = np.insert(phi, 0, np.pi)
            phi = np.append(phi, -np.pi)

    f = ax.fill(phi, theta, color=face_color, alpha=face_alpha, **kwargs)
    refs.append(f)

    return refs


def galactic_plane(ax, flipped=True, fermi_quat=None, outer_color='dimgray',
                   inner_color='black', line_alpha=0.5, center_alpha=0.75,
                   **kwargs):
    """Plot the galactic plane on the sky
        
    Inputs:
    -----------
    ax: matplotlib.axes
        Plot axes object
    outer_color: str, optional
        The color of the outer line
    inner_color: str, optional
        The color of the inner line
    line_alpha: float, optional
        The alpha of the line
    center_alpha: float, optional
        The alpha of the center
    fermi_quat: np.array, optional
        If set, rotate the galactic plane into the Fermi frame
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    **kwargs: optional
        Other plotting options
    """
    fermi = False
    # if a quaternion is sent, then plot in the Fermi frame
    if fermi_quat is not None:
        flipped = False
        fermi = True

    # longitude and latitude arrays
    lon_array = np.arange(0, 360, dtype=float)
    lat = np.zeros_like(lon_array)
    gc = SkyCoord(l=lon_array * u.degree, b=lat * u.degree, frame='galactic')
    ra, dec = gc.icrs.ra.deg, gc.icrs.dec.deg
    # plot in Fermi frame
    if fermi_quat is not None:
        ra, dec = radec_to_spacecraft(ra, dec, fermi_quat)

    line1 = sky_line(ra, dec, ax, color=outer_color, linewidth=3,
                     alpha=line_alpha, flipped=flipped, fermi=fermi)
    line2 = sky_line(ra, dec, ax, color=inner_color, linewidth=1,
                     alpha=line_alpha, flipped=flipped, fermi=fermi)

    # plot Galactic center
    pt1 = sky_point(ra[0], dec[0], ax, marker='o', c=outer_color, s=100,
                    alpha=center_alpha, edgecolor=None, flipped=flipped,
                    fermi=fermi)
    pt2 = sky_point(ra[0], dec[0], ax, marker='o', c=inner_color, s=20,
                    alpha=center_alpha, edgecolor=None, flipped=flipped,
                    fermi=fermi)

    return [line1, line2, pt1, pt2]


def sky_heatmap(x, y, heatmap, ax, cmap='RdPu', norm=None, flipped=True,
                fermi=False, **kwargs):
    """Plot a heatmap on the sky as a colormap gradient
       
    Inputs:
    -----------
    x: np.array
        The azimuthal coordinate array of the heatmap grid
    y: np.array
        The polar coordinate array of the heatmap grid
    heatmap: np.array
        The heatmap array, of shape (num_x, num_y)
    ax: matplotlib.axes
        Plot axes object
    cmap: str, optional
        The colormap. Default is 'RdPu'
    norm: matplotlib.colors.Normalize or similar, optional
        The normalization used to scale the colormapping to the heatmap values.
        This can be initialized by Normalize, LogNorm, SymLogNorm, PowerNorm,
        or some custom normalization.  Default is PowerNorm(gamma=0.3).
    flipped: bool, optional
        If True, the azimuthal axis is flipped, following equatorial convention
    fermi: bool, optional
        If True, plot in Fermi spacecraft coordinates, else plot in equatorial.
        Default is False.
    **kwargs: optional
        Other plotting options
    
    Returns:
    --------
    image: matplotlib.collections.QuadMesh
        The reference to the heatmap plot object
    """
    theta = np.deg2rad(y)
    phi = np.deg2rad(x - 180.0)
    if fermi:
        flipped = False
        theta = np.pi / 2.0 - theta
        phi -= np.pi
        phi[phi < -np.pi] += 2 * np.pi
        shift = int(phi.shape[0] / 2.0)
        phi = np.roll(phi, shift)
        heatmap = np.roll(heatmap, shift, axis=1)

    if flipped:
        phi = -phi

    image = ax.pcolormesh(phi, theta, heatmap, rasterized=True, cmap=cmap,
                          norm=norm)
    return image