import numpy as np

from matplotlib import pyplot as plt
%matplotlib inline

"""
Explanation: Contents
This notebook covers the basics of creating an event list object and carrying out various operations such as simulating time and energies, joining, storing and retrieving event lists.
Setup
Import some useful libraries.
End of explanation
"""


from stingray import EventList, Lightcurve

"""
Explanation: Import some relevant stingray classes.
End of explanation
"""


times = [0.5, 1.1, 2.2, 3.7]
mjdref=58000.

"""
Explanation: Creating EventList from Photon Arrival Times
Given photon arrival times, an eventlist object can be created. Times are assumed to be seconds from a reference MJD, that can optionally be specified with the mjdref keyword and attribute.
End of explanation
"""


ev = EventList(times, mjdref=mjdref)
ev.time

"""
Explanation: Create event list object by passing arrival times as argument.
End of explanation
"""


energy = [0., 3., 4., 20.]
ev = EventList(times, energy=energy, mjdref=mjdref)

"""
Explanation: One can add all sorts of data to the EventList object, it is very flexible. In general, we suggest to stick with easily interpretable attributes, like energy or pi.
ev.energy = [0., 3., 4., 20.]
is the same as
End of explanation
"""


gti = [[0, 4]]
ev = EventList(times, gti=gti, energy=energy, mjdref=mjdref)

"""
Explanation: It is always recommended to specify the good time intervals (GTIs) of the event list, as the time intervals where the instrument was fully operational. If not specified, GTIs are defined automatically as the first and the last event time.
End of explanation
"""


table = ev.to_astropy_table()
table

"""
Explanation: Roundtrip to Astropy-compatible formats
EventList has the following methods that allow an easy roundtrip to Astropy objects: to_astropy_table, to_astropy_timeseries, from_astropy_table, from_astropy_timeseries
This allows a better interoperability with the Astropy ecosystem. 
In this roundtrip, a Table or Timeseries object is created, having as columns time and all other attributes of the same size (e.g. pi, energy), and the rest of the attributes (e.g. gti, mjdref) in the table's metadata.
End of explanation
"""


timeseries = ev.to_astropy_timeseries()
timeseries

table.meta

"""
Explanation: When converting to Timeseries, times are transformed into astropy.time.TimeDelta objects.
End of explanation
"""


table_ev = EventList.from_astropy_table(table)
table_ts = EventList.from_astropy_timeseries(timeseries)


table_ev.time, table_ts.time

"""
Explanation: Of course, these objects can be converted back to event lists. The user should be careful in defining the proper column names and metadata so that the final object is a valid EventList
End of explanation
"""


ev = EventList.read('events.fits', 'ogip')

"""
Explanation: Loading and writing EventList objects
We made it possible to save and load data in a number of different formats.
The general syntax is
```
ev = EventList.read(filename, format)
ev.write(filename, format)
```
There are three main blocks of formats that might be useful:


(read-only) HEASoft-compatible formats -> read event data from HEASOFT-supported missions


pickle: reading and saving EventLists from/to Python pickle objects


Any format compatible with astropy.table.Table objects.


Loading an EventList from an X-ray observation in HEASoft-compatible format
Loading event data from HEASoft-supported missions in FITS format is easy. It's sufficient to use the read method with hea or, equivalently, ogip, as format. 
Beware: please use hea or ogip, not fits! It would make the roundrip to Astropy tables more complicated, as Astropy supports a generic FITS writer which is not necessarily compatible with HEASoft.
End of explanation
"""


ev.time[:10]

ev.mjdref

ev.gti

"""
Explanation: Times are saved to the time attribute, GTIs to the gti attribute, MJDREF to the mjdref attribute, etc.
End of explanation
"""


ev.write("events.p", "pickle")
ev2 = EventList.read("events.p", "pickle")

np.allclose(ev2.time, ev.time), np.allclose(ev2.gti, ev.gti)

"""
Explanation: Roundtrip to pickle objects
It is possible to save and load eventlist objects using pickle.
End of explanation
"""


ev.write("events.hdf5", "hdf5")
ev3 = EventList.read("events.hdf5", "hdf5")
ev3.time[:10]

# Try the round trip again to verify that everything works

ev.write("events.ecsv", "ascii.ecsv")
ev4 = EventList.read("events.ecsv", "ascii.ecsv")
!cat events.ecsv
ev4.time[:10]

