#!/usr/bin/env python
# coding: utf-8

# # AAS233 WWT Workshop Notebook (January 2019; Seattle, WA)
# 
# *Please note: these notebooks are preserved for posterity but not kept up-to-date, so you might run into issues with older notebooks.*

# ## Import and initialize statements from 'Get started'
# 
# *(Select a cell and press `ctrl` + `Enter` to run it.)*

# In[ ]:


import time

from astropy.coordinates import SkyCoord
from astropy import units as u
from astropy.table import Table, Column
from astropy.time import Time, TimeDelta # note capitalization

import numpy as np

from astroquery.vizier import Vizier
from astroquery.nasa_exoplanet_archive import NasaExoplanetArchive

from pywwt.jupyter import WWTJupyterWidget


# In[ ]:


wwt = WWTJupyterWidget()
wwt


# *WWT works under HTTP, while Binder uses HTTPS. As such, in order to access all features of pywwt, Binder users will need to disable HTTPS for this session. Firefox users can do this with the following steps:*
# - *Click the lock icon to the left of the address bar,*
# - *Select the arrow next to “Connection,”*
# - *Then, press the “Disable protection for now” button. (Protection is disabled for this page and only during this session, not for your whole browser. We are working on upgrading WWT so this step is no longer necessary.)*

# ## Plotting K2’s footprint and exoplanet yield

# **First, we use a class from the `astroquery` package to query NASA's Exoplanet Archive and return an `astropy` table that contains information about every confirmed exoplanet so far. It's a large table, so we'll only keep the columns necessary to run this example.**
# 
# *(If the `astroquery` call doesn't load, uncomment and run the following cell to download the table from a `pywwt`-affiliated GitHub repository.)*

# In[ ]:


# xarch = NasaExoplanetArchive.get_confirmed_planets_table()
# xarch.keep_columns(['pl_hostname', 'pl_letter', 'st_dist', 'ra', 'dec'])


# In[ ]:


xarch = Table.read('../data/k2_table.ecsv', format='ascii.ecsv')


# *(Open a new cell above this one and see what the table looks like by entering `xarch` and running it.)*

# **Next, we slice the table twice so it only contains entries of 'b' planets observed by K2.** 
# 
# K2 planets and stars are named using a consistent system: 
# 
# ```K2 + 'system number' + 'letter'```. 
# 
# The letter for a star is typically 'a', the first exoplanet confirmed in the system is 'b', and so on.
# 
# We are plotting stars, but since the table only contains exoplanets (letters from 'b' onward), we will slice the table so only planets with the `b` suffix remain. That gives us one entry per system.
# 
# **`k2_pl` searches the `xarch` column that lists star names and only keeps those with K2 at the beginning. Then, `k2_st` searches the suffix of the planet name and only keeps `b` planets.**
# 
# *(If you aren't familiar with list comprehensions, see [Note 1](aas233_tutorial.ipynb#Note-1).)*

# In[ ]:


# select only k2 planets
k2_pl = xarch[ [row.startswith('K2') for row in xarch['pl_hostname']] ]

# select only 'b' planets from k2
k2_st = k2_pl[ ['b' in row for row in k2_pl['pl_letter']] ]


# **After that, we load the footprint of the K2 campaign fields on the sky using pywwt's `add_fov` method.**
# 
# *(Loading K2's footprint usually takes a minute, so please be patient with this step.)*

# In[ ]:


field = wwt.add_fov(wwt.instruments.k2)


# **Move the viewer down here and watch what happens next.**

# **Once the footprint has loaded, we finally create the data layer that will plot our points in the viewer.**
# 
# `pywwt` reads the locations from the keyword arguments `lon_att` and `lat_att`, and we feed it the names of the columns in `k2_st` that contain the stellar right ascensions and declinations. We can also set other attributes like color and size in the function call.

# In[ ]:


lay = wwt.layers.add_table_layer(table=k2_st, frame='Sky',
                                 lon_att='ra', lat_att='dec',
                                 color='#c4d600', size_scale=40)


# Do you see the points? If not, try making them larger.

# In[ ]:


lay.size_scale *= 2


