# file: lotsofplots3.py
# Modified version of LOSC tutorial lotsofplots.py from
# https://losc.ligo.org/tutorial05/
# Small adjustments by Eric Myers <EricMyers47@gmail.com> on 02Aug2018
######################################################################
 
import numpy as np
import h5py
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import readligo as rl

#---------------------
# Read in strain data
#---------------------
#fileName = 'L-L1_LOSC_4_V1-1126318080-4096.hdf5'
fileName = '../LIGOdata/H-H2_LOSC_4_V1-842690560-4096.hdf5'
strain, time, channel_dict = rl.loaddata(fileName)
print("strain contains ",len(strain)," samples")
print("time contains ", len(time)," items")
ts = time[1] - time[0] #-- Time between samples
fs = int(1.0 / ts)          #-- Sampling frequency

#---------------------------------------------------------
# Find a good segment, get first 16 seconds of data
#---------------------------------------------------------
segList = rl.dq_channel_to_seglist(channel_dict['DEFAULT'], fs)
length = 16  # seconds
strain_seg = strain[segList[0]][0:(length*fs)]
time_seg = time[segList[0]][0:(length*fs)]

fig = plt.figure(figsize=(8.5,9))
fig.subplots_adjust(wspace=0.3, hspace=0.5)

#---------------------
# 1) Plot the time series
#----------------------

plot321 = plt.subplot(321)
plot321.title.set_text('1) Strain Time Series')
plt.plot(time_seg - time_seg[0], strain_seg)
plt.xlabel('Time since GPS ' + str(time_seg[0]))
plt.ylabel('Strain')


#------------------------------------------
# 2) Apply a Blackman Window, and plot the FFT
#------------------------------------------
window = np.blackman(strain_seg.size)
windowed_strain = strain_seg*window
freq = np.fft.rfftfreq(len(windowed_strain))*fs
## -EAM divide by EACH frequency, not the sampling frequency
freq_domain = np.fft.rfft(windowed_strain) / (freq+0.0000000001)


plot322 = plt.subplot(322)
plot322.title.set_text('2) Blackman Windowed FFT')
plt.loglog( freq, abs(freq_domain) )
plt.axis([10, fs/2.0, 1e-24, 1e-18])
plt.grid(True)
plt.xlabel('Freq (Hz)')
plt.ylabel('Strain / Hz')


#----------------------------------
# 3) Make PSD for first chunk of data
#----------------------------------
plot323 = plt.subplot(323)
plot323.title.set_text('3) Power Spectral Density (PSD)')
                       
Pxx, freqs = mlab.psd(strain_seg, Fs = fs, NFFT=fs)
plt.loglog(freqs, Pxx)
plt.axis([10, 2000, 1e-46, 1e-36])
plt.grid(True)
plt.ylabel('PSD')
plt.xlabel('Freq (Hz)')

#-------------------------
# 4) Plot the ASD
#-------------------------------
plot324 = plt.subplot(324)
plot324.title.set_text('4) Amplitude Spectral Density (ASD)')
## -EAM divide by freq of bin to get units of strain/sqrt(Hz)
plt.loglog(freqs, np.sqrt(Pxx/(freqs+0.00001)))
plt.axis([10, 2000, 1e-24, 1e-18])
plt.grid(True)
plt.xlabel('Freq (Hz)')
plt.ylabel('ASD [Strain / Hz$^{1/2}$]')

#--------------------
# 5) Make a spectrogram
#-------------------
NFFT = 1024
window = np.blackman(NFFT)

plot325 = plt.subplot(325)
plot325.title.set_text('5) Spectrogram')
spec_power, freqs, bins, im = plt.specgram(strain_seg, NFFT=NFFT, Fs=fs, 
                                   window=window)
plt.xlabel('Time (s)')
plt.ylabel('Freq (Hz)')

#------------------------------------------
# Renormalize by average power in freq. bin
#-----------------------------------------
med_power = np.zeros(freqs.shape)
norm_spec_power = np.zeros(spec_power.shape)
index = 0
for row in spec_power:
    med_power[index] = np.median(row)
    norm_spec_power[index] = row / med_power[index]
    index += 1

plot326 = plt.subplot(326)
plot326.pcolormesh(bins, freqs, np.log10(norm_spec_power))
plot326.title.set_text('6) Spectrogram rescaled by average power')
plt.xlabel('Time (s)')
plt.ylabel('Freq (Hz)')

plt.savefig("lotsofplots3.png")
print("Exit plot window to exit script")
plt.show()
print("DONE.")