"""
Explanation: Roundtrip to Astropy-compatible formats
If the read and write methods receive a format which is not hea, ogip, or pickle, the event list is transformed into an Astropy Table object with the methods described above, and the readers and writers from the Table class are used instead. This allows to extend the save/load operations to a large number of formats, including hdf5 and enhanced CSV (ascii.ecsv).
Note that columns coming from the EVENTS (or equivalent) fits extension, those having the same length as time, when converting to astropy tables they become columns of the table. All the others, including gti, are treated as metadata.
Care should be used in selecting formats that preserve metadata. For example, simple CSV format loses all metadata, including MJDREF, GTIs etc.
End of explanation
"""


times = np.arange(3)
counts = np.floor(np.random.rand(3)*5)
lc = Lightcurve(times, counts, skip_checks=True, dt=1.)

lc.time, lc.counts

"""
Explanation: Transforming a Lightcurve into an EventList.
Event lists can be obtained from light curves, where the standard followed is as follows: as many events are created as the counts in the lightcurve at the time specified by time bins.
To demonstrate this, let us define a light curve.
End of explanation
"""


ev = EventList.from_lc(lc)
ev.time

"""
Explanation: Now, eventlist can be loaded by calling static from_lc() method.
End of explanation
"""


ev = EventList()
ev.simulate_times(lc)
ev.time

"""
Explanation: Simulating EventList from Lightcurve
An arguably better way is having proper random events, reproducing the initial light curve within the errors. Stingray does this by using the inverse CDF method, using the light curve as a binned probability distribution.
Please note that in this case we will have to create the EventList object before (in technical terms, simulate_times is not a static method.). See simulation tutorial for more details.
End of explanation
"""


dt = 1.
times = np.arange(50)
counts = np.floor(np.random.rand(50)*50000)
lc = Lightcurve(times, counts, skip_checks=True, dt=1.)

"""
Explanation: Creating a light curve from an EventList object
After simulating event list, the original light curve can be recovered. Let's demonstrate by creating a light curve.
End of explanation
"""


ev = EventList()
ev = ev.from_lc(lc)

ev.gti

"""
Explanation: Simulate an event list.
End of explanation
"""


lc_new = ev.to_lc(dt=1)

"""
Explanation: Recover original light curve curve using to_lc() method. Here, dt defines time resolution, tstart the starting time, and tseg the total time duration.
End of explanation
"""


plt.plot(lc.time, lc.counts,'r-', lc_new.counts, 'g-', drawstyle="steps-mid")
plt.xlabel('Times')
plt.ylabel('Counts')

"""
Explanation: Let us verify that this has worked properly, by comparing the input and output light curves
End of explanation
"""


plt.plot(lc.time, lc.counts - lc_new.counts, 'g-', drawstyle="steps-mid")
plt.xlabel('Times')
plt.ylabel('Counts')

"""
Explanation: ... and their difference
End of explanation
"""


spectrum = [[1, 2, 3, 4, 5, 6],[1000, 2040, 1000, 3000, 4020, 2070]]

ev = EventList(time=np.sort(np.random.uniform(0, 1000, 12)))
ev.simulate_energies(spectrum)

ev.energy

"""
Explanation: As can be seen from the figure above, the recovered light curve is aligned with the original light curve.
Simulating Energies
In order to simulate photon energies, a spectral distribution needs to be passed.
The spectrum input is a two-dimensional array, with the energies in keV in the first dimension and the number of counts in the second. The count array will be normalized before the simulation: the raw counts do not matter, but only the ratio of the counts in each bin to the total.
Again, the energies are simulated using an inverse CDF method.
End of explanation
"""


ev1 = EventList(time=[1,2,3], gti=[[0.5, 3.5]])
ev2 = EventList(time=[4,5], gti=[[3.5, 5.5]])
ev = ev1.join(ev2)
ev.time, ev.gti

ev1 = EventList(time=[1,2,3], gti=[[0.5, 3.5]])
ev2 = EventList(time=[1.2, 3.3, 5.6], gti=[[0.6, 7.8]])
ev = ev1.join(ev2)
ev.time, ev.gti

"""
Explanation: Joining EventLists
Two event lists can also be joined together. If the GTI do not overlap, the event times and GTIs are appended. Otherwise, the GTIs are crossed (i.e., only the overlapping parts are saved) and the events merged together.
End of explanation
"""