# If you want to see the points and footprint without a sky background, `pywwt` also has a black layer available.

# In[ ]:


wwt.foreground_opacity = 0
wwt.background = wwt.imagery.other.black


# **It is also possible to visualize data layers in `pywwt`'s 3D solar system mode. To do so, we'll switch the view.**
# 
# We will also remove the K2 footprint, as `add_fov` is designed to work in 2D.

# In[ ]:


field.remove()
wwt.set_view('solar system')


# **Bring the viewer down here, zoom out, and you'll see K2's exoplanet-hosting stars in 3D. All of them are currently the same distance from Earth -- let's fix that.**
# 
# Data layers also have attributes related to distance and altitude. We set `alt_type` to the correct type for our scenario, and let `pywwt` know which column of the table contains distance with `alt_att`.

# In[ ]:


lay.alt_type = 'distance'
lay.alt_att = 'st_dist'


# **Zoom out again and see the stars at their proper distances.**

# ## Plotting and watching gamma ray bursts (GRBs)

# **In this example (adapted from [a notebook](https://nbviewer.jupyter.org/github/marksubbarao/pyWWT_AAS225/blob/master/Fermi%20Gamma%20Ray%20Bursts%20.ipynb) by Mark SubbaRao), we use the Fermi Gamma-ray Burst Monitor's GRB catalog to get information about GRB events from 2008 to 2014.**
# 
# The goal here is not just to plot the locations of these events, but also to use `pywwt`'s time controls to bring this time series data set to life. We will loop through time, refresh the data layer as the GRBs "occur," and watch as they pop up in the viewer.

# **First, we'll retrieve the data from a `pywwt`-affiliated repository on GitHub.**
# 
# *(This table, like the exoplanet one, can also be also be loaded via `astroquery`, but the process of manipulating it to fit our needs takes more time than we have in the tutorial. If you're curious about it, see [Note 2](aas233_tutorial.ipynb#Note-2).)*

# In[ ]:


bursts = Table.read('../data/grb_table.ecsv', format='ascii.ecsv')


# *(Again, open a new cell above this one and see what the table looks like by entering `bursts` and running it.)*

# **Let's refresh the viewer by re-executing the `wwt = WWTJupyterWidget()` cell. After that, since we're working with gamma ray bursts, let's switch to a background layer in that region of the electromagnetic spectrum.**

# In[ ]:


wwt.foreground_opacity = 0
wwt.background = wwt.imagery.gamma.fermi


# **Now, let's create the data layer.**
# 
# This time, notice that we assign the 'cmap' column to be tracked by `pywwt`. This is because we will have the points change color both during and after the GRB events occur.

# In[ ]:


lay = wwt.layers.add_table_layer(table=bursts, frame='Sky',
                                 lon_att='ra', lat_att='dec',
                                 size_scale=90, cmap_att='cmap')


# **We need to define some variables before running the time loop.**
# - *`start`* and *`finish`* are `astropy` Time objects and define the beginning and end of the loop through time. They're set to be the times just before the first and after the last GRB events, respectively.
# - *`step`* is an `astropy` TimeDelta object and can be added to a Time object to create a new time. We will "step" the loop forward in intervals of 10 days. It needs a unit from `astropy`'s Units object, too.
# - *`now`* is the Time object that will be added to by `step` in each iteration of the loop. It represents the current time.

# In[ ]:


start = Time(bursts[0]['iso_times'], format='jd')
finish = Time(bursts[-1]['iso_times'], format='jd')
#start = Time('2008-07-11', format='iso')
#finish = Time('2008-08-01', format='iso')

step = TimeDelta(10 * u.day)
now = start + step


# **Next, let's bring the viewer down here and center it on the earliest GRB in the data set.**
# 
# That way, we'll see colors begin to change immediately.

# In[ ]:


wwt.center_on_coordinates(SkyCoord(bursts[0]['ra'], bursts[0]['dec'], 
                                   unit=u.deg))


