%matplotlib inline

# load auxiliary libraries
import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits

# import stingray
import stingray

plt.style.use('seaborn-talk')

"""
Explanation: Dynamical Power Spectra (on real data)
End of explanation
"""


f = fits.open('emr_cleaned.fits')

"""
Explanation: All starts with a lightcurve..
Open the event file with astropy.io.fits
End of explanation
"""


dt = f[1].header['TIMEDEL']

"""
Explanation: The time resolution is stored in the header of the first extension under the Keyword TIMEDEL
End of explanation
"""


toa = f[1].data['Time']

"""
Explanation: The collumn TIME of the first extension stores the time of each event
End of explanation
"""


lc = stingray.Lightcurve.make_lightcurve(toa=toa, dt=dt)

lc.plot()

"""
Explanation: Let's create a Lightcurve from the Events time of arrival witha a given time resolution
End of explanation
"""


dynspec = stingray.DynamicalPowerspectrum(lc=lc, segment_size=16, norm='leahy')

"""
Explanation: DynamicPowerspectrum
Let's create a dynamic powerspectrum with the a segment size of 16s and the powers with a "leahy" normalization
End of explanation
"""


dynspec.dyn_ps

"""
Explanation: The dyn_ps attribute stores the power matrix, each column corresponds to the powerspectrum of each segment of the light curve
End of explanation
"""


extent = min(dynspec.time), max(dynspec.time), max(dynspec.freq), min(dynspec.freq)

plt.imshow(dynspec.dyn_ps, origin="lower left", aspect="auto", vmin=1.98, vmax=3.0,
           interpolation="none", extent=extent)
plt.colorbar()
plt.ylim(700,850)

print("The dynamical powerspectrun has {} frequency bins and {} time bins".format(len(dynspec.freq), len(dynspec.time)))

"""
Explanation: To plot the DynamicalPowerspectrum matrix, we use the attributes time  and freq to set the extend of the image axis. have a look at the documentation of matplotlib's imshow().
End of explanation
"""


print("The current frequency resolution is {}".format(dynspec.df))

"""
Explanation: # Rebinning in Frequency
End of explanation
"""


dynspec.rebin_frequency(df_new=2.0, method="average")

print("The new frequency resolution is {}".format(dynspec.df))

"""
Explanation: Let's rebin to a frequency resolution of 2 Hz and using the average of the power
End of explanation
"""


extent = min(dynspec.time), max(dynspec.time), min(dynspec.freq), max(dynspec.freq)
plt.imshow(dynspec.dyn_ps, origin="lower", aspect="auto", vmin=1.98, vmax=3.0,
           interpolation="none", extent=extent)
plt.colorbar()
plt.ylim(500, 1000)

extent = min(dynspec.time), max(dynspec.time), min(dynspec.freq), max(dynspec.freq)
plt.imshow(dynspec.dyn_ps, origin="lower", aspect="auto", vmin=2.0, vmax=3.0,
           interpolation="none", extent=extent)
plt.colorbar()
plt.ylim(700,850)

"""
Explanation: Let's see how the Dynamical Powerspectrum looks now
End of explanation
"""


print("The current time resolution is {}".format(dynspec.dt))

"""
Explanation: Rebin time
Let's try to improve the visualization by rebinnin our matrix in the time axis
End of explanation
"""


dynspec.rebin_time(dt_new=64.0, method="average")

print("The new time resolution is {}".format(dynspec.dt))

extent = min(dynspec.time), max(dynspec.time), min(dynspec.freq), max(dynspec.freq)
plt.imshow(dynspec.dyn_ps, origin="lower", aspect="auto", vmin=2.0, vmax=3.0,
           interpolation="none", extent=extent)
plt.colorbar()
plt.ylim(700,850)

"""
Explanation: Let's rebin to a time resolution of 64 s
End of explanation
"""


tracing = dynspec.trace_maximum(min_freq=755, max_freq=782)

"""
Explanation: Trace maximun
Let's use the method trace_maximum() to find the index of the maximum on each powerspectrum in a certain frequency range. For example, between 755 and 782Hz)
End of explanation
"""


plt.plot(dynspec.time, dynspec.freq[tracing], color='red', alpha=1)
plt.show()

"""
Explanation: This is how the trace function looks like
End of explanation
"""


extent = min(dynspec.time), max(dynspec.time), min(dynspec.freq), max(dynspec.freq)
plt.imshow(dynspec.dyn_ps, origin="lower", aspect="auto", vmin=2.0, vmax=3.0,
           interpolation="none", extent=extent, alpha=0.7)
plt.colorbar()
plt.ylim(740,800)
plt.plot(dynspec.time, dynspec.freq[tracing], color='red', lw=3, alpha=1)
plt.show()

"""
Explanation: Let's plot it on top of the dynamic spectrum
End of explanation
"""


