import readligo as rl
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
plt.ion()

# -- Set the times marking the start and stop of data downloaded.
dstart = 843894784
dstop = dstart + 4096

# -----------------------------------------
# Plot ASDs for data which passes CBC CAT 2
# ----------------------------------------

# -- Create a CBC High mass post CAT 2 segment list
seglist = rl.getsegs(dstart, dstop, 'H1', flag='CBCHIGH_CAT2')

# -- Load the data for each of the first 3 segments
for (start, stop) in seglist[0:3]:
    strain, meta, dq = rl.getstrain(start, stop, 'H1')

    #-- For each segment, plot the ASD    
    fs = int(1.0/meta['dt'])
    Pxx, freqs = mlab.psd(strain, Fs = fs, NFFT=4*fs)
    plt.loglog(freqs, np.sqrt(Pxx), label=str(start))

# -- Make the plot look nice
plt.axis([10, 2000, 1e-23, 1e-18])
plt.grid('on')
plt.xlabel('Freq (Hz)')
plt.ylabel('Strain / Hz$^{1/2}$')
plt.legend()

# -----------------------------------------------
# Now let's try finding a burst hardware injection
# ------------------------------------------------

# -- Define a segment list of hardware injection times
seglist = rl.getsegs(dstart, dstop, 'H1', flag='HW')

# -- As before, load the data for each segment
for (start, stop) in seglist:
    strain, meta, dq = rl.getstrain(start, stop, 'H1')

    # -- Make a spectragram of the data
    fs = int(1.0/meta['dt'])
    NFFT = fs//2       
    window = np.blackman(NFFT)
    spec_power, freqs, bins = mlab.specgram(strain, NFFT=NFFT, Fs=fs, 
                                               window=window)
    
    # -- Normalize each frequency bin to the median power in that bin
    for n in range(spec_power.shape[0]):
        spec_power[n] = spec_power[n] / np.median(spec_power[n])

    # -- Plot the spectragram
    plt.figure()
    ax = plt.subplot(111)
    ax.pcolormesh(bins, freqs, np.log10(spec_power))
    plt.xlabel('Time since GPS {0} (s)'.format(start))
    plt.ylabel('Freq (Hz)')
    plt.axis([0, stop-start, 0, fs//2])
plt.show()