# Basic notebook imports
%matplotlib inline

import matplotlib
import pylab as plt
import numpy as np
import healpy as hp

"""
Explanation: <center>
   Go back to the Index
</center>
Chapter 3: Skymap and HEALPix
In this chapter we introduce the interface between skymap and HEALPix (Gorski et al. 2005). The ability to accurately draw HEALPix maps with minimal overhead is one of the most powerful (and useful) feature of skymap.
End of explanation
"""


import skymap
from skymap import Skymap,McBrydeSkymap,OrthoSkymap
from skymap import SurveySkymap,SurveyMcBryde,SurveyOrtho
from skymap import DESSkymap,BlissSkymap

SKYMAPS = [Skymap,McBrydeSkymap,OrthoSkymap]
SURVEYS = [SurveySkymap,SurveyMcBryde,SurveyOrtho]

NSIDE = 8

"""
Explanation: We start by importing several classes from the skymap module and setting a few constants that we will use in this example. By this point, these modules and classes should be familiar.
End of explanation
"""


hpxmap = np.arange(hp.nside2npix(NSIDE))
fig,axes = plt.subplots(1,3,figsize=(20,4))
for i,cls in enumerate(SKYMAPS):
    plt.sca(axes[i])
    m = cls()
    m.draw_hpxmap(hpxmap,xsize=200)
    plt.title('HEALPix Map (%s)'%cls.__name__)

"""
Explanation: We start off by creating and drawing a simple, low-resolution HEALPix map in ring format.
End of explanation
"""


hpxmap = np.arange(hp.nside2npix(NSIDE))
fig,axes = plt.subplots(1,3,figsize=(20,4))
for i,proj in enumerate(['sinu','kav7','vandg']):
    plt.sca(axes[i])
    smap = Skymap(projection=proj,lon_0=0)
    im,lon,lat,values = smap.draw_hpxmap(hpxmap,xsize=200)
    plt.title('HEALPix Map (%s)'%proj)

"""
Explanation: This is nothing special compared to what you can do with healpy.visfunc, but remember that skymap gives you access to all the power of basemap.
End of explanation
"""


plt.figure()
plt.pcolormesh(lon,lat,values)

"""
Explanation: The draw_hpxmap function returns the image, longitude and latitude of each pixel, and the values of each pixel. You can take that output and draw it yourself if you like.
End of explanation
"""


pix = hpxmap = np.arange(525,535)
fig,axes = plt.subplots(1,3,figsize=(20,4))
for i,cls in enumerate(SKYMAPS):
    plt.sca(axes[i])
    m = cls()
    m.draw_hpxmap(hpxmap,pix,NSIDE,xsize=200)
    plt.title('Partial HEALPix Map (%s)'%cls.__name__)

"""
Explanation: The draw_hpxmap function can also take a partial HEALPix map where the pixels and values are specified explicitly.
End of explanation
"""


pix = np.arange(525,535)

fig,axes = plt.subplots(1,3,figsize=(20,4))
plt.sca(axes[0])
smap = McBrydeSkymap()
hpxmap = hp.UNSEEN * np.ones(hp.nside2npix(NSIDE))
hpxmap[pix] = pix
smap.draw_hpxmap(hpxmap,xsize=400)

plt.sca(axes[1])
smap = McBrydeSkymap()
hpxmap = np.nan * np.ones(hp.nside2npix(NSIDE))
hpxmap[pix] = pix
smap.draw_hpxmap(hpxmap,xsize=400)

plt.sca(axes[2])
smap = McBrydeSkymap()
hpxmap = np.arange(hp.nside2npix(NSIDE))
hpxmap = np.ma.array(hpxmap, mask=~np.in1d(hpxmap,pix))
out = smap.draw_hpxmap(hpxmap,xsize=200)

"""
Explanation: Notice that in the above maps all pixels that are not explicitly specified are masked by draw_hpxmap. The same behavior can be accomplished by setting HEALPix values to np.NaN, hp.UNSEEN or by using the numpy.MaskedArray module.
End of explanation
"""


# These are random, non-uniform points
size = int(1e5)
lon = np.random.uniform(0,360,size=size)
lat = np.random.uniform(-90,90,size=size)

fig,axes = plt.subplots(1,3,figsize=(20,4))
for i,cls in enumerate(SKYMAPS):
    plt.sca(axes[i])
    smap = cls()                                                  
    hpxmap,(im,lon,lat,values) = smap.draw_hpxbin(lon,lat,nside=16,xsize=200)
    plt.title('HEALPix Binning (%s)'%cls.__name__)

"""
Explanation: We have also implemented a HEALPix analog to the matplotlib.hexbin function called hpxbin (this function returns the healpix map in addition to image and associated pixel values).
End of explanation
"""


import skymap.healpix

ra,dec = 45,-45
fig,axes = plt.subplots(1,3,figsize=(15,4))

for i,nside in enumerate([512,4096,4096*2**5]):
    radius = np.degrees(50*hp.nside2resol(nside))
    pixels = skymap.healpix.ang2disc(nside,ra,dec,radius)
    values = pixels
    
    plt.sca(axes[i])
    # Use the Cassini projection (because we can)                                
    m = Skymap(projection='cass', lon_0=ra, lat_0=dec, celestial=False,
               llcrnrlon=ra+2*radius,urcrnrlon=ra-1.5*radius,
               llcrnrlat=dec-1.2*radius,urcrnrlat=dec+1.5*radius,
               parallels=False, meridians=False)

    m.draw_hpxmap(values,pixels,nside=nside,xsize=200)
    m.draw_parallels(np.linspace(dec-2*radius,dec+2*radius,5),
                     labelstyle='+/-',labels=[1,0,0,0],fmt = '%.2f')
    m.draw_meridians(np.linspace(ra-2*radius,ra+2*radius,5),
                     labelstyle='+/-',labels=[0,0,0,1],fmt = '%.2f')

    plt.title('HEALPix Zoom (nside=%i)'%nside)

"""
Explanation: Because skymap can draw partial HEALPix maps with explicit pixel labeling, it is possible to draw small regions of the sky much more efficiently the healpy. The true power comes in the ability to plot ultra-high resolution HEALPix maps (nside > 1e5). This would take an enormous amount of memory if you needed to create the entire map. In the plots below notice the change in the size of the axes ranges.
End of explanation
"""


