r''' Example scipt to plot a LIGO skymap with neutrinos overlayed.
Will also plot angular error contours and confidence level containing
LIGO PDF.

Author: Raamis Hussain
Date  : 05/21/19
'''

import numpy  as np
import healpy as hp
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import meander #pip install this to draw contours around skymap

from matplotlib import cm

##############################  Some plotting functions we will use ##################################
def compute_ang_err(ra, dec, sigma):
    r''' Compute circular uncertainty contour around point
    on a healpix map.
    Parameters:
    -----------
    ra: float
        Right ascension of center point
    dec: float
        declination of center point
    sigma:
        Angular uncertainty (radius of circle to draw
        around (ra,dec))
    Returns:
    --------
    Theta: array
        theta values of contour
    Phi: array
        phi values of contour
    '''

    dec = np.pi/2 - dec
    sigma = np.rad2deg(sigma)
    delta, step, bins = 0, 0, 0
    delta= sigma/180.0*np.pi
    step = 1./np.sin(delta)/20.
    bins = int(360./step)
    Theta = np.zeros(bins+1, dtype=np.double)
    Phi = np.zeros(bins+1, dtype=np.double)
    # define the contour
    for j in range(0,bins) :
            phi = j*step/180.*np.pi
            vx = np.cos(phi)*np.sin(ra)*np.sin(delta) + np.cos(ra)*(np.cos(delta)*np.sin(dec) + np.cos(dec)*np.sin(delta)*np.sin(phi))
            vy = np.cos(delta)*np.sin(dec)*np.sin(ra) + np.sin(delta)*(-np.cos(ra)*np.cos(phi) + np.cos(dec)*np.sin(ra)*np.sin(phi))
            vz = np.cos(dec)*np.cos(delta) - np.sin(dec)*np.sin(delta)*np.sin(phi)
            idx = hp.vec2pix(nside, vx, vy, vz)
            DEC, RA = hp.pix2ang(nside, idx)
            Theta[j] = DEC
            Phi[j] = RA
    Theta[bins] = Theta[0]
    Phi[bins] = Phi[0]

    return Theta, Phi

def compute_contours(proportions,samples):
    r''' Plot containment contour around desired level.
    E.g 90% containment of a PDF on a healpix map

    Parameters:
    -----------
    proportions: list
        list of containment level to make contours for.
        E.g [0.68,0.9]
    samples: array
        array of values read in from healpix map
        E.g samples = hp.read_map(file)
    Returns:
    --------
    theta_list: list
        List of arrays containing theta values for desired contours
    phi_list: list
        List of arrays containing phi values for desired contours
    '''

    levels = []
    sorted_samples = list(reversed(list(sorted(samples))))
    nside = hp.pixelfunc.get_nside(samples)
    sample_points = np.array(hp.pix2ang(nside,np.arange(len(samples)))).T
    for proportion in proportions:
        level_index = (np.cumsum(sorted_samples) > proportion).tolist().index(True)
        level = (sorted_samples[level_index] + (sorted_samples[level_index+1] if level_index+1 < len(samples) else 0)) / 2.0
        levels.append(level)
    contours_by_level = meander.spherical_contours(sample_points, samples, levels)

    theta_list = []; phi_list=[]
    for contours in contours_by_level:
        for contour in contours:
            theta, phi = contour.T
            phi[phi<0] += 2.0*np.pi
            theta_list.append(theta)
            phi_list.append(phi)

    return theta_list, phi_list
######################################################################################################

# Fits file containing Gravitational wave skymap
fitsFile = 'https://gracedb.ligo.org/api/superevents/S190521r/files/bayestar.fits.gz'

### Read map and get probabilities
probs = hp.read_map(fitsFile)
nside = hp.pixelfunc.get_nside(probs)

# Choose color map and set background to white
cmap = cm.YlOrRd
cmap.set_under("w")

# Plot GW skymap in Mollweide projection
hp.mollview(probs,cbar=True,unit=r'Probability',min=0,max=3e-5,rot=180,cmap=cmap)
hp.graticule() # Set grid lines

# Let's generate 5 random neutrinos on the sky and plot them
theta= np.random.uniform(0,np.pi,size=5)
phi = np.random.uniform(0,2*np.pi,size=5)
hp.projscatter(theta,phi,c='b',marker='x',label='GFU Event')

# Now let's assign some angular error to each event. Generally these
# neutrinos have ~1 deg angular error (39% containment) but we will make it
# larger here for visual effect
sigma = np.random.uniform(np.radians(2),np.radians(5),size=5)

#plot events on sky with error contours
for i in range(theta.size):
    my_contour = compute_ang_err(phi[i],np.pi/2-theta[i],sigma[i])
    hp.projplot(my_contour[0], my_contour[1], linewidth=1.75, color="gray",
                linestyle='solid',coord='C')

# Draw containment contour around GW skymap
probs = hp.pixelfunc.ud_grade(probs, 64) #reduce nside to make it faster
probs = probs/np.sum(probs)
pixels = np.arange(probs.size)
sample_points = np.array(hp.pix2ang(nside,pixels)).T

### plot 90% containment contour of GW PDF
levels = [0.90]
theta_contour, phi_contour = compute_contours(levels,probs)
hp.projplot(theta_contour[0],phi_contour[0],linewidth=1,c='k',label='GW (90% C.L)')
for i in range(1,len(theta_contour)):
    hp.projplot(theta_contour[i],phi_contour[i],linewidth=1,c='k')

#Add labels to plot
plt.text(2.0,0., r"$0^\circ$", ha="left", va="center")
plt.text(1.9,0.45, r"$30^\circ$", ha="left", va="center")
plt.text(1.4,0.8, r"$60^\circ$", ha="left", va="center")
plt.text(1.9,-0.45, r"$-30^\circ$", ha="left", va="center")
plt.text(1.4,-0.8, r"$-60^\circ$", ha="left", va="center")
plt.text(1.333, -0.15, r"$60^\circ$", ha="center", va="center")
plt.text(.666, -0.15, r"$120^\circ$", ha="center", va="center")
plt.text(0.0, -0.15, r"$180^\circ$", ha="center", va="center")
plt.text(-.666, -0.15, r"$240^\circ$", ha="center", va="center")
plt.text(-1.333, -0.15, r"$300^\circ$", ha="center", va="center")
plt.text(-2.0, -0.15, r"$360^\circ$", ha="center", va="center")

plt.title('GW Skymap',fontsize=15)
plt.legend(loc=1,bbox_to_anchor=(1.08,1.09),prop={'size': 14})
plt.savefig('GW_skymap.png',bbox_inches='tight')