# **Finally, it's time to write and run the loop.**
# 
# In each iteration, here is the strategy:
# - Set the time in the viewer to `now`;
# - `active`: Compare `now` to the values in the 'iso_times' column of `bursts` to find which entries had an event in the period between the previous iteration and now ([Note 3.1](aas233_tutorial.ipynb#Note-3));
# - `last_active`: Do the same to find which entries had an event in the `step`-day period before the current period ([Note 3.2](aas233_tutorial.ipynb#Note-3));
# - Turn `burst` entries with events in that period white (active);
# - Turn `burst` entries occuring just before that period green (passed);
# - Update the layer in `pywwt` with the new version of `burst` where the colors are changed;
# - Use the `sleep` method of the `time` package to pause the loop so the data layer has time to finish updating;
# - Add `step` to `now` and repeat until we've reached the time of `finish`.

# In[ ]:


while now <= finish:
    wwt.set_current_time(now)
        
    active = np.intersect1d(
        np.where(bursts['iso_times'] > (now - step).jd),
        np.where(bursts['iso_times'] <= now.jd))
    
    passed = np.intersect1d(
        np.where(bursts['iso_times'] > (now - 2 * step).jd),
        np.where(bursts['iso_times'] <= (now - step).jd))
    
    if len(active) > 0:
        bursts[active[0] : active[-1] + 1]['cmap'] = '#ffffff'
    if len(passed) > 0:
        bursts[passed[0] : passed[-1] + 1]['cmap'] = '#00ff00'
    
    lay.update_data(bursts)
    
    #print(now)
    time.sleep(1)
    now += step


# ***Is it working? If so, congratulations, enjoy the view, and stay in touch! Find us at the AAS booth, work with us all day tomorrow at Hack Together Day, and contribute ideas, suggestions, and fixes at our [GitHub repository](https://github.com/WorldWideTelescope/pywwt).***

# #### Note 1
# *The constructions used to build `k2_pl` and `k2_st` are called "list comprehensions," and they are shortcuts for building lists. The regular, equivalent method to create `k2_st`, looks like this:*
# 
# ```
# index_list = []
# for i, row in enumerate(xarch, 0):
#     if row['pl_hostname'].startswith('K2'):
#         index_list.append(i)
# 
# k2_pl = xarch[index_list]
# ```

# #### Note 2
# 
# *For my own reference as well as yours, this is how I used astroquery to build the table of GRB events.*
# 
# ```
# # retrieve the list of GRBs
# v = Vizier()
# v.ROW_LIMIT = -1
# cats = v.get_catalogs('J/ApJS/223/28N/GBM')
# 
# # take the table and cut unnecessary columns
# bursts = cats[0]
# bursts.keep_columns(["GRB", "RAJ2000", "DEJ2000", "Time",
#                      "Trig", "Fl.w", "Fl.n"])
# 
# # rename some columns
# bursts.rename_column('RAJ2000','ra')
# bursts.rename_column('DEJ2000','dec')
# bursts.rename_column('Trig', 'obs_time')
# 
# # only keep entries with nonzero 'GRB' values 
# # (we need that for the loop)
# bursts = bursts[bursts["GRB"].nonzero()[0]]
# 
# # add a column with times readable by pywwt
# grb_times = []
# for i, name in enumerate(bursts['GRB']):
#     # name format: 080714B
#     # year (2 digits), month (2 digits), day (2 digits), letter
#     # style is changed to '2008-07-14 02:04:12.053', 
#     # ...then converted to julian date
# 
#     formatted = Time('20' + name[:2] + '-' 
#                  + name[2:4] + '-' + name[4:6] + ' ' 
#                  + bursts['obs_time'][i], format='iso').jd
#     grb_times.append(formatted)
# 
# bursts.add_column(Column(grb_times), name='iso_times')
# 
# # add a colormap column to the table
# cmap = ['#000000'] * len(bursts)
# bursts.add_column(Column(cmap), name='cmap')
# ```

# #### Note 3
# 1. i.e. between `now - step` and `now`. e.g. between July 11 and July 21 if `now == Time('2008-07-21')` and `step == TimeDelta(10 * u.day)`;
# 2. Given the variables in 1, the period here would be between between July 1 and July 11.

# ## Credits
# 
# This notebook was prepared by O. Justin Otor, with contributions from Thomas Robitaille and Peter K. G. Williams